From 898d5fad02ec299161f0f8d5c8a60e52a7897445 Mon Sep 17 00:00:00 2001 From: drowzeys <264866615+drowzeys@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:47:37 -0400 Subject: [PATCH 1/3] lfm3_omni: LFM3-Omni-2B on Metal: speech in, interleaved text and 24 kHz speech out A standalone family in Drama 3's shape: the four Liquid stacks of the bf16 release (speech encoder and projector, LFM3 backbone, depth decoder, Liquid codec decoder) over one metallib, the log-mel front end on the device, each codebook's pick on the GPU, and the codec streamed between backbone rows. `zig build lfm3-omni-generate` makes one turn from a WAV, writes and scores stage dumps against a reference's, and serves POST /v1/audio/chat (server-sent text deltas and 16-bit 24 kHz audio) and GET /health, one request at a time. Products of many rows take the tensor unit through the prefill GEMM's blocks; other chips a simdgroup-matrix product. Host tests cover the mel front end, the prompt, the streamed overlap-add, the synthesis basis and the Metal argument layouts. --- zig/build/lfm3_omni_metal.zig | 51 ++ zig/build/metal_targets.zig | 1 + zig/kernels/metal/lfm3_omni/om_common.h | 56 +++ zig/kernels/metal/lfm3_omni/omni.metal | 474 +++++++++++++++++++ zig/kernels/metal/lfm3_omni/omni_gemm.metal | 235 +++++++++ zig/kernels/metal/lfm3_omni/omni_nax.metal | 39 ++ zig/kernels/metal/lfm3_omni/omni_rows.metal | 271 +++++++++++ zig/src/families/lfm3_omni/common.zig | 231 +++++++++ zig/src/families/lfm3_omni/config.zig | 147 ++++++ zig/src/families/lfm3_omni/generate.zig | 330 +++++++++++++ zig/src/families/lfm3_omni/mel.zig | 108 +++++ zig/src/families/lfm3_omni/metal/engine.zig | 468 ++++++++++++++++++ zig/src/families/lfm3_omni/metal/gpu.zig | 286 +++++++++++ zig/src/families/lfm3_omni/metal/kernels.zig | 296 ++++++++++++ zig/src/families/lfm3_omni/metal/load.zig | 159 +++++++ zig/src/families/lfm3_omni/metal/reply.zig | 203 ++++++++ zig/src/families/lfm3_omni/metal/stack.zig | 310 ++++++++++++ zig/src/families/lfm3_omni/omni_metal.zig | 20 + zig/src/families/lfm3_omni/prompt.zig | 27 ++ zig/src/families/lfm3_omni/serve.zig | 237 ++++++++++ zig/src/families/lfm3_omni/spec.zig | 45 ++ 21 files changed, 3994 insertions(+) create mode 100644 zig/build/lfm3_omni_metal.zig create mode 100644 zig/kernels/metal/lfm3_omni/om_common.h create mode 100644 zig/kernels/metal/lfm3_omni/omni.metal create mode 100644 zig/kernels/metal/lfm3_omni/omni_gemm.metal create mode 100644 zig/kernels/metal/lfm3_omni/omni_nax.metal create mode 100644 zig/kernels/metal/lfm3_omni/omni_rows.metal create mode 100644 zig/src/families/lfm3_omni/common.zig create mode 100644 zig/src/families/lfm3_omni/config.zig create mode 100644 zig/src/families/lfm3_omni/generate.zig create mode 100644 zig/src/families/lfm3_omni/mel.zig create mode 100644 zig/src/families/lfm3_omni/metal/engine.zig create mode 100644 zig/src/families/lfm3_omni/metal/gpu.zig create mode 100644 zig/src/families/lfm3_omni/metal/kernels.zig create mode 100644 zig/src/families/lfm3_omni/metal/load.zig create mode 100644 zig/src/families/lfm3_omni/metal/reply.zig create mode 100644 zig/src/families/lfm3_omni/metal/stack.zig create mode 100644 zig/src/families/lfm3_omni/omni_metal.zig create mode 100644 zig/src/families/lfm3_omni/prompt.zig create mode 100644 zig/src/families/lfm3_omni/serve.zig create mode 100644 zig/src/families/lfm3_omni/spec.zig diff --git a/zig/build/lfm3_omni_metal.zig b/zig/build/lfm3_omni_metal.zig new file mode 100644 index 0000000000..549271ae93 --- /dev/null +++ b/zig/build/lfm3_omni_metal.zig @@ -0,0 +1,51 @@ +//! macOS: LFM3-Omni's metallib, the family module over Metal, `zig build lfm3-omni-generate` and its host tests. +const std = @import("std"); + +const dir = "zig/kernels/metal/lfm3_omni"; + +/// The kernels with safe math and precise fp32 functions (the CUDA port's rounding op for op) and the tensor-unit ones. +fn metallib(b: *std.Build, metal_std: []const u8) std.Build.LazyPath { + const link = b.addSystemCommand(&.{ "xcrun", "-sdk", "macosx", "metallib", "-o" }); + const out = link.addOutputFileArg("lfm3_omni.metallib"); + for ([_][]const u8{ "omni", "omni_gemm", "omni_rows", "omni_nax" }) |name| { + const cc = b.addSystemCommand(&.{ "xcrun", "-sdk", "macosx", "metal", b.fmt("-std={s}", .{metal_std}), "-fmetal-math-mode=safe", "-fmetal-math-fp32-functions=precise", "-c" }); + cc.addFileArg(b.path(b.fmt("{s}/{s}.metal", .{ dir, name }))); + cc.addFileInput(b.path(dir ++ "/om_common.h")); + cc.addFileInput(b.path("zig/kernels/metal/nax.h")); + cc.addFileInput(b.path("zig/kernels/metal/prefill/gemm_nax.metal")); + cc.addArg("-o"); + link.addFileArg(cc.addOutputFileArg(b.fmt("{s}.air", .{name}))); + } + return out; +} + +pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, metal_std: []const u8, test_step: *std.Build.Step) void { + const lib = metallib(b, metal_std); + const wf = b.addWriteFiles(); + _ = wf.addCopyFile(lib, "lfm3_omni.metallib"); + const bytes = b.createModule(.{ .root_source_file = wf.add("metallib.zig", "pub const bytes = @embedFile(\"lfm3_omni.metallib\");\n") }); + const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + core.addImport("tokenizer", tokenizer); + // the WAV reader / resampler is Drama 3's + const wav = b.createModule(.{ .root_source_file = b.path("zig/src/families/drama3/wav.zig"), .target = target, .optimize = optimize }); + const family = b.createModule(.{ + .root_source_file = b.path("zig/src/families/lfm3_omni/omni_metal.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "core", .module = core }, .{ .name = "wav", .module = wav }, .{ .name = "lfm3_omni_metallib", .module = bytes } }, + }); + const tool = b.createModule(.{ + .root_source_file = b.path("zig/src/families/lfm3_omni/generate.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "core", .module = core }, .{ .name = "lfm3_omni", .module = family } }, + }); + const exe = b.addExecutable(.{ .name = "lfm3-omni-generate", .root_module = tool }); + b.step("lfm3-omni-generate", "LFM3-Omni speech in, text and speech out on Metal, with its stage dumps").dependOn(&b.addInstallArtifact(exe, .{}).step); + const host = b.addRunArtifact(b.addTest(.{ .root_module = family })); + b.step("test-lfm3-omni-host", "LFM3-Omni's mel front end, prompt, overlap-add and Metal argument layouts on the host").dependOn(&host.step); + test_step.dependOn(&host.step); +} diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 8a16d7432e..8416f974ae 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -168,6 +168,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s }); @import("drama3_metal.zig").addNative(b, target, optimize, metal, engines, metal_std); @import("drama3_metal.zig").targets(b, target, optimize, metal, metal_std, test_step); + @import("lfm3_omni_metal.zig").targets(b, target, optimize, metal, metal_std, test_step); const programs = [_]struct { name: []const u8, path: []const u8, about: []const u8 }{ .{ .name = "tf-metal-tests", .path = "zig/tests/gpu_tests.zig", .about = "GPU runtime tests" }, diff --git a/zig/kernels/metal/lfm3_omni/om_common.h b/zig/kernels/metal/lfm3_omni/om_common.h new file mode 100644 index 0000000000..4e8c46e1a1 --- /dev/null +++ b/zig/kernels/metal/lfm3_omni/om_common.h @@ -0,0 +1,56 @@ +// LFM3-Omni's Metal helpers: bf16 as ushort bits (round to nearest even), erf for the exact GELU, group reductions. +#pragma once +#include +using namespace metal; + +typedef ushort bf; + +inline float f(bf x) { return as_type(uint(x) << 16); } + +inline bf b(float x) { + const uint u = as_type(x); + if ((u & 0x7FFFFFFFu) > 0x7F800000u) return bf((u >> 16) | 0x40u); + return bf((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16); +} + +inline float rb(float x) { return f(b(x)); } + +inline float sigm(float x) { return 1.0f / (1.0f + precise::exp(-x)); } + +// erf for the exact GELU: Numerical Recipes' erfcc Chebyshev fit, and the Taylor series near zero where it cancels. +inline float erf_f(float xf) { + const float x = fabs(xf); + float r; + if (x < 0.25f) { + const float t = x * x; + r = 1.1283791670955126f * x * (1.0f - t * (1.0f / 3.0f - t * (0.1f - t * (1.0f / 42.0f - t * (1.0f / 216.0f))))); + } else { + const float t = 1.0f / (1.0f + 0.5f * x); + const float e = t * precise::exp(-x * x - 1.26551223f + t * (1.00002368f + t * (0.37409196f + t * (0.09678418f + t * (-0.18628806f + t * (0.27886807f + t * (-1.13520398f + t * (1.48851587f + t * (-0.82215223f + t * 0.17087277f))))))))); + r = 1.0f - e; + } + return xf < 0.0f ? -r : r; +} + +template +inline float group_sum(float v, threadgroup float* red, uint tid) { + v = simd_sum(v); + threadgroup_barrier(mem_flags::mem_threadgroup); + if ((tid & 31) == 0) red[tid >> 5] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (uint i = 0; i < NT / 32; ++i) t += red[i]; + return t; +} + +template +inline float group_max(float v, threadgroup float* red, uint tid) { + v = simd_max(v); + threadgroup_barrier(mem_flags::mem_threadgroup); + if ((tid & 31) == 0) red[tid >> 5] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = -INFINITY; + for (uint i = 0; i < NT / 32; ++i) t = max(t, red[i]); + return t; +} + diff --git a/zig/kernels/metal/lfm3_omni/omni.metal b/zig/kernels/metal/lfm3_omni/omni.metal new file mode 100644 index 0000000000..01d20e7ee6 --- /dev/null +++ b/zig/kernels/metal/lfm3_omni/omni.metal @@ -0,0 +1,474 @@ +// LFM3-Omni's Metal ops for its four Liquid stacks, rounding to bf16 wherever the bf16 reference rounds. +#include "om_common.h" + +// ---- norms + +struct NormArgs { + int D, ld, ldy; + float eps; + int mode, ldw, has_r, has_pw; +}; + +// RMSNorm a group a row, X = bf16(X + R) first (R normed by PW if given); mode 0 rounds x * rsqrt, mode 1 once. +template +inline void rmsnorm_body(device bf* X, device const bf* R, device const bf* Wb, device bf* Y, device const bf* PW, constant NormArgs& a, + uint row, uint tid, threadgroup float* red) { + device const bf* W = Wb + (long)row * a.ldw; + device bf* x = X + (long)row * a.ld; + device bf* y = Y + (long)row * a.ldy; + device const bf* r = R + (long)row * a.ld; + float inv_r = 0.0f; + if (a.has_r && a.has_pw) { + float sr = 0.0f; + for (int i = tid; i < a.D; i += NT) sr += f(r[i]) * f(r[i]); + inv_r = precise::rsqrt(group_sum(sr, red, tid) / float(a.D) + a.eps); + } + float ss = 0.0f; + for (int i = tid; i < a.D; i += NT) { + float v = f(x[i]); + if (a.has_r) { + const float add = a.has_pw ? rb(f(PW[i]) * rb(f(r[i]) * inv_r)) : f(r[i]); + v = rb(v + add); + x[i] = b(v); + } + ss += v * v; + } + const float inv = precise::rsqrt(group_sum(ss, red, tid) / float(a.D) + a.eps); + for (int i = tid; i < a.D; i += NT) { + const float v = f(x[i]) * inv; + y[i] = a.mode == 0 ? b(f(W[i]) * rb(v)) : b(v * f(W[i])); + } +} + +#define OM_RMSNORM(NAME, NT) \ + [[max_total_threads_per_threadgroup(NT)]] kernel void NAME( \ + device bf* X [[buffer(0)]], device const bf* R [[buffer(1)]], device const bf* Wb [[buffer(2)]], device bf* Y [[buffer(3)]], \ + device const bf* PW [[buffer(4)]], constant NormArgs& a [[buffer(5)]], uint row [[threadgroup_position_in_grid]], \ + uint tid [[thread_index_in_threadgroup]]) { \ + threadgroup float red[NT / 32]; \ + rmsnorm_body(X, R, Wb, Y, PW, a, row, tid, red); \ + } +OM_RMSNORM(om_rmsnorm, 256) +OM_RMSNORM(om_rmsnorm_1k, 1024) + +struct HeadNormArgs { + int T, H, hd, ld; + float eps; +}; + +// Per-head RMSNorm (mode 0 rounding) in place: one simdgroup a (row, head). +kernel void om_head_norm(device bf* X [[buffer(0)]], device const bf* W [[buffer(1)]], constant HeadNormArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const int gw = gid >> 5, lane = gid & 31; + if (gw >= a.T * a.H) return; + device bf* x = X + (long)(gw / a.H) * a.ld + (gw % a.H) * a.hd; + float ss = 0.0f; + for (int i = lane; i < a.hd; i += 32) ss += f(x[i]) * f(x[i]); + ss = simd_sum(ss); + const float inv = precise::rsqrt(ss / float(a.hd) + a.eps); + for (int i = lane; i < a.hd; i += 32) x[i] = b(f(W[i]) * rb(f(x[i]) * inv)); +} + +struct LayerNormArgs { + int D; + float eps; +}; + +// LayerNorm with bias over rows of D in fp32, one rounding (a group a row). +[[max_total_threads_per_threadgroup(256)]] kernel void om_layernorm(device const bf* X [[buffer(0)]], device const bf* W [[buffer(1)]], + device const bf* Bb [[buffer(2)]], device bf* Y [[buffer(3)]], + constant LayerNormArgs& a [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + threadgroup float red[8]; + device const bf* x = X + (long)row * a.D; + float s = 0.0f; + for (int i = tid; i < a.D; i += 256) s += f(x[i]); + const float mean = group_sum<256>(s, red, tid) / float(a.D); + float v = 0.0f; + for (int i = tid; i < a.D; i += 256) { + const float d = f(x[i]) - mean; + v += d * d; + } + const float rstd = precise::rsqrt(group_sum<256>(v, red, tid) / float(a.D) + a.eps); + for (int i = tid; i < a.D; i += 256) Y[(long)row * a.D + i] = b((f(x[i]) - mean) * rstd * f(W[i]) + f(Bb[i])); +} + +// ---- rotary, caches, positions + +struct RopeArgs { + int T, H, hd, ld, pos0, mode; +}; + +// Rotate-half RoPE in place at pos0 + t; mode 0 fp32 cos/sin, 1 bf16 cos/sin, 2 every product and sum rounded. +kernel void om_rope(device bf* X [[buffer(0)]], device const float* INV [[buffer(1)]], constant RopeArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const int half_ = a.hd >> 1; + const long i = gid; + if (i >= (long)a.T * a.H * half_) return; + const int p = int(i % half_), h = int((i / half_) % a.H), t = int(i / ((long)half_ * a.H)); + device bf* x = X + (long)t * a.ld + h * a.hd; + const float fr = float(a.pos0 + t) * INV[p]; + float c = precise::cos(fr), s = precise::sin(fr); + if (a.mode != 0) { + c = rb(c); + s = rb(s); + } + const float xa = f(x[p]), z = f(x[p + half_]); + if (a.mode == 2) { + x[p] = b(rb(xa * c) + rb(-z * s)); + x[p + half_] = b(rb(z * c) + rb(xa * s)); + } else { + x[p] = b(xa * c - z * s); + x[p + half_] = b(z * c + xa * s); + } +} + +struct CopyRowsArgs { + int T, n, lds, ldd; +}; + +kernel void om_copy_rows(device const bf* S [[buffer(0)]], device bf* D [[buffer(1)]], constant CopyRowsArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.T * a.n) return; + const int t = int(i / a.n), c = int(i % a.n); + D[(long)t * a.ldd + c] = S[(long)t * a.lds + c]; +} + +// n 32-bit words from the argument bytes (small host writes in order with the dispatches). +kernel void om_set_words(device uint* D [[buffer(0)]], constant uint* V [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) D[gid] = V[gid]; +} + +kernel void om_zero_words(device uint* D [[buffer(0)]], constant uint& n [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + if (gid < n) D[gid] = 0; +} + +kernel void om_copy_words(device const uint* S [[buffer(0)]], device uint* D [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) D[gid] = S[gid]; +} + +// ---- Liquid's gated short conv + +struct ConvArgs { + int ld, ldy, T, D, K; +}; + +// y = C * conv(B * x) over T rows of B | C | x, a thread a channel, STATE [D, K-1] updated. +kernel void om_short_conv(device const bf* BCx [[buffer(0)]], device const bf* Wc [[buffer(1)]], device bf* STATE [[buffer(2)]], + device bf* Y [[buffer(3)]], constant ConvArgs& a [[buffer(4)]], uint d [[thread_position_in_grid]]) { + if (int(d) >= a.D) return; + const int K = a.K, D = a.D; + float hist[8]; + float w[8]; + for (int j = 0; j < K; ++j) w[j] = f(Wc[d * K + j]); + for (int j = 0; j < K - 1; ++j) hist[j] = f(STATE[d * (K - 1) + j]); + for (int t = 0; t < a.T; ++t) { + device const bf* row = BCx + (long)t * a.ld; + const float bx = rb(f(row[d]) * f(row[2 * D + d])); + float acc = 0.0f; + for (int j = 0; j < K - 1; ++j) acc += w[j] * hist[j]; + acc += w[K - 1] * bx; + for (int j = 0; j + 1 < K - 1; ++j) hist[j] = hist[j + 1]; + if (K > 1) hist[K - 2] = bx; + Y[(long)t * a.ldy + d] = b(f(row[D + d]) * rb(acc)); + } + for (int j = 0; j < K - 1; ++j) STATE[d * (K - 1) + j] = b(hist[j]); +} + +// As om_short_conv for many rows, a thread a (row, channel); the state is left alone (om_conv_state follows). +kernel void om_short_conv_rows(device const bf* BCx [[buffer(0)]], device const bf* Wc [[buffer(1)]], device const bf* STATE [[buffer(2)]], + device bf* Y [[buffer(3)]], constant ConvArgs& a [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.T * a.D) return; + const int K = a.K, D = a.D; + const int t = int(i / D), d = int(i % D); + float acc = 0.0f; + for (int j = 0; j < K; ++j) { + const int src = t - (K - 1) + j; + float v; + if (src >= 0) { + device const bf* row = BCx + (long)src * a.ld; + v = rb(f(row[d]) * f(row[2 * D + d])); + } else { + v = f(STATE[d * (K - 1) + (K - 1 + src)]); + } + acc += f(Wc[d * K + j]) * v; + } + Y[(long)t * a.ldy + d] = b(f(BCx[(long)t * a.ld + D + d]) * rb(acc)); +} + +kernel void om_conv_state(device const bf* BCx [[buffer(0)]], device bf* STATE [[buffer(1)]], constant ConvArgs& a [[buffer(2)]], + uint d [[thread_position_in_grid]]) { + if (int(d) >= a.D) return; + const int K = a.K, D = a.D, T = a.T; + for (int j = 0; j < K - 1; ++j) { + const int src = T - (K - 1) + j; + if (src >= 0) { + device const bf* row = BCx + (long)src * a.ld; + STATE[d * (K - 1) + j] = b(f(row[d]) * f(row[2 * D + d])); + } else { + STATE[d * (K - 1) + j] = STATE[d * (K - 1) + j + T]; + } + } +} + +// ---- attention + +struct AttnArgs { + int ldq, ldkv, ldo, T, H, Hkv, hd, pos0, kv_len, window, causal; + float scale; + int has_gate, ldg; +}; + +// Prompt attention: 16 rows of a head a group over 32-key tiles, online softmax, then LFM3's gate G if given. +#define AT_Q 16 +#define AT_K 32 +[[max_total_threads_per_threadgroup(128)]] kernel void om_attention_tiles(device const bf* Q [[buffer(0)]], device const bf* Kc [[buffer(1)]], + device const bf* Vc [[buffer(2)]], device bf* O [[buffer(3)]], + device const bf* G [[buffer(4)]], constant AttnArgs& a [[buffer(5)]], + uint2 tg [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + threadgroup float qs[AT_Q * 128]; + threadgroup float ss[AT_Q * AT_K]; + threadgroup bf ks[AT_K * 130]; + threadgroup bf vs[AT_K * 130]; + const int hd = a.hd, hdp = hd + 2; + const int h = tg.y, kh = h / (a.H / a.Hkv); + const int r0 = tg.x * AT_Q; + const int warp = tid >> 5, lane = tid & 31; + const int dpl = hd >> 5; + for (int i = tid; i < AT_Q * hd; i += 128) { + const int r = i / hd, c = i % hd; + qs[i] = r0 + r < a.T ? f(Q[(long)(r0 + r) * a.ldq + h * hd + c]) : 0.0f; + } + const int last = min(a.T, r0 + AT_Q) - 1; + int hi = a.causal ? a.pos0 + last + 1 : a.kv_len; + hi = min(hi, a.kv_len); + const int lo = (a.causal && a.window > 0) ? max(0, a.pos0 + r0 + 1 - a.window) : 0; + float m[4], l[4], acc[4][4]; + for (int x = 0; x < 4; ++x) { + m[x] = -INFINITY; + l[x] = 0.0f; + for (int c = 0; c < 4; ++c) acc[x][c] = 0.0f; + } + for (int k0 = lo; k0 < hi; k0 += AT_K) { + const int nk = min(AT_K, hi - k0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = tid; i < nk * hd; i += 128) { + const int k = i / hd, c = i % hd; + const long g = (long)(k0 + k) * a.ldkv + kh * hd + c; + ks[k * hdp + c] = Kc[g]; + vs[k * hdp + c] = Vc[g]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = tid; i < AT_Q * AT_K; i += 128) { + const int r = i / AT_K, k = i % AT_K; + const int pos = a.pos0 + r0 + r, kp = k0 + k; + float s = -INFINITY; + if (k < nk && r0 + r < a.T && (!a.causal || kp <= pos) && (!(a.causal && a.window > 0) || kp > pos - a.window)) { + s = 0.0f; + threadgroup const float* qq = qs + r * hd; + threadgroup const bf* kr = ks + k * hdp; + for (int c = 0; c < hd; c += 2) s += qq[c] * f(kr[c]) + qq[c + 1] * f(kr[c + 1]); + s *= a.scale; + } + ss[i] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int x = 0; x < 4; ++x) { + const int r = warp * 4 + x; + threadgroup float* sr = ss + r * AT_K; + const float mx = simd_max(sr[lane]); + const float mn = max(m[x], mx); + if (mn == -INFINITY) continue; + const float scl = precise::exp(m[x] - mn); + const float p0 = precise::exp(sr[lane] - mn); + sr[lane] = p0; + const float ps = simd_sum(p0); + l[x] = l[x] * scl + ps; + m[x] = mn; + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int c = 0; c < 4; ++c) acc[x][c] *= scl; + // only the row's visible keys: a masked key's zero weight times a non-finite value is NaN + const int pos = a.pos0 + r0 + r; + const int k_lo = (a.causal && a.window > 0) ? max(0, pos - a.window + 1 - k0) : 0; + const int k_hi = a.causal ? min(nk, pos - k0 + 1) : nk; + for (int k = k_lo; k < k_hi; ++k) { + const float p = sr[k]; + threadgroup const bf* vr = vs + k * hdp; + for (int c = 0; c < 4; ++c) + if (c < dpl) acc[x][c] += p * f(vr[lane + 32 * c]); + } + } + } + for (int x = 0; x < 4; ++x) { + const int r = r0 + warp * 4 + x; + if (r >= a.T) continue; + const float inv = 1.0f / l[x]; + const float g = a.has_gate ? rb(sigm(f(G[(long)r * a.ldg + h]))) : 1.0f; + for (int c = 0; c < 4; ++c) + if (c < dpl) { + const bf o = b(acc[x][c] * inv); + O[(long)r * a.ldo + h * hd + lane + 32 * c] = a.has_gate ? b(f(o) * g) : o; + } + } +} + +// ---- elementwise + +struct SwigluArgs { + int T, F, inter; +}; + +// y = bf16(bf16(silu(a)) * c) from rows of 2F: a | c halves, or (inter) pairs a_j, c_j side by side. +kernel void om_swiglu(device const bf* AC [[buffer(0)]], device bf* Y [[buffer(1)]], constant SwigluArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.T * a.F) return; + const long t = i / a.F; + const int c = int(i % a.F); + const long ia = a.inter ? t * 2 * a.F + 2 * c : t * 2 * a.F + c; + const long ic = a.inter ? ia + 1 : ia + a.F; + const float x = f(AC[ia]); + Y[i] = b(rb(x / (1.0f + precise::exp(-x))) * f(AC[ic])); +} + +struct BiasActArgs { + int rows, cols, per_row, act, has_bias; +}; + +// Y = act(bf16(A + bias)) from an fp32 product; bias per column (per_row 0) or per row (1); act 1 SiLU, 2 exact GELU. +kernel void om_bias_act(device const float* A [[buffer(0)]], device const bf* Bias [[buffer(1)]], device bf* Y [[buffer(2)]], + constant BiasActArgs& a [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.rows * a.cols) return; + const int r = int(i / a.cols), c = int(i % a.cols); + float v = A[i] + (a.has_bias ? f(Bias[a.per_row ? r : c]) : 0.0f); + v = rb(v); + if (a.act == 1) v = v / (1.0f + precise::exp(-v)); + else if (a.act == 2) v = 0.5f * v * (1.0f + erf_f(v * 0.70710678118654752f)); + Y[i] = b(v); +} + +struct EmbedArgs { + int T, D, NQ, size; +}; + +// Y[t] = TABLE[IDS[t]]; ids < 0 leave the row alone. +kernel void om_embed(device const bf* TABLE [[buffer(0)]], device const int* IDS [[buffer(1)]], device bf* Y [[buffer(2)]], + constant EmbedArgs& a [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.T * a.D) return; + const int t = int(i / a.D), c = int(i % a.D); + const int id = IDS[t]; + if (id >= 0) Y[i] = TABLE[(long)id * a.D + c]; +} + +// A speech frame's input row: the fp32 sum over the codebooks of TABLE[code + q * size], one rounding. +kernel void om_embed_codes(device const bf* TABLE [[buffer(0)]], device const int* CODES [[buffer(1)]], device bf* Y [[buffer(2)]], + constant EmbedArgs& a [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.T * a.D) return; + const int t = int(i / a.D), c = int(i % a.D); + float s = 0.0f; + for (int q = 0; q < a.NQ; ++q) s += f(TABLE[((long)CODES[t * a.NQ + q] + (long)q * a.size) * a.D + c]); + Y[i] = b(s); +} + +// The depth decoder's slots: X[q] = bf16(DIN[q] + (MASKED[q] ? MASK : TABLES[q][CODES[q]])). +kernel void om_depth_in(device const bf* DIN [[buffer(0)]], device const bf* MASK [[buffer(1)]], device const bf* TABLES [[buffer(2)]], + device const int* CODES [[buffer(3)]], device const int* MASKED [[buffer(4)]], device bf* X [[buffer(5)]], + constant EmbedArgs& a [[buffer(6)]], uint gid [[thread_position_in_grid]]) { + const int i = gid; + if (i >= a.NQ * a.D) return; + const int q = i / a.D, c = i % a.D; + const float e = MASKED[q] ? f(MASK[c]) : f(TABLES[((long)q * a.size + CODES[q]) * a.D + c]); + X[i] = b(f(DIN[i]) + e); +} + +struct Conv2dArgs { + int CIN, H, W, C, H2, W2, silu; +}; + +// Stride-2 3x3 conv (padding 1, bias, rounded) over [C, H, W]; CIN 1 reads plane 0, else depthwise; then SiLU. +kernel void om_conv2d_s2(device const bf* X [[buffer(0)]], device const bf* Wt [[buffer(1)]], device const bf* Bias [[buffer(2)]], + device bf* Y [[buffer(3)]], constant Conv2dArgs& a [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.C * a.H2 * a.W2) return; + const int c = int(i / ((long)a.H2 * a.W2)), h = int((i / a.W2) % a.H2), w = int(i % a.W2); + device const bf* x = X + (long)(a.CIN == 1 ? 0 : c) * a.H * a.W; + float acc = 0.0f; + for (int kh = 0; kh < 3; ++kh) { + const int ih = 2 * h - 1 + kh; + if (ih < 0 || ih >= a.H) continue; + for (int kw = 0; kw < 3; ++kw) { + const int iw = 2 * w - 1 + kw; + if (iw < 0 || iw >= a.W) continue; + acc += f(Wt[c * 9 + kh * 3 + kw]) * f(x[(long)ih * a.W + iw]); + } + } + float v = rb(acc + f(Bias[c])); + if (a.silu) v = rb(v / (1.0f + precise::exp(-v))); + Y[i] = b(v); +} + +kernel void om_f32_bf16(device const float* X [[buffer(0)]], device bf* Y [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) Y[gid] = b(X[gid]); +} + +kernel void om_bf16_f32(device const bf* X [[buffer(0)]], device float* Y [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) Y[gid] = f(X[gid]); +} + +struct FsqArgs { + int n, NQ; +}; + +// The codec's finite scalar quantizer read back: CODES [n, NQ] to latents [n, NQ * 4], levels 8 7 6 6. +kernel void om_fsq_decode(device const int* CODES [[buffer(0)]], device bf* L [[buffer(1)]], constant FsqArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const int i = gid; + if (i >= a.n * a.NQ) return; + const int t = i / a.NQ, q = i % a.NQ; + const int levels[4] = {8, 7, 6, 6}; + const int basis[4] = {1, 8, 56, 336}; + const int idx = CODES[t * a.NQ + q]; + for (int j = 0; j < 4; ++j) { + const int hf = levels[j] / 2; + const int digit = (idx / basis[j]) % levels[j] - hf; + L[(long)t * a.NQ * 4 + q * 4 + j] = b(float(digit) / float(hf)); + } +} + +struct RepeatArgs { + int n, D, rep; +}; + +kernel void om_repeat_rows(device const bf* X [[buffer(0)]], device bf* Y [[buffer(1)]], constant RepeatArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.n * a.rep * a.D) return; + const long r = i / a.D; + Y[i] = X[(r / a.rep) * a.D + i % a.D]; +} + +struct PolarArgs { + int T, NB; +}; + +// The codec head's rows (log-magnitude | phase) to a real | imaginary spectrum in fp32. +kernel void om_polar(device const bf* X [[buffer(0)]], device float* S [[buffer(1)]], constant PolarArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.T * a.NB) return; + const long t = i / a.NB; + const int k = int(i % a.NB); + const float m = precise::exp(f(X[t * 2 * a.NB + k])), ph = f(X[t * 2 * a.NB + a.NB + k]); + S[t * 2 * a.NB + k] = m * precise::cos(ph); + S[t * 2 * a.NB + a.NB + k] = m * precise::sin(ph); +} diff --git a/zig/kernels/metal/lfm3_omni/omni_gemm.metal b/zig/kernels/metal/lfm3_omni/omni_gemm.metal new file mode 100644 index 0000000000..52c66a4373 --- /dev/null +++ b/zig/kernels/metal/lfm3_omni/omni_gemm.metal @@ -0,0 +1,235 @@ +// LFM3-Omni's simdgroup products for chips without tensor units, the depth decoder's picks and the log-mel front end. +#include "om_common.h" + +inline void better(thread float& best, thread int& at, float v, int i) { + if (v > best || (v == best && i < at)) { + best = v; + at = i; + } +} + +// ---- C = A B (either transposed) on the simdgroup matrix units from fp32 tiles, fp32 sums; a batch on z + +struct GemmArgs { + int M, N, K, lda, ldb, ldc, ta, tb; + long sa, sb, sc; +}; + +template +inline float ld_in(device const TI* p, long i); +template <> +inline float ld_in(device const bf* p, long i) { + return f(p[i]); +} +template <> +inline float ld_in(device const float* p, long i) { + return p[i]; +} + +template +inline void st_out(device TO* p, long i, float v); +template <> +inline void st_out(device bf* p, long i, float v) { + p[i] = b(v); +} +template <> +inline void st_out(device float* p, long i, float v) { + p[i] = v; +} + +template +inline void gemm_body(device const TI* A, device const TI* B, device TO* C, constant GemmArgs& a, uint3 tg, uint tid, uint sg, + uint lane, threadgroup float* As, threadgroup float* Bs, threadgroup float* Cs) { + constexpr int BM = 64, BN = 64, BK = 16; + A += (long)tg.z * a.sa; + B += (long)tg.z * a.sb; + C += (long)tg.z * a.sc; + const int m0 = tg.y * BM, n0 = tg.x * BN; + const int wm = (sg >> 1) * 32, wn = (sg & 1) * 32; + simdgroup_float8x8 acc[4][4]; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) acc[i][j] = simdgroup_float8x8(0.0f); + for (int k0 = 0; k0 < a.K; k0 += BK) { + // A tile [BM][BK], B tile [BK][BN]: 1024 values each, 8 a thread + for (int i = tid; i < BM * BK; i += 128) { + int r, k; + if (a.ta) { + r = i % BM; + k = i / BM; + } else { + r = i / BK; + k = i % BK; + } + const int gm = m0 + r, gk = k0 + k; + float v = 0.0f; + if (gm < a.M && gk < a.K) v = ld_in(A, a.ta ? (long)gk * a.lda + gm : (long)gm * a.lda + gk); + As[r * BK + k] = v; + } + for (int i = tid; i < BK * BN; i += 128) { + int k, c; + if (a.tb) { + k = i % BK; + c = i / BK; + } else { + k = i / BN; + c = i % BN; + } + const int gk = k0 + k, gn = n0 + c; + float v = 0.0f; + if (gn < a.N && gk < a.K) v = ld_in(B, a.tb ? (long)gn * a.ldb + gk : (long)gk * a.ldb + gn); + Bs[k * BN + c] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int kk = 0; kk < BK; kk += 8) { + simdgroup_float8x8 am[4], bm[4]; + for (int i = 0; i < 4; ++i) simdgroup_load(am[i], As + (wm + i * 8) * BK + kk, BK); + for (int j = 0; j < 4; ++j) simdgroup_load(bm[j], Bs + kk * BN + wn + j * 8, BN); + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_multiply_accumulate(acc[i][j], am[i], bm[j], acc[i][j]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + threadgroup float* cs = Cs + sg * 64; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) { + simdgroup_store(acc[i][j], cs, 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int e = lane; e < 64; e += 32) { + const int gm = m0 + wm + i * 8 + e / 8, gn = n0 + wn + j * 8 + e % 8; + if (gm < a.M && gn < a.N) st_out(C, (long)gm * a.ldc + gn, cs[e]); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } +} + +#define OM_GEMM(NAME, TI, TO) \ + [[max_total_threads_per_threadgroup(128)]] kernel void NAME(device const TI* A [[buffer(0)]], device const TI* B [[buffer(1)]], \ + device TO* C [[buffer(2)]], constant GemmArgs& a [[buffer(3)]], \ + uint3 tg [[threadgroup_position_in_grid]], \ + uint tid [[thread_index_in_threadgroup]], \ + uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]]) { \ + threadgroup float As[64 * 16]; \ + threadgroup float Bs[16 * 64]; \ + threadgroup float Cs[4 * 64]; \ + gemm_body(A, B, C, a, tg, tid, sg, lane, As, Bs, Cs); \ + } +OM_GEMM(om_gemm_bb, bf, bf) +OM_GEMM(om_gemm_bf, bf, float) +OM_GEMM(om_gemm_ff, float, float) + +// ---- the depth decoder's picks on the device + +inline ulong mix64(ulong z) { + z += 0x9E3779B97F4A7C15ul; + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ul; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBul; + return z ^ (z >> 31); +} + +struct PickArgs { + int p0, hn; + float penalty, temp; + ulong seed; + int frame, NQ, size, pad; +}; + +// Codebook ORDER[p0 + group]'s code after HIST's repetition penalty: greedy (lower id on ties) or a hashed Gumbel max. +[[max_total_threads_per_threadgroup(256)]] kernel void om_depth_pick(device const bf* LOG [[buffer(0)]], device const int* ORDER [[buffer(1)]], + device const int* HIST [[buffer(2)]], device int* CODES [[buffer(3)]], + device int* MASKED [[buffer(4)]], constant PickArgs& a [[buffer(5)]], + uint blk [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + threadgroup uchar seen[2048]; + threadgroup float bv[8]; + threadgroup int bi[8]; + const int q = ORDER[a.p0 + blk]; + for (int i = tid; i < a.size; i += 256) seen[i] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int j = tid; j < a.hn; j += 256) seen[HIST[j * a.NQ + q]] = 1; + threadgroup_barrier(mem_flags::mem_threadgroup); + float best = -INFINITY; + int at = 0x7fffffff; + for (int i = tid; i < a.size; i += 256) { + float v = f(LOG[(long)q * a.size + i]); + if (seen[i]) v = v < 0.0f ? v * a.penalty : v / a.penalty; + if (a.temp > 0.0f) { + const ulong h = mix64(a.seed ^ mix64(((ulong)a.frame << 32) ^ ((ulong)q << 16) ^ (ulong)i)); + const float u = (float(h >> 40) + 0.5f) * (1.0f / 16777216.0f); + v = v / a.temp - precise::log(-precise::log(u)); + } + better(best, at, v, i); + } + for (ushort o = 16; o > 0; o >>= 1) better(best, at, simd_shuffle_xor(best, o), simd_shuffle_xor(at, o)); + if ((tid & 31) == 0) { + bv[tid >> 5] = best; + bi[tid >> 5] = at; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + best = bv[0]; + at = bi[0]; + for (int w = 1; w < 8; ++w) better(best, at, bv[w], bi[w]); + CODES[q] = at; + MASKED[q] = 0; + } +} + +struct FrameEndArgs { + int NQ, eos, slot, push; +}; + +// A finished frame: any end-of-speech code makes it all EOS, else with push it joins HIST's ring at row `slot`. +kernel void om_frame_end(device int* CODES [[buffer(0)]], device int* HIST [[buffer(1)]], constant FrameEndArgs& a [[buffer(2)]], + uint tid [[thread_position_in_grid]]) { + if (tid != 0) return; + bool any = false; + for (int q = 0; q < a.NQ; ++q) any |= CODES[q] == a.eos; + if (any) + for (int q = 0; q < a.NQ; ++q) CODES[q] = a.eos; + else if (a.push) + for (int q = 0; q < a.NQ; ++q) HIST[a.slot * a.NQ + q] = CODES[q]; +} + +// ---- the log-mel front end: frames of the padded clip, |DFT|^2 by product, mel by product, log + +struct MelArgs { + int len, nf, nfft, hop; +}; + +kernel void om_mel_frames(device const float* X [[buffer(0)]], device const float* WIN [[buffer(1)]], device float* FR [[buffer(2)]], + constant MelArgs& a [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.nf * a.nfft) return; + const int t = int(i / a.nfft), n = int(i % a.nfft); + const int s = t * a.hop + n; + FR[i] = s < a.len ? X[s] * WIN[n] : 0.0f; +} + +struct PowerArgs { + int nf, bins; +}; + +kernel void om_power(device const float* S [[buffer(0)]], device float* P [[buffer(1)]], constant PowerArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.nf * a.bins) return; + const long t = i / a.bins; + const int k = int(i % a.bins); + const float re = S[t * 2 * a.bins + k], im = S[t * 2 * a.bins + a.bins + k]; + const float m = precise::sqrt(re * re + im * im); + P[i] = m * m; +} + +struct LogmelArgs { + int nf, mels; + float floor_; +}; + +kernel void om_logmel(device const float* M [[buffer(0)]], device float* OUT [[buffer(1)]], constant LogmelArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const long i = gid; + if (i >= (long)a.nf * a.mels) return; + const int t = int(i / a.mels), m = int(i % a.mels); + OUT[(long)m * a.nf + t] = (max(precise::log10(M[i]), a.floor_) + 4.0f) / 4.0f; +} diff --git a/zig/kernels/metal/lfm3_omni/omni_nax.metal b/zig/kernels/metal/lfm3_omni/omni_nax.metal new file mode 100644 index 0000000000..80d9f37f81 --- /dev/null +++ b/zig/kernels/metal/lfm3_omni/omni_nax.metal @@ -0,0 +1,39 @@ +// LFM3-Omni's bf16 products of many rows on the tensor unit, on the prefill GEMM's blocks; fp32 sums, one rounding. +#include "../prefill/gemm_nax.metal" + +struct NaxArgs { + int M, N, K, lda, ldb, ldd, kblocks, pad; + long sa, sb, sd; +}; + +template +inline void om_nax(const device bfloat* A, const device bfloat* B, device O* D, constant NaxArgs& a, uint3 tg, uint sg, + uint lane) { + constexpr int BM = 64, BN = 128, BK = 256, WN = 4; + const int row = int(tg.y) * BM + 32 * int(sg / WN), col = int(tg.x) * BN + 32 * int(sg % WN); + A += a.sa * tg.z + (TA ? long(row) : long(row) * a.lda); + B += a.sb * tg.z + (TB ? long(col) * a.ldb : long(col)); + D += a.sd * tg.z + long(row) * a.ldd + col; + const int rows = min(32, a.M - row), cols = min(32, a.N - col); + const short2 home = tfp::frag_home(ushort(lane)); + tfp::frag acc[2][2]; + if (rows == 32 && cols == 32) { + tfp::gemm_block(acc, A, B, a.lda, a.ldb, a.K, a.kblocks, rows, cols, home); + tfp::gemm_store(acc, D, a.ldd, rows, cols, home); + } else { + tfp::gemm_block(acc, A, B, a.lda, a.ldb, a.K, a.kblocks, rows, cols, home); + if (rows > 0 && cols > 0) tfp::gemm_store(acc, D, a.ldd, rows, cols, home); + } +} + +#define OM_NAX(NAME, O, TA, TB) \ + [[max_total_threads_per_threadgroup(256)]] kernel void NAME( \ + const device bfloat* A [[buffer(0)]], const device bfloat* B [[buffer(1)]], device O* D [[buffer(2)]], \ + constant NaxArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], \ + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { \ + om_nax(A, B, D, a, tg, sg, lane); \ + } +OM_NAX(om_nax_bb_nt, bfloat, false, true) +OM_NAX(om_nax_bf_nt, float, false, true) +OM_NAX(om_nax_bf_nn, float, false, false) +OM_NAX(om_nax_bf_tt, float, true, true) diff --git a/zig/kernels/metal/lfm3_omni/omni_rows.metal b/zig/kernels/metal/lfm3_omni/omni_rows.metal new file mode 100644 index 0000000000..b4898d783c --- /dev/null +++ b/zig/kernels/metal/lfm3_omni/omni_rows.metal @@ -0,0 +1,271 @@ +// LFM3-Omni's few-row path, each kernel rounding where the separate ops it fuses round (om_head_norm, om_rope, ...). +#include "om_common.h" + +struct RowsArgs { + int K, N, T, ldx, ldy, has_y, has_yf, epi; + long sw, sx, sy; +}; + +inline float4 bf4(uint2 w) { + return float4(as_type(w.x << 16), as_type(w.x & 0xFFFF0000u), as_type(w.y << 16), as_type(w.y & 0xFFFF0000u)); +} + +// y[r, n] = x[r] . W[n] for T <= R rows, 4 weight rows a simdgroup; epi 0 bf16 Y and / or fp32 YF, epi 2 SwiGLU pairs. +template +inline void rows_body(device const bf* W, device const bf* X, device bf* Y, device float* YF, constant RowsArgs& a, uint3 tg, + uint sg, uint lane) { + W += (long)tg.z * a.sw; + X += (long)tg.z * a.sx; + Y += (long)tg.z * a.sy; + const int nr = min(R, a.T); + const int K = a.K; + const int n0 = (tg.x * 4 + sg) * 4; + if (n0 >= a.N) return; + device const uint4* w[4]; + for (int j = 0; j < 4; ++j) w[j] = (device const uint4*)(W + (long)min(n0 + j, a.N - 1) * K); + float acc[R][4]; + for (int r = 0; r < R; ++r) + for (int j = 0; j < 4; ++j) acc[r][j] = 0.0f; + for (int c = lane; c < K / 8; c += 32) { + uint4 wv[4]; + for (int j = 0; j < 4; ++j) wv[j] = w[j][c]; + for (int r = 0; r < R; ++r) { + if (r >= nr) break; + const uint4 xv = ((device const uint4*)(X + (long)r * a.ldx))[c]; + const float4 x0 = bf4(xv.xy), x1 = bf4(xv.zw); + for (int j = 0; j < 4; ++j) { + const float4 w0 = bf4(wv[j].xy), w1 = bf4(wv[j].zw); + float s = acc[r][j]; + s = fma(w0.x, x0.x, s); + s = fma(w0.y, x0.y, s); + s = fma(w0.z, x0.z, s); + s = fma(w0.w, x0.w, s); + s = fma(w1.x, x1.x, s); + s = fma(w1.y, x1.y, s); + s = fma(w1.z, x1.z, s); + s = fma(w1.w, x1.w, s); + acc[r][j] = s; + } + } + } + for (int r = 0; r < R; ++r) { + if (r >= nr) break; + float v[4]; + for (int j = 0; j < 4; ++j) v[j] = simd_sum(acc[r][j]); + if (lane != 0) continue; + if (a.epi == 2) { + for (int p = 0; p < 2; ++p) { + const float x = rb(v[2 * p]), c = rb(v[2 * p + 1]); + Y[(long)r * a.ldy + n0 / 2 + p] = b(rb(x / (1.0f + precise::exp(-x))) * c); + } + continue; + } + for (int j = 0; j < 4; ++j) { + if (n0 + j >= a.N) break; + if (a.has_y) Y[(long)r * a.ldy + n0 + j] = b(v[j]); + if (a.has_yf) YF[(long)r * a.N + n0 + j] = v[j]; + } + } +} + +#define OM_ROWS(R) \ + [[max_total_threads_per_threadgroup(128)]] kernel void om_rows_##R( \ + device const bf* W [[buffer(0)]], device const bf* X [[buffer(1)]], device bf* Y [[buffer(2)]], device float* YF [[buffer(3)]], \ + constant RowsArgs& a [[buffer(4)]], uint3 tg [[threadgroup_position_in_grid]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]]) { \ + rows_body(W, X, Y, YF, a, tg, sg, lane); \ + } +OM_ROWS(1) +OM_ROWS(4) +OM_ROWS(8) +OM_ROWS(12) +OM_ROWS(16) + +struct QkArgs { + int T, H, Hkv, hd, ld, pos0, mode, nk; + float eps; +}; + +// q | k | v | gate rows: q and k heads normed and rotated as om_rope, k and v written to cache row pos0 + t. +kernel void om_qk_prep(device bf* P [[buffer(0)]], device const bf* QN [[buffer(1)]], device const bf* KN [[buffer(2)]], + device const float* INV [[buffer(3)]], device bf* KC [[buffer(4)]], device bf* VC [[buffer(5)]], + constant QkArgs& a [[buffer(6)]], uint g [[threadgroup_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + threadgroup float sv[128]; + const int heads = a.H + a.Hkv; + const int t = g / heads, h = g % heads; + const bool q = h < a.H; + const int kh = h - a.H; + device bf* x = P + (long)t * a.ld + (q ? h * a.hd : a.H * a.hd + kh * a.hd); + device const bf* w = q ? QN : KN; + float ss = 0.0f; + for (int i = lane; i < a.hd; i += 32) ss += f(x[i]) * f(x[i]); + ss = simd_sum(ss); + const float inv = precise::rsqrt(ss / float(a.hd) + a.eps); + for (int i = lane; i < a.hd; i += 32) sv[i] = rb(f(w[i]) * rb(f(x[i]) * inv)); + simdgroup_barrier(mem_flags::mem_threadgroup); + const int half_ = a.hd >> 1; + device bf* out = q ? x : KC + (long)(a.pos0 + t) * a.nk + kh * a.hd; + for (int p = lane; p < half_; p += 32) { + const float fr = float(a.pos0 + t) * INV[p]; + float c = precise::cos(fr), s = precise::sin(fr); + if (a.mode != 0) { + c = rb(c); + s = rb(s); + } + const float xa = sv[p], z = sv[p + half_]; + if (a.mode == 2) { + out[p] = b(rb(xa * c) + rb(-z * s)); + out[p + half_] = b(rb(z * c) + rb(xa * s)); + } else { + out[p] = b(xa * c - z * s); + out[p + half_] = b(z * c + xa * s); + } + } + if (!q) { + device const bf* v = P + (long)t * a.ld + a.H * a.hd + a.nk + kh * a.hd; + for (int i = lane; i < a.hd; i += 32) VC[(long)(a.pos0 + t) * a.nk + kh * a.hd + i] = v[i]; + } +} + +struct AttendArgs { + int ldkv, H, Hkv, hd, pos, kv_len, window, causal, has_gate, gate_off; + float scale; +}; + +#define OM_SPAN 4096 + +inline float dot8(uint4 k, threadgroup const float* q) { + const float4 k0 = bf4(k.xy), k1 = bf4(k.zw); + return k0.x * q[0] + k0.y * q[1] + k0.z * q[2] + k0.w * q[3] + k1.x * q[4] + k1.y * q[5] + k1.z * q[6] + k1.w * q[7]; +} + +// One query row over the cache, a group of 256 a head, then LFM3's sigmoid gate; hd % 8 == 0, hd <= 128. +[[max_total_threads_per_threadgroup(256)]] kernel void om_attend(device const bf* Q [[buffer(0)]], device const bf* Kc [[buffer(1)]], + device const bf* Vc [[buffer(2)]], device bf* O [[buffer(3)]], + constant AttendArgs& a [[buffer(4)]], uint h [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + threadgroup float sc[OM_SPAN]; + threadgroup float qs[128]; + threadgroup float part[2048]; + threadgroup float red[8]; + const int kh = h / (a.H / a.Hkv); + const int hi = a.causal ? min(a.pos + 1, a.kv_len) : a.kv_len; + const int lo = (a.causal && a.window > 0) ? max(0, a.pos + 1 - a.window) : 0; + const int chunks = a.hd / 8; + if (int(tid) < a.hd) qs[tid] = f(Q[h * a.hd + tid]); + threadgroup_barrier(mem_flags::mem_threadgroup); + float m = -INFINITY; + for (int k = lo + tid; k < hi; k += 256) { + device const uint4* kr = (device const uint4*)(Kc + (long)k * a.ldkv + kh * a.hd); + float s = 0.0f; + for (int c = 0; c < chunks; ++c) s += dot8(kr[c], qs + 8 * c); + s *= a.scale; + sc[k - lo] = s; + m = max(m, s); + } + m = group_max<256>(m, red, tid); + float l = 0.0f; + for (int k = lo + tid; k < hi; k += 256) { + const float e = precise::exp(sc[k - lo] - m); + sc[k - lo] = e; + l += e; + } + l = group_sum<256>(l, red, tid); + const int groups = 256 / chunks; + const int dc = tid % chunks, kg = tid / chunks; + float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + if (kg < groups) { + for (int k = lo + kg; k < hi; k += groups) { + const float p = sc[k - lo]; + const uint4 v = ((device const uint4*)(Vc + (long)k * a.ldkv + kh * a.hd))[dc]; + const float4 v0 = bf4(v.xy), v1 = bf4(v.zw); + acc[0] += p * v0.x; + acc[1] += p * v0.y; + acc[2] += p * v0.z; + acc[3] += p * v0.w; + acc[4] += p * v1.x; + acc[5] += p * v1.y; + acc[6] += p * v1.z; + acc[7] += p * v1.w; + } + for (int j = 0; j < 8; ++j) part[kg * a.hd + dc * 8 + j] = acc[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (int(tid) < a.hd) { + float o = 0.0f; + for (int g = 0; g < groups; ++g) o += part[g * a.hd + tid]; + bf ob = b(o * (1.0f / l)); + if (a.has_gate) ob = b(f(ob) * rb(sigm(f(Q[a.gate_off + h])))); + O[h * a.hd + tid] = ob; + } +} + +struct ArgmaxArgs2 { + int N, ns, per; +}; + +// Stage 1 of the greedy pick: each group's best of `per` logits (ids in SUP excluded), as (value, index) pairs. +[[max_total_threads_per_threadgroup(256)]] kernel void om_argmax_part(device const bf* X [[buffer(0)]], device const int* SUP [[buffer(1)]], + device float* PV [[buffer(2)]], device int* PI [[buffer(3)]], + constant ArgmaxArgs2& a [[buffer(4)]], uint g [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + threadgroup float bv[8]; + threadgroup int bi[8]; + float best = -INFINITY; + int at = 0x7fffffff; + const int lo = g * a.per, hi = min(a.N, lo + a.per); + for (int i = lo + tid; i < hi; i += 256) { + bool skip = false; + for (int j = 0; j < a.ns; ++j) skip |= SUP[j] == i; + if (skip) continue; + const float v = f(X[i]); + if (v > best || (v == best && i < at)) { + best = v; + at = i; + } + } + for (ushort o = 16; o > 0; o >>= 1) { + const float ov = simd_shuffle_xor(best, o); + const int oi = simd_shuffle_xor(at, o); + if (ov > best || (ov == best && oi < at)) { + best = ov; + at = oi; + } + } + if ((tid & 31) == 0) { + bv[tid >> 5] = best; + bi[tid >> 5] = at; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + for (int w = 1; w < 8; ++w) + if (bv[w] > best || (bv[w] == best && bi[w] < at)) { + best = bv[w]; + at = bi[w]; + } + PV[g] = best; + PI[g] = at; + } +} + +// Stage 2: the first index of the largest over the n parts. +kernel void om_argmax_final(device const float* PV [[buffer(0)]], device const int* PI [[buffer(1)]], device int* OUT [[buffer(2)]], + constant int& n [[buffer(3)]], uint lane [[thread_index_in_threadgroup]]) { + float best = -INFINITY; + int at = 0x7fffffff; + for (int i = lane; i < n; i += 32) { + if (PV[i] > best || (PV[i] == best && PI[i] < at)) { + best = PV[i]; + at = PI[i]; + } + } + for (ushort o = 16; o > 0; o >>= 1) { + const float ov = simd_shuffle_xor(best, o); + const int oi = simd_shuffle_xor(at, o); + if (ov > best || (ov == best && oi < at)) { + best = ov; + at = oi; + } + } + if (lane == 0) OUT[0] = at; +} diff --git a/zig/src/families/lfm3_omni/common.zig b/zig/src/families/lfm3_omni/common.zig new file mode 100644 index 0000000000..71c404295c --- /dev/null +++ b/zig/src/families/lfm3_omni/common.zig @@ -0,0 +1,231 @@ +//! The reply loop's backend-neutral parts: sampling, sink, statistics, the codec's synthesis and the host's code pick. +const std = @import("std"); +const cfg = @import("config.zig"); + +pub const Sampling = struct { + /// 0 picks the most likely code; the model card's setting is 0.4. + audio_temperature: f32 = 0.4, + audio_top_k: usize = 0, + audio_repetition_penalty: f32 = 1.2, + audio_repetition_window: usize = 64, + /// Reveal the codebooks in order 0..11 instead of a random order (the reference's randperm). + identity_order: bool = false, + seed: u64 = 0, + /// Generated slots (a text token or a speech frame each). + max_new: usize = 512, +}; + +/// Where a reply goes as it is made. +pub const Sink = struct { + ctx: *anyopaque, + text: ?*const fn (*anyopaque, u32) anyerror!void = null, + frame: ?*const fn (*anyopaque, *const [cfg.codebooks]i32) anyerror!void = null, + pcm: ?*const fn (*anyopaque, []const f32) anyerror!void = null, + /// After every backbone row (the engine's `lastHidden` reads it). + hidden: ?*const fn (*anyopaque) anyerror!void = null, +}; + +/// When the codec runs while speech frames arrive: `first` frames for the first audio, then every `chunk`. +pub const Streaming = struct { + first: usize = 1, + chunk: usize = 8, + /// Decode audio at all (frames only otherwise). + audio: bool = true, +}; + +/// Teacher forcing: the reference reply's slots and frames are fed whatever the engine picks; the picks are scored. +pub const Forced = struct { + seq: []const u32, + frames: []const [cfg.codebooks]i32, + text_hits: usize = 0, + text_total: usize = 0, + code_hits: usize = 0, + code_total: usize = 0, + frame_hits: usize = 0, + frame_total: usize = 0, +}; + +pub const Stats = struct { + mel_ms: f64 = 0, + encode_ms: f64 = 0, + prefill_ms: f64 = 0, + first_text_ms: f64 = 0, + first_frame_ms: f64 = 0, + first_pcm_ms: f64 = 0, + total_ms: f64 = 0, + step_ms: f64 = 0, + depth_ms: f64 = 0, + codec_ms: f64 = 0, + text_tokens: usize = 0, + frames: usize = 0, + samples: usize = 0, + prompt_rows: usize = 0, + audio_tokens: usize = 0, +}; + +pub fn nowNs(io: std.Io) i128 { + return @intCast(std.Io.Clock.awake.now(io).toNanoseconds()); +} + +pub fn msSince(io: std.Io, t0: i128) f64 { + return @as(f64, @floatFromInt(nowNs(io) - t0)) / 1e6; +} + +pub const Mode = enum { text, audio, interleaved }; + +/// Frames the device's repetition window holds (the reference's default window). +pub const hist_cap = 64; + +/// The codec's spectrum to windowed frames, [bins re | bins im] x n_fft: irfft (norm "backward") times the window. +pub fn synthesisBasis(gpa: std.mem.Allocator, window: []const f32) ![]f32 { + const n = cfg.n_fft; + const nb = cfg.spec_bins; + const m = try gpa.alloc(f32, 2 * nb * n); + for (0..nb) |k| { + const c: f64 = if (k == 0 or k == nb - 1) 1.0 else 2.0; + for (0..n) |t| { + const a = 2.0 * std.math.pi * @as(f64, @floatFromInt((k * t) % n)) / @as(f64, @floatFromInt(n)); + const w: f64 = window[t]; + m[k * n + t] = @floatCast(c * @cos(a) / @as(f64, @floatFromInt(n)) * w); + m[(nb + k) * n + t] = @floatCast(-c * @sin(a) / @as(f64, @floatFromInt(n)) * w); + } + } + return m; +} + +/// The codec's overlap-add on the host: frames in, finished samples out (held back where later frames still add). +pub const OverlapAdd = struct { + gpa: std.mem.Allocator, + window_sq: [cfg.n_fft]f32, + acc: std.ArrayList(f32) = .empty, + env: std.ArrayList(f32) = .empty, + rows: usize = 0, + emitted: usize = cfg.istft_trim, + out: std.ArrayList(f32) = .empty, + + pub fn init(gpa: std.mem.Allocator, window: []const f32) OverlapAdd { + var o: OverlapAdd = .{ .gpa = gpa, .window_sq = undefined }; + for (&o.window_sq, window) |*q, w| q.* = w * w; + return o; + } + + pub fn deinit(o: *OverlapAdd) void { + o.acc.deinit(o.gpa); + o.env.deinit(o.gpa); + o.out.deinit(o.gpa); + } + + /// Adds windowed frames [n, n_fft]; returns the samples now final (valid until the next call). + pub fn add(o: *OverlapAdd, frames: []const f32, final: bool) ![]const f32 { + const n = frames.len / cfg.n_fft; + for (0..n) |j| { + const base = (o.rows + j) * cfg.hop; + const need = base + cfg.n_fft; + while (o.acc.items.len < need) { + try o.acc.append(o.gpa, 0); + try o.env.append(o.gpa, 0); + } + const fr = frames[j * cfg.n_fft ..][0..cfg.n_fft]; + for (0..cfg.n_fft) |t| { + o.acc.items[base + t] += fr[t]; + o.env.items[base + t] += o.window_sq[t]; + } + } + o.rows += n; + const end = if (final) o.rows * cfg.hop + cfg.istft_trim else o.rows * cfg.hop; + o.out.clearRetainingCapacity(); + while (o.emitted < end and o.emitted < o.acc.items.len) : (o.emitted += 1) try o.out.append(o.gpa, o.acc.items[o.emitted] / o.env.items[o.emitted]); + return o.out.items; + } +}; + +pub fn bf16(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// One codebook's code: the repetition penalty over its recent codes, then greedy or tempered sampling. +pub fn pickCode(logits: []f32, q: usize, smp: Sampling, rng: std.Random, history: []const [cfg.codebooks]i32) i32 { + const p = smp.audio_repetition_penalty; + if (p != 1.0) { + var seen: [cfg.codebook_size]bool = @splat(false); + for (history) |fr| { + const c: usize = @intCast(fr[q]); + if (seen[c]) continue; + seen[c] = true; + logits[c] = if (logits[c] < 0) logits[c] * p else logits[c] / p; + } + } + if (smp.audio_temperature == 0) { + var best: usize = 0; + for (logits, 0..) |v, i| if (v > logits[best]) { + best = i; + }; + return @intCast(best); + } + const t = smp.audio_temperature; + var m: f32 = -std.math.inf(f32); + for (logits) |*v| { + v.* /= t; + m = @max(m, v.*); + } + if (smp.audio_top_k > 0 and smp.audio_top_k < logits.len) { + // the k-th largest as a floor + var tmp: [cfg.codebook_size]f32 = undefined; + @memcpy(tmp[0..logits.len], logits); + std.mem.sort(f32, tmp[0..logits.len], {}, std.sort.desc(f32)); + const floor = tmp[smp.audio_top_k - 1]; + for (logits) |*v| if (v.* < floor) { + v.* = -std.math.inf(f32); + }; + } + var sum: f64 = 0; + for (logits) |*v| { + v.* = @exp(v.* - m); + sum += v.*; + } + var u = rng.float(f64) * sum; + for (logits, 0..) |v, i| { + u -= v; + if (u <= 0) return @intCast(i); + } + return @intCast(logits.len - 1); +} + +test "overlap-add: streamed chunks give the samples of one whole decode" { + const gpa = std.testing.allocator; + var window: [cfg.n_fft]f32 = undefined; + for (&window, 0..) |*w, i| w.* = @floatCast(0.5 - 0.5 * @cos(2.0 * std.math.pi * @as(f64, @floatFromInt(i)) / cfg.n_fft)); + const rows = 37; + const frames = try gpa.alloc(f32, rows * cfg.n_fft); + defer gpa.free(frames); + var prng = std.Random.DefaultPrng.init(7); + for (frames) |*x| x.* = prng.random().float(f32) - 0.5; + var whole = OverlapAdd.init(gpa, &window); + defer whole.deinit(); + const a = try gpa.dupe(f32, try whole.add(frames, true)); + defer gpa.free(a); + try std.testing.expectEqual(@as(usize, rows * cfg.hop), a.len); + var parts = OverlapAdd.init(gpa, &window); + defer parts.deinit(); + var b: std.ArrayList(f32) = .empty; + defer b.deinit(gpa); + var at: usize = 0; + for ([_]usize{ 6, 6, 12, 1, 12 }) |n| { + try b.appendSlice(gpa, try parts.add(frames[at * cfg.n_fft .. (at + n) * cfg.n_fft], at + n == rows)); + at += n; + } + try std.testing.expectEqualSlices(f32, a, b.items); +} + +test "synthesis basis: a cosine bin comes back as the windowed cosine" { + const gpa = std.testing.allocator; + var window: [cfg.n_fft]f32 = @splat(1); + const m = try synthesisBasis(gpa, &window); + defer gpa.free(m); + // the spectrum of cos(2 pi 3 t / n) is n / 2 at bin 3 (real) + for (0..cfg.n_fft) |t| { + const want = @cos(2.0 * std.math.pi * 3.0 * @as(f64, @floatFromInt(t)) / cfg.n_fft); + const got = @as(f64, m[3 * cfg.n_fft + t]) * cfg.n_fft / 2.0; + try std.testing.expectApproxEqAbs(want, got, 1e-5); + } +} diff --git a/zig/src/families/lfm3_omni/config.zig b/zig/src/families/lfm3_omni/config.zig new file mode 100644 index 0000000000..3f16136d10 --- /dev/null +++ b/zig/src/families/lfm3_omni/config.zig @@ -0,0 +1,147 @@ +//! LFM3-Omni-2B's shapes and token ids (the released config), checked against the checkpoint's config.json at load. +const std = @import("std"); +const stack = @import("spec.zig"); + +pub const hidden = 2048; +pub const vocab = 128000; +pub const codebooks = 12; +pub const codebook_size = 2017; // 2016 codes and the end-of-speech code +pub const audio_eos_code = 2016; +pub const maskgit_schedule = [_]usize{ 4, 4, 4 }; +pub const interleaved_n_text = 6; +pub const interleaved_n_audio = 12; + +pub const bos_token_id = 124894; +pub const eos_token_id = 124900; // <|im_end|> +pub const text_end_token_id = 125014; +pub const audio_start_token_id = 125012; +pub const mixed_start_token_id = 125015; +pub const audio_token_id = 125031; +pub const audio_out_token_id = 125032; +pub const image_token_id = 125030; +pub const video_token_id = 125029; +pub const vision_start_token_id = 125009; +pub const vision_end_token_id = 125010; + +// speech encoder +pub const mel_bins = 128; +pub const conv_channels = 256; +pub const enc_d = 768; +// depth decoder +pub const depth_d = 1024; +// codec +pub const codec_d = 512; +pub const quantizer_dim = 256; +pub const frame_stack = 6; +pub const n_fft = 1280; +pub const hop = 320; +pub const spec_bins = n_fft / 2 + 1; +pub const istft_trim = (n_fft - hop) / 2; +pub const sample_rate = 24000; +pub const frame_samples = hop * frame_stack; + +const C = stack.Kind.conv; +const F = stack.Kind.full; +const S = stack.Kind.sliding; + +pub const encoder_kinds = [_]stack.Kind{ C, C, F, C, C, F, C, C, F, C, C, F, C, C, F, F, F }; +pub const backbone_kinds = [_]stack.Kind{ C, C, F, C, S, C, F, C, S, C, F, C, S, C, F, C, S, C, F, C, S, C, F, C }; +pub const depth_kinds = [_]stack.Kind{ F, F, F }; +pub const codec_kinds = [_]stack.Kind{ C, C, S, C, S, C, S, C }; + +/// The speech encoder: LFM2 layers (3-tap conv, no output gate), causal attention over the last 32 tokens. +pub fn encoderSpec(max_tokens: usize) stack.Spec { + return .{ .name = "speech encoder", .prefix = "model.audio_encoder.lfm.", .d = enc_d, .ffn = 1792, .heads = 8, .kv_heads = 4, .head_dim = 96, .conv_k = 3, .kinds = &encoder_kinds, .window = 32, .window_all = true, .sandwich = false, .lfm3 = false, .rope_full = .{ .theta = 1e6 }, .rope_mode = 2, .max_rows = max_tokens, .max_ctx = max_tokens }; +} + +pub fn backboneSpec(rows: usize, ctx: usize) stack.Spec { + return .{ .name = "backbone", .prefix = "model.language_model.", .d = hidden, .ffn = 7680, .heads = 16, .kv_heads = 4, .head_dim = 128, .conv_k = 7, .kinds = &backbone_kinds, .window = 256, .sandwich = true, .lfm3 = true, .rope_full = .{ .theta = 1e6, .fraction = 0.5 }, .rope_sliding = .{ .theta = 10000 }, .rope_mode = 1, .max_rows = rows, .max_ctx = ctx }; +} + +/// The depth decoder: three bidirectional layers over a frame's 12 codebook slots. +pub fn depthSpec() stack.Spec { + return .{ .name = "depth decoder", .prefix = "model.depth_decoder.model.", .d = depth_d, .ffn = 2816, .heads = 32, .kv_heads = 8, .head_dim = 32, .conv_k = 7, .kinds = &depth_kinds, .causal = false, .sandwich = false, .lfm3 = true, .rope_full = .{ .theta = 1e6 }, .rope_mode = 1, .max_rows = codebooks, .max_ctx = codebooks }; +} + +/// The Liquid codec's decoder at 75 rows a second: sliding attention over the last 30 rows, RoPE in fp32. +pub fn codecSpec(rows: usize, ctx: usize) stack.Spec { + return .{ .name = "codec decoder", .prefix = "decoder.lfm.", .d = codec_d, .ffn = 2304, .heads = 8, .kv_heads = 2, .head_dim = 64, .conv_k = 7, .kinds = &codec_kinds, .window = 30, .sandwich = true, .lfm3 = true, .rope_full = .{ .theta = 1e6 }, .rope_sliding = .{ .theta = 1e6 }, .rope_mode = 0, .max_rows = rows, .max_ctx = ctx }; +} + +fn get(o: std.json.ObjectMap, key: []const u8) !std.json.Value { + return o.get(key) orelse { + std.log.err("lfm3_omni: config.json has no {s}", .{key}); + return error.BadConfig; + }; +} + +fn int(o: std.json.ObjectMap, key: []const u8, want: i64) !void { + const v = try get(o, key); + if (v != .integer or v.integer != want) { + std.log.err("lfm3_omni: config {s} is not {d}; only LFM3-Omni-2B's shapes are supported", .{ key, want }); + return error.UnsupportedConfig; + } +} + +fn kinds(o: std.json.ObjectMap, want: []const stack.Kind) !void { + const v = try get(o, "layer_types"); + if (v != .array or v.array.items.len != want.len) return error.UnsupportedConfig; + for (v.array.items, want) |x, k| { + if (x != .string) return error.UnsupportedConfig; + const name = switch (k) { + .conv => "conv", + .full => "full_attention", + .sliding => "sliding_attention", + }; + if (!std.mem.eql(u8, x.string, name)) return error.UnsupportedConfig; + } +} + +fn object(v: std.json.Value) !std.json.ObjectMap { + if (v != .object) return error.BadConfig; + return v.object; +} + +/// Refuses a checkpoint whose config differs from LFM3-Omni-2B's in anything this port hard-codes. +pub fn check(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(text); + var tree = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer tree.deinit(); + const root = try object(tree.value); + const mt = try get(root, "model_type"); + if (mt != .string or !std.mem.eql(u8, mt.string, "lfm3_omni")) return error.NotLfm3Omni; + try int(root, "num_codebooks", codebooks); + try int(root, "codebook_size", codebook_size); + try int(root, "audio_eos_code", audio_eos_code); + try int(root, "audio_token_id", audio_token_id); + try int(root, "audio_out_token_id", audio_out_token_id); + try int(root, "mixed_start_token_id", mixed_start_token_id); + try int(root, "text_end_token_id", text_end_token_id); + try int(root, "interleaved_n_text", interleaved_n_text); + try int(root, "interleaved_n_audio", interleaved_n_audio); + const tc = try object(try get(root, "text_config")); + try int(tc, "hidden_size", hidden); + try int(tc, "num_hidden_layers", backbone_kinds.len); + try int(tc, "intermediate_size", 7680); + try int(tc, "num_attention_heads", 16); + try int(tc, "num_key_value_heads", 4); + try int(tc, "sliding_window", 256); + try int(tc, "vocab_size", vocab); + try kinds(tc, &backbone_kinds); + const ac = try object(try get(root, "audio_encoder_config")); + try int(ac, "hidden_size", enc_d); + try int(ac, "num_mel_bins", mel_bins); + try int(ac, "sliding_window", 32); + try int(ac, "conv_channels", conv_channels); + try kinds(try object(try get(ac, "text_config")), &encoder_kinds); + const dc = try object(try get(root, "depth_config")); + try int(dc, "hidden_size", depth_d); + try int(dc, "num_hidden_layers", depth_kinds.len); + try int(dc, "head_dim", 32); + const ms = try get(root, "maskgit_schedule"); + if (ms != .array or ms.array.items.len != maskgit_schedule.len) return error.UnsupportedConfig; + for (ms.array.items, maskgit_schedule) |x, w| if (x != .integer or x.integer != w) return error.UnsupportedConfig; +} diff --git a/zig/src/families/lfm3_omni/generate.zig b/zig/src/families/lfm3_omni/generate.zig new file mode 100644 index 0000000000..fc167ea238 --- /dev/null +++ b/zig/src/families/lfm3_omni/generate.zig @@ -0,0 +1,330 @@ +//! lfm3-omni-generate: LFM3-Omni-2B speech in, text and speech out, its server, stage dumps and teacher forcing. +const std = @import("std"); +const core = @import("core"); +const om = @import("lfm3_omni"); + +const cfg = om.config; +const Engine = om.Engine; + +const usage = + \\usage: lfm3-omni-generate --model DIR --serve HOST:PORT [--system TEXT | --system-file FILE] [sampling flags] + \\ lfm3-omni-generate --model DIR --wav IN.wav [--out REPLY.wav] [--system TEXT | --system-file FILE] + \\ [--audio-temperature 0.4] [--audio-repetition-penalty 1.2] [--audio-top-k 0] + \\ [--identity-order 1] [--seed 0] [--max-new 512] [--first-chunk 1] [--chunk 8] + \\ [--dump DIR] [--force DIR] [--repeat 1] + \\ + \\--serve: POST /v1/audio/chat {"audio": base64 WAV, "system"?, "stream"? (true), "audio_temperature"?, "seed"?, + \\ "max_new"?} answers with server-sent events (text deltas, base64 pcm16 24 kHz audio, then done) or one JSON. + \\The input is any PCM or float WAV, resampled to 16 kHz; the reply is 16-bit PCM mono at 24 kHz. + \\--dump DIR (an existing directory) writes the stages as raw little-endian arrays: mel.f32 [128, frames], + \\ enc.f32 and audio_emb.f32 [tokens, ...], input_ids.i32, prefill_hidden.f32 [rows, 2048], first_logits.f32, + \\ depth_logits0.f32 [12, 2017], gen_seq.i32, gen_codes.i32 [frames, 12], step_hidden.f32, codec_lin.f32 + \\ [6 * frames, 1282] and reply.f32 (the whole reply through the codec at once). + \\--codec-ref DIR: with --dump, codec_lin.f32 and reply.f32 decode DIR's gen_codes.i32 ([12, frames]) instead. + \\--depth-check DIR: the depth decoder alone on DIR's depth_h.f32 [n, 2048] and depth_codes.i32 [n, 12]: three + \\ passes a row (identity order), every pass's logits written to depth_check.f32 [n, 3, 12, 2017]; no reply. + \\--graphs 0 launches the backbone's rows and the depth decoder's frames one kernel at a time (no CUDA graphs). + \\--host-mel 1 computes the log-mel features on the host (a float64 DFT) instead of the device. + \\--force DIR reads a reference's gen_seq.i32 and gen_codes.i32 ([12, frames]) and feeds them back, scoring the + \\ engine's own picks (greedy text, its audio sampling settings) at every slot. + \\The last stderr line is "RESULT lfm3-omni OK ..." with the timings. + \\ +; + +const Args = struct { + model: []const u8 = "", + wav: []const u8 = "", + out: ?[]const u8 = null, + system: []const u8 = "Respond with interleaved text and audio.", + smp: om.engine.Sampling = .{}, + stream: om.engine.Streaming = .{}, + dump: ?[]const u8 = null, + force: ?[]const u8 = null, + codec_ref: ?[]const u8 = null, + depth_check: ?[]const u8 = null, + repeat: usize = 1, + host_mel: bool = false, + serve: ?[]const u8 = null, + graphs: bool = true, +}; + +pub fn main(init: std.process.Init) !u8 { + const arena = init.arena.allocator(); + const gpa = std.heap.c_allocator; + const io = init.io; + const argv = try init.minimal.args.toSlice(arena); + var a: Args = .{}; + var i: usize = 1; + while (i < argv.len) : (i += 2) { + if (i + 1 >= argv.len) { + std.debug.print("{s}", .{usage}); + return 2; + } + const k = argv[i]; + const v = argv[i + 1]; + if (eq(k, "--model")) a.model = v else if (eq(k, "--wav")) a.wav = v else if (eq(k, "--out")) a.out = v else if (eq(k, "--system")) a.system = v else if (eq(k, "--system-file")) { + a.system = std.mem.trim(u8, try std.Io.Dir.cwd().readFileAlloc(io, v, arena, .limited(1 << 20)), " \n\r\t"); + } else if (eq(k, "--audio-temperature")) a.smp.audio_temperature = try std.fmt.parseFloat(f32, v) else if (eq(k, "--audio-repetition-penalty")) a.smp.audio_repetition_penalty = try std.fmt.parseFloat(f32, v) else if (eq(k, "--audio-top-k")) a.smp.audio_top_k = try std.fmt.parseInt(usize, v, 10) else if (eq(k, "--identity-order")) a.smp.identity_order = !eq(v, "0") else if (eq(k, "--seed")) a.smp.seed = try std.fmt.parseInt(u64, v, 10) else if (eq(k, "--max-new")) a.smp.max_new = try std.fmt.parseInt(usize, v, 10) else if (eq(k, "--first-chunk")) a.stream.first = try std.fmt.parseInt(usize, v, 10) else if (eq(k, "--chunk")) a.stream.chunk = try std.fmt.parseInt(usize, v, 10) else if (eq(k, "--dump")) a.dump = v else if (eq(k, "--force")) a.force = v else if (eq(k, "--codec-ref")) a.codec_ref = v else if (eq(k, "--depth-check")) a.depth_check = v else if (eq(k, "--serve")) a.serve = v else if (eq(k, "--graphs")) a.graphs = !eq(v, "0") else if (eq(k, "--host-mel")) a.host_mel = !eq(v, "0") else if (eq(k, "--repeat")) a.repeat = try std.fmt.parseInt(usize, v, 10) else { + std.debug.print("unknown flag {s}\n{s}", .{ k, usage }); + return 2; + } + } + if (a.depth_check) |dc| { + const e = try Engine.create(gpa, io, a.model, .{}); + defer e.destroy(); + try depthCheck(arena, io, e, dc); + return 0; + } + if (a.serve) |hp| { + if (a.model.len == 0) { + std.debug.print("{s}", .{usage}); + return 2; + } + const colon = std.mem.lastIndexOfScalar(u8, hp, ':') orelse return error.ServeIsHostColonPort; + const tpath0 = try std.fs.path.join(arena, &.{ a.model, "tokenizer.json" }); + var tok0 = try core.tokenizer.loadTokenizer(io, gpa, tpath0); + defer tok0.deinit(); + const e = try Engine.create(gpa, io, a.model, .{ .graphs = a.graphs }); + defer e.destroy(); + try warm(gpa, io, e, &tok0, a); + try om.serve.run(gpa, io, e, &tok0, .{ .host = hp[0..colon], .port = try std.fmt.parseInt(u16, hp[colon + 1 ..], 10), .system = a.system, .smp = a.smp, .stream = a.stream }); + return 0; + } + if (a.model.len == 0 or a.wav.len == 0) { + std.debug.print("{s}", .{usage}); + return 2; + } + const tpath = try std.fs.path.join(arena, &.{ a.model, "tokenizer.json" }); + var tok = try core.tokenizer.loadTokenizer(io, gpa, tpath); + defer tok.deinit(); + const t_load = om.engine.nowNs(io); + const e = try Engine.create(gpa, io, a.model, .{ .graphs = a.graphs }); + defer e.destroy(); + std.log.info("lfm3-omni: loaded in {d:.1} s, {d:.2} GB on the device", .{ @as(f64, @floatFromInt(om.engine.nowNs(io) - t_load)) / 1e9, @as(f64, @floatFromInt(e.deviceBytes())) / 1e9 }); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, a.wav, arena, .limited(1 << 30)); + const au = try om.wav.parse(arena, bytes); + const pcm = try om.wav.resample(arena, au.samples, au.rate, om.mel.sample_rate); + var r: usize = 0; + while (r < a.repeat) : (r += 1) try run(gpa, io, a, e, &tok, pcm, r + 1 == a.repeat); + return 0; +} + +/// One short turn on a second of quiet noise, so the first request finds cuBLAS's plans and the kernels loaded. +fn warm(gpa: std.mem.Allocator, io: std.Io, e: *Engine, tok: *const core.tokenizer.Tokenizer, a: Args) !void { + var ar_state = std.heap.ArenaAllocator.init(gpa); + defer ar_state.deinit(); + const ar = ar_state.allocator(); + const pcm = try ar.alloc(f32, om.mel.sample_rate); + var prng = std.Random.DefaultPrng.init(1); + for (pcm) |*x| x.* = (prng.random().float(f32) - 0.5) * 0.002; + const nf = try e.featuresDevice(pcm); + const n = try e.encode(null, nf); + const p = try om.prompt.build(ar, tok, a.system, n); + try e.prefill(p.ids, p.audio_at, null); + var st: om.engine.Stats = .{}; + var smp = a.smp; + smp.max_new = 24; + var sink_ctx: u8 = 0; + try e.reply(smp, .{}, .{ .ctx = &sink_ctx }, null, &st, om.engine.nowNs(io)); +} + +fn eq(x: []const u8, y: []const u8) bool { + return std.mem.eql(u8, x, y); +} + +const Out = struct { + gpa: std.mem.Allocator, + e: *Engine, + text: std.ArrayList(u32) = .empty, + seq: std.ArrayList(i32) = .empty, + frames: std.ArrayList([cfg.codebooks]i32) = .empty, + pcm: std.ArrayList(f32) = .empty, + hidden: ?*std.ArrayList(f32) = null, + + fn onText(ctx: *anyopaque, id: u32) anyerror!void { + const o: *Out = @ptrCast(@alignCast(ctx)); + try o.text.append(o.gpa, id); + try o.seq.append(o.gpa, @intCast(id)); + } + fn onFrame(ctx: *anyopaque, codes: *const [cfg.codebooks]i32) anyerror!void { + const o: *Out = @ptrCast(@alignCast(ctx)); + try o.frames.append(o.gpa, codes.*); + try o.seq.append(o.gpa, cfg.audio_out_token_id); + } + fn onPcm(ctx: *anyopaque, x: []const f32) anyerror!void { + const o: *Out = @ptrCast(@alignCast(ctx)); + try o.pcm.appendSlice(o.gpa, x); + } + fn onHidden(ctx: *anyopaque) anyerror!void { + const o: *Out = @ptrCast(@alignCast(ctx)); + if (o.hidden) |h| { + var row: [cfg.hidden]f32 = undefined; + try o.e.lastHidden(&row); + try h.appendSlice(o.gpa, &row); + } + } + fn deinit(o: *Out) void { + o.text.deinit(o.gpa); + o.seq.deinit(o.gpa); + o.frames.deinit(o.gpa); + o.pcm.deinit(o.gpa); + } +}; + +fn writeRaw(io: std.Io, dir: []const u8, name: []const u8, data: []const u8) !void { + var buf: [1024]u8 = undefined; + const p = try std.fmt.bufPrint(&buf, "{s}/{s}", .{ dir, name }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = p, .data = data }); +} + +fn readRaw(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) ![]u8 { + var buf: [1024]u8 = undefined; + const p = try std.fmt.bufPrint(&buf, "{s}/{s}", .{ dir, name }); + return std.Io.Dir.cwd().readFileAlloc(io, p, gpa, .limited(1 << 30)); +} + +fn run(gpa: std.mem.Allocator, io: std.Io, a: Args, e: *Engine, tok: *const core.tokenizer.Tokenizer, pcm: []const f32, last: bool) !void { + var ar_state = std.heap.ArenaAllocator.init(gpa); + defer ar_state.deinit(); + const ar = ar_state.allocator(); + var st: om.engine.Stats = .{}; + const t_req = om.engine.nowNs(io); + // speech in + var nf = om.mel.frames(pcm.len); + const feats = try ar.alloc(f32, cfg.mel_bins * nf); + if (a.host_mel) e.features(pcm, feats) else nf = try e.featuresDevice(pcm); + st.mel_ms = @as(f64, @floatFromInt(om.engine.nowNs(io) - t_req)) / 1e6; + const t_enc = om.engine.nowNs(io); + const n_audio = try e.encode(if (a.host_mel) feats else null, nf); + const dump = if (last) a.dump else null; + if (dump) |d| { + try e.stream.synchronize(); + if (!a.host_mel) try e.d2h(std.mem.sliceAsBytes(feats), e.mel_f32); + try writeRaw(io, d, "mel.f32", std.mem.sliceAsBytes(feats)); + try writeRaw(io, d, "enc.f32", std.mem.sliceAsBytes(try deviceRows(ar, e, e.enc_out, n_audio * cfg.enc_d))); + try writeRaw(io, d, "audio_emb.f32", std.mem.sliceAsBytes(try deviceRows(ar, e, e.audio_emb, n_audio * cfg.hidden))); + } + st.encode_ms = @as(f64, @floatFromInt(om.engine.nowNs(io) - t_enc)) / 1e6; + st.audio_tokens = n_audio; + const p = try om.prompt.build(ar, tok, a.system, n_audio); + st.prompt_rows = p.ids.len; + const t_pre = om.engine.nowNs(io); + var all_hidden: ?[]f32 = null; + if (dump != null) all_hidden = try ar.alloc(f32, p.ids.len * cfg.hidden); + try e.prefill(p.ids, p.audio_at, all_hidden); + st.prefill_ms = @as(f64, @floatFromInt(om.engine.nowNs(io) - t_pre)) / 1e6; + var hidden: std.ArrayList(f32) = .empty; + var out: Out = .{ .gpa = gpa, .e = e }; + defer out.deinit(); + if (dump) |d| { + try writeRaw(io, d, "input_ids.i32", std.mem.sliceAsBytes(p.ids)); + try writeRaw(io, d, "prefill_hidden.f32", std.mem.sliceAsBytes(all_hidden.?)); + const lg = try ar.alloc(f32, cfg.vocab); + try e.textLogits(lg); + try writeRaw(io, d, "first_logits.f32", std.mem.sliceAsBytes(lg)); + out.hidden = &hidden; + try Out.onHidden(&out); + } + var forced: ?om.engine.Forced = null; + if (a.force) |fd| { + const seq_b = try readRaw(ar, io, fd, "gen_seq.i32"); + const seq: []const u32 = @alignCast(std.mem.bytesAsSlice(u32, seq_b)); + forced = .{ .seq = seq, .frames = try readFrames(ar, io, fd) }; + } + const sink: om.engine.Sink = .{ .ctx = &out, .text = Out.onText, .frame = Out.onFrame, .pcm = Out.onPcm, .hidden = if (dump != null) Out.onHidden else null }; + try e.reply(a.smp, a.stream, sink, if (forced) |*f| f else null, &st, t_req); + // the reply's text (special tokens left out) + var words: std.ArrayList(u32) = .empty; + for (out.text.items) |id| if (id < cfg.bos_token_id) try words.append(ar, id); + const text = try tok.decode(ar, words.items, true); + if (a.out) |o| { + if (last) try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = o, .data = try om.wav.encode16(ar, out.pcm.items, cfg.sample_rate) }); + } + if (dump) |d| { + try writeRaw(io, d, "gen_seq.i32", std.mem.sliceAsBytes(out.seq.items)); + try writeRaw(io, d, "gen_codes.i32", std.mem.sliceAsBytes(out.frames.items)); + try writeRaw(io, d, "step_hidden.f32", std.mem.sliceAsBytes(hidden.items)); + try writeRaw(io, d, "reply_stream.f32", std.mem.sliceAsBytes(out.pcm.items)); + var frames: []const [cfg.codebooks]i32 = out.frames.items; + if (a.codec_ref) |cr| frames = try readFrames(ar, io, cr); + try dumpDepthAndCodec(ar, io, e, d, p.ids, p.audio_at, frames, a.stream); + hidden.deinit(gpa); + } + if (forced) |f| std.debug.print("FORCED text {d}/{d} codes {d}/{d} frames {d}/{d}\n", .{ f.text_hits, f.text_total, f.code_hits, f.code_total, f.frame_hits, f.frame_total }); + std.debug.print("TEXT {s}\n", .{text}); + std.debug.print("RESULT lfm3-omni OK input_s={d:.2} audio_tokens={d} prompt_rows={d} mel_ms={d:.1} encode_ms={d:.1} prefill_ms={d:.1} first_text_ms={d:.1} first_frame_ms={d:.1} first_pcm_ms={d:.1} text_tokens={d} frames={d} speech_s={d:.2} step_ms={d:.1} depth_ms={d:.1} codec_ms={d:.1} total_ms={d:.1}\n", .{ @as(f64, @floatFromInt(pcm.len)) / 16000.0, st.audio_tokens, st.prompt_rows, st.mel_ms, st.encode_ms, st.prefill_ms, st.first_text_ms, st.first_frame_ms, st.first_pcm_ms, st.text_tokens, st.frames, @as(f64, @floatFromInt(st.samples)) / 24000.0, st.step_ms, st.depth_ms, st.codec_ms, st.total_ms }); +} + +/// The depth decoder on given hidden states: per row, the passes of the identity order with the given codes revealed. +fn depthCheck(ar: std.mem.Allocator, io: std.Io, e: *Engine, dir: []const u8) !void { + const hb = try readRaw(ar, io, dir, "depth_h.f32"); + const cb = try readRaw(ar, io, dir, "depth_codes.i32"); + const h: []const f32 = @alignCast(std.mem.bytesAsSlice(f32, hb)); + const c: []const i32 = @alignCast(std.mem.bytesAsSlice(i32, cb)); + const n = h.len / cfg.hidden; + var out: std.ArrayList(f32) = .empty; + for (0..n) |r| { + try e.setHidden(h[r * cfg.hidden ..][0..cfg.hidden]); + try e.depthStart(); + var codes: [cfg.codebooks]i32 = @splat(0); + var masked: [cfg.codebooks]bool = @splat(true); + var at: usize = 0; + for (cfg.maskgit_schedule) |k| { + try out.appendSlice(ar, try e.depthPass(&codes, &masked)); + for (at..at + k) |q| { + codes[q] = c[r * cfg.codebooks + q]; + masked[q] = false; + } + at += k; + } + } + try writeRaw(io, dir, "depth_check.f32", std.mem.sliceAsBytes(out.items)); + std.debug.print("RESULT lfm3-omni depth-check OK rows={d}\n", .{n}); +} + +/// A reference's gen_codes.i32 ([12, frames]) as frames. +fn readFrames(ar: std.mem.Allocator, io: std.Io, dir: []const u8) ![]const [cfg.codebooks]i32 { + const codes_b = try readRaw(ar, io, dir, "gen_codes.i32"); + const cm: []const i32 = @alignCast(std.mem.bytesAsSlice(i32, codes_b)); + const n = cm.len / cfg.codebooks; + const frames = try ar.alloc([cfg.codebooks]i32, n); + for (frames, 0..) |*fr, t| for (fr, 0..) |*c, q| { + c.* = cm[q * n + t]; + }; + return frames; +} + +fn deviceRows(ar: std.mem.Allocator, e: *Engine, src: u64, n: usize) ![]f32 { + const hb = try ar.alloc(u16, n); + try e.d2h(std.mem.sliceAsBytes(hb), src); + const out = try ar.alloc(f32, n); + for (out, hb) |*o, h| o.* = @bitCast(@as(u32, h) << 16); + return out; +} + +/// The first frame's depth logits from the prompt (every codebook masked) and the whole reply through the codec. +fn dumpDepthAndCodec(ar: std.mem.Allocator, io: std.Io, e: *Engine, d: []const u8, ids: []const u32, audio_at: usize, frames: []const [cfg.codebooks]i32, sm: om.engine.Streaming) !void { + // the reference scores the prompt without the interleaved marker + try e.prefill(ids[0 .. ids.len - 1], audio_at, null); + try e.depthStart(); + const codes: [cfg.codebooks]i32 = @splat(0); + const masked: [cfg.codebooks]bool = @splat(true); + const lg = try e.depthPass(&codes, &masked); + try writeRaw(io, d, "depth_logits0.f32", std.mem.sliceAsBytes(lg)); + try e.codecReset(); + var lin: std.ArrayList(f32) = .empty; + var ola = om.engine.OverlapAdd.init(ar, &e.window); + var wav: std.ArrayList(f32) = .empty; + var at: usize = 0; + // in the streaming chunks the server uses: the first `first` frames, then `chunk` at a time + while (at < frames.len) { + const n = @min(if (at == 0) sm.first else sm.chunk, frames.len - at); + const l = try ar.alloc(f32, n * cfg.frame_stack * 2 * cfg.spec_bins); + const fr = try e.codecChunk(frames[at .. at + n], l); + try lin.appendSlice(ar, l); + try wav.appendSlice(ar, try ola.add(fr, at + n == frames.len)); + at += n; + } + try writeRaw(io, d, "codec_lin.f32", std.mem.sliceAsBytes(lin.items)); + try writeRaw(io, d, "reply.f32", std.mem.sliceAsBytes(wav.items)); +} diff --git a/zig/src/families/lfm3_omni/mel.zig b/zig/src/families/lfm3_omni/mel.zig new file mode 100644 index 0000000000..14d5634269 --- /dev/null +++ b/zig/src/families/lfm3_omni/mel.zig @@ -0,0 +1,108 @@ +//! LFM3-Omni's host log-mel: 400-point Hann STFT every 160 samples, 128 Slaney bins, (max(log10, -12) + 4) / 4. +const std = @import("std"); + +pub const sample_rate = 16000; +pub const n_fft = 400; +pub const hop = 160; +pub const bins = n_fft / 2 + 1; +pub const mels = 128; +const log_floor = -12.0; + +/// Frames for `n` samples once padded to whole hops (at least one window). +pub fn frames(n: usize) usize { + return 1 + (padded(n) - n_fft) / hop; +} + +fn padded(n: usize) usize { + if (n < n_fft) return n_fft; + return n + (hop - (n - n_fft) % hop) % hop; +} + +/// The speech encoder's tokens for `f` frames: one per 8 frames, rounded up. +pub fn tokens(f: usize) usize { + return (f + 7) / 8; +} + +fn hzToMel(f: f64) f64 { + const logstep = 27.0 / @log(6.4); + return if (f >= 1000.0) 15.0 + @log(f / 1000.0) * logstep else 3.0 * f / 200.0; +} + +fn melToHz(m: f64) f64 { + const logstep = @log(6.4) / 27.0; + return if (m >= 15.0) 1000.0 * @exp(logstep * (m - 15.0)) else 200.0 * m / 3.0; +} + +/// transformers' mel_filter_bank(201, 128, 0, 8000, 16000, norm="slaney", mel_scale="slaney"), [bins][mels]. +pub fn filterBank(out: *[bins][mels]f64) void { + var hz: [mels + 2]f64 = undefined; + const lo = hzToMel(0); + const hi = hzToMel(sample_rate / 2); + for (&hz, 0..) |*h, i| h.* = melToHz(lo + (hi - lo) * @as(f64, @floatFromInt(i)) / @as(f64, mels + 1)); + for (out, 0..) |*row, k| { + const fk = @as(f64, @floatFromInt(k)) * (sample_rate / 2) / @as(f64, bins - 1); + for (row, 0..) |*v, m| { + const down = -(hz[m] - fk) / (hz[m + 1] - hz[m]); + const up = (hz[m + 2] - fk) / (hz[m + 2] - hz[m + 1]); + v.* = @max(0.0, @min(down, up)) * (2.0 / (hz[m + 2] - hz[m])); + } + } +} + +pub const Mel = struct { + filters: [bins][mels]f32, + cos: [n_fft][bins]f64, + sin: [n_fft][bins]f64, + window: [n_fft]f64, + + pub fn init(m: *Mel) void { + var fb: [bins][mels]f64 = undefined; + filterBank(&fb); + for (&m.filters, fb) |*r, s| for (r, s) |*v, x| { + v.* = @floatCast(x); + }; + for (&m.window, 0..) |*w, i| w.* = 0.5 - 0.5 * @cos(2.0 * std.math.pi * @as(f64, @floatFromInt(i)) / n_fft); + for (0..n_fft) |n| for (0..bins) |k| { + const a = 2.0 * std.math.pi * @as(f64, @floatFromInt((n * k) % n_fft)) / n_fft; + m.cos[n][k] = @cos(a); + m.sin[n][k] = @sin(a); + }; + } + + /// Features of `x` (16 kHz mono) into `out` [mels, frames(x.len)], row-major. + pub fn compute(m: *const Mel, x: []const f32, out: []f32) void { + const nf = frames(x.len); + std.debug.assert(out.len == mels * nf); + var power: [bins]f32 = undefined; + for (0..nf) |t| { + var re: [bins]f64 = @splat(0); + var im: [bins]f64 = @splat(0); + for (0..n_fft) |n| { + const i = t * hop + n; + const v = (if (i < x.len) @as(f64, x[i]) else 0.0) * m.window[n]; + if (v == 0) continue; + for (0..bins) |k| { + re[k] += v * m.cos[n][k]; + im[k] -= v * m.sin[n][k]; + } + } + for (&power, re, im) |*p, r, i| { + const a: f32 = @floatCast(@sqrt(r * r + i * i)); + p.* = a * a; + } + for (0..mels) |j| { + var s: f32 = 0; + for (0..bins) |k| s += m.filters[k][j] * power[k]; + const l = @max(@log10(s), log_floor); + out[j * nf + t] = (l + 4.0) / 4.0; + } + } + } +}; + +test "frames and tokens" { + try std.testing.expectEqual(@as(usize, 1), frames(10)); + try std.testing.expectEqual(@as(usize, 2), frames(401)); + try std.testing.expectEqual(@as(usize, 422), frames(400 + 421 * 160)); + try std.testing.expectEqual(@as(usize, 53), tokens(422)); +} diff --git a/zig/src/families/lfm3_omni/metal/engine.zig b/zig/src/families/lfm3_omni/metal/engine.zig new file mode 100644 index 0000000000..195e47697d --- /dev/null +++ b/zig/src/families/lfm3_omni/metal/engine.zig @@ -0,0 +1,468 @@ +//! LFM3-Omni-2B on Metal in bf16 as released, one request at a time, with the CUDA engine's public calls. +const std = @import("std"); +const metallib = @import("lfm3_omni_metallib"); +const gpu = @import("gpu.zig"); +const kern = @import("kernels.zig"); +const stack = @import("stack.zig"); +const load = @import("load.zig"); +const melmod = @import("../mel.zig"); +const cfg = @import("../config.zig"); +const common = @import("../common.zig"); + +const Stack = stack.Stack; +const Arena = stack.Arena; +const Gpu = gpu.Gpu; + +pub const Options = struct { + /// Longest input clip (seconds of 16 kHz audio) the encoder's buffers take. + max_input_s: usize = 60, + /// Backbone rows: the window (prompt and reply) and a prompt chunk. + max_ctx: usize = 4096, + prefill_rows: usize = 512, + /// Longest reply in speech frames (12.5 a second) and the codec's largest chunk. + max_frames: usize = 1024, + codec_chunk: usize = 64, + /// Accepted for the CUDA engine's flags; Metal encodes each row's dispatches into one command buffer. + graphs: bool = true, +}; + +pub const Sampling = common.Sampling; +pub const Sink = common.Sink; +pub const Streaming = common.Streaming; +pub const Forced = common.Forced; +pub const Stats = common.Stats; +pub const OverlapAdd = common.OverlapAdd; +pub const nowNs = common.nowNs; +const bf16 = common.bf16; +const pickCode = common.pickCode; + +const Downsample = struct { + conv_in_w: u64, + conv_in_b: u64, + conv_mid_w: u64, + conv_mid_b: u64, + lin_mid_w: u64, + lin_mid_b: u64, + conv_end_w: u64, + conv_end_b: u64, + lin_end_w: u64, + lin_end_b: u64, + lin_out_w: u64, + lin_out_b: u64, +}; + +/// The CUDA engine's stream, here the device's open command buffer: `synchronize` waits for the GPU. +pub const Stream = struct { + g: *Gpu, + pub fn synchronize(s: Stream) !void { + try s.g.flush(); + } +}; + +pub const Engine = struct { + gpa: std.mem.Allocator, + io: std.Io, + o: Options, + g: *Gpu, + stream: Stream, + k: kern.Kernels, + a: Arena, + mel: *melmod.Mel, + window: [cfg.n_fft]f32, + + enc: Stack, + bb: Stack, + depth: Stack, + codec: Stack, + ds: Downsample, + // projector + ln_w: u64, + ln_b: u64, + l1_w: u64, + l1_b: u64, + l2_w: u64, + l2_b: u64, + // backbone tables + embed: u64, + cb_embed: u64, + // depth decoder + depth_w: u64, + depth_b: u64, + mask_emb: u64, + depth_tables: u64, // [12 * 2017, 1024] + depth_norms: u64, // [12, 1024] + // codec + po_w: u64, + po_b: u64, + fq_w: u64, + fq_b: u64, + lin_w: u64, + lin_b: u64, + basis: u64, // f32 [1282, 1280] + + // the log-mel front end on the device + mel_pcm: u64, + mel_win: u64, + mel_dft: u64, + mel_fb: u64, + mel_fr: u64, + mel_s: u64, + mel_p: u64, + mel_m: u64, + // the device's frame: order | codes | masked (i32 [3 * 12]) and the repetition window's frames (a ring) + dframe: u64, + hist: u64, + // scratch + mel_f32: u64, + mel_bf: u64, + a1: u64, + a2: u64, + a3: u64, + acc: u64, + enc_x: u64, + enc_out: u64, + ln_out: u64, + h1: u64, + audio_emb: u64, + audio_tokens: usize = 0, + /// The newest backbone row's hidden state (bf16 on the device). + last: u64 = 0, + ids: u64, + bb_x: u64, + bb_out: u64, + logits: u64, + pick: u64, + sup: u64, + amax_ws: u64, // the greedy pick's parts + din_f32: u64, + din: u64, + dx: u64, + dout: u64, + dz: u64, + dlog: u64, + dlog_f32: u64, + dcodes: u64, + dmask: u64, + c_codes: u64, + c_lat: u64, + c_p1: u64, + c_p2: u64, + c_up: u64, + c_out: u64, + c_f32: u64, + c_lin: u64, + c_spec: u64, + c_frames: u64, + host_frames: []f32, + host_logits: []f32, + /// LFM3_OMNI_PROFILE=1: every phase waited for alone, its GPU time summed (logged after each reply). + prof: bool = false, + prof_s: [4]f64 = @splat(0), + prof_n: [4]usize = @splat(0), + + pub const Phase = enum { row, frame, text, codec }; + + /// Ends a phase begun when the device's GPU time was `t0` (profile only). + pub fn phase(e: *Engine, p: Phase, t0: f64) !void { + if (!e.prof) return; + try e.g.flush(); + e.prof_s[@backingInt(p)] += e.g.gpu_s - t0; + e.prof_n[@backingInt(p)] += 1; + } + + pub fn create(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, o: Options) !*Engine { + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = undefined; + e.gpa = gpa; + e.io = io; + e.o = o; + e.audio_tokens = 0; + e.last = 0; + e.prof = if (std.c.getenv("LFM3_OMNI_PROFILE")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else false; + e.prof_s = @splat(0); + e.prof_n = @splat(0); + try cfg.check(gpa, io, model_dir); + e.g = try Gpu.open(gpa, metallib.bytes); + errdefer e.g.close(); + e.stream = .{ .g = e.g }; + e.k = .{ .g = e.g }; + e.a = .{ .g = e.g }; + e.mel = try gpa.create(melmod.Mel); + errdefer gpa.destroy(e.mel); + e.mel.init(); + try load.model(e, model_dir); + try load.scratch(e); + return e; + } + + pub fn destroy(e: *Engine) void { + e.g.flush() catch {}; + e.gpa.free(e.host_frames); + e.gpa.free(e.host_logits); + for ([_]*Stack{ &e.enc, &e.bb, &e.depth, &e.codec }) |s| e.gpa.free(s.layers); + e.g.close(); + e.gpa.destroy(e.mel); + e.gpa.destroy(e); + } + + pub fn deviceBytes(e: *const Engine) usize { + return e.g.bytes; + } + + + // ---- host copies (unified memory: the GPU's work so far finishes first) + + pub fn h2d(e: *Engine, dst: u64, bytes: []const u8) !void { + if (bytes.len % 4 == 0 and bytes.len <= 4096 and bytes.len > 0) return e.g.set(dst, bytes); + try e.g.put(dst, bytes); + } + + pub fn d2h(e: *Engine, out: []u8, src: u64) !void { + try e.g.get(src, out); + } + + fn d2d(e: *Engine, dst: u64, src: u64, n: usize) void { + e.g.copy(src, dst, n); + } + + // ---- speech in + + /// Log-mel features of 16 kHz `pcm` into `out` (host, [128, frames]). + pub fn features(e: *Engine, pcm: []const f32, out: []f32) void { + e.mel.compute(pcm, out); + } + + /// Log-mel features of 16 kHz `pcm` on the device (into `mel_f32`, [128, frames]); returns the frames. + pub fn featuresDevice(e: *Engine, pcm: []const f32) !usize { + const k = &e.k; + const nf = melmod.frames(pcm.len); + if (nf > melmod.frames(e.o.max_input_s * melmod.sample_rate)) return error.InputTooLong; + try e.g.put(e.mel_pcm, std.mem.sliceAsBytes(pcm)); + k.melFrames(e.mel_pcm, pcm.len, e.mel_win, e.mel_fr, nf, melmod.n_fft, melmod.hop); + k.matmul(false, false, nf, 2 * melmod.bins, melmod.n_fft, e.mel_fr, e.mel_dft, e.mel_s, .f32, .f32); + k.power(e.mel_s, e.mel_p, nf, melmod.bins); + k.matmul(false, false, nf, melmod.mels, melmod.bins, e.mel_p, e.mel_fb, e.mel_m, .f32, .f32); + k.logmel(e.mel_m, e.mel_f32, nf, melmod.mels, -12.0); + return nf; + } + + /// The speech encoder and projector over features [128, nf] (null: `featuresDevice`'s), a row of `audio_emb` per 8. + pub fn encode(e: *Engine, feats: ?[]const f32, nf: usize) !usize { + const k = &e.k; + const C = cfg.conv_channels; + if (nf > melmod.frames(e.o.max_input_s * melmod.sample_rate)) return error.InputTooLong; + if (feats) |fh| try e.g.put(e.mel_f32, std.mem.sliceAsBytes(fh)); + k.f32ToBf16(e.mel_f32, e.mel_bf, cfg.mel_bins * nf); + const w1 = (nf + 1) / 2; + const w2 = (w1 + 1) / 2; + const w3 = (w2 + 1) / 2; + k.conv2dS2(e.mel_bf, 1, cfg.mel_bins, nf, e.ds.conv_in_w, e.ds.conv_in_b, e.a1, C, 64, w1, true); + k.conv2dS2(e.a1, C, 64, w1, e.ds.conv_mid_w, e.ds.conv_mid_b, e.a2, C, 32, w2, false); + k.matmul(false, false, C, 32 * w2, C, e.ds.lin_mid_w, e.a2, e.acc, .bf16, .f32); + k.biasAct(e.acc, e.ds.lin_mid_b, e.a3, C, 32 * w2, true, .silu); + k.conv2dS2(e.a3, C, 32, w2, e.ds.conv_end_w, e.ds.conv_end_b, e.a2, C, 16, w3, false); + k.matmul(false, false, C, 16 * w3, C, e.ds.lin_end_w, e.a2, e.acc, .bf16, .f32); + k.biasAct(e.acc, e.ds.lin_end_b, e.a1, C, 16 * w3, true, .silu); + // [C * 16, w3] rows to tokens: Z [w3, 768] = A^T Wo^T + k.matmul(true, true, w3, cfg.enc_d, C * 16, e.a1, e.ds.lin_out_w, e.acc, .bf16, .f32); + k.biasAct(e.acc, e.ds.lin_out_b, e.enc_x, w3, cfg.enc_d, false, .none); + e.enc.reset(); + try e.enc.forward(k, e.enc_x, w3, e.enc_out); + k.layernorm(e.enc_out, e.ln_w, e.ln_b, e.ln_out, cfg.enc_d, w3, 1e-5); + k.linearF32(e.l1_w, e.ln_out, cfg.enc_d, cfg.enc_d, cfg.hidden, e.acc, w3); + k.biasAct(e.acc, e.l1_b, e.h1, w3, cfg.hidden, false, .gelu); + k.linearF32(e.l2_w, e.h1, cfg.hidden, cfg.hidden, cfg.hidden, e.acc, w3); + k.biasAct(e.acc, e.l2_b, e.audio_emb, w3, cfg.hidden, false, .none); + e.audio_tokens = w3; + return w3; + } + + // ---- the backbone + + /// Prompt rows from an empty window, the audio rows from `audio_at` from the encoder; `hidden_out` gets every row. + pub fn prefill(e: *Engine, ids: []const u32, audio_at: usize, hidden_out: ?[]f32) !void { + const k = &e.k; + const R = e.o.prefill_rows; + e.bb.reset(); + var at: usize = 0; + const ibuf = try e.gpa.alloc(i32, R); + defer e.gpa.free(ibuf); + const hb = try e.gpa.alloc(u16, R * cfg.hidden); + defer e.gpa.free(hb); + while (at < ids.len) { + const t = @min(R, ids.len - at); + for (0..t) |j| { + const g = at + j; + const audio = g >= audio_at and g < audio_at + e.audio_tokens; + ibuf[j] = if (audio) -1 else @intCast(ids[g]); + } + // a chunk's ids in order with the previous chunk's work + var w: usize = 0; + while (w < t) : (w += 1024) e.g.set(e.ids + w * 4, std.mem.sliceAsBytes(ibuf[w..@min(t, w + 1024)])); + k.embed(e.embed, e.ids, e.bb_x, t, cfg.hidden); + const lo = @max(at, audio_at); + const hi = @min(at + t, audio_at + e.audio_tokens); + if (hi > lo) e.d2d(e.bb_x + (lo - at) * cfg.hidden * 2, e.audio_emb + (lo - audio_at) * cfg.hidden * 2, (hi - lo) * cfg.hidden * 2); + try e.bb.forward(k, e.bb_x, t, e.bb_out); + if (hidden_out) |ho| { + try e.d2h(std.mem.sliceAsBytes(hb[0 .. t * cfg.hidden]), e.bb_out); + for (0..t * cfg.hidden) |j| ho[at * cfg.hidden + j] = bf16(hb[j]); + } + e.last = e.bb_out + (t - 1) * cfg.hidden * 2; + at += t; + } + } + + /// One backbone row from the embedding in `bb_x` row 0; the state lands in `bb_out` row 0. Committed, not waited. + pub fn stepRow(e: *Engine) !void { + const t0 = e.g.gpu_s; + e.last = e.bb_out; + try e.bb.forward(&e.k, e.bb_x, 1, e.bb_out); + try e.g.commit(); + try e.phase(.row, t0); + } + + /// The greedy text token from the last hidden state (placeholders excluded), left on the device in `pick`. + pub fn textPick(e: *Engine) !u32 { + const t0 = e.g.gpu_s; + defer e.phase(.text, t0) catch {}; + e.k.linear(e.embed, e.last, cfg.hidden, cfg.hidden, cfg.vocab, e.logits, cfg.vocab, 1); + e.k.argmax(e.logits, cfg.vocab, e.sup, 4, e.pick, e.amax_ws); + var v: i32 = 0; + try e.d2h(std.mem.asBytes(&v), e.pick); + return @intCast(v); + } + + /// Text logits (fp32 on the host) from the last hidden state. + pub fn textLogits(e: *Engine, out: []f32) !void { + e.k.linear(e.embed, e.last, cfg.hidden, cfg.hidden, cfg.vocab, e.logits, cfg.vocab, 1); + const hb = try e.gpa.alloc(u16, cfg.vocab); + defer e.gpa.free(hb); + try e.d2h(std.mem.sliceAsBytes(hb), e.logits); + for (out, hb) |*o, h| o.* = bf16(h); + } + + /// Sets the newest hidden state from the host (bf16-exact fp32 values), for checks of the heads alone. + pub fn setHidden(e: *Engine, h: []const f32) !void { + var hb: [cfg.hidden]u16 = undefined; + for (&hb, h) |*o, v| o.* = @intCast(@as(u32, @bitCast(v)) >> 16); + try e.h2d(e.bb_out, std.mem.sliceAsBytes(&hb)); + e.last = e.bb_out; + } + + pub fn lastHidden(e: *Engine, out: []f32) !void { + var hb: [cfg.hidden]u16 = undefined; + try e.d2h(std.mem.sliceAsBytes(&hb), e.last); + for (out, hb) |*o, h| o.* = bf16(h); + } + + // ---- the depth decoder + + /// The depth decoder's slots from the last hidden state (run once a frame). + pub fn depthStart(e: *Engine) !void { + e.k.linearF32(e.depth_w, e.last, cfg.hidden, cfg.hidden, cfg.codebooks * cfg.depth_d, e.din_f32, 1); + e.k.biasAct(e.din_f32, e.depth_b, e.din, 1, cfg.codebooks * cfg.depth_d, false, .none); + } + + /// One pass's logits [12, 2017] in `dlog` given the codes and the masked set at `codes` / `masked` on the device. + fn depthLogits(e: *Engine, codes: u64, masked: u64) !void { + const k = &e.k; + const NQ = cfg.codebooks; + k.depthIn(e.din, e.mask_emb, e.depth_tables, codes, masked, e.dx, NQ, cfg.depth_d, cfg.codebook_size); + e.depth.reset(); + try e.depth.forward(k, e.dx, NQ, e.dout); + k.rmsnormRows(e.dout, 0, e.depth_norms, cfg.depth_d, e.dz, cfg.depth_d, cfg.depth_d, cfg.depth_d, NQ, 1e-5, .head); + k.linearBatched(e.depth_tables, cfg.codebook_size * cfg.depth_d, e.dz, cfg.depth_d, cfg.depth_d, cfg.codebook_size, e.dlog, cfg.codebook_size, 1, NQ); + } + + /// One pass: every codebook's logits (host, fp32 [12, 2017]) given the known codes and the masked set. + pub fn depthPass(e: *Engine, codes: *const [cfg.codebooks]i32, masked: *const [cfg.codebooks]bool) ![]f32 { + var m: [cfg.codebooks]i32 = undefined; + for (&m, masked) |*x, y| x.* = @intFromBool(y); + e.g.set(e.dcodes, std.mem.sliceAsBytes(codes)); + e.g.set(e.dmask, std.mem.sliceAsBytes(&m)); + try e.depthLogits(e.dcodes, e.dmask); + e.k.bf16ToF32(e.dlog, e.dlog_f32, cfg.codebooks * cfg.codebook_size); + try e.d2h(std.mem.sliceAsBytes(e.host_logits), e.dlog_f32); + return e.host_logits; + } + + /// A speech frame's codes from the last hidden state: the codebooks revealed over the mask schedule's passes. + pub fn sampleFrame(e: *Engine, smp: Sampling, rng: std.Random, history: []const [cfg.codebooks]i32) ![cfg.codebooks]i32 { + try e.depthStart(); + var codes: [cfg.codebooks]i32 = @splat(0); + var masked: [cfg.codebooks]bool = @splat(true); + var order: [cfg.codebooks]usize = undefined; + for (&order, 0..) |*x, i| x.* = i; + if (!smp.identity_order) rng.shuffle(usize, &order); + var at: usize = 0; + for (cfg.maskgit_schedule) |n| { + const logits = try e.depthPass(&codes, &masked); + for (order[at .. at + n]) |q| codes[q] = pickCode(logits[q * cfg.codebook_size ..][0..cfg.codebook_size], q, smp, rng, history); + for (order[at .. at + n]) |q| masked[q] = false; + at += n; + } + return codes; + } + + /// A speech frame picked on the device into `dframe`'s codes, pushed to `hist` at `slot` unless it ends speech. + pub fn frameDevice(e: *Engine, smp: Sampling, order: *const [cfg.codebooks]i32, frame: usize, hn: usize, slot: usize) !void { + const NQ = cfg.codebooks; + const k = &e.k; + var init: [3 * NQ]i32 = undefined; + @memcpy(init[0..NQ], order); + @memset(init[NQ .. 2 * NQ], 0); + @memset(init[2 * NQ .. 3 * NQ], 1); + e.g.set(e.dframe, std.mem.sliceAsBytes(&init)); + const codes = e.dframe + NQ * 4; + const masked = e.dframe + 2 * NQ * 4; + try e.depthStart(); + var at: usize = 0; + for (cfg.maskgit_schedule) |n| { + try e.depthLogits(codes, masked); + k.depthPick(.{ .logits = e.dlog, .order = e.dframe, .p0 = at, .n = n, .hist = e.hist, .hn = hn, .penalty = smp.audio_repetition_penalty, .temperature = smp.audio_temperature, .seed = smp.seed, .frame = frame, .codes = codes, .masked = masked, .nq = NQ, .size = cfg.codebook_size }); + at += n; + } + k.frameEnd(codes, e.hist, NQ, cfg.audio_eos_code, slot, true); + } + + // ---- the codec + + /// Decodes `frames` (each 12 codes) after those already decoded; windowed frames for the overlap-add (host). + pub fn codecChunk(e: *Engine, frames: []const [cfg.codebooks]i32, lin_out: ?[]f32) ![]const f32 { + const k = &e.k; + const n = frames.len; + if (n == 0) return &.{}; + if (n > e.o.codec_chunk) return error.CodecChunkTooLong; + try e.h2d(e.c_codes, std.mem.sliceAsBytes(frames)); + k.fsqDecode(e.c_codes, e.c_lat, n, cfg.codebooks); + k.linearF32(e.po_w, e.c_lat, cfg.codebooks * 4, cfg.codebooks * 4, cfg.quantizer_dim, e.c_f32, n); + k.biasAct(e.c_f32, e.po_b, e.c_p1, n, cfg.quantizer_dim, false, .none); + k.linearF32(e.fq_w, e.c_p1, cfg.quantizer_dim, cfg.quantizer_dim, cfg.codec_d, e.c_f32, n); + k.biasAct(e.c_f32, e.fq_b, e.c_p2, n, cfg.codec_d, false, .none); + const r = cfg.frame_stack * n; + k.repeatRows(e.c_p2, e.c_up, n, cfg.codec_d, cfg.frame_stack); + try e.codec.forward(k, e.c_up, r, e.c_out); + k.linearF32(e.lin_w, e.c_out, cfg.codec_d, cfg.codec_d, 2 * cfg.spec_bins, e.c_f32, r); + k.biasAct(e.c_f32, e.lin_b, e.c_lin, r, 2 * cfg.spec_bins, false, .none); + if (lin_out) |lo| { + const hb = try e.gpa.alloc(u16, r * 2 * cfg.spec_bins); + defer e.gpa.free(hb); + try e.d2h(std.mem.sliceAsBytes(hb), e.c_lin); + for (lo[0..hb.len], hb) |*o, h| o.* = bf16(h); + } + k.polar(e.c_lin, e.c_spec, r, cfg.spec_bins); + k.matmul(false, false, r, cfg.n_fft, 2 * cfg.spec_bins, e.c_spec, e.basis, e.c_frames, .f32, .f32); + const out = e.host_frames[0 .. r * cfg.n_fft]; + try e.d2h(std.mem.sliceAsBytes(out), e.c_frames); + return out; + } + + pub fn codecReset(e: *Engine) !void { + e.codec.reset(); + } + + /// A reply after `prefill`, until <|im_end|> in text or `max_new` slots (reply.zig). + pub const reply = @import("reply.zig").reply; +}; diff --git a/zig/src/families/lfm3_omni/metal/gpu.zig b/zig/src/families/lfm3_omni/metal/gpu.zig new file mode 100644 index 0000000000..508dac0330 --- /dev/null +++ b/zig/src/families/lfm3_omni/metal/gpu.zig @@ -0,0 +1,286 @@ +//! LFM3-Omni's Metal device: one queue, buffers as u64 (slot << 40 | offset), one serial encoder until a flush. +const std = @import("std"); +const mtl = @import("metal"); + +/// Every pipeline of lfm3_omni.metallib, by its kernel name. +pub const P = enum { + om_rmsnorm, + om_rmsnorm_1k, + om_head_norm, + om_layernorm, + om_rope, + om_copy_rows, + om_set_words, + om_zero_words, + om_copy_words, + om_short_conv, + om_short_conv_rows, + om_conv_state, + om_attention_tiles, + om_attend, + om_qk_prep, + om_swiglu, + om_bias_act, + om_embed, + om_embed_codes, + om_depth_in, + om_conv2d_s2, + om_f32_bf16, + om_bf16_f32, + om_fsq_decode, + om_repeat_rows, + om_polar, + om_argmax_part, + om_argmax_final, + om_rows_1, + om_rows_4, + om_rows_8, + om_rows_12, + om_rows_16, + om_gemm_bb, + om_gemm_bf, + om_gemm_ff, + om_depth_pick, + om_frame_end, + om_mel_frames, + om_power, + om_logmel, +}; + +const n_pipes = @typeInfo(P).@"enum".field_names.len; + +/// The tensor-unit products (omni_nax.metal), built only on a GPU that has the units (M5 and later). +pub const Nax = enum { om_nax_bb_nt, om_nax_bf_nt, om_nax_bf_nn, om_nax_bf_tt }; +const n_nax = @typeInfo(Nax).@"enum".field_names.len; +const slot_shift = 40; + +pub const Ref = struct { buf: mtl.Buffer, off: usize }; + +pub const Gpu = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + lib: mtl.Library, + pipes: [n_pipes]mtl.Pipeline, + /// Null without tensor units (the simdgroup-matrix product takes every product then). + nax: [n_nax]?mtl.Pipeline = @splat(null), + bufs: std.ArrayList(mtl.Buffer), // by slot; slot 0 is a small dummy, bound where a kernel's optional input is null + bytes: usize = 0, + pool: ?mtl.objc.Pool = null, + cb: ?mtl.CommandBuffer = null, + enc: ?mtl.ComputeEncoder = null, + pending: ?mtl.CommandBuffer = null, // committed, not yet waited for + /// Dispatches encoded since the last commit, and in all (for the profile). + encoded: usize = 0, + dispatches: usize = 0, + /// GPU seconds of the command buffers waited for (`flush`), for the profile. + gpu_s: f64 = 0, + + pub fn open(gpa: std.mem.Allocator, metallib: []const u8) !*Gpu { + const g = try gpa.create(Gpu); + errdefer gpa.destroy(g); + g.* = .{ .gpa = gpa, .device = try mtl.Device.init(), .queue = undefined, .lib = undefined, .pipes = undefined, .bufs = .empty }; + errdefer g.device.deinit(); + g.queue = try g.device.queue(); + errdefer g.queue.deinit(); + g.lib = try mtl.Library.fromBytes(g.device, metallib); + errdefer g.lib.deinit(); + var made: usize = 0; + errdefer for (g.pipes[0..made]) |pipe| pipe.deinit(); + inline for (@typeInfo(P).@"enum".field_names, 0..) |name, i| { + g.pipes[i] = try mtl.Pipeline.init(g.device, g.lib, name, false); + made += 1; + } + errdefer for (g.nax) |n| if (n) |pipe| pipe.deinit(); + // LFM3_OMNI_NAX=0 keeps every product on the simdgroup-matrix kernels (an A/B of the two) + const off = if (std.c.getenv("LFM3_OMNI_NAX")) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; + if (g.device.tensorUnits() and !off) inline for (@typeInfo(Nax).@"enum".field_names, 0..) |name, i| { + g.nax[i] = mtl.Pipeline.initOptional(g.device, g.lib, name, false) catch null; + }; + errdefer g.bufs.deinit(gpa); + const dummy = try g.device.buffer(256, mtl.ResourceOptions.shared); + try g.bufs.append(gpa, dummy); + return g; + } + + pub fn close(g: *Gpu) void { + g.flush() catch {}; + for (g.bufs.items) |b| b.deinit(); + g.bufs.deinit(g.gpa); + for (&g.pipes) |*p| p.deinit(); + for (g.nax) |n| if (n) |pipe| pipe.deinit(); + g.lib.deinit(); + g.queue.deinit(); + g.device.deinit(); + g.gpa.destroy(g); + } + + /// A new zeroed shared buffer of `bytes`, by address. + pub fn alloc(g: *Gpu, bytes: usize) !u64 { + const len = std.mem.alignForward(usize, @max(bytes, 256), 256); + if (len > g.device.maxWorkingSet() -| g.device.allocated()) return error.OutOfMemory; + const b = try g.device.buffer(len, mtl.ResourceOptions.shared); + errdefer b.deinit(); + @memset(b.contents()[0..len], 0); + try g.bufs.append(g.gpa, b); + g.bytes += len; + return @as(u64, g.bufs.items.len - 1) << slot_shift; + } + + pub fn ref(g: *const Gpu, addr: u64) Ref { + return .{ .buf = g.bufs.items[@intCast(addr >> slot_shift)], .off = @intCast(addr & ((@as(u64, 1) << slot_shift) - 1)) }; + } + + /// The host view of `len` bytes at `addr` (unified memory; the GPU's writes are there after a flush). + pub fn host(g: *const Gpu, addr: u64, len: usize) []u8 { + const r = g.ref(addr); + return r.buf.contents()[r.off..][0..len]; + } + + /// The encoder of the open command buffer (made on first use). + pub fn encoder(g: *Gpu) mtl.ComputeEncoder { + if (g.enc) |e| return e; + if (g.pool == null) g.pool = mtl.objc.Pool.push(); + const cb = g.queue.commandBufferUnretained(); + g.cb = cb; + const e = cb.compute(.serial); + g.enc = e; + return e; + } + + /// One dispatch: buffers bound at 0.. (address 0 the dummy), then `args` (extern struct, scalar or void) after. + pub fn go(g: *Gpu, p: P, groups: [3]usize, threads: usize, bufs: []const u64, args: anytype) void { + const e = g.encoder(); + e.setPipeline(g.pipes[@backingInt(p)]); + for (bufs, 0..) |a, i| { + const r = g.ref(a); + e.setBuffer(r.buf, r.off, i); + } + if (@TypeOf(args) != void) e.setBytes(std.mem.asBytes(&args), bufs.len); + e.dispatchGroups(mtl.Size.of(groups[0], groups[1], groups[2]), mtl.Size.of(threads, 1, 1)); + g.encoded += 1; + g.dispatches += 1; + } + + /// `go` with `tg_bytes` of threadgroup memory at threadgroup index 0. + pub fn goTg(g: *Gpu, p: P, groups: [3]usize, threads: usize, bufs: []const u64, args: anytype, tg_bytes: usize) void { + const e = g.encoder(); + e.setPipeline(g.pipes[@backingInt(p)]); + for (bufs, 0..) |a, i| { + const r = g.ref(a); + e.setBuffer(r.buf, r.off, i); + } + e.setBytes(std.mem.asBytes(&args), bufs.len); + e.setThreadgroupMemory(std.mem.alignForward(usize, @max(tg_bytes, 16), 16), 0); + e.dispatchGroups(mtl.Size.of(groups[0], groups[1], groups[2]), mtl.Size.of(threads, 1, 1)); + g.encoded += 1; + g.dispatches += 1; + } + + /// Whether every tensor-unit product built. + pub fn hasNax(g: *const Gpu) bool { + for (g.nax) |n| if (n == null) return false; + return true; + } + + /// A tensor-unit product's dispatch (as `go`). + pub fn goNax(g: *Gpu, p: Nax, groups: [3]usize, threads: usize, bufs: []const u64, args: anytype) void { + const e = g.encoder(); + e.setPipeline(g.nax[@backingInt(p)].?); + for (bufs, 0..) |a, i| { + const r = g.ref(a); + e.setBuffer(r.buf, r.off, i); + } + e.setBytes(std.mem.asBytes(&args), bufs.len); + e.dispatchGroups(mtl.Size.of(groups[0], groups[1], groups[2]), mtl.Size.of(threads, 1, 1)); + g.encoded += 1; + g.dispatches += 1; + } + + /// One dispatch with a raw byte argument (at most 4 KB) after the buffers, and a scalar after it. + pub fn goBytes(g: *Gpu, p: P, groups: [3]usize, threads: usize, bufs: []const u64, bytes: []const u8, scalar: anytype) void { + const e = g.encoder(); + e.setPipeline(g.pipes[@backingInt(p)]); + for (bufs, 0..) |a, i| { + const r = g.ref(a); + e.setBuffer(r.buf, r.off, i); + } + e.setBytes(bytes, bufs.len); + e.setBytes(std.mem.asBytes(&scalar), bufs.len + 1); + e.dispatchGroups(mtl.Size.of(groups[0], groups[1], groups[2]), mtl.Size.of(threads, 1, 1)); + g.encoded += 1; + g.dispatches += 1; + } + + /// Commit what is encoded without waiting (the previous commit is waited for first). + pub fn commit(g: *Gpu) !void { + const e = g.enc orelse return; + e.end(); + const cb = g.cb.?.retain(); + cb.commit(); + g.enc = null; + g.cb = null; + g.encoded = 0; + defer { + if (g.pool) |p| p.pop(); + g.pool = null; + } + g.waitPending() catch |err| { + cb.wait(); + cb.release(); + return err; + }; + g.pending = cb; + } + + fn waitPending(g: *Gpu) !void { + const cb = g.pending orelse return; + g.pending = null; + defer cb.release(); + cb.wait(); + g.gpu_s += cb.gpuSeconds(); + if (cb.failure()) |msg| { + std.log.err("lfm3_omni metal: command buffer failed: {s}", .{msg}); + return error.CommandBufferFailed; + } + } + + /// Commit and wait until the GPU is done. + pub fn flush(g: *Gpu) !void { + try g.commit(); + try g.waitPending(); + } + + /// Host bytes into a buffer after the GPU is done with it (a flush first). + pub fn put(g: *Gpu, dst: u64, bytes: []const u8) !void { + try g.flush(); + @memcpy(g.host(dst, bytes.len), bytes); + } + + /// Buffer bytes to the host once the GPU's work so far is done. + pub fn get(g: *Gpu, src: u64, out: []u8) !void { + try g.flush(); + @memcpy(out, g.host(src, out.len)); + } + + /// Small host bytes (whole 32-bit words, at most 4 KB) written in order with the dispatches, without a flush. + pub fn set(g: *Gpu, dst: u64, bytes: []const u8) void { + std.debug.assert(bytes.len % 4 == 0 and bytes.len <= 4096 and bytes.len > 0); + const n: u32 = @intCast(bytes.len / 4); + g.goBytes(.om_set_words, .{ (n + 255) / 256, 1, 1 }, 256, &.{dst}, bytes, n); + } + + /// `bytes` (whole words) zeroed on the GPU, in order with the dispatches. + pub fn zero(g: *Gpu, dst: u64, bytes: usize) void { + std.debug.assert(bytes % 4 == 0); + const n: u32 = @intCast(bytes / 4); + g.go(.om_zero_words, .{ (n + 255) / 256, 1, 1 }, 256, &.{dst}, n); + } + + /// `bytes` copied on the GPU, in order with the dispatches around it. + pub fn copy(g: *Gpu, src: u64, dst: u64, bytes: usize) void { + std.debug.assert(bytes % 4 == 0); + const n: u32 = @intCast(bytes / 4); + g.go(.om_copy_words, .{ (n + 255) / 256, 1, 1 }, 256, &.{ src, dst }, n); + } +}; diff --git a/zig/src/families/lfm3_omni/metal/kernels.zig b/zig/src/families/lfm3_omni/metal/kernels.zig new file mode 100644 index 0000000000..60c0a53689 --- /dev/null +++ b/zig/src/families/lfm3_omni/metal/kernels.zig @@ -0,0 +1,296 @@ +//! LFM3-Omni's Metal launches, each binding its buffers and argument struct as zig/kernels/metal/lfm3_omni declares. +const std = @import("std"); +const gpu = @import("gpu.zig"); + +const Gpu = gpu.Gpu; + +pub const Act = enum(i32) { none = 0, silu = 1, gelu = 2 }; +pub const Norm = enum(i32) { lfm = 0, head = 1 }; +pub const DataType = enum { bf16, f32 }; + +/// Rows the few-row product takes; more go through the tiled product. +pub const gemv_rows = 16; + +fn i(x: usize) i32 { + return @intCast(x); +} + +fn blocks(n: usize, per: usize) usize { + return (n + per - 1) / per; +} + +pub const Kernels = struct { + g: *Gpu, + + fn go(k: *const Kernels, p: gpu.P, n: usize, per: usize, bufs: []const u64, args: anytype) void { + k.g.go(p, .{ blocks(n, per), 1, 1 }, per, bufs, args); + } + + // ---- matrix products + + /// Whether `t` rows of `kdim` (rows `ldx` apart) take the few-row product. + pub fn few(t: usize, kdim: usize, ldx: usize, x: u64) bool { + return t <= gemv_rows and kdim % 8 == 0 and ldx % 8 == 0 and x % 16 == 0; + } + + /// The few-row product's result: bf16 / fp32 rows, or SwiGLU of interleaved (a, c) column pairs. + pub const Epi = enum(i32) { plain = 0, swiglu = 2 }; + + pub const RowsArgs = extern struct { K: i32, N: i32, T: i32, ldx: i32, ldy: i32, has_y: i32, has_yf: i32, epi: i32, sw: i64, sx: i64, sy: i64 }; + + /// The few-row product into bf16 y and / or fp32 yf; with `.swiglu` columns are (a, c) pairs, y n / 2 wide. + pub fn product(k: *const Kernels, w: u64, x: u64, ldx: usize, kdim: usize, n: usize, t: usize, y: u64, ldy: usize, yf: u64, batch: usize, sw: usize, sx: usize, sy: usize, epi: Epi) void { + const p: gpu.P = if (t <= 1) .om_rows_1 else if (t <= 4) .om_rows_4 else if (t <= 8) .om_rows_8 else if (t <= 12) .om_rows_12 else .om_rows_16; + k.g.go(p, .{ blocks(n, 16), 1, batch }, 128, &.{ w, x, y, yf }, RowsArgs{ .K = i(kdim), .N = i(n), .T = i(t), .ldx = i(ldx), .ldy = i(ldy), .has_y = @intFromBool(y != 0), .has_yf = @intFromBool(yf != 0), .epi = @backingInt(epi), .sw = @intCast(sw), .sx = @intCast(sx), .sy = @intCast(sy) }); + } + + pub fn gemv(k: *const Kernels, w: u64, x: u64, ldx: usize, kdim: usize, n: usize, t: usize, y: u64, ldy: usize, yf: u64, batch: usize, sw: usize, sx: usize, sy: usize) void { + k.product(w, x, ldx, kdim, n, t, y, ldy, yf, batch, sw, sx, sy, .plain); + } + + /// C [m, n] = A [m, kd] B [kd, n], A stored [kd, m] when `ta`, B [n, kd] when `tb`; `batch` products. + pub fn gemm(k: *const Kernels, ta: bool, tb: bool, m: usize, n: usize, kd: usize, a: u64, lda: usize, b: u64, ldb: usize, c: u64, ldc: usize, in_t: DataType, out_t: DataType, batch: usize, sa: usize, sb: usize, sc: usize) void { + const A = extern struct { M: i32, N: i32, K: i32, lda: i32, ldb: i32, ldc: i32, ta: i32, tb: i32, sa: i64, sb: i64, sc: i64 }; + const p: gpu.P = switch (in_t) { + .bf16 => if (out_t == .bf16) .om_gemm_bb else .om_gemm_bf, + .f32 => .om_gemm_ff, + }; + std.debug.assert(!(in_t == .f32 and out_t == .bf16)); + if (in_t == .bf16 and k.g.hasNax()) { + const nax: ?gpu.Nax = if (!ta and tb) (if (out_t == .bf16) .om_nax_bb_nt else .om_nax_bf_nt) else if (out_t == .f32 and !ta and !tb) .om_nax_bf_nn else if (out_t == .f32 and ta and tb) .om_nax_bf_tt else null; + if (nax) |np| { + const N = extern struct { M: i32, N: i32, K: i32, lda: i32, ldb: i32, ldd: i32, kblocks: i32, pad: i32 = 0, sa: i64, sb: i64, sd: i64 }; + return k.g.goNax(np, .{ blocks(n, 128), blocks(m, 64), batch }, 256, &.{ a, b, c }, N{ .M = i(m), .N = i(n), .K = i(kd), .lda = i(lda), .ldb = i(ldb), .ldd = i(ldc), .kblocks = i(kd / 256), .sa = @intCast(sa), .sb = @intCast(sb), .sd = @intCast(sc) }); + } + } + k.g.go(p, .{ blocks(n, 64), blocks(m, 64), batch }, 128, &.{ a, b, c }, A{ .M = i(m), .N = i(n), .K = i(kd), .lda = i(lda), .ldb = i(ldb), .ldc = i(ldc), .ta = @intFromBool(ta), .tb = @intFromBool(tb), .sa = @intCast(sa), .sb = @intCast(sb), .sc = @intCast(sc) }); + } + + /// y [t, n] = x [t, k] w^T with w [n, k]; x rows `ldx` apart, y rows `ldy` apart; bf16 out, fp32 accumulate. + pub fn linear(k: *const Kernels, w: u64, x: u64, ldx: usize, kdim: usize, n: usize, y: u64, ldy: usize, t: usize) void { + if (few(t, kdim, ldx, x)) return k.gemv(w, x, ldx, kdim, n, t, y, ldy, 0, 1, 0, 0, 0); + k.gemm(false, true, t, n, kdim, x, ldx, w, kdim, y, ldy, .bf16, .bf16, 1, 0, 0, 0); + } + + /// As `linear`, the result in fp32 (for a bias added before the one rounding). + pub fn linearF32(k: *const Kernels, w: u64, x: u64, ldx: usize, kdim: usize, n: usize, y: u64, t: usize) void { + if (few(t, kdim, ldx, x)) return k.gemv(w, x, ldx, kdim, n, t, 0, n, y, 1, 0, 0, 0); + k.gemm(false, true, t, n, kdim, x, ldx, w, kdim, y, n, .bf16, .f32, 1, 0, 0, 0); + } + + /// Row-major C [m, n] = A [m, kd] B [kd, n] with either transposed in memory (`ta`: A stored [kd, m]). + pub fn matmul(k: *const Kernels, ta: bool, tb: bool, m: usize, n: usize, kd: usize, a: u64, b: u64, c: u64, in_t: DataType, out_t: DataType) void { + k.gemm(ta, tb, m, n, kd, a, if (ta) m else kd, b, if (tb) kd else n, c, n, in_t, out_t, 1, 0, 0, 0); + } + + /// `batch` products y_j [t, n] = x_j [t, kd] w_j^T, the operands `sw`, `sx`, `sy` elements apart; bf16 out. + pub fn linearBatched(k: *const Kernels, w: u64, sw: usize, x: u64, sx: usize, kd: usize, n: usize, y: u64, sy: usize, t: usize, batch: usize) void { + if (few(t, kd, kd, x) and sx % 8 == 0) return k.gemv(w, x, kd, kd, n, t, y, n, 0, batch, sw, sx, sy); + k.gemm(false, true, t, n, kd, x, kd, w, kd, y, n, .bf16, .bf16, batch, sx, sw, sy); + } + + // ---- elementwise and row kernels + + pub fn rmsnorm(k: *const Kernels, x: u64, r: u64, w: u64, y: u64, d: usize, ld: usize, ldy: usize, rows: usize, eps: f32, mode: Norm) void { + k.rmsnormPost(x, r, 0, w, 0, y, d, ld, ldy, rows, eps, mode); + } + + /// As `rmsnorm` with a weight row a row (`ldw` elements apart). + pub fn rmsnormRows(k: *const Kernels, x: u64, r: u64, w: u64, ldw: usize, y: u64, d: usize, ld: usize, ldy: usize, rows: usize, eps: f32, mode: Norm) void { + k.rmsnormPost(x, r, 0, w, ldw, y, d, ld, ldy, rows, eps, mode); + } + + /// As `rmsnormRows`, r first normed by `pw` (a sandwich layer's post norm) when given. + pub fn rmsnormPost(k: *const Kernels, x: u64, r: u64, pw: u64, w: u64, ldw: usize, y: u64, d: usize, ld: usize, ldy: usize, rows: usize, eps: f32, mode: Norm) void { + const A = extern struct { D: i32, ld: i32, ldy: i32, eps: f32, mode: i32, ldw: i32, has_r: i32, has_pw: i32 }; + // a few rows: a wide group each (the row's latency is all there is) + const wide = rows <= gemv_rows and d >= 1024; + k.g.go(if (wide) .om_rmsnorm_1k else .om_rmsnorm, .{ rows, 1, 1 }, if (wide) 1024 else 256, &.{ x, r, w, y, pw }, A{ .D = i(d), .ld = i(ld), .ldy = i(ldy), .eps = eps, .mode = @backingInt(mode), .ldw = i(ldw), .has_r = @intFromBool(r != 0), .has_pw = @intFromBool(pw != 0) }); + } + + /// q and k heads of `proj` normed and rotated, k and v written to cache row pos0 + t. + pub fn qkPrep(k: *const Kernels, proj: u64, qn: u64, kn: u64, inv: u64, kc: u64, vc: u64, t: usize, h: usize, hkv: usize, hd: usize, ld: usize, pos0: usize, mode: u8, nk: usize, eps: f32) void { + const A = extern struct { T: i32, H: i32, Hkv: i32, hd: i32, ld: i32, pos0: i32, mode: i32, nk: i32, eps: f32 }; + std.debug.assert(hd <= 128); + k.g.go(.om_qk_prep, .{ t * (h + hkv), 1, 1 }, 32, &.{ proj, qn, kn, inv, kc, vc }, A{ .T = i(t), .H = i(h), .Hkv = i(hkv), .hd = i(hd), .ld = i(ld), .pos0 = i(pos0), .mode = mode, .nk = i(nk), .eps = eps }); + } + + pub fn headNorm(k: *const Kernels, x: u64, w: u64, t: usize, h: usize, hd: usize, ld: usize, eps: f32) void { + const A = extern struct { T: i32, H: i32, hd: i32, ld: i32, eps: f32 }; + k.go(.om_head_norm, t * h * 32, 256, &.{ x, w }, A{ .T = i(t), .H = i(h), .hd = i(hd), .ld = i(ld), .eps = eps }); + } + + pub fn rope(k: *const Kernels, x: u64, inv: u64, t: usize, h: usize, hd: usize, ld: usize, pos0: usize, mode: u8) void { + const A = extern struct { T: i32, H: i32, hd: i32, ld: i32, pos0: i32, mode: i32 }; + k.go(.om_rope, t * h * hd / 2, 256, &.{ x, inv }, A{ .T = i(t), .H = i(h), .hd = i(hd), .ld = i(ld), .pos0 = i(pos0), .mode = mode }); + } + + pub fn copyRows(k: *const Kernels, src: u64, dst: u64, t: usize, n: usize, lds: usize, ldd: usize) void { + const A = extern struct { T: i32, n: i32, lds: i32, ldd: i32 }; + k.go(.om_copy_rows, t * n, 256, &.{ src, dst }, A{ .T = i(t), .n = i(n), .lds = i(lds), .ldd = i(ldd) }); + } + + pub fn shortConv(k: *const Kernels, bcx: u64, ld: usize, w: u64, state: u64, y: u64, ldy: usize, t: usize, d: usize, kw: usize) void { + const A = extern struct { ld: i32, ldy: i32, T: i32, D: i32, K: i32 }; + const a: A = .{ .ld = i(ld), .ldy = i(ldy), .T = i(t), .D = i(d), .K = i(kw) }; + if (t > 4) { + k.go(.om_short_conv_rows, t * d, 256, &.{ bcx, w, state, y }, a); + return k.go(.om_conv_state, d, 128, &.{ bcx, state }, a); + } + k.go(.om_short_conv, d, 128, &.{ bcx, w, state, y }, a); + } + + pub const Attn = struct { + q: u64, + ldq: usize, + kc: u64, + vc: u64, + ldkv: usize, + o: u64, + ldo: usize, + t: usize, + heads: usize, + kv_heads: usize, + hd: usize, + pos0: usize, + kv_len: usize, + window: usize, + causal: bool, + /// LFM3's per-head gate logits (rows `ldg` apart), applied to the output; 0: none. + gate: u64 = 0, + ldg: usize = 0, + }; + + pub fn attention(k: *const Kernels, p: Attn) !void { + if (p.hd % 32 != 0 or p.hd > 128) return error.UnsupportedHeadDim; + if (p.t == 1 and p.kv_len > 4096) return error.AttentionSpanTooLong; + if (p.t == 1) { + const A = extern struct { ldkv: i32, H: i32, Hkv: i32, hd: i32, pos: i32, kv_len: i32, window: i32, causal: i32, has_gate: i32, gate_off: i32, scale: f32 }; + const span = if (p.causal and p.window > 0) @min(p.window, p.kv_len) else p.kv_len; + if (span > 4096) return error.AttentionSpanTooLong; + // the gate's logits sit in the query's row (q | k | v | gate) + const off: usize = if (p.gate != 0) (p.gate - p.q) / 2 else 0; + return k.g.go(.om_attend, .{ p.heads, 1, 1 }, 256, &.{ p.q, p.kc, p.vc, p.o }, A{ .ldkv = i(p.ldkv), .H = i(p.heads), .Hkv = i(p.kv_heads), .hd = i(p.hd), .pos = i(p.pos0), .kv_len = i(p.kv_len), .window = i(p.window), .causal = @intFromBool(p.causal), .has_gate = @intFromBool(p.gate != 0), .gate_off = i(off), .scale = 1.0 / @sqrt(@as(f32, @floatFromInt(p.hd))) }); + } + const A = extern struct { ldq: i32, ldkv: i32, ldo: i32, T: i32, H: i32, Hkv: i32, hd: i32, pos0: i32, kv_len: i32, window: i32, causal: i32, scale: f32, has_gate: i32, ldg: i32 }; + const a: A = .{ .ldq = i(p.ldq), .ldkv = i(p.ldkv), .ldo = i(p.ldo), .T = i(p.t), .H = i(p.heads), .Hkv = i(p.kv_heads), .hd = i(p.hd), .pos0 = i(p.pos0), .kv_len = i(p.kv_len), .window = i(p.window), .causal = @intFromBool(p.causal), .scale = 1.0 / @sqrt(@as(f32, @floatFromInt(p.hd))), .has_gate = @intFromBool(p.gate != 0), .ldg = i(p.ldg) }; + k.g.go(.om_attention_tiles, .{ blocks(p.t, 16), p.heads, 1 }, 128, &.{ p.q, p.kc, p.vc, p.o, p.gate }, a); + } + + /// SwiGLU over rows of 2f: halves, or interleaved (a, c) pairs. + pub fn swiglu(k: *const Kernels, ac: u64, y: u64, t: usize, f: usize, inter: bool) void { + const A = extern struct { T: i32, F: i32, inter: i32 }; + k.go(.om_swiglu, t * f, 256, &.{ ac, y }, A{ .T = i(t), .F = i(f), .inter = @intFromBool(inter) }); + } + + pub fn biasAct(k: *const Kernels, acc: u64, bias: u64, y: u64, rows: usize, cols: usize, per_row: bool, act: Act) void { + const A = extern struct { rows: i32, cols: i32, per_row: i32, act: i32, has_bias: i32 }; + k.go(.om_bias_act, rows * cols, 256, &.{ acc, bias, y }, A{ .rows = i(rows), .cols = i(cols), .per_row = @intFromBool(per_row), .act = @backingInt(act), .has_bias = @intFromBool(bias != 0) }); + } + + pub fn layernorm(k: *const Kernels, x: u64, w: u64, bias: u64, y: u64, d: usize, rows: usize, eps: f32) void { + const A = extern struct { D: i32, eps: f32 }; + k.g.go(.om_layernorm, .{ rows, 1, 1 }, 256, &.{ x, w, bias, y }, A{ .D = i(d), .eps = eps }); + } + + const EmbedArgs = extern struct { T: i32, D: i32, NQ: i32, size: i32 }; + + pub fn embed(k: *const Kernels, table: u64, ids: u64, y: u64, t: usize, d: usize) void { + k.go(.om_embed, t * d, 256, &.{ table, ids, y }, EmbedArgs{ .T = i(t), .D = i(d), .NQ = 0, .size = 0 }); + } + + pub fn embedCodes(k: *const Kernels, table: u64, codes: u64, y: u64, t: usize, d: usize, nq: usize, size: usize) void { + k.go(.om_embed_codes, t * d, 256, &.{ table, codes, y }, EmbedArgs{ .T = i(t), .D = i(d), .NQ = i(nq), .size = i(size) }); + } + + pub fn depthIn(k: *const Kernels, din: u64, mask: u64, tables: u64, codes: u64, masked: u64, x: u64, nq: usize, d: usize, size: usize) void { + k.go(.om_depth_in, nq * d, 256, &.{ din, mask, tables, codes, masked, x }, EmbedArgs{ .T = 0, .D = i(d), .NQ = i(nq), .size = i(size) }); + } + + pub fn conv2dS2(k: *const Kernels, x: u64, cin: usize, h: usize, w: usize, wt: u64, bias: u64, y: u64, c: usize, h2: usize, w2: usize, silu: bool) void { + const A = extern struct { CIN: i32, H: i32, W: i32, C: i32, H2: i32, W2: i32, silu: i32 }; + k.go(.om_conv2d_s2, c * h2 * w2, 256, &.{ x, wt, bias, y }, A{ .CIN = i(cin), .H = i(h), .W = i(w), .C = i(c), .H2 = i(h2), .W2 = i(w2), .silu = @intFromBool(silu) }); + } + + pub fn f32ToBf16(k: *const Kernels, x: u64, y: u64, n: usize) void { + k.go(.om_f32_bf16, n, 256, &.{ x, y }, @as(u32, @intCast(n))); + } + + pub fn bf16ToF32(k: *const Kernels, x: u64, y: u64, n: usize) void { + k.go(.om_bf16_f32, n, 256, &.{ x, y }, @as(u32, @intCast(n))); + } + + pub fn fsqDecode(k: *const Kernels, codes: u64, l: u64, n: usize, nq: usize) void { + const A = extern struct { n: i32, NQ: i32 }; + k.go(.om_fsq_decode, n * nq, 128, &.{ codes, l }, A{ .n = i(n), .NQ = i(nq) }); + } + + pub fn repeatRows(k: *const Kernels, x: u64, y: u64, n: usize, d: usize, rep: usize) void { + const A = extern struct { n: i32, D: i32, rep: i32 }; + k.go(.om_repeat_rows, n * rep * d, 256, &.{ x, y }, A{ .n = i(n), .D = i(d), .rep = i(rep) }); + } + + pub fn polar(k: *const Kernels, x: u64, spec: u64, t: usize, nb: usize) void { + const A = extern struct { T: i32, NB: i32 }; + k.go(.om_polar, t * nb, 256, &.{ x, spec }, A{ .T = i(t), .NB = i(nb) }); + } + + /// The first index of the largest of n bf16 logits, `sup` excluded, in two stages through `ws` (8 bytes a part). + pub fn argmax(k: *const Kernels, x: u64, n: usize, sup: u64, ns: usize, out: u64, ws: u64) void { + const A = extern struct { N: i32, ns: i32, per: i32 }; + const per = 1024; + const parts = blocks(n, per); + k.g.go(.om_argmax_part, .{ parts, 1, 1 }, 256, &.{ x, sup, ws, ws + parts * 4 }, A{ .N = i(n), .ns = i(ns), .per = per }); + k.g.go(.om_argmax_final, .{ 1, 1, 1 }, 32, &.{ ws, ws + parts * 4, out }, @as(i32, i(parts))); + } + + pub const Pick = struct { + logits: u64, + order: u64, + p0: usize, + n: usize, + hist: u64, + hn: usize, + penalty: f32, + temperature: f32, + seed: u64, + frame: usize, + codes: u64, + masked: u64, + nq: usize, + size: usize, + }; + + pub fn depthPick(k: *const Kernels, p: Pick) void { + const A = extern struct { p0: i32, hn: i32, penalty: f32, temp: f32, seed: u64, frame: i32, NQ: i32, size: i32, pad: i32 = 0 }; + std.debug.assert(p.size <= 2048); + k.g.go(.om_depth_pick, .{ p.n, 1, 1 }, 256, &.{ p.logits, p.order, p.hist, p.codes, p.masked }, A{ .p0 = i(p.p0), .hn = i(p.hn), .penalty = p.penalty, .temp = p.temperature, .seed = p.seed, .frame = i(p.frame), .NQ = i(p.nq), .size = i(p.size) }); + } + + /// The frame's end: all end-of-speech if any code is; else, with `push`, the frame joins the window at `slot`. + pub fn frameEnd(k: *const Kernels, codes: u64, hist: u64, nq: usize, eos: usize, slot: usize, push: bool) void { + const A = extern struct { NQ: i32, eos: i32, slot: i32, push: i32 }; + k.g.go(.om_frame_end, .{ 1, 1, 1 }, 32, &.{ codes, hist }, A{ .NQ = i(nq), .eos = i(eos), .slot = i(slot), .push = @intFromBool(push) }); + } + + pub fn melFrames(k: *const Kernels, x: u64, len: usize, win: u64, fr: u64, nf: usize, nfft: usize, hop: usize) void { + const A = extern struct { len: i32, nf: i32, nfft: i32, hop: i32 }; + k.go(.om_mel_frames, nf * nfft, 256, &.{ x, win, fr }, A{ .len = i(len), .nf = i(nf), .nfft = i(nfft), .hop = i(hop) }); + } + + pub fn power(k: *const Kernels, spec: u64, p: u64, nf: usize, nb: usize) void { + const A = extern struct { nf: i32, bins: i32 }; + k.go(.om_power, nf * nb, 256, &.{ spec, p }, A{ .nf = i(nf), .bins = i(nb) }); + } + + pub fn logmel(k: *const Kernels, m: u64, out: u64, nf: usize, mels: usize, floor: f32) void { + const A = extern struct { nf: i32, mels: i32, floor_: f32 }; + k.go(.om_logmel, nf * mels, 256, &.{ m, out }, A{ .nf = i(nf), .mels = i(mels), .floor_ = floor }); + } +}; + +test "argument structs match the kernels' layouts" { + try std.testing.expectEqual(@as(usize, 56), @sizeOf(Kernels.RowsArgs)); + const P = extern struct { p0: i32, hn: i32, penalty: f32, temp: f32, seed: u64, frame: i32, NQ: i32, size: i32, pad: i32 = 0 }; + try std.testing.expectEqual(@as(usize, 40), @sizeOf(P)); +} diff --git a/zig/src/families/lfm3_omni/metal/load.zig b/zig/src/families/lfm3_omni/metal/load.zig new file mode 100644 index 0000000000..4c98920731 --- /dev/null +++ b/zig/src/families/lfm3_omni/metal/load.zig @@ -0,0 +1,159 @@ +//! The Metal engine's weights from the checkpoint and its device buffers, sized by its options. +const std = @import("std"); +const st = @import("core").safetensors; +const stack = @import("stack.zig"); +const melmod = @import("../mel.zig"); +const cfg = @import("../config.zig"); +const common = @import("../common.zig"); +const Engine = @import("engine.zig").Engine; + +const Arena = stack.Arena; +const hist_cap = common.hist_cap; + +pub fn model(e: *Engine, dir: []const u8) !void { + const gpa = e.gpa; + const mpath = try std.fs.path.join(gpa, &.{ dir, "model.safetensors" }); + defer gpa.free(mpath); + var f = try st.File.openPrefixes(gpa, e.io, mpath, &.{ "model.audio_encoder.", "model.audio_projector.", "model.language_model.", "model.codebook_embedding.", "model.depth_decoder." }); + defer f.close(e.io); + const a = &e.a; + const max_tokens = melmod.tokens(melmod.frames(e.o.max_input_s * melmod.sample_rate)); + e.enc = try stack.load(gpa, a, &f, cfg.encoderSpec(max_tokens)); + e.bb = try stack.load(gpa, a, &f, cfg.backboneSpec(e.o.prefill_rows, e.o.max_ctx)); + e.depth = try stack.load(gpa, a, &f, cfg.depthSpec()); + const p = "model.audio_encoder.downsample."; + const C = cfg.conv_channels; + e.ds = .{ + .conv_in_w = try up(a, &f, p ++ "conv_in.weight", &.{ C, 1, 3, 3 }), + .conv_in_b = try up(a, &f, p ++ "conv_in.bias", &.{C}), + .conv_mid_w = try up(a, &f, p ++ "conv_mid.weight", &.{ C, 1, 3, 3 }), + .conv_mid_b = try up(a, &f, p ++ "conv_mid.bias", &.{C}), + .lin_mid_w = try up(a, &f, p ++ "lin_mid.weight", &.{ C, C, 1, 1 }), + .lin_mid_b = try up(a, &f, p ++ "lin_mid.bias", &.{C}), + .conv_end_w = try up(a, &f, p ++ "conv_end.weight", &.{ C, 1, 3, 3 }), + .conv_end_b = try up(a, &f, p ++ "conv_end.bias", &.{C}), + .lin_end_w = try up(a, &f, p ++ "lin_end.weight", &.{ C, C, 1, 1 }), + .lin_end_b = try up(a, &f, p ++ "lin_end.bias", &.{C}), + .lin_out_w = try up(a, &f, p ++ "lin_out.weight", &.{ cfg.enc_d, C * cfg.mel_bins / 8, 1 }), + .lin_out_b = try up(a, &f, p ++ "lin_out.bias", &.{cfg.enc_d}), + }; + const q = "model.audio_projector."; + e.ln_w = try up(a, &f, q ++ "layer_norm.weight", &.{cfg.enc_d}); + e.ln_b = try up(a, &f, q ++ "layer_norm.bias", &.{cfg.enc_d}); + e.l1_w = try up(a, &f, q ++ "linear_1.weight", &.{ cfg.hidden, cfg.enc_d }); + e.l1_b = try up(a, &f, q ++ "linear_1.bias", &.{cfg.hidden}); + e.l2_w = try up(a, &f, q ++ "linear_2.weight", &.{ cfg.hidden, cfg.hidden }); + e.l2_b = try up(a, &f, q ++ "linear_2.bias", &.{cfg.hidden}); + e.embed = try up(a, &f, "model.language_model.embed_tokens.weight", &.{ cfg.vocab, cfg.hidden }); + e.cb_embed = try up(a, &f, "model.codebook_embedding.weight", &.{ cfg.codebooks * cfg.codebook_size, cfg.hidden }); + const dd = "model.depth_decoder."; + e.depth_w = try up(a, &f, dd ++ "depth_linear.weight", &.{ cfg.codebooks * cfg.depth_d, cfg.hidden }); + e.depth_b = try up(a, &f, dd ++ "depth_linear.bias", &.{cfg.codebooks * cfg.depth_d}); + e.mask_emb = try up(a, &f, dd ++ "mask_embedding", &.{cfg.depth_d}); + const tbytes = cfg.codebook_size * cfg.depth_d * 2; + e.depth_tables = try a.alloc(cfg.codebooks * tbytes); + e.depth_norms = try a.alloc(cfg.codebooks * cfg.depth_d * 2); + var name: [128]u8 = undefined; + for (0..cfg.codebooks) |i| { + a.write(e.depth_tables, i * tbytes, try stack.tensor(&f, try std.fmt.bufPrint(&name, dd ++ "depth_embeddings.{d}.embedding.weight", .{i}), &.{ cfg.codebook_size, cfg.depth_d })); + a.write(e.depth_norms, i * cfg.depth_d * 2, try stack.tensor(&f, try std.fmt.bufPrint(&name, dd ++ "depth_embeddings.{d}.embedding_norm.weight", .{i}), &.{cfg.depth_d})); + } + f.release(); + + const cpath = try std.fs.path.join(gpa, &.{ dir, "audio_tokenizer", "model.safetensors" }); + defer gpa.free(cpath); + var c = try st.File.openPrefixes(gpa, e.io, cpath, &.{ "decoder.", "quantizer.", "from_quantizer." }); + defer c.close(e.io); + e.codec = try stack.load(gpa, a, &c, cfg.codecSpec(6 * e.o.codec_chunk, 6 * e.o.max_frames + 6 * e.o.codec_chunk)); + e.po_w = try up(a, &c, "quantizer.project_out.weight", &.{ cfg.quantizer_dim, cfg.codebooks * 4 }); + e.po_b = try up(a, &c, "quantizer.project_out.bias", &.{cfg.quantizer_dim}); + e.fq_w = try up(a, &c, "from_quantizer.weight", &.{ cfg.codec_d, cfg.quantizer_dim }); + e.fq_b = try up(a, &c, "from_quantizer.bias", &.{cfg.codec_d}); + e.lin_w = try up(a, &c, "decoder.lin.weight", &.{ 2 * cfg.spec_bins, cfg.codec_d }); + e.lin_b = try up(a, &c, "decoder.lin.bias", &.{2 * cfg.spec_bins}); + const wt = c.get("decoder.istft.window") orelse return error.MissingTensor; + if (!wt.is(.f32, &.{cfg.n_fft})) return error.BadTensor; + @memcpy(std.mem.asBytes(&e.window), wt.bytes); + const basis = try common.synthesisBasis(gpa, &e.window); + defer gpa.free(basis); + e.basis = try a.upload(std.mem.sliceAsBytes(basis)); + c.release(); +} + +fn up(a: *Arena, f: *const st.File, name: []const u8, shape: []const usize) !u64 { + return a.upload(try stack.tensor(f, name, shape)); +} + +pub fn scratch(e: *Engine) !void { + const a = &e.a; + const F = melmod.frames(e.o.max_input_s * melmod.sample_rate); + const T = melmod.tokens(F); + const C = cfg.conv_channels; + const w1 = (F + 1) / 2; + const w2 = (w1 + 1) / 2; + e.mel_f32 = try a.alloc(cfg.mel_bins * F * 4); + e.mel_pcm = try a.alloc((e.o.max_input_s * melmod.sample_rate + melmod.n_fft) * 4); + e.mel_fr = try a.alloc(F * melmod.n_fft * 4); + e.mel_s = try a.alloc(F * 2 * melmod.bins * 4); + e.mel_p = try a.alloc(F * melmod.bins * 4); + e.mel_m = try a.alloc(F * melmod.mels * 4); + { + const gpa = e.gpa; + const win = try gpa.alloc(f32, melmod.n_fft); + defer gpa.free(win); + for (win, e.mel.window) |*w, x| w.* = @floatCast(x); + e.mel_win = try a.upload(std.mem.sliceAsBytes(win)); + const dft = try gpa.alloc(f32, melmod.n_fft * 2 * melmod.bins); + defer gpa.free(dft); + for (0..melmod.n_fft) |n| for (0..melmod.bins) |kk| { + dft[n * 2 * melmod.bins + kk] = @floatCast(e.mel.cos[n][kk]); + dft[n * 2 * melmod.bins + melmod.bins + kk] = @floatCast(e.mel.sin[n][kk]); + }; + e.mel_dft = try a.upload(std.mem.sliceAsBytes(dft)); + e.mel_fb = try a.upload(std.mem.sliceAsBytes(&e.mel.filters)); + } + e.dframe = try a.alloc((3 * cfg.codebooks + 2) * 4); + e.hist = try a.zeroed(hist_cap * cfg.codebooks * 4); + e.mel_bf = try a.alloc(cfg.mel_bins * F * 2); + e.a1 = try a.alloc(C * 64 * w1 * 2); + e.a2 = try a.alloc(C * 32 * w2 * 2); + e.a3 = try a.alloc(C * 32 * w2 * 2); + e.acc = try a.alloc(@max(C * 32 * w2, T * cfg.hidden) * 4); + e.enc_x = try a.alloc(T * cfg.enc_d * 2); + e.enc_out = try a.alloc(T * cfg.enc_d * 2); + e.ln_out = try a.alloc(T * cfg.enc_d * 2); + e.h1 = try a.alloc(T * cfg.hidden * 2); + e.audio_emb = try a.alloc(T * cfg.hidden * 2); + const R = e.o.prefill_rows; + e.ids = try a.alloc(R * 4); + e.bb_x = try a.alloc(R * cfg.hidden * 2); + e.bb_out = try a.alloc(R * cfg.hidden * 2); + e.logits = try a.alloc(cfg.vocab * 2); + e.pick = try a.alloc(4); + e.amax_ws = try a.alloc(8 * (cfg.vocab / 1024 + 1)); + const sup = [_]i32{ cfg.audio_token_id, cfg.audio_out_token_id, cfg.image_token_id, cfg.video_token_id }; + e.sup = try a.upload(std.mem.sliceAsBytes(&sup)); + const NQ = cfg.codebooks; + e.din_f32 = try a.alloc(NQ * cfg.depth_d * 4); + e.din = try a.alloc(NQ * cfg.depth_d * 2); + e.dx = try a.alloc(NQ * cfg.depth_d * 2); + e.dout = try a.alloc(NQ * cfg.depth_d * 2); + e.dz = try a.alloc(NQ * cfg.depth_d * 2); + e.dlog = try a.alloc(NQ * cfg.codebook_size * 2); + e.dlog_f32 = try a.alloc(NQ * cfg.codebook_size * 4); + e.dcodes = try a.alloc(NQ * 4); + e.dmask = try a.alloc(NQ * 4); + const CC = e.o.codec_chunk; + e.c_codes = try a.alloc(NQ * CC * 4); + e.c_lat = try a.alloc(CC * NQ * 4 * 2); + e.c_p1 = try a.alloc(CC * cfg.quantizer_dim * 2); + e.c_p2 = try a.alloc(CC * cfg.codec_d * 2); + e.c_up = try a.alloc(6 * CC * cfg.codec_d * 2); + e.c_out = try a.alloc(6 * CC * cfg.codec_d * 2); + e.c_f32 = try a.alloc(6 * CC * 2 * cfg.spec_bins * 4); + e.c_lin = try a.alloc(6 * CC * 2 * cfg.spec_bins * 2); + e.c_spec = try a.alloc(6 * CC * 2 * cfg.spec_bins * 4); + e.c_frames = try a.alloc(6 * CC * cfg.n_fft * 4); + e.host_frames = try e.gpa.alloc(f32, 6 * CC * cfg.n_fft); + e.host_logits = try e.gpa.alloc(f32, NQ * cfg.codebook_size); +} diff --git a/zig/src/families/lfm3_omni/metal/reply.zig b/zig/src/families/lfm3_omni/metal/reply.zig new file mode 100644 index 0000000000..94f9523aad --- /dev/null +++ b/zig/src/families/lfm3_omni/metal/reply.zig @@ -0,0 +1,203 @@ +//! A reply on Metal: text slots and speech frames in the model's interleaved schedule, the codec streamed between rows. +const std = @import("std"); +const cfg = @import("../config.zig"); +const common = @import("../common.zig"); +const Engine = @import("engine.zig").Engine; + +const Sampling = common.Sampling; +const Sink = common.Sink; +const Streaming = common.Streaming; +const Forced = common.Forced; +const Stats = common.Stats; +const OverlapAdd = common.OverlapAdd; +const nowNs = common.nowNs; +const msSince = common.msSince; +const Mode = common.Mode; +const hist_cap = common.hist_cap; + +/// Generates a reply after `prefill` into `sink` as it is made, the audio through the codec per `stream`. +pub fn reply(e: *Engine, smp: Sampling, stream: Streaming, sink: Sink, forced: ?*Forced, st_: *Stats, t_req: i128) !void { + var prng = std.Random.DefaultPrng.init(smp.seed); + const rng = prng.random(); + const io = e.io; + var ola = OverlapAdd.init(e.gpa, &e.window); + defer ola.deinit(); + if (stream.audio) try e.codecReset(); + var history: std.ArrayList([cfg.codebooks]i32) = .empty; + defer history.deinit(e.gpa); + var pending: std.ArrayList([cfg.codebooks]i32) = .empty; + defer pending.deinit(e.gpa); + var decoded: usize = 0; + var mode: Mode = .interleaved; + var speaking = false; + var text_done = false; + var slots_left: usize = cfg.interleaved_n_text; + var fi: usize = 0; + var si: usize = 1; + var step: usize = 0; + // the device path samples frames on the GPU (one copy back a frame); teacher forcing and top-k use the host's + const device = forced == null and smp.audio_top_k == 0; + var hist_n: usize = 0; + while (step < smp.max_new) : (step += 1) { + const interleaving = mode == .interleaved; + const last_step = step + 1 == smp.max_new; + var finished = false; + if (mode == .audio or (interleaving and speaking)) { + const t0 = nowNs(io); + var codes: [cfg.codebooks]i32 = undefined; + if (device) { + var order: [cfg.codebooks]i32 = undefined; + var perm: [cfg.codebooks]usize = undefined; + for (&perm, 0..) |*x, i| x.* = i; + if (!smp.identity_order) rng.shuffle(usize, &perm); + for (&order, perm) |*o, x| o.* = @intCast(x); + const g0 = e.g.gpu_s; + try e.frameDevice(smp, &order, step, @min(hist_n, hist_cap), hist_n % hist_cap); + try e.phase(.frame, g0); + try e.d2h(std.mem.sliceAsBytes(&codes), e.dframe + cfg.codebooks * 4); + } else { + const lo = history.items.len -| smp.audio_repetition_window; + codes = try e.sampleFrame(smp, rng, history.items[lo..]); + } + st_.depth_ms += msSince(io, t0); + if (forced) |fc| { + var mine_eos = false; + for (codes) |c| mine_eos = mine_eos or c == cfg.audio_eos_code; + if (si < fc.seq.len and fc.seq[si] == cfg.audio_out_token_id and fi < fc.frames.len) { + const ref = fc.frames[fi]; + var all = true; + for (codes, ref) |c, r| { + fc.code_hits += @intFromBool(c == r); + all = all and c == r; + } + fc.code_total += cfg.codebooks; + fc.frame_hits += @intFromBool(all); + codes = ref; + fi += 1; + si += 1; + } else { + fc.frame_hits += @intFromBool(mine_eos); + codes = @splat(cfg.audio_eos_code); + } + fc.frame_total += 1; + } + var eos = false; + for (codes) |c| eos = eos or c == cfg.audio_eos_code; + if (eos) { + codes = @splat(cfg.audio_eos_code); + mode = .text; + } else { + // the device's frame already joined the window (om_frame_end) + if (device) hist_n += 1 else try history.append(e.gpa, codes); + st_.frames += 1; + if (st_.frames == 1) st_.first_frame_ms = msSince(io, t_req); + if (sink.frame) |f| try f(sink.ctx, &codes); + try pending.append(e.gpa, codes); + } + if (!device) e.g.set(e.dframe + cfg.codebooks * 4, std.mem.sliceAsBytes(&codes)); + e.k.embedCodes(e.cb_embed, e.dframe + cfg.codebooks * 4, e.bb_x, 1, cfg.hidden, cfg.codebooks, cfg.codebook_size); + } else { + var tok = try e.textPick(); + if (forced) |fc| { + if (si >= fc.seq.len) break; + const ref = fc.seq[si]; + si += 1; + fc.text_hits += @intFromBool(tok == ref); + fc.text_total += 1; + tok = ref; + } + st_.text_tokens += 1; + if (st_.text_tokens == 1) st_.first_text_ms = msSince(io, t_req); + if (sink.text) |f| try f(sink.ctx, tok); + if (tok == cfg.text_end_token_id) { + if (interleaving) { + text_done = true; + speaking = true; + } else finished = true; + } else if (tok == cfg.eos_token_id) { + finished = true; + } else if (interleaving) {} else if (tok == cfg.audio_start_token_id) { + mode = .audio; + } else if (tok == cfg.mixed_start_token_id) { + mode = .interleaved; + slots_left = cfg.interleaved_n_text; + speaking = false; + text_done = false; + } + if (forced != null) { + const id: i32 = @intCast(tok); + e.g.set(e.pick, std.mem.asBytes(&id)); + } + e.k.embed(e.embed, e.pick, e.bb_x, 1, cfg.hidden); + } + if (interleaving and !text_done) { + slots_left -= 1; + if (slots_left == 0) { + speaking = !speaking; + slots_left = if (speaking) cfg.interleaved_n_audio else cfg.interleaved_n_text; + } + } + if (forced) |fc| if (si >= fc.seq.len and mode != .audio and !(mode == .interleaved and speaking)) { + finished = true; + }; + // the codec before the next backbone row: the first frames as soon as they exist, then whole chunks + var frames_out: []const f32 = &.{}; + if (stream.audio) { + const want: usize = if (decoded == 0) stream.first else stream.chunk; + if (pending.items.len >= want or (pending.items.len > 0 and (finished or mode == .text))) { + const t2 = nowNs(io); + const g0 = e.g.gpu_s; + frames_out = try e.codecChunk(pending.items, null); + try e.phase(.codec, g0); + decoded += pending.items.len; + pending.clearRetainingCapacity(); + st_.codec_ms += msSince(io, t2); + } + } + const go_on = !(finished or last_step); + const t1 = nowNs(io); + if (go_on) try e.stepRow(); + // the overlap-add on the host while the row runs + if (frames_out.len > 0) { + const pcm = try ola.add(frames_out, false); + if (pcm.len > 0) { + if (st_.samples == 0) st_.first_pcm_ms = msSince(io, t_req); + st_.samples += pcm.len; + if (sink.pcm) |f| try f(sink.ctx, pcm); + } + } + if (!go_on) break; + if (sink.hidden) |f| try f(sink.ctx); + st_.step_ms += msSince(io, t1); + } + if (stream.audio) try flushCodec(e, &ola, &pending, &decoded, sink, st_, t_req, true); + st_.total_ms = msSince(io, t_req); + if (e.prof) { + for (0..4) |i| std.log.info("lfm3-omni profile: {s} {d} x {d:.3} ms GPU", .{ @tagName(@as(Engine.Phase, @fromBackingInt(@intCast(i)))), e.prof_n[i], if (e.prof_n[i] > 0) e.prof_s[i] * 1e3 / @as(f64, @floatFromInt(e.prof_n[i])) else 0 }); + std.log.info("lfm3-omni profile: {d} dispatches", .{e.g.dispatches}); + e.prof_s = @splat(0); + e.prof_n = @splat(0); + e.g.dispatches = 0; + } +} + +fn flushCodec(e: *Engine, ola: *OverlapAdd, pending: *std.ArrayList([cfg.codebooks]i32), decoded: *usize, sink: Sink, st_: *Stats, t_req: i128, final: bool) !void { + const t0 = nowNs(e.io); + var at: usize = 0; + while (at < pending.items.len or (final and at == 0)) { + const n = @min(e.o.codec_chunk, pending.items.len - at); + const frames = try e.codecChunk(pending.items[at .. at + n], null); + const last = final and at + n >= pending.items.len; + const pcm = try ola.add(frames, last); + if (pcm.len > 0) { + if (st_.samples == 0) st_.first_pcm_ms = msSince(e.io, t_req); + st_.samples += pcm.len; + if (sink.pcm) |f| try f(sink.ctx, pcm); + } + decoded.* += n; + at += n; + if (n == 0) break; + } + pending.clearRetainingCapacity(); + st_.codec_ms += msSince(e.io, t0); +} diff --git a/zig/src/families/lfm3_omni/metal/stack.zig b/zig/src/families/lfm3_omni/metal/stack.zig new file mode 100644 index 0000000000..c006084993 --- /dev/null +++ b/zig/src/families/lfm3_omni/metal/stack.zig @@ -0,0 +1,310 @@ +//! One Liquid stack on Metal (gated short conv or attention layers, then SwiGLU) with its caches. +const std = @import("std"); +const st = @import("core").safetensors; +const kern = @import("kernels.zig"); +const gpu = @import("gpu.zig"); +const spec = @import("../spec.zig"); + +pub const Kind = spec.Kind; +pub const Rope = spec.Rope; +pub const Spec = spec.Spec; + +pub const Layer = struct { + kind: Kind, + op_norm: u64, + ffn_norm: u64, + post_op_norm: u64 = 0, + post_ffn_norm: u64 = 0, + // conv + in_proj: u64 = 0, // [3D, D] + conv_w: u64 = 0, // [D, K] + out_proj: u64 = 0, // [D, D] + state: u64 = 0, // [D, K-1] + // attention + qkvg: u64 = 0, // [q | k | v | gate, D] + q_norm: u64 = 0, + k_norm: u64 = 0, + o_proj: u64 = 0, // [D, H * hd] + kc: u64 = 0, // [max_ctx, Hkv * hd] + vc: u64 = 0, + // feed-forward + w13: u64 = 0, // [2F, D] + w2: u64 = 0, // [D, F] +}; + +pub const Stack = struct { + s: Spec, + layers: []Layer, + final_norm: u64, + inv_full: u64, + inv_sliding: u64, + // scratch + xn: u64, + proj: u64, + y: u64, + tmp: u64, + states: u64, + states_bytes: usize, + /// Rows already in the caches (the next row's position). + pos: usize = 0, + g: *gpu.Gpu, + + pub fn nq(x: *const Stack) usize { + return x.s.heads * x.s.head_dim; + } + pub fn nk(x: *const Stack) usize { + return x.s.kv_heads * x.s.head_dim; + } + fn qkvgCols(x: *const Stack) usize { + return x.nq() + 2 * x.nk() + if (x.s.lfm3) x.s.heads else 0; + } + + /// Clears the conv states and the cache position (in order with the dispatches; the caches' rows are overwritten). + pub fn reset(x: *Stack) void { + x.pos = 0; + if (x.states_bytes > 0) x.g.zero(x.states, x.states_bytes); + } + + /// Runs `t` rows of the residual stream `x` in place at the cache position, then the final norm into `out`. + pub fn forward(x_: *Stack, k: *const kern.Kernels, x: u64, t: usize, out: u64) !void { + const sp = x_.s; + if (t == 0) return; + if (t > sp.max_rows) return error.TooManyRows; + if (x_.pos + t > sp.max_ctx) return error.ContextFull; + if (t <= kern.gemv_rows and kern.Kernels.few(t, sp.d, sp.d, x)) try x_.forwardFew(k, x, t, out) else try x_.forwardRows(k, x, t, out); + x_.pos += t; + } + + fn attend(x_: *const Stack, k: *const kern.Kernels, l: *const Layer, t: usize) !void { + const sp = x_.s; + const nq_ = x_.nq(); + const nk_ = x_.nk(); + const cols = x_.qkvgCols(); + const windowed = l.kind == .sliding or (sp.window_all and sp.window > 0); + try k.attention(.{ .q = x_.proj, .ldq = cols, .kc = l.kc, .vc = l.vc, .ldkv = nk_, .o = x_.y, .ldo = nq_, .t = t, .heads = sp.heads, .kv_heads = sp.kv_heads, .hd = sp.head_dim, .pos0 = x_.pos, .kv_len = x_.pos + t, .window = if (windowed) sp.window else 0, .causal = sp.causal, .gate = if (sp.lfm3) x_.proj + (nq_ + 2 * nk_) * 2 else 0, .ldg = cols }); + } + + /// Up to 16 rows, fused: heads' norm, rotary and cache write; the gate in attention; SwiGLU in the w1 | w3 product. + fn forwardFew(x_: *Stack, k: *const kern.Kernels, x: u64, t: usize, out: u64) !void { + const sp = x_.s; + const D = sp.d; + const nq_ = x_.nq(); + const cols = x_.qkvgCols(); + var pending: u64 = 0; // the last block's output, added to the residual by the next norm + var pending_pw: u64 = 0; // its post norm (sandwich layers) + for (x_.layers) |*l| { + k.rmsnormPost(x, pending, pending_pw, l.op_norm, 0, x_.xn, D, D, D, t, sp.eps, .lfm); + switch (l.kind) { + .conv => { + k.product(l.in_proj, x_.xn, D, D, 3 * D, t, x_.proj, 3 * D, 0, 1, 0, 0, 0, .plain); + k.shortConv(x_.proj, 3 * D, l.conv_w, l.state, x_.y, D, t, D, sp.conv_k); + k.linear(l.out_proj, x_.y, D, D, D, x_.tmp, D, t); + }, + .full, .sliding => { + k.product(l.qkvg, x_.xn, D, D, cols, t, x_.proj, cols, 0, 1, 0, 0, 0, .plain); + const inv = if (l.kind == .sliding) x_.inv_sliding else x_.inv_full; + k.qkPrep(x_.proj, l.q_norm, l.k_norm, inv, l.kc, l.vc, t, sp.heads, sp.kv_heads, sp.head_dim, cols, x_.pos, sp.rope_mode, x_.nk(), sp.eps); + try x_.attend(k, l, t); + k.linear(l.o_proj, x_.y, nq_, nq_, D, x_.tmp, D, t); + }, + } + k.rmsnormPost(x, x_.tmp, if (sp.sandwich) l.post_op_norm else 0, l.ffn_norm, 0, x_.xn, D, D, D, t, sp.eps, .lfm); + k.product(l.w13, x_.xn, D, D, 2 * sp.ffn, t, x_.y, sp.ffn, 0, 1, 0, 0, 0, .swiglu); + k.linear(l.w2, x_.y, sp.ffn, sp.ffn, D, x_.tmp, D, t); + pending = x_.tmp; + pending_pw = if (sp.sandwich) l.post_ffn_norm else 0; + } + k.rmsnormPost(x, pending, pending_pw, x_.final_norm, 0, out, D, D, D, t, sp.eps, .lfm); + } + + /// Many rows (prompt chunks, the speech encoder, codec chunks): one op a dispatch. + fn forwardRows(x_: *Stack, k: *const kern.Kernels, x: u64, t: usize, out: u64) !void { + const sp = x_.s; + const D = sp.d; + const nq_ = x_.nq(); + const nk_ = x_.nk(); + const cols = x_.qkvgCols(); + var pending: u64 = 0; // the last block's output, added to the residual by the next norm + for (x_.layers) |*l| { + k.rmsnorm(x, pending, l.op_norm, x_.xn, D, D, D, t, sp.eps, .lfm); + switch (l.kind) { + .conv => { + k.linear(l.in_proj, x_.xn, D, D, 3 * D, x_.proj, 3 * D, t); + k.shortConv(x_.proj, 3 * D, l.conv_w, l.state, x_.y, D, t, D, sp.conv_k); + k.linear(l.out_proj, x_.y, D, D, D, x_.tmp, D, t); + }, + .full, .sliding => { + k.linear(l.qkvg, x_.xn, D, D, cols, x_.proj, cols, t); + k.headNorm(x_.proj, l.q_norm, t, sp.heads, sp.head_dim, cols, sp.eps); + k.headNorm(x_.proj + nq_ * 2, l.k_norm, t, sp.kv_heads, sp.head_dim, cols, sp.eps); + const inv = if (l.kind == .sliding) x_.inv_sliding else x_.inv_full; + k.rope(x_.proj, inv, t, sp.heads, sp.head_dim, cols, x_.pos, sp.rope_mode); + k.rope(x_.proj + nq_ * 2, inv, t, sp.kv_heads, sp.head_dim, cols, x_.pos, sp.rope_mode); + k.copyRows(x_.proj + nq_ * 2, l.kc + x_.pos * nk_ * 2, t, nk_, cols, nk_); + k.copyRows(x_.proj + (nq_ + nk_) * 2, l.vc + x_.pos * nk_ * 2, t, nk_, cols, nk_); + try x_.attend(k, l, t); + k.linear(l.o_proj, x_.y, nq_, nq_, D, x_.tmp, D, t); + }, + } + if (sp.sandwich) k.rmsnorm(x_.tmp, 0, l.post_op_norm, x_.tmp, D, D, D, t, sp.eps, .lfm); + k.rmsnorm(x, x_.tmp, l.ffn_norm, x_.xn, D, D, D, t, sp.eps, .lfm); + k.linear(l.w13, x_.xn, D, D, 2 * sp.ffn, x_.proj, 2 * sp.ffn, t); + k.swiglu(x_.proj, x_.y, t, sp.ffn, true); + k.linear(l.w2, x_.y, sp.ffn, sp.ffn, D, x_.tmp, D, t); + if (sp.sandwich) k.rmsnorm(x_.tmp, 0, l.post_ffn_norm, x_.tmp, D, D, D, t, sp.eps, .lfm); + pending = x_.tmp; + } + k.rmsnorm(x, pending, x_.final_norm, out, D, D, D, t, sp.eps, .lfm); + } +}; + +/// Device buffers a loader hands out (freed with the device). +pub const Arena = struct { + g: *gpu.Gpu, + + pub fn alloc(a: *Arena, n: usize) !u64 { + return a.g.alloc(n); + } + + pub fn zeroed(a: *Arena, n: usize) !u64 { + return a.g.alloc(n); + } + + pub fn upload(a: *Arena, bytes: []const u8) !u64 { + const p = try a.g.alloc(bytes.len); + @memcpy(a.g.host(p, bytes.len), bytes); + return p; + } + + /// Bytes at `off` of a buffer made here (before any GPU work reads it). + pub fn write(a: *Arena, p: u64, off: usize, bytes: []const u8) void { + @memcpy(a.g.host(p + off, bytes.len), bytes); + } +}; + +/// A bf16 tensor of the given shape, or a named refusal. +pub fn tensor(f: *const st.File, name: []const u8, shape: []const usize) ![]const u8 { + const t = f.get(name) orelse { + std.log.err("lfm3_omni: {s} is missing", .{name}); + return error.MissingTensor; + }; + if (!t.is(.bf16, shape)) { + std.log.err("lfm3_omni: {s} is {s} {any}, expected bf16 {any}", .{ name, @tagName(t.dtype), t.shape[0..t.rank], shape }); + return error.BadTensor; + } + return t.bytes; +} + +fn one(a: *Arena, f: *const st.File, name: []const u8, shape: []const usize) !u64 { + return a.upload(try tensor(f, name, shape)); +} + +/// Tensors of `k` columns stacked by rows into one matrix. +fn stacked(a: *Arena, f: *const st.File, names: []const []const u8, rows: []const usize, k: usize) !u64 { + var total: usize = 0; + for (rows) |r| total += r; + const p = try a.alloc(total * k * 2); + var at: usize = 0; + for (names, rows) |n, r| { + a.write(p, at * k * 2, try tensor(f, n, &.{ r, k })); + at += r; + } + return p; +} + +/// Two [n, k] tensors as one [2n, k] matrix of row pairs (w1[j], w3[j]): SwiGLU's operands side by side. +fn interleaved(a: *Arena, f: *const st.File, n1: []const u8, n3: []const u8, n: usize, k: usize) !u64 { + const p = try a.alloc(2 * n * k * 2); + const t1 = try tensor(f, n1, &.{ n, k }); + const t3 = try tensor(f, n3, &.{ n, k }); + for (0..n) |j| { + a.write(p, (2 * j) * k * 2, t1[j * k * 2 ..][0 .. k * 2]); + a.write(p, (2 * j + 1) * k * 2, t3[j * k * 2 ..][0 .. k * 2]); + } + return p; +} + +fn invFreq(gpa: std.mem.Allocator, a: *Arena, r: Rope, hd: usize) !u64 { + const v = try gpa.alloc(f32, hd / 2); + defer gpa.free(v); + spec.invFreq(v, r, hd); + return a.upload(std.mem.sliceAsBytes(v)); +} + +/// Loads a stack's weights from `f` and allocates its caches and scratch in `a`. +pub fn load(gpa: std.mem.Allocator, a: *Arena, f: *const st.File, sp: Spec) !Stack { + const D = sp.d; + const hd = sp.head_dim; + var x: Stack = undefined; + x.s = sp; + x.g = a.g; + x.pos = 0; + x.layers = try gpa.alloc(Layer, sp.kinds.len); + errdefer gpa.free(x.layers); + var nbuf: [5][256]u8 = undefined; + const P = sp.prefix; + var conv_layers: usize = 0; + for (sp.kinds) |kd| conv_layers += @intFromBool(kd == .conv); + x.states_bytes = conv_layers * D * (sp.conv_k - 1) * 2; + x.states = if (x.states_bytes > 0) try a.zeroed(x.states_bytes) else 0; + var conv_i: usize = 0; + for (sp.kinds, 0..) |kd, li| { + var l: Layer = .{ .kind = kd, .op_norm = 0, .ffn_norm = 0 }; + const n0 = try std.fmt.bufPrint(&nbuf[0], "{s}layers.{d}.", .{ P, li }); + l.op_norm = try one(a, f, try cat(&nbuf[1], n0, "operator_norm.weight"), &.{D}); + l.ffn_norm = try one(a, f, try cat(&nbuf[1], n0, "ffn_norm.weight"), &.{D}); + if (sp.sandwich) { + l.post_op_norm = try one(a, f, try cat(&nbuf[1], n0, "post_operator_norm.weight"), &.{D}); + l.post_ffn_norm = try one(a, f, try cat(&nbuf[1], n0, "post_ffn_norm.weight"), &.{D}); + } + switch (kd) { + .conv => { + l.in_proj = try one(a, f, try cat(&nbuf[1], n0, "conv.in_proj.weight"), &.{ 3 * D, D }); + l.conv_w = try one(a, f, try cat(&nbuf[1], n0, "conv.conv.weight"), &.{ D, 1, sp.conv_k }); + l.out_proj = try one(a, f, try cat(&nbuf[1], n0, "conv.out_proj.weight"), &.{ D, D }); + l.state = x.states + conv_i * D * (sp.conv_k - 1) * 2; + conv_i += 1; + }, + .full, .sliding => { + const nq_ = sp.heads * hd; + const nk_ = sp.kv_heads * hd; + const q = try cat(&nbuf[1], n0, "self_attn.q_proj.weight"); + const kk = try cat(&nbuf[2], n0, "self_attn.k_proj.weight"); + const v = try cat(&nbuf[3], n0, "self_attn.v_proj.weight"); + if (sp.lfm3) { + const g = try cat(&nbuf[4], n0, "self_attn.gate_proj.weight"); + l.qkvg = try stacked(a, f, &.{ q, kk, v, g }, &.{ nq_, nk_, nk_, sp.heads }, D); + l.q_norm = try one(a, f, try cat(&nbuf[1], n0, "self_attn.q_norm.weight"), &.{hd}); + l.k_norm = try one(a, f, try cat(&nbuf[1], n0, "self_attn.k_norm.weight"), &.{hd}); + l.o_proj = try one(a, f, try cat(&nbuf[1], n0, "self_attn.o_proj.weight"), &.{ D, nq_ }); + } else { + l.qkvg = try stacked(a, f, &.{ q, kk, v }, &.{ nq_, nk_, nk_ }, D); + l.q_norm = try one(a, f, try cat(&nbuf[1], n0, "self_attn.q_layernorm.weight"), &.{hd}); + l.k_norm = try one(a, f, try cat(&nbuf[1], n0, "self_attn.k_layernorm.weight"), &.{hd}); + l.o_proj = try one(a, f, try cat(&nbuf[1], n0, "self_attn.out_proj.weight"), &.{ D, nq_ }); + } + l.kc = try a.alloc(sp.max_ctx * nk_ * 2); + l.vc = try a.alloc(sp.max_ctx * nk_ * 2); + }, + } + const w1 = try cat(&nbuf[1], n0, "feed_forward.w1.weight"); + const w3 = try cat(&nbuf[2], n0, "feed_forward.w3.weight"); + l.w13 = try interleaved(a, f, w1, w3, sp.ffn, D); + l.w2 = try one(a, f, try cat(&nbuf[1], n0, "feed_forward.w2.weight"), &.{ D, sp.ffn }); + x.layers[li] = l; + } + x.final_norm = try one(a, f, try cat(&nbuf[1], P, "embedding_norm.weight"), &.{D}); + x.inv_full = try invFreq(gpa, a, sp.rope_full, hd); + x.inv_sliding = try invFreq(gpa, a, sp.rope_sliding, hd); + const R = sp.max_rows; + const cols = @max(3 * D, @max(sp.heads * hd + 2 * sp.kv_heads * hd + sp.heads, 2 * sp.ffn)); + x.xn = try a.alloc(R * D * 2); + x.proj = try a.alloc(R * cols * 2); + x.y = try a.alloc(R * @max(D, @max(sp.ffn, sp.heads * hd)) * 2); + x.tmp = try a.alloc(R * D * 2); + return x; +} + +fn cat(buf: []u8, a: []const u8, b: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}{s}", .{ a, b }); +} diff --git a/zig/src/families/lfm3_omni/omni_metal.zig b/zig/src/families/lfm3_omni/omni_metal.zig new file mode 100644 index 0000000000..03160e63a1 --- /dev/null +++ b/zig/src/families/lfm3_omni/omni_metal.zig @@ -0,0 +1,20 @@ +//! LFM3-Omni-2B on Metal, with the module surface generate.zig and serve.zig build against on either backend. +pub const config = @import("config.zig"); +pub const mel = @import("mel.zig"); +pub const common = @import("common.zig"); +pub const spec = @import("spec.zig"); +pub const stack = @import("metal/stack.zig"); +pub const kernels = @import("metal/kernels.zig"); +pub const engine = @import("metal/engine.zig"); +pub const prompt = @import("prompt.zig"); +pub const serve = @import("serve.zig"); +pub const wav = @import("wav"); + +pub const Engine = engine.Engine; + +test { + _ = mel; + _ = prompt; + _ = common; + _ = kernels; +} diff --git a/zig/src/families/lfm3_omni/prompt.zig b/zig/src/families/lfm3_omni/prompt.zig new file mode 100644 index 0000000000..2c34c67b36 --- /dev/null +++ b/zig/src/families/lfm3_omni/prompt.zig @@ -0,0 +1,27 @@ +//! LFM3-Omni's chat prompt for a system text and one spoken turn, ending on the interleaved generation marker. +const std = @import("std"); +const tokenizer = @import("core").tokenizer; +const cfg = @import("config.zig"); + +pub const Prompt = struct { + ids: []u32, + /// The first of the contiguous `<|audio|>` rows the encoder's rows replace. + audio_at: usize, +}; + +/// The system turn, a user turn of `n_audio` `<|audio|>` rows, then the assistant's empty think and `<|mixed_start|>`. +pub fn build(gpa: std.mem.Allocator, tok: *const tokenizer.Tokenizer, system: []const u8, n_audio: usize) !Prompt { + const head = try std.fmt.allocPrint(gpa, "<|startoftext|><|im_start|>system\n{s}<|im_end|>\n<|im_start|>user\n<|vision_start|>", .{system}); + defer gpa.free(head); + const a = try tok.encode(gpa, head); + defer gpa.free(a); + const b = try tok.encode(gpa, "<|vision_end|><|im_end|>\n<|im_start|>assistant\n"); + defer gpa.free(b); + const ids = try gpa.alloc(u32, a.len + n_audio + b.len + 1); + @memcpy(ids[0..a.len], a); + @memset(ids[a.len..][0..n_audio], cfg.audio_token_id); + @memcpy(ids[a.len + n_audio ..][0..b.len], b); + ids[ids.len - 1] = cfg.mixed_start_token_id; + if (a[a.len - 1] != cfg.vision_start_token_id or b[0] != cfg.vision_end_token_id) return error.TokenizerSpecialTokens; + return .{ .ids = ids, .audio_at = a.len }; +} diff --git a/zig/src/families/lfm3_omni/serve.zig b/zig/src/families/lfm3_omni/serve.zig new file mode 100644 index 0000000000..a6700b3454 --- /dev/null +++ b/zig/src/families/lfm3_omni/serve.zig @@ -0,0 +1,237 @@ +//! HTTP/1.1 for one LFM3-Omni engine, one request at a time: POST /v1/audio/chat (SSE or whole) and GET /health. +const std = @import("std"); +const core = @import("core"); +const engine = @import("common.zig"); +const cfg = @import("config.zig"); +const mel = @import("mel.zig"); +const prompt = @import("prompt.zig"); +const wav = @import("wav"); + +extern "c" fn socket(domain: c_int, ty: c_int, proto: c_int) c_int; +extern "c" fn setsockopt(fd: c_int, level: c_int, name: c_int, val: *const anyopaque, len: u32) c_int; +extern "c" fn bind(fd: c_int, addr: *const anyopaque, len: u32) c_int; +extern "c" fn listen(fd: c_int, n: c_int) c_int; +extern "c" fn accept(fd: c_int, addr: ?*anyopaque, len: ?*u32) c_int; +extern "c" fn recv(fd: c_int, buf: [*]u8, n: usize, flags: c_int) isize; +extern "c" fn send(fd: c_int, buf: [*]const u8, n: usize, flags: c_int) isize; +extern "c" fn close(fd: c_int) c_int; +extern "c" fn inet_pton(af: c_int, src: [*:0]const u8, dst: *anyopaque) c_int; + +const SockAddrIn = extern struct { family: u16, port: u16, addr: u32, zero: [8]u8 = @splat(0) }; +const msg_nosignal = 0x4000; +const max_body = 64 << 20; + +pub const Options = struct { + host: []const u8 = "127.0.0.1", + port: u16 = 8441, + /// The system text a request without one gets. + system: []const u8 = "Respond with interleaved text and audio.", + smp: engine.Sampling = .{}, + stream: engine.Streaming = .{}, +}; + +const Conn = struct { + fd: c_int, + closed: bool = false, + + fn writeAll(c: *Conn, b: []const u8) !void { + var at: usize = 0; + while (at < b.len) { + const n = send(c.fd, b.ptr + at, b.len - at, msg_nosignal); + if (n <= 0) { + c.closed = true; + return error.ClientGone; + } + at += @intCast(n); + } + } + + /// One chunk of a chunked body. + fn chunk(c: *Conn, b: []const u8) !void { + var h: [32]u8 = undefined; + try c.writeAll(try std.fmt.bufPrint(&h, "{x}\r\n", .{b.len})); + try c.writeAll(b); + try c.writeAll("\r\n"); + } +}; + +pub fn run(gpa: std.mem.Allocator, io: std.Io, e: anytype, tok: *const core.tokenizer.Tokenizer, o: Options) !void { + const fd = socket(2, 1, 0); + if (fd < 0) return error.Socket; + defer _ = close(fd); + const one: c_int = 1; + _ = setsockopt(fd, 1, 2, &one, @sizeOf(c_int)); + var addr: SockAddrIn = .{ .family = 2, .port = std.mem.nativeToBig(u16, o.port), .addr = 0 }; + const hz = try gpa.dupeSentinel(u8, o.host, 0); + defer gpa.free(hz); + if (inet_pton(2, hz.ptr, &addr.addr) != 1) return error.BadHost; + if (bind(fd, &addr, @sizeOf(SockAddrIn)) != 0) return error.Bind; + if (listen(fd, 16) != 0) return error.Listen; + std.log.info("lfm3-omni: serving POST /v1/audio/chat on {s}:{d}", .{ o.host, o.port }); + while (true) { + const cfd = accept(fd, null, null); + if (cfd < 0) continue; + var c: Conn = .{ .fd = cfd }; + handle(gpa, io, e, tok, o, &c) catch |err| { + std.log.warn("lfm3-omni: request failed: {s}", .{@errorName(err)}); + if (!c.closed) respond(&c, "500 Internal Server Error", "application/json", "{\"error\":\"internal error\"}") catch {}; + }; + _ = close(cfd); + } +} + +fn respond(c: *Conn, status: []const u8, ctype: []const u8, body: []const u8) !void { + var h: [256]u8 = undefined; + try c.writeAll(try std.fmt.bufPrint(&h, "HTTP/1.1 {s}\r\nContent-Type: {s}\r\nContent-Length: {d}\r\nConnection: close\r\n\r\n", .{ status, ctype, body.len })); + try c.writeAll(body); +} + +fn badRequest(c: *Conn, a: std.mem.Allocator, msg: []const u8) !void { + try respond(c, "400 Bad Request", "application/json", try std.fmt.allocPrint(a, "{{\"error\":{s}}}", .{try jstr(a, msg)})); +} + +/// A JSON string literal. +fn jstr(a: std.mem.Allocator, s: []const u8) ![]const u8 { + return std.json.Stringify.valueAlloc(a, s, .{}); +} + +/// The request head and body (Content-Length only). +fn readRequest(a: std.mem.Allocator, c: *Conn) !struct { method: []const u8, path: []const u8, body: []const u8 } { + var buf: std.ArrayList(u8) = .empty; + var head_end: ?usize = null; + var tmp: [65536]u8 = undefined; + while (head_end == null) { + const n = recv(c.fd, &tmp, tmp.len, 0); + if (n <= 0) return error.ClientGone; + try buf.appendSlice(a, tmp[0..@intCast(n)]); + head_end = std.mem.indexOf(u8, buf.items, "\r\n\r\n"); + if (head_end == null and buf.items.len > 1 << 16) return error.HeadTooLong; + } + const head = buf.items[0..head_end.?]; + var lines = std.mem.splitSequence(u8, head, "\r\n"); + const first = lines.next() orelse return error.BadRequest; + var parts = std.mem.splitScalar(u8, first, ' '); + const method = parts.next() orelse return error.BadRequest; + const path = parts.next() orelse return error.BadRequest; + var len: usize = 0; + while (lines.next()) |l| { + const colon = std.mem.indexOfScalar(u8, l, ':') orelse continue; + if (std.ascii.eqlIgnoreCase(std.mem.trim(u8, l[0..colon], " "), "content-length")) + len = try std.fmt.parseInt(usize, std.mem.trim(u8, l[colon + 1 ..], " "), 10); + } + if (len > max_body) return error.BodyTooLarge; + const start = head_end.? + 4; + while (buf.items.len < start + len) { + const n = recv(c.fd, &tmp, tmp.len, 0); + if (n <= 0) return error.ClientGone; + try buf.appendSlice(a, tmp[0..@intCast(n)]); + } + return .{ .method = method, .path = path, .body = buf.items[start .. start + len] }; +} + +const Turn = struct { + a: std.mem.Allocator, + c: *Conn, + tok: *const core.tokenizer.Tokenizer, + stream: bool, + text_ids: std.ArrayList(u32) = .empty, + sent_text: usize = 0, + pcm: std.ArrayList(f32) = .empty, + buf: std.ArrayList(u8) = .empty, + first_audio_ms: f64 = 0, + t0: i128, + io: std.Io, + + fn event(t: *Turn, json_text: []const u8) !void { + t.buf.clearRetainingCapacity(); + try t.buf.appendSlice(t.a, "data: "); + try t.buf.appendSlice(t.a, json_text); + try t.buf.appendSlice(t.a, "\n\n"); + try t.c.chunk(t.buf.items); + } + + /// The reply's words so far (special tokens left out) past what was sent, whole UTF-8 sequences only. + fn onText(ctx: *anyopaque, id: u32) anyerror!void { + const t: *Turn = @ptrCast(@alignCast(ctx)); + if (id >= cfg.bos_token_id) return; + try t.text_ids.append(t.a, id); + if (!t.stream) return; + const all = try t.tok.decode(t.a, t.text_ids.items, true); + var end = all.len; + while (end > t.sent_text and !std.unicode.utf8ValidateSlice(all[t.sent_text..end])) end -= 1; + if (end <= t.sent_text) return; + try t.event(try std.fmt.allocPrint(t.a, "{{\"type\":\"text\",\"delta\":{s}}}", .{try jstr(t.a, all[t.sent_text..end])})); + t.sent_text = end; + } + + fn onPcm(ctx: *anyopaque, x: []const f32) anyerror!void { + const t: *Turn = @ptrCast(@alignCast(ctx)); + if (t.pcm.items.len == 0 and t.first_audio_ms == 0) t.first_audio_ms = @as(f64, @floatFromInt(engine.nowNs(t.io) - t.t0)) / 1e6; + if (!t.stream) return t.pcm.appendSlice(t.a, x); + const raw = try t.a.alloc(u8, x.len * 2); + wav.pcm16(x, raw); + const b64 = try t.a.alloc(u8, std.base64.standard.Encoder.calcSize(raw.len)); + _ = std.base64.standard.Encoder.encode(b64, raw); + const j = try std.fmt.allocPrint(t.a, "{{\"type\":\"audio\",\"format\":\"pcm16\",\"sample_rate\":24000,\"data\":\"{s}\"}}", .{b64}); + try t.event(j); + } +}; + +fn num(v: ?std.json.Value) ?f64 { + const x = v orelse return null; + return switch (x) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => null, + }; +} + +fn handle(gpa: std.mem.Allocator, io: std.Io, e: anytype, tok: *const core.tokenizer.Tokenizer, o: Options, c: *Conn) !void { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const req = try readRequest(a, c); + if (std.mem.eql(u8, req.method, "GET") and std.mem.eql(u8, req.path, "/health")) return respond(c, "200 OK", "application/json", "{\"ok\":true,\"model\":\"lfm3-omni\"}"); + if (!std.mem.eql(u8, req.method, "POST") or !std.mem.eql(u8, req.path, "/v1/audio/chat")) return respond(c, "404 Not Found", "application/json", "{\"error\":\"POST /v1/audio/chat or GET /health\"}"); + const t0 = engine.nowNs(io); + const parsed = std.json.parseFromSlice(std.json.Value, a, req.body, .{}) catch return badRequest(c, a, "the body is not JSON"); + if (parsed.value != .object) return badRequest(c, a, "the body is not a JSON object"); + const obj = parsed.value.object; + const audio_b64 = if (obj.get("audio")) |v| (if (v == .string) v.string else return badRequest(c, a, "audio: a base64 WAV")) else return badRequest(c, a, "audio is required (a base64 WAV)"); + const wav_bytes = try a.alloc(u8, std.base64.standard.Decoder.calcSizeForSlice(audio_b64) catch return badRequest(c, a, "audio is not base64")); + std.base64.standard.Decoder.decode(wav_bytes, audio_b64) catch return badRequest(c, a, "audio is not base64"); + const au = wav.parse(a, wav_bytes) catch return badRequest(c, a, "audio is not a PCM or float WAV"); + const pcm = try wav.resample(a, au.samples, au.rate, mel.sample_rate); + if (pcm.len == 0) return badRequest(c, a, "the audio is empty"); + if (pcm.len > e.o.max_input_s * mel.sample_rate) return badRequest(c, a, "the audio is longer than the server takes"); + const system = if (obj.get("system")) |v| (if (v == .string) v.string else o.system) else o.system; + const stream_out = if (obj.get("stream")) |v| (v == .bool and v.bool) else true; + var smp = o.smp; + if (num(obj.get("audio_temperature"))) |x| smp.audio_temperature = @floatCast(x); + if (num(obj.get("audio_repetition_penalty"))) |x| smp.audio_repetition_penalty = @floatCast(x); + if (num(obj.get("seed"))) |x| smp.seed = @intFromFloat(x); + if (num(obj.get("max_new"))) |x| smp.max_new = @min(@as(usize, @intFromFloat(x)), e.o.max_frames); + smp.audio_top_k = 0; + + var st: engine.Stats = .{}; + var turn: Turn = .{ .a = a, .c = c, .tok = tok, .stream = stream_out, .t0 = t0, .io = io }; + const nf = try e.featuresDevice(pcm); + const n_audio = try e.encode(null, nf); + const p = try prompt.build(a, tok, system, n_audio); + if (p.ids.len + smp.max_new + 1 > e.o.max_ctx) return badRequest(c, a, "the prompt and reply do not fit the window"); + try e.prefill(p.ids, p.audio_at, null); + if (stream_out) try c.writeAll("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n"); + try e.reply(smp, o.stream, .{ .ctx = &turn, .text = Turn.onText, .pcm = Turn.onPcm }, null, &st, t0); + const text = try tok.decode(a, turn.text_ids.items, true); + const stats = try std.fmt.allocPrint(a, "\"text\":{s},\"frames\":{d},\"speech_s\":{d:.3},\"audio_tokens\":{d},\"prompt_rows\":{d},\"first_text_ms\":{d:.1},\"first_audio_ms\":{d:.1},\"total_ms\":{d:.1}", .{ try jstr(a, text), st.frames, @as(f64, @floatFromInt(st.samples)) / cfg.sample_rate, n_audio, p.ids.len, st.first_text_ms, turn.first_audio_ms, st.total_ms }); + if (!stream_out) { + const wb = try wav.encode16(a, turn.pcm.items, cfg.sample_rate); + const b64 = try a.alloc(u8, std.base64.standard.Encoder.calcSize(wb.len)); + _ = std.base64.standard.Encoder.encode(b64, wb); + return respond(c, "200 OK", "application/json", try std.fmt.allocPrint(a, "{{{s},\"audio\":\"{s}\"}}", .{ stats, b64 })); + } + try turn.event(try std.fmt.allocPrint(a, "{{\"type\":\"done\",{s}}}", .{stats})); + try turn.event("[DONE]"); + try c.chunk(""); + std.log.info("lfm3-omni: {d} audio tokens, {d} frames, first audio {d:.1} ms, total {d:.1} ms", .{ n_audio, st.frames, turn.first_audio_ms, st.total_ms }); +} diff --git a/zig/src/families/lfm3_omni/spec.zig b/zig/src/families/lfm3_omni/spec.zig new file mode 100644 index 0000000000..7f3722d4c9 --- /dev/null +++ b/zig/src/families/lfm3_omni/spec.zig @@ -0,0 +1,45 @@ +//! A Liquid stack's description (LFM2 or LFM3 layers), shared by the CUDA and Metal stacks. +const std = @import("std"); + +pub const Kind = enum { conv, full, sliding }; + +pub const Rope = struct { + theta: f32, + /// Rotated share of the head (LFM3's zero-tail partial rotation keeps the full head's frequency progression). + fraction: f32 = 1.0, +}; + +pub const Spec = struct { + name: []const u8, + /// Weight names: "{prefix}layers.{i}...." and "{prefix}embedding_norm.weight". + prefix: []const u8, + d: usize, + ffn: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + conv_k: usize, + kinds: []const Kind, + /// Keys a sliding layer (and, in the speech encoder, every attention layer) sees, its own row included. + window: usize = 0, + /// The speech encoder masks every attention layer to `window`. + window_all: bool = false, + causal: bool = true, + sandwich: bool, + /// LFM3 attention: a sigmoid output gate per head and q_norm / k_norm / o_proj names; LFM2: q_layernorm / out_proj. + lfm3: bool, + rope_full: Rope, + rope_sliding: Rope = .{ .theta = 10000 }, + /// 0: fp32 cos/sin; 1: bf16 cos/sin, fp32 products; 2: bf16 arithmetic. + rope_mode: u8, + eps: f32 = 1e-5, + /// Rows one forward takes (scratch) and rows the caches hold. + max_rows: usize, + max_ctx: usize, +}; + +/// Inverse frequencies of `hd / 2` pairs, zero past the rotated share, as the reference computes them in fp32. +pub fn invFreq(out: []f32, r: Rope, hd: usize) void { + const rot: usize = @as(usize, @intFromFloat(r.fraction * @as(f32, @floatFromInt(hd)))) / 2; + for (out, 0..) |*x, i| x.* = if (i < rot) 1.0 / std.math.pow(f32, r.theta, @as(f32, @floatFromInt(2 * i)) / @as(f32, @floatFromInt(hd))) else 0; +} From ade7077162cfd315e1073ac813937f545ac5aa41 Mon Sep 17 00:00:00 2001 From: drowzeys <264866615+drowzeys@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:47:37 -0400 Subject: [PATCH 2/3] kernel contract: every LFM3-Omni Metal pipeline, against float64 at its bf16 points, a control each Registers all 45 pipelines of lfm3_omni.metallib (69 registrations over the stacks' variants) in tf-kernel-contract: gates 1-4 against a float64 replay of each kernel's rounding points with the spread its fp32 steps allow, stated rules for the picks and the word moves, and a planted source change per registration. Gate 5 records the launches with their M1/M2 limits from --offline. --- zig/kernels/metal/sources.zig | 10 + .../core/contract/lfm3_omni_metal_records.zig | 54 ++ zig/src/core/contract/registry.zig | 2 +- zig/tests/contract/metal_limits.zig | 4 + zig/tests/contract/metal_main.zig | 2 + zig/tests/contract/metal_omni.zig | 85 ++++ zig/tests/contract/metal_omni_attn.zig | 221 +++++++++ zig/tests/contract/metal_omni_common.zig | 441 +++++++++++++++++ zig/tests/contract/metal_omni_conv.zig | 467 ++++++++++++++++++ zig/tests/contract/metal_omni_elem.zig | 418 ++++++++++++++++ zig/tests/contract/metal_omni_gemm.zig | 291 +++++++++++ zig/tests/contract/metal_omni_norms.zig | 432 ++++++++++++++++ zig/tests/contract/metal_omni_pick.zig | 307 ++++++++++++ 13 files changed, 2733 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/contract/lfm3_omni_metal_records.zig create mode 100644 zig/tests/contract/metal_omni.zig create mode 100644 zig/tests/contract/metal_omni_attn.zig create mode 100644 zig/tests/contract/metal_omni_common.zig create mode 100644 zig/tests/contract/metal_omni_conv.zig create mode 100644 zig/tests/contract/metal_omni_elem.zig create mode 100644 zig/tests/contract/metal_omni_gemm.zig create mode 100644 zig/tests/contract/metal_omni_norms.zig create mode 100644 zig/tests/contract/metal_omni_pick.zig diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 6e06cb42f9..482e8552da 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -200,6 +200,16 @@ pub const drama3 = struct { pub const codec = @embedFile("drama3/codec.metal"); }; +/// LFM3-Omni's kernels (lfm3_omni.metallib is built ahead of time); the kernel contract compiles them from this text. +pub const lfm3_omni = struct { + pub const common = @embedFile("lfm3_omni/om_common.h"); + pub const omni = @embedFile("lfm3_omni/omni.metal"); + pub const gemm = @embedFile("lfm3_omni/omni_gemm.metal"); + pub const rows = @embedFile("lfm3_omni/omni_rows.metal"); + pub const nax = @embedFile("lfm3_omni/omni_nax.metal"); + pub const gemm_nax = @embedFile("prefill/gemm_nax.metal"); +}; + pub const core_weight_only_body = @embedFile("core/weight_only.h"); pub const core_weight_only = @embedFile("core/weight_only.metal"); pub const core_hc_activation_body = @embedFile("core/hc_activation.h"); diff --git a/zig/src/core/contract/lfm3_omni_metal_records.zig b/zig/src/core/contract/lfm3_omni_metal_records.zig new file mode 100644 index 0000000000..4a8c5e9c97 --- /dev/null +++ b/zig/src/core/contract/lfm3_omni_metal_records.zig @@ -0,0 +1,54 @@ +//! LFM3-Omni's Metal launches (lfm3_omni.metallib), as kernels.zig and gpu.zig launch them. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "lfm3 omni omni om_rmsnorm", .threads = 256, .tg_memory = 32, .reserved = 256 }, + .{ .name = "lfm3 omni omni om_rmsnorm_1k", .threads = 1024, .tg_memory = 128, .reserved = 1024 }, + .{ .name = "lfm3 omni omni om_head_norm", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_layernorm", .threads = 256, .tg_memory = 32, .reserved = 256 }, + .{ .name = "lfm3 omni omni om_rope", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_copy_rows", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_set_words", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_zero_words", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_copy_words", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_short_conv", .threads = 128, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_short_conv_rows", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_conv_state", .threads = 128, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_attention_tiles", .threads = 128, .tg_memory = 26880, .reserved = 128 }, + .{ .name = "lfm3 omni omni om_swiglu", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_bias_act", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_embed", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_embed_codes", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_depth_in", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_conv2d_s2", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_f32_bf16", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_bf16_f32", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_fsq_decode", .threads = 128, .tg_memory = 0, .m1 = 928, .m2 = 928 }, + .{ .name = "lfm3 omni omni om_repeat_rows", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni omni om_polar", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni gemm om_gemm_bb", .threads = 128, .tg_memory = 9216, .reserved = 128 }, + .{ .name = "lfm3 omni gemm om_gemm_bf", .threads = 128, .tg_memory = 9216, .reserved = 128 }, + .{ .name = "lfm3 omni gemm om_gemm_ff", .threads = 128, .tg_memory = 9216, .reserved = 128 }, + .{ .name = "lfm3 omni gemm om_depth_pick", .threads = 256, .tg_memory = 2112, .reserved = 256 }, + .{ .name = "lfm3 omni gemm om_frame_end", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni gemm om_mel_frames", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni gemm om_power", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni gemm om_logmel", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni rows om_rows_1", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "lfm3 omni rows om_rows_4", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "lfm3 omni rows om_rows_8", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "lfm3 omni rows om_rows_12", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "lfm3 omni rows om_rows_16", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "lfm3 omni rows om_qk_prep", .threads = 32, .tg_memory = 512, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni rows om_attend", .threads = 256, .tg_memory = 25120, .reserved = 256 }, + .{ .name = "lfm3 omni rows om_argmax_part", .threads = 256, .tg_memory = 64, .reserved = 256 }, + .{ .name = "lfm3 omni rows om_argmax_final", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, + .{ .name = "lfm3 omni nax om_nax_bb_nt", .threads = 256, .tg_memory = 0, .reserved = 256, .offline = false }, + .{ .name = "lfm3 omni nax om_nax_bf_nt", .threads = 256, .tg_memory = 0, .reserved = 256, .offline = false }, + .{ .name = "lfm3 omni nax om_nax_bf_nn", .threads = 256, .tg_memory = 0, .reserved = 256, .offline = false }, + .{ .name = "lfm3 omni nax om_nax_bf_tt", .threads = 256, .tg_memory = 0, .reserved = 256, .offline = false }, +}; + +test "the LFM3-Omni launches fit every Apple family" { + try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true)); +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 92e6a40b76..80fdc86155 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -572,7 +572,7 @@ test { _ = paging; } /// The Metal collector keeps measured slots and the restored vision declarations. -pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ +pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("lfm3_omni_metal_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ @import("moss_image_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 23b6f6ac8b..eecd638e26 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -61,6 +61,10 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La if (std.mem.startsWith(u8, d.name, "metal lfm2 ")) return .{ .decl = d, .function = fun, .here = sources.lfm2, .pre_m5 = sources.lfm2 }; if (std.mem.startsWith(u8, d.name, "nemotron paged ")) return pagedLaunch(a, device, d, fun); if (std.mem.startsWith(u8, d.name, "nemotron attention ")) return attentionLaunch(a, device, d, fun); + if (std.mem.startsWith(u8, d.name, "lfm3 omni ")) { + const body = try @import("metal_omni.zig").limitsSource(a, d.name); + return .{ .decl = d, .function = fun, .here = body, .pre_m5 = body }; + } if (std.mem.startsWith(u8, d.name, "drama3 ")) { const text_d3 = try @import("metal_drama3.zig").source(a, d.name); return .{ .decl = d, .function = fun, .here = text_d3, .pre_m5 = text_d3 }; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 34567cef0f..864f2bd47c 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -403,6 +403,8 @@ pub fn main(init: std.process.Init) !u8 { defer native_gdn.destroy(); const moss_image = try @import("metal_moss_image.zig").Suite.create(a, device, queue, &kernels, o.only); defer moss_image.destroy(); + const omni = try @import("metal_omni.zig").Suite.create(a, device, queue, &kernels, o.only); + defer omni.destroy(); var t: g.Tally = .{}; const t0 = std.Io.Clock.awake.now(init.io); if (o.gates) { diff --git a/zig/tests/contract/metal_omni.zig b/zig/tests/contract/metal_omni.zig new file mode 100644 index 0000000000..26120fd4b6 --- /dev/null +++ b/zig/tests/contract/metal_omni.zig @@ -0,0 +1,85 @@ +//! LFM3-Omni's Metal kernels under the contract: every pipeline of lfm3_omni.metallib, each with a planted control. +const std = @import("std"); +const mtl = @import("metal"); +const k = @import("metal_omni_common.zig"); +const rig_ = @import("rig.zig"); + +const g = k.g; + +/// Every registration: norms and rotary, convs and gathers, elementwise, picks, attention, products. +pub const all = @import("metal_omni_norms.zig").entries ++ @import("metal_omni_conv.zig").entries ++ @import("metal_omni_elem.zig").entries ++ @import("metal_omni_pick.zig").entries ++ @import("metal_omni_attn.zig").entries ++ @import("metal_omni_gemm.zig").entries; + +/// A gate 5 label's source ("lfm3 omni "), its includes inlined. +pub fn limitsSource(a: std.mem.Allocator, label: []const u8) ![]u8 { + var it = std.mem.tokenizeScalar(u8, label, ' '); + _ = it.next(); + _ = it.next(); + const word = it.next() orelse return error.UnknownLaunch; + const file = std.meta.stringToEnum(k.File, word) orelse return error.UnknownLaunch; + return k.text(a, file); +} + +pub const Suite = struct { + a: std.mem.Allocator, + probes: std.ArrayList(*k.Probe) = .empty, + libs: std.ArrayList(mtl.Library) = .empty, + names: std.ArrayList([]u8) = .empty, + + fn keep(o: ?[]const u8, name: []const u8) bool { + return if (o) |needle| std.mem.indexOf(u8, name, needle) != null else true; + } + + fn add(s: *Suite, device: mtl.Device, queue: mtl.Queue, lib: mtl.Library, e: k.Entry, name: []const u8, control: ?g.Gate, out: *std.ArrayList(g.Kernel)) !void { + const pipe = try mtl.Pipeline.init(device, lib, e.pipe, false); + errdefer pipe.deinit(); + const p = try s.a.create(k.Probe); + errdefer s.a.destroy(p); + var spec = e.spec; + spec.name = name; + p.* = .{ .a = s.a, .rig = .{ .a = s.a, .device = device, .queue = queue }, .pipe = pipe, .def = e.def, .spec_ = spec, .variant = e.variant }; + try s.probes.append(s.a, p); + try out.append(s.a, p.kernel(control)); + } + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, out: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + errdefer s.destroy(); + const tensor = device.tensorUnits(); + var served: [4]?mtl.Library = @splat(null); + defer for (served) |l| if (l) |x| x.deinit(); + for (all) |e| { + if (e.file == .nax and !tensor) { + if (keep(only, e.spec.name)) std.debug.print("UNSUPPORTED {s}: tensor units required; the simdgroup products serve this GPU\n", .{e.spec.name}); + continue; + } + if (keep(only, e.spec.name)) { + const i = @intFromEnum(e.file); + if (served[i] == null) served[i] = try k.compile(a, device, e.file, null); + try s.add(device, queue, served[i].?, e, e.spec.name, null, out); + } + const ctl = e.control orelse continue; + const name = try std.fmt.allocPrint(a, "{s} control: {s}", .{ e.spec.name, e.what }); + try s.names.append(a, name); + if (!keep(only, name)) continue; + const lib = try k.compile(a, device, e.file, ctl); + try s.libs.append(a, lib); + try s.add(device, queue, lib, e, name, .shapes, out); + } + return s; + } + + pub fn destroy(s: *Suite) void { + for (s.probes.items) |p| { + p.pipe.deinit(); + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + for (s.libs.items) |l| l.deinit(); + s.libs.deinit(s.a); + for (s.names.items) |n| s.a.free(n); + s.names.deinit(s.a); + s.a.destroy(s); + } +}; diff --git a/zig/tests/contract/metal_omni_attn.zig b/zig/tests/contract/metal_omni_attn.zig new file mode 100644 index 0000000000..0e28ee1666 --- /dev/null +++ b/zig/tests/contract/metal_omni_attn.zig @@ -0,0 +1,221 @@ +//! LFM3-Omni's Metal attention: prompt rows in tiles (om_attention_tiles) and one row over the cache (om_attend). +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const k = @import("metal_omni_common.zig"); + +const c = k.c; +const R = k.R; +const Probe = k.Probe; + +/// A stack's attention: heads, kv heads, head width, window (0: none), causal, and LFM3's gate. +const Cfg = struct { h: usize, hkv: usize, hd: usize, window: usize, causal: bool, gate: bool }; +const cfgs = [_]Cfg{ + .{ .h = 16, .hkv = 4, .hd = 128, .window = 0, .causal = true, .gate = true }, + .{ .h = 16, .hkv = 4, .hd = 128, .window = 256, .causal = true, .gate = true }, + .{ .h = 8, .hkv = 4, .hd = 96, .window = 32, .causal = true, .gate = false }, + .{ .h = 32, .hkv = 8, .hd = 32, .window = 0, .causal = false, .gate = true }, + .{ .h = 8, .hkv = 2, .hd = 64, .window = 30, .causal = true, .gate = true }, +}; + +fn tiled(p: *const Probe) bool { + return p.variant & 0x10 == 0; +} + +fn cfgOf(p: *const Probe) Cfg { + return cfgs[p.variant & 0xf]; +} + +/// The cache's rows: the rows before (shape k) and the new rows (shape rows). +fn kvLen(p: *const Probe) usize { + return p.shape.k + p.shape.rows; +} + +/// A query row's width: its heads, and om_attend's gate logits after them (the projection's row). +fn ldq(p: *const Probe) usize { + const f = cfgOf(p); + return f.h * f.hd + if (tiled(p)) 0 else f.h; +} + +/// Shape rows: new query rows (om_attend runs one launch a row, as decode steps), n the head width, k rows before. +fn attnPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + const f = cfgOf(p); + if (s.rows == 0 or s.n != f.hd or (!f.causal and s.k != 0) or kvLen(p) > 4096) return false; + const nkv = kvLen(p) * f.hkv * f.hd; + const q = try p.alloc(u16, 0, s.rows * ldq(p)); + const kc = try p.alloc(u16, 1, nkv); + const vc = try p.alloc(u16, 2, nkv); + const g = try p.alloc(u16, 3, s.rows * f.h); + _ = try p.alloc(u16, 4, s.rows * f.h * f.hd); + if (p.fill == .stress) { + // equal scores (q = 0); over the last row's keys, integers whose partials grow large then cancel to a few + @memset(q, c.toBf16(30)); + for (0..s.rows) |r| @memset(q[r * ldq(p) ..][0 .. f.h * f.hd], 0); + k.randomBf16(rng, kc, -1, 1); + const n = kvLen(p); + const lo = if (f.causal and f.window > 0) n -| f.window else 0; + const half = (n - lo) / 2; + for (0..n) |jj| { + const j = jj -| lo; + // the second half's order turned by one, so 32 lanes' strided partials do not cancel early + const at = if (j < half) j else (j - half + 1) % @max(half, 1); + const pos: f32 = if (at % 2 == 0) 127 else 4064; + const res: f32 = if (at < 8 and at % 2 == 0) 1 else 0; + const v: f32 = if (jj < lo) 0 else if (j < half) pos else if (j < 2 * half) res - pos else 0; + @memset(vc[jj * f.hkv * f.hd ..][0 .. f.hkv * f.hd], c.toBf16(v)); + } + @memset(g, c.toBf16(30)); + } else { + k.randomBf16(rng, q, -1, 1); + k.randomBf16(rng, kc, -1, 1); + k.randomBf16(rng, vc, -2, 2); + k.randomBf16(rng, g, -3, 3); + } + if (p.fill == .nonfinite) { + // the last key's values: NaN, +Inf, -Inf in three columns and both infinities (two keys) in a fourth + const last = (kvLen(p) - 1) * f.hkv * f.hd; + const prev = if (kvLen(p) > 1) last - f.hkv * f.hd else last; + for ([_]usize{ last, last + 1, last + 2, last + 3, prev + 3 }, [_]u16{ k.nan16, k.inf16, k.ninf16, k.inf16, k.ninf16 }) |at, v| vc[at] = v; + // a NaN query in the first row's second kv group + q[f.h / f.hkv * f.hd + 3] = k.nan16; + } + for (0..4) |i| try p.in(i); + try p.out(4, false); + return true; +} + +const TileArgs = extern struct { ldq: i32, ldkv: i32, ldo: i32, T: i32, H: i32, Hkv: i32, hd: i32, pos0: i32, kv_len: i32, window: i32, causal: i32, scale: f32, has_gate: i32, ldg: i32 }; +const AttendArgs = extern struct { ldkv: i32, H: i32, Hkv: i32, hd: i32, pos: i32, kv_len: i32, window: i32, causal: i32, has_gate: i32, gate_off: i32, scale: f32 }; + +fn scale(f: Cfg) f32 { + return 1.0 / @sqrt(@as(f32, @floatFromInt(f.hd))); +} + +fn attnLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const f = cfgOf(p); + if (tiled(p)) { + p.bind(e, &.{ 0, 1, 2, 4, if (f.gate) 3 else null }, TileArgs{ .ldq = k.i32_(f.h * f.hd), .ldkv = k.i32_(f.hkv * f.hd), .ldo = k.i32_(f.h * f.hd), .T = k.i32_(s.rows), .H = k.i32_(f.h), .Hkv = k.i32_(f.hkv), .hd = k.i32_(f.hd), .pos0 = k.i32_(s.k), .kv_len = k.i32_(kvLen(p)), .window = k.i32_(f.window), .causal = @intFromBool(f.causal), .scale = scale(f), .has_gate = @intFromBool(f.gate), .ldg = k.i32_(f.h) }); + return Probe.go(e, .{ k.blocks(s.rows, 16), f.h, 1 }, 128); + } + e.setPipeline(p.pipe); + for (0..s.rows) |r| { + const qr = Probe.at(p.refs[0], r * ldq(p) * 2); + const o = Probe.at(p.refs[4], r * f.h * f.hd * 2); + e.setBuffer(qr.buf, qr.off, 0); + e.setBuffer(p.refs[1].buf, p.refs[1].off, 1); + e.setBuffer(p.refs[2].buf, p.refs[2].off, 2); + e.setBuffer(o.buf, o.off, 3); + const a = AttendArgs{ .ldkv = k.i32_(f.hkv * f.hd), .H = k.i32_(f.h), .Hkv = k.i32_(f.hkv), .hd = k.i32_(f.hd), .pos = k.i32_(s.k + r), .kv_len = k.i32_(s.k + r + 1), .window = k.i32_(f.window), .causal = @intFromBool(f.causal), .has_gate = @intFromBool(f.gate), .gate_off = k.i32_(f.h * f.hd), .scale = scale(f) }; + e.setBytes(std.mem.asBytes(&a), 4); + Probe.go(e, .{ f.h, 1, 1 }, 256); + } +} + +fn attnCount(p: *Probe) usize { + return p.shape.rows * cfgOf(p).h * cfgOf(p).hd; +} + +/// Softmax(q k^T scale) v in float64 a (row, head), bounded by the fp32 steps' spread; then bf16 and the gate. +fn attnEval(p: *Probe, mode: k.Mode, out: []R) !void { + const s = p.shape; + const f = cfgOf(p); + const sc: f64 = scale(f); + const q = p.get(u16, 0); + const kc = p.get(u16, 1); + const vc = p.get(u16, 2); + const ldkv = f.hkv * f.hd; + const n_kv = kvLen(p); + const buf = try p.a.alloc(f64, 4 * n_kv); + defer p.a.free(buf); + const sv = buf[0..n_kv]; + const ones = buf[n_kv .. 2 * n_kv]; + const col = buf[2 * n_kv .. 3 * n_kv]; + const unit = c.unit(.f32); + for (0..s.rows) |r| for (0..f.h) |h| { + const kh = h / (f.h / f.hkv); + const pos = s.k + r; + const hi = if (f.causal) pos + 1 else n_kv; + const lo = if (f.causal and f.window > 0) (pos + 1) -| f.window else 0; + const qr = q[r * ldq(p) + h * f.hd ..][0..f.hd]; + var m: f64 = -std.math.inf(f64); + var nan = false; + var e_max: f64 = 0; + for (lo..hi) |j| { + var dot: f64 = 0; + var mag: f64 = 0; + for (qr, kc[j * ldkv + kh * f.hd ..][0..f.hd]) |a, b| { + const t = @as(f64, c.fromBf16(a)) * c.fromBf16(b); + dot += t; + mag += @abs(t); + } + sv[j - lo] = dot * sc; + e_max = @max(e_max, c.gamma(@floatFromInt(f.hd + 2), unit) * mag * sc + 2 * unit * @abs(sv[j - lo])); + if (std.math.isNan(sv[j - lo])) nan = true; + m = @max(m, sv[j - lo]); + } + if (nan) m = std.math.nan(f64); + const n = hi - lo; + const tiles: f64 = if (tiled(p)) @floatFromInt((n + 31) / 32 + 1) else 1; + const e_p = 2 * e_max + (tiles + 2) * 8 * 0x1p-23; + const e_sum = c.gamma(@floatFromInt(n + 16), unit); + var l: f64 = 0; + for (sv[0..n]) |*x| { + x.* = @exp(x.* - m); + l += x.*; + } + for (0..f.hd) |e| { + var o: f64 = 0; + var mag: f64 = 0; + for (sv[0..n], lo..) |w, j| { + const v: f64 = c.fromBf16(vc[j * ldkv + kh * f.hd + e]); + o += w * v; + mag += w * @abs(v); + col[j - lo] = v; + } + var val = o / l; + var d = mag / l * (2 * e_p + 2 * e_sum) + @abs(val) * 4 * unit; + switch (mode) { + .ref => {}, + .exact => d = @abs(val) * 4 * unit, + .sim => |x| { + @memset(ones[0..n], 1); + val = contract.rows.sumIn(x.acc, x.order, col[0..n]) / contract.rows.sumIn(x.acc, x.order, ones[0..n]); + d = 0; + }, + } + var res = (R{ .v = val, .lo = val - d, .hi = val + d }).bf(); + if (f.gate) { + const gl = if (tiled(p)) p.get(u16, 3)[r * f.h + h] else q[r * ldq(p) + f.h * f.hd + h]; + res = res.mul(k.sigmR(R.bf16(gl)).bf()).bf(); + } + out[(r * f.h + h) * f.hd + e] = res; + } + }; +} + +const attn_def: k.Def = .{ .prep = attnPrep, .launch = attnLaunch, .count = attnCount, .eval = attnEval }; + +fn attnDepth(s: c.Shape) usize { + return s.n + s.k + s.rows; +} + +fn spec(name: []const u8, served: []const c.Shape, tile: c.Tile) c.Spec { + return .{ .name = name, .accum = .f32, .depth = attnDepth, .ops = 8, .out = .bf16, .tiles = .{ tile, .{}, .{} }, .served = served }; +} + +const mask = "if (k < nk && r0 + r < a.T && (!a.causal || kp <= pos) && (!(a.causal && a.window > 0) || kp > pos - a.window)) {"; +const gate = "if (a.has_gate) ob = b(f(ob) * rb(sigm(f(Q[a.gate_off + h]))));"; +const t16: c.Tile = .{ .size = 16 }; + +pub const entries = [_]k.Entry{ + .{ .def = &attn_def, .spec = spec("metal lfm3 omni om_attention_tiles full, gated (backbone)", &.{ .{ .rows = 64, .n = 128, .k = 37 }, .{ .rows = 5, .n = 128, .k = 0 } }, t16), .file = .omni, .pipe = "om_attention_tiles", .variant = 0, .control = .{ mask, "if (k < nk && r0 + r < a.T && (!a.causal || kp <= pos + 1) && (!(a.causal && a.window > 0) || kp > pos - a.window)) {" }, .what = "sees one key past the row" }, + .{ .def = &attn_def, .spec = spec("metal lfm3 omni om_attention_tiles sliding 256 (backbone)", &.{.{ .rows = 40, .n = 128, .k = 300 }}, t16), .file = .omni, .pipe = "om_attention_tiles", .variant = 1, .control = .{ mask, "if (k < nk && r0 + r < a.T && (!a.causal || kp <= pos) && (!(a.causal && a.window > 0) || kp >= pos - a.window)) {" }, .what = "keeps one key past the window" }, + .{ .def = &attn_def, .spec = spec("metal lfm3 omni om_attention_tiles window 32 (encoder)", &.{ .{ .rows = 94, .n = 96, .k = 0 }, .{ .rows = 20, .n = 96, .k = 0 } }, t16), .file = .omni, .pipe = "om_attention_tiles", .variant = 2, .control = .{ "if (c < dpl) acc[x][c] += p * f(vr[lane + 32 * c]);", "if (c < dpl) acc[x][c] += rb(p) * f(vr[lane + 32 * c]);" }, .what = "adds bf16(p)" }, + .{ .def = &attn_def, .spec = spec("metal lfm3 omni om_attention_tiles bidirectional, gated (depth decoder)", &.{.{ .rows = 12, .n = 32, .k = 0 }}, .{}), .file = .omni, .pipe = "om_attention_tiles", .variant = 3, .control = .{ "const float g = a.has_gate ? rb(sigm(f(G[(long)r * a.ldg + h]))) : 1.0f;", "const float g = a.has_gate ? (sigm(f(G[(long)r * a.ldg + h]))) : 1.0f;" }, .what = "drops bf16(sigmoid(gate))" }, + .{ .def = &attn_def, .spec = spec("metal lfm3 omni om_attention_tiles sliding 30, gated (codec)", &.{ .{ .rows = 48, .n = 64, .k = 50 }, .{ .rows = 6, .n = 64, .k = 50 } }, t16), .file = .omni, .pipe = "om_attention_tiles", .variant = 4, .control = .{ "const bf o = b(acc[x][c] * inv);", "const bf o = b(acc[x][c] * rb(inv));" }, .what = "adds bf16(1 / l)" }, + .{ .def = &attn_def, .spec = spec("metal lfm3 omni om_attend full, gated (backbone)", &.{ .{ .rows = 1, .n = 128, .k = 37 }, .{ .rows = 3, .n = 128, .k = 299 } }, .{}), .file = .rows, .pipe = "om_attend", .variant = 0x10, .control = .{ gate, "if (a.has_gate) ob = b(f(ob) * (sigm(f(Q[a.gate_off + h]))));" }, .what = "drops bf16(sigmoid(gate))" }, + .{ .def = &attn_def, .spec = spec("metal lfm3 omni om_attend sliding 256 (backbone)", &.{.{ .rows = 2, .n = 128, .k = 300 }}, .{}), .file = .rows, .pipe = "om_attend", .variant = 0x11, .control = .{ "const int lo = (a.causal && a.window > 0) ? max(0, a.pos + 1 - a.window) : 0;", "const int lo = (a.causal && a.window > 0) ? max(0, a.pos - a.window) : 0;" }, .what = "keeps one key past the window" }, + .{ .def = &attn_def, .spec = spec("metal lfm3 omni om_attend sliding 30, gated (codec)", &.{.{ .rows = 1, .n = 64, .k = 50 }}, .{}), .file = .rows, .pipe = "om_attend", .variant = 0x14, .control = .{ "bf ob = b(o * (1.0f / l));", "bf ob = b(o * rb(1.0f / l));" }, .what = "adds bf16(1 / l)" }, +}; diff --git a/zig/tests/contract/metal_omni_common.zig b/zig/tests/contract/metal_omni_common.zig new file mode 100644 index 0000000000..52f0172e15 --- /dev/null +++ b/zig/tests/contract/metal_omni_common.zig @@ -0,0 +1,441 @@ +//! LFM3-Omni's Metal kernel tests' shared parts: the sources, float64 values with their fp32 spread, and the probe. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const sources = @import("kernel_sources"); +const rig_ = @import("rig.zig"); + +pub const c = contract.spec; +pub const g = contract.gates; +pub const Ref = rig_.Ref; + +/// The metallib's four sources; a pipeline belongs to one. +pub const File = enum { omni, gemm, rows, nax }; + +/// A source's text as the metallib compiles it, its includes inlined. +pub fn text(a: std.mem.Allocator, file: File) ![]u8 { + const s = sources.lfm3_omni; + const common = try std.mem.replaceOwned(u8, a, s.common, "#pragma once\n", ""); + defer a.free(common); + return switch (file) { + .omni, .gemm, .rows => std.mem.replaceOwned(u8, a, switch (file) { + .omni => s.omni, + .gemm => s.gemm, + else => s.rows, + }, "#include \"om_common.h\"", common), + .nax => blk: { + const gemm = try std.mem.replaceOwned(u8, a, s.gemm_nax, "#include \"../nax.h\"", sources.nax); + defer a.free(gemm); + break :blk std.mem.replaceOwned(u8, a, s.nax, "#include \"../prefill/gemm_nax.metal\"", gemm); + }, + }; +} + +/// The metallib's options: safe math and precise fp32 functions, the OS's Metal language. +pub fn options() mtl.CompileOptions { + var o = mtl.CompileOptions.mlx(); + o.math_mode = .safe; + o.math_functions = .precise; + return o; +} + +/// A source compiled with `from` replaced once by `to` (a planted control), or as served. +pub fn compile(a: std.mem.Allocator, device: mtl.Device, file: File, edit: ?[2][]const u8) !mtl.Library { + const served = try text(a, file); + defer a.free(served); + const t = if (edit) |e| blk: { + if (std.mem.count(u8, served, e[0]) != 1) { + std.debug.print("control text not found once in {t}: {s}\n", .{ file, e[0] }); + return error.ControlTextNotFoundOnce; + } + break :blk try std.mem.replaceOwned(u8, a, served, e[0], e[1]); + } else try a.dupe(u8, served); + defer a.free(t); + return mtl.Library.fromSource(device, t, options()); +} + +// Metal's error bounds for the precise functions the kernels call, in fp32 ulps. +pub const div_ulps = 2.5; +pub const rsqrt_ulps = 2; +pub const sqrt_ulps = 1; +pub const fn_ulps = 4; + +fn bfv(v: f64) f64 { + return c.fromBf16(c.toBf16(@floatCast(v))); +} + +/// A float64 value at the kernel's bf16 points, and the range its fp32 steps may land in. +pub const R = struct { + v: f64, + lo: f64, + hi: f64, + + pub fn of(x: f64) R { + return .{ .v = x, .lo = x, .hi = x }; + } + pub fn bf16(h: u16) R { + return of(c.fromBf16(h)); + } + /// Rounded to bf16 (nearest even), both ends with it. + pub fn bf(x: R) R { + return .{ .v = bfv(x.v), .lo = bfv(x.lo), .hi = bfv(x.hi) }; + } + /// An fp32 step good to `ulps`. + pub fn fp(x: R, ulps: f64) R { + const u = ulps * 0x1p-23; + const lo = if (std.math.isFinite(x.lo)) x.lo - u * @abs(x.lo) - 0x1p-149 else x.lo; + const hi = if (std.math.isFinite(x.hi)) x.hi + u * @abs(x.hi) + 0x1p-149 else x.hi; + return .{ .v = x.v, .lo = lo, .hi = hi }; + } + pub fn add(x: R, y: R) R { + return (R{ .v = x.v + y.v, .lo = x.lo + y.lo, .hi = x.hi + y.hi }).fp(0.5); + } + pub fn sub(x: R, y: R) R { + return x.add(y.neg()); + } + pub fn neg(x: R) R { + return .{ .v = -x.v, .lo = -x.hi, .hi = -x.lo }; + } + fn span(x: R, y: R, v: f64, p: [4]f64) R { + _ = x; + _ = y; + return .{ .v = v, .lo = @min(@min(p[0], p[1]), @min(p[2], p[3])), .hi = @max(@max(p[0], p[1]), @max(p[2], p[3])) }; + } + pub fn mulExact(x: R, y: R) R { + return span(x, y, x.v * y.v, .{ x.lo * y.lo, x.lo * y.hi, x.hi * y.lo, x.hi * y.hi }); + } + pub fn mul(x: R, y: R) R { + return x.mulExact(y).fp(0.5); + } + /// x / y with y's range on one side of zero. + pub fn div(x: R, y: R) R { + return span(x, y, x.v / y.v, .{ x.lo / y.lo, x.lo / y.hi, x.hi / y.lo, x.hi / y.hi }).fp(div_ulps); + } + /// f of an increasing function, good to `ulps`. + pub fn up(x: R, comptime f: fn (f64) f64, ulps: f64) R { + return (R{ .v = f(x.v), .lo = f(x.lo), .hi = f(x.hi) }).fp(ulps); + } + /// f with |f'| <= l over the range, good to `ulps`. + pub fn lip(x: R, comptime f: fn (f64) f64, l: f64, ulps: f64) R { + const v = f(x.v); + const d = l * @max(x.v - x.lo, x.hi - x.v); + return (R{ .v = v, .lo = v - d, .hi = v + d }).fp(ulps); + } + pub fn rsqrt(x: R) R { + return (R{ .v = 1 / @sqrt(x.v), .lo = 1 / @sqrt(x.hi), .hi = 1 / @sqrt(x.lo) }).fp(rsqrt_ulps); + } +}; + +pub fn exp(x: f64) f64 { + return @exp(x); +} +pub fn sigmoid(x: f64) f64 { + return 1 / (1 + @exp(-x)); +} +pub fn silu(x: f64) f64 { + return x / (1 + @exp(-x)); +} + +/// exp(-x) as fp32 gives it: past fp32's largest value it is +Inf. +fn expNeg(x: R) R { + var e = x.neg().up(exp, fn_ulps); + const top: f64 = std.math.floatMax(f32); + if (e.v > top) e.v = std.math.inf(f64); + if (e.lo > top) e.lo = std.math.inf(f64); + if (e.hi > top) e.hi = std.math.inf(f64); + return e; +} + +/// SiLU as the kernels round it: x / (1 + exp(-x)), the exp and the division fp32 steps. +pub fn siluR(x: R) R { + return x.div(R.of(1).add(expNeg(x))); +} + +/// sigmoid as om_common.h's sigm. +pub fn sigmR(x: R) R { + return R.of(1).div(R.of(1).add(expNeg(x))); +} + +/// erf by its Taylor series (|x| < 4), else +-1 less erfc's asymptote: an independent float64 erf. +pub fn erf(x: f64) f64 { + const ax = @abs(x); + if (std.math.isNan(x)) return x; + var r: f64 = 1; + if (ax < 4) { + var term = ax; + var s: f64 = ax; + var n: f64 = 0; + while (@abs(term) > 1e-18 * @abs(s)) { + n += 1; + term *= -ax * ax / n; + s += term / (2 * n + 1); + } + r = s * 2 / @sqrt(std.math.pi); + } else r = 1 - @exp(-ax * ax) / (ax * @sqrt(std.math.pi)); + return if (x < 0) -r else r; +} + +/// The exact GELU as om_bias_act rounds it: 0.5 v (1 + erf(v / sqrt 2)), erf_f good to 2e-7. +pub fn geluR(v: R) R { + const z = v.mul(R.of(0.70710678118654752)); + var e = z.lip(erf, 1.1283791670955126, fn_ulps); + e.lo -= 2e-7; + e.hi += 2e-7; + return R.of(0.5).mul(v).mul(R.of(1).add(e)); +} + +/// How a replay runs its sums: as float64 with fp32's spread, exactly (gate 4's inputs), or in an accumulator. +pub const Mode = union(enum) { ref, exact, sim: struct { acc: c.Accum, order: g.Order } }; + +/// A sum of `terms` (`depth` fp32 roundings deep at most) under `mode`. +pub fn sum(mode: Mode, buf: []f64, terms: []const R, depth: usize) R { + switch (mode) { + .sim => |s| { + for (buf[0..terms.len], terms) |*b, t| b.* = t.v; + return R.of(contract.rows.sumIn(s.acc, s.order, buf[0..terms.len])); + }, + else => {}, + } + var r = R.of(0); + var m: f64 = 0; + for (terms) |t| { + r.v += t.v; + r.lo += t.lo; + r.hi += t.hi; + m += @max(@abs(t.lo), @abs(t.hi)); + } + if (mode == .ref) { + const e = c.gamma(@floatFromInt(depth), c.unit(.f32)) * m; + r.lo -= e; + r.hi += e; + } + return r; +} + +pub const nan16: u16 = 0x7fc0; +pub const inf16: u16 = 0x7f80; +pub const ninf16: u16 = 0xff80; + +pub fn randomBf16(rng: std.Random, out: []u16, lo: f32, hi: f32) void { + for (out) |*v| v.* = c.toBf16(lo + rng.float(f32) * (hi - lo)); +} + +/// Integers in [-bound, bound] as bf16 (exact while bound <= 256). +pub fn intsBf16(rng: std.Random, out: []u16, bound: u32) void { + for (out) |*v| v.* = c.toBf16(@as(f32, @floatFromInt(rng.uintLessThan(u32, 2 * bound + 1))) - @as(f32, @floatFromInt(bound))); +} + +pub fn noSums(_: c.Shape) usize { + return 0; +} + +/// One kernel's test: its pipeline, host inputs and device buffers; `Def` says how to fill, launch and replay it. +pub const Probe = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + def: *const Def, + spec_: c.Spec, + /// The variant a def serves (a mode, an epilogue, rows a launch). + variant: u32 = 0, + shape: c.Shape = undefined, + fill: g.Fill = .random, + host: [8][]align(8) u8 = @splat(&.{}), + refs: [8]Ref = undefined, + /// In-place outputs: their bytes before every run. + keep: [3]?struct { at: usize, ref: Ref } = @splat(null), + dummy: Ref = undefined, + lo: []f64 = &.{}, + hi: []f64 = &.{}, + n_out: usize = 0, + scratch: []f64 = &.{}, + + pub fn kernel(p: *Probe, control: ?g.Gate) g.Kernel { + return .{ .spec = p.spec_, .ctx = p, .vt = if (p.def.check != null) &vt_checked else &vt, .control = control }; + } + const vt = g.hooks(Probe); + const vt_checked = blk: { + var v = g.hooks(Probe); + v.check = checkHook; + break :blk v; + }; + fn checkHook(ctx: *anyopaque, gate: g.Gate) anyerror!?[]const u8 { + const p: *Probe = @ptrCast(@alignCast(ctx)); + return p.def.check.?(p, gate, try p.rig.outputBytes()); + } + + pub fn deinit(p: *Probe) void { + p.release(); + p.a.free(p.lo); + p.a.free(p.hi); + p.a.free(p.scratch); + p.rig.deinit(); + } + + fn release(p: *Probe) void { + for (&p.host) |*h| { + p.a.free(h.*); + h.* = &.{}; + } + p.keep = @splat(null); + } + + /// Host bytes for input slot `i`, typed. + pub fn alloc(p: *Probe, comptime T: type, i: usize, n: usize) ![]T { + const bytes = try p.a.alignedAlloc(u8, .@"8", n * @sizeOf(T)); + @memset(bytes, 0); + p.host[i] = bytes; + return @alignCast(std.mem.bytesAsSlice(T, bytes)); + } + pub fn get(p: *const Probe, comptime T: type, i: usize) []T { + return @alignCast(std.mem.bytesAsSlice(T, p.host[i])); + } + /// Slot `i` bound as an input. + pub fn in(p: *Probe, i: usize) !void { + p.refs[i] = try p.rig.input(p.host[i]); + } + /// Slot `i` bound as an output (poisoned), or in place (its host bytes restored before each run). + pub fn out(p: *Probe, i: usize, in_place: bool) !void { + p.refs[i] = try p.rig.output(p.host[i].len); + if (!in_place) return; + for (&p.keep) |*k| if (k.* == null) { + k.* = .{ .at = i, .ref = p.refs[i] }; + return; + }; + return error.TooManyInPlace; + } + + pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + p.rig.reset(); + p.release(); + p.shape = shape; + p.fill = fill; + var prng = std.Random.DefaultPrng.init(seed); + if (!try p.def.prep(p, prng.random())) return false; + p.dummy = try p.rig.scratch(64); + return true; + } + + pub fn run(p: *Probe) !void { + for (p.keep) |k| if (k) |x| @memcpy(p.rig.bytes(x.ref, p.host[x.at].len), p.host[x.at]); + const pass = p.rig.begin(); + p.def.launch(p, pass.e); + try p.rig.finish(pass); + } + + /// Binds slot refs (null: the dummy) at 0.., then the argument struct. + pub fn bind(p: *Probe, e: mtl.ComputeEncoder, slots: []const ?usize, args: anytype) void { + e.setPipeline(p.pipe); + for (slots, 0..) |s, i| { + const r = if (s) |x| p.refs[x] else p.dummy; + e.setBuffer(r.buf, r.off, i); + } + e.setBytes(std.mem.asBytes(&args), slots.len); + } + + /// A ref `bytes` further into its buffer. + pub fn at(r: Ref, bytes: usize) Ref { + return .{ .buf = r.buf, .off = r.off + bytes }; + } + + pub fn go(e: mtl.ComputeEncoder, groups: [3]usize, threads: usize) void { + e.dispatchGroups(mtl.Size.of(groups[0], groups[1], groups[2]), mtl.Size.of(threads, 1, 1)); + } + + pub fn count(p: *Probe) usize { + return p.def.count(p); + } + + pub fn replay(p: *Probe, mode: Mode, out_: []R) !void { + if (p.scratch.len < 1 << 16) { + p.a.free(p.scratch); + p.scratch = try p.a.alloc(f64, 1 << 16); + } + try p.def.eval(p, mode, out_); + } + + pub fn reference(p: *Probe, want: []f64, mag: []f64) !void { + const n = p.count(); + const rs = try p.a.alloc(R, n); + defer p.a.free(rs); + try p.replay(if (p.fill == .stress) .exact else .ref, rs); + p.a.free(p.lo); + p.a.free(p.hi); + p.lo = try p.a.alloc(f64, n); + p.hi = try p.a.alloc(f64, n); + for (rs, want, mag, p.lo, p.hi) |r, *w, *m, *l, *h| { + w.* = r.v; + m.* = if (std.math.isFinite(r.v)) @abs(r.v) else 0; + l.* = r.lo; + h.* = r.hi; + } + } + + /// The allowance about the float64 value: its range's farther end (a non-finite value is checked by class). + pub fn bounds(p: *Probe, want: []const f64, _: []const f64, allow: []f64) !void { + for (allow, want, p.lo, p.hi) |*l, w, lo, hi| { + const d = @max(@abs(w - lo), @abs(hi - w)); + l.* = if (std.math.isFinite(w) and std.math.isFinite(d)) d + 0x1p-140 else 0; + } + } + + pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out_: []f64) !void { + const rs = try p.a.alloc(R, out_.len); + defer p.a.free(rs); + try p.replay(.{ .sim = .{ .acc = acc, .order = order } }, rs); + for (out_, rs) |*o, r| o.* = r.v; + } +}; + +/// A kernel's test: inputs for `p.shape` (false: refused), its launch, its outputs and their float64 replay. +pub const Def = struct { + prep: *const fn (*Probe, std.Random) anyerror!bool, + launch: *const fn (*Probe, mtl.ComputeEncoder) void, + count: *const fn (*Probe) usize, + eval: *const fn (*Probe, Mode, []R) anyerror!void, + check: ?*const fn (*Probe, g.Gate, []const u8) anyerror!?[]const u8 = null, +}; + +/// An entry: the def, its spec, the pipeline and source it runs, its variant, and the control's text change if any. +pub const Entry = struct { + def: *const Def, + spec: c.Spec, + file: File, + pipe: []const u8, + variant: u32 = 0, + control: ?[2][]const u8 = null, + what: []const u8 = "", +}; + +pub fn blocks(n: usize, per: usize) usize { + return (n + per - 1) / per; +} + +pub fn i32_(x: usize) i32 { + return @intCast(x); +} + +/// A stated-rule kernel's check: every output inside its replay's range, a non-finite one of the replay's class. +pub fn checkReplay(p: *Probe, _: g.Gate, out: []const u8) !?[]const u8 { + const n = p.count(); + const rs = try p.a.alloc(R, n); + defer p.a.free(rs); + try p.replay(.ref, rs); + for (rs, 0..) |r, i| { + const got = c.read(p.spec_.out, out, i); + const ok = if (std.math.isFinite(r.v)) got >= r.lo and got <= r.hi else c.classOf(got) == c.classOf(r.v); + if (!ok) { + std.debug.print(" output {d}: {e} where the replay gives {e} in [{e}, {e}]\n", .{ i, got, r.v, r.lo, r.hi }); + return "an output leaves its replay's range"; + } + } + return null; +} + +/// A raw kernel's check: its output bytes equal `want` (the replay's bytes) exactly. +pub fn sameBytes(got: []const u8, want: []const u8) ?[]const u8 { + if (std.mem.eql(u8, got, want)) return null; + const at = std.mem.indexOfDiff(u8, got, want) orelse got.len; + std.debug.print(" byte {d} differs\n", .{at}); + return "its bytes differ from the replay's"; +} diff --git a/zig/tests/contract/metal_omni_conv.zig b/zig/tests/contract/metal_omni_conv.zig new file mode 100644 index 0000000000..1746adb81b --- /dev/null +++ b/zig/tests/contract/metal_omni_conv.zig @@ -0,0 +1,467 @@ +//! LFM3-Omni's Metal convolutions and gathers: the gated short conv, its state, the encoder's strided conv, embeddings. +const std = @import("std"); +const mtl = @import("metal"); +const k = @import("metal_omni_common.zig"); + +const c = k.c; +const R = k.R; +const Probe = k.Probe; + +fn plantRows(x: []u16, rows: usize, d: usize) void { + const at = [_][2]usize{ .{ 0, d / 3 }, .{ 1, d / 2 }, .{ 2, d - 1 }, .{ 3, 0 }, .{ 3, d - 1 } }; + const put = [_]u16{ k.nan16, k.inf16, k.ninf16, k.inf16, k.ninf16 }; + for (at, put) |p, v| if (p[0] < rows) { + x[p[0] * d + p[1]] = v; + }; +} + +/// An odd integer in [47, 61], whose products with another pass 2^11 (fp16's last exact integer). +fn big(rng: std.Random, signed: bool) f32 { + const v: f32 = @floatFromInt(47 + 2 * rng.uintLessThan(u32, 8)); + return if (signed and rng.boolean()) -v else v; +} + +/// An odd integer in [1, top] (top odd), with a random sign when `signed`. +fn odd(rng: std.Random, top: u32, signed: bool) f32 { + const v: f32 = @floatFromInt(2 * rng.uintLessThan(u32, (top + 1) / 2) + 1); + return if (signed and rng.boolean()) -v else v; +} + +// --- the gated short conv: om_short_conv (a few rows, the state updated), om_short_conv_rows and om_conv_state + +const ConvArgs = extern struct { ld: i32, ldy: i32, T: i32, D: i32, K: i32 }; + +/// B | C | x rows, taps, state; stress: B = C = 1, x and state h, taps (w, -w - 2) pairs then K: partials past 2^11. +fn convInputs(p: *Probe, rng: std.Random) !void { + const s = p.shape; + const d = s.n; + const kw = s.k; + const bcx = try p.alloc(u16, 0, s.rows * 3 * d); + const w = try p.alloc(u16, 1, d * kw); + const state = try p.alloc(u16, 2, d * (kw - 1)); + if (p.fill == .stress) { + for (0..d) |ch| { + const h = c.toBf16(big(rng, true)); + for (0..s.rows) |t| { + bcx[t * 3 * d + ch] = c.toBf16(1); + bcx[t * 3 * d + d + ch] = c.toBf16(1); + bcx[t * 3 * d + 2 * d + ch] = h; + } + for (0..kw - 1) |j| state[ch * (kw - 1) + j] = h; + var j: usize = 0; + while (j + 1 < kw) : (j += 2) { + const v = big(rng, false); + w[ch * kw + j] = c.toBf16(v); + w[ch * kw + j + 1] = c.toBf16(-v - 2); + } + w[ch * kw + kw - 1] = c.toBf16(@floatFromInt(kw)); + } + return; + } + k.randomBf16(rng, bcx, -2, 2); + k.randomBf16(rng, w, -1, 1); + k.randomBf16(rng, state, -2, 2); + // the last rows, which the state keeps + const r4 = @min(s.rows, 4); + if (p.fill == .nonfinite) plantRows(bcx[(s.rows - r4) * 3 * d ..], r4, 3 * d); +} + +fn convPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.k < 2 or s.k > 8 or s.k % 2 == 0 or s.n == 0) return false; + try convInputs(p, rng); + _ = try p.alloc(u16, 3, s.rows * s.n); + try p.in(0); + try p.in(1); + if (p.variant == 0) try p.out(2, true) else try p.in(2); + try p.out(3, false); + return true; +} + +fn convLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const a = ConvArgs{ .ld = k.i32_(3 * s.n), .ldy = k.i32_(s.n), .T = k.i32_(s.rows), .D = k.i32_(s.n), .K = k.i32_(s.k) }; + p.bind(e, &.{ 0, 1, 2, 3 }, a); + if (p.variant == 0) Probe.go(e, .{ k.blocks(s.n, 128), 1, 1 }, 128) else Probe.go(e, .{ k.blocks(s.rows * s.n, 256), 1, 1 }, 256); +} + +fn convCount(p: *Probe) usize { + return p.shape.rows * p.shape.n + if (p.variant == 0) p.shape.n * (p.shape.k - 1) else 0; +} + +fn bx(p: *Probe, t: usize, ch: usize) R { + const d = p.shape.n; + const row = p.get(u16, 0)[t * 3 * d ..]; + return R.bf16(row[ch]).mul(R.bf16(row[2 * d + ch])).bf(); +} + +/// Both forms: the taps over the state then B x rows, summed in order; y = bf16(C bf16(acc)). +fn convEval(p: *Probe, mode: k.Mode, out: []R) !void { + const s = p.shape; + const d = s.n; + const kw = s.k; + const w = p.get(u16, 1); + const state = p.get(u16, 2); + const ys = if (p.variant == 0) out[d * (kw - 1) ..] else out; + var terms: [8]R = undefined; + for (0..d) |ch| { + for (0..s.rows) |t| { + for (0..kw) |j| { + const src = @as(isize, @intCast(t + j)) - @as(isize, @intCast(kw - 1)); + const v = if (src >= 0) bx(p, @intCast(src), ch) else R.bf16(state[ch * (kw - 1) + @as(usize, @intCast(@as(isize, @intCast(kw - 1)) + src))]); + terms[j] = R.bf16(w[ch * kw + j]).mulExact(v); + } + const acc = k.sum(mode, p.scratch, terms[0..kw], kw + 1).bf(); + ys[t * d + ch] = R.bf16(p.get(u16, 0)[t * 3 * d + d + ch]).mul(acc).bf(); + } + if (p.variant == 0) for (0..kw - 1) |j| { + const src = @as(isize, @intCast(s.rows + j)) - @as(isize, @intCast(kw - 1)); + out[ch * (kw - 1) + j] = if (src >= 0) bx(p, @intCast(src), ch) else R.bf16(state[ch * (kw - 1) + j + s.rows]); + }; + } +} + +const conv_def: k.Def = .{ .prep = convPrep, .launch = convLaunch, .count = convCount, .eval = convEval }; + +fn convDepth(s: c.Shape) usize { + return s.k + 1; +} + +fn stateDef() k.Def { + return .{ .prep = statePrep, .launch = stateLaunch, .count = stateCount, .eval = stateEval }; +} + +fn statePrep(p: *Probe, rng: std.Random) !bool { + if (p.shape.k < 2 or p.shape.k > 8) return false; + try convInputs(p, rng); + try p.in(0); + try p.out(2, true); + return true; +} + +fn stateLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + p.bind(e, &.{ 0, 2 }, ConvArgs{ .ld = k.i32_(3 * s.n), .ldy = k.i32_(s.n), .T = k.i32_(s.rows), .D = k.i32_(s.n), .K = k.i32_(s.k) }); + Probe.go(e, .{ k.blocks(s.n, 128), 1, 1 }, 128); +} + +fn stateCount(p: *Probe) usize { + return p.shape.n * (p.shape.k - 1); +} + +fn stateEval(p: *Probe, _: k.Mode, out: []R) !void { + const s = p.shape; + const kw = s.k; + const state = p.get(u16, 2); + for (0..s.n) |ch| for (0..kw - 1) |j| { + const src = @as(isize, @intCast(s.rows + j)) - @as(isize, @intCast(kw - 1)); + out[ch * (kw - 1) + j] = if (src >= 0) bx(p, @intCast(src), ch) else R.bf16(state[ch * (kw - 1) + j + s.rows]); + }; +} + +const state_def = stateDef(); + +// --- om_conv2d_s2: the encoder's stride-2 3x3 convs; shape rows C, n H, k W; variant 0 conv_in (+SiLU), 1 depthwise + +const Conv2dArgs = extern struct { CIN: i32, H: i32, W: i32, C: i32, H2: i32, W2: i32, silu: i32 }; + +fn c2Prep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.rows == 0 or s.n < 2 or s.k < 2) return false; + const planes: usize = if (p.variant == 0) 1 else s.rows; + const x = try p.alloc(u16, 0, planes * s.n * s.k); + const w = try p.alloc(u16, 1, s.rows * 9); + const bias = try p.alloc(u16, 2, s.rows); + _ = try p.alloc(u16, 3, s.rows * ((s.n + 1) / 2) * ((s.k + 1) / 2)); + if (p.fill == .stress) { + // every plane one odd integer, taps (w, -w - 2) pairs then 9: interior sums come to the plane's value + for (0..planes) |pl| @memset(x[pl * s.n * s.k ..][0 .. s.n * s.k], c.toBf16(big(rng, true))); + for (0..s.rows) |ch| { + var j: usize = 0; + while (j < 8) : (j += 2) { + const v = big(rng, false); + w[ch * 9 + j] = c.toBf16(v); + w[ch * 9 + j + 1] = c.toBf16(-v - 2); + } + w[ch * 9 + 8] = c.toBf16(9); + } + @memset(bias, 0); + } else { + k.randomBf16(rng, x, -2, 2); + k.randomBf16(rng, w, -1, 1); + k.randomBf16(rng, bias, -0.5, 0.5); + if (p.fill == .nonfinite) plantRows(x, @min(planes, 4) * s.n, s.k); + } + for (0..3) |i| try p.in(i); + try p.out(3, false); + return true; +} + +fn c2Launch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const h2 = (s.n + 1) / 2; + const w2 = (s.k + 1) / 2; + p.bind(e, &.{ 0, 1, 2, 3 }, Conv2dArgs{ .CIN = if (p.variant == 0) 1 else k.i32_(s.rows), .H = k.i32_(s.n), .W = k.i32_(s.k), .C = k.i32_(s.rows), .H2 = k.i32_(h2), .W2 = k.i32_(w2), .silu = @intFromBool(p.variant == 0) }); + Probe.go(e, .{ k.blocks(s.rows * h2 * w2, 256), 1, 1 }, 256); +} + +fn c2Count(p: *Probe) usize { + return p.shape.rows * ((p.shape.n + 1) / 2) * ((p.shape.k + 1) / 2); +} + +fn c2Eval(p: *Probe, mode: k.Mode, out: []R) !void { + const s = p.shape; + const h2 = (s.n + 1) / 2; + const w2 = (s.k + 1) / 2; + const x = p.get(u16, 0); + const w = p.get(u16, 1); + var terms: [9]R = undefined; + for (0..s.rows) |ch| for (0..h2) |h| for (0..w2) |col| { + const plane = x[(if (p.variant == 0) 0 else ch) * s.n * s.k ..]; + var n: usize = 0; + for (0..3) |kh| for (0..3) |kc| { + const ih = @as(isize, @intCast(2 * h + kh)) - 1; + const iw = @as(isize, @intCast(2 * col + kc)) - 1; + if (ih < 0 or ih >= s.n or iw < 0 or iw >= s.k) continue; + terms[n] = R.bf16(w[ch * 9 + kh * 3 + kc]).mulExact(R.bf16(plane[@as(usize, @intCast(ih)) * s.k + @as(usize, @intCast(iw))])); + n += 1; + }; + var v = k.sum(mode, p.scratch, terms[0..n], 10).add(R.bf16(p.get(u16, 2)[ch])).bf(); + if (p.variant == 0) v = k.siluR(v).bf(); + out[(ch * h2 + h) * w2 + col] = v; + }; +} + +const c2_def: k.Def = .{ .prep = c2Prep, .launch = c2Launch, .count = c2Count, .eval = c2Eval }; + +fn depth10(_: c.Shape) usize { + return 10; +} + +// --- gathers: om_embed (ids < 0 leave a row), om_embed_codes (12 codebooks summed), om_depth_in + +const EmbedArgs = extern struct { T: i32, D: i32, NQ: i32, size: i32 }; +const nq = 12; +const size = 2017; +const embed_rows = 4096; + +fn embedPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0 or s.rows == 0) return false; + const table = try p.alloc(u16, 0, embed_rows * s.n); + const ids = try p.alloc(i32, 1, s.rows); + const y = try p.alloc(u16, 2, s.rows * s.n); + k.randomBf16(rng, table, -1, 1); + k.randomBf16(rng, y, -1, 1); + for (ids, 0..) |*id, t| id.* = if (t % 5 == 3) -1 else @intCast(rng.uintLessThan(u32, embed_rows)); + if (p.fill == .nonfinite) { + ids[0] = 7; + ids[1] = 9; + plantRows(table[7 * s.n ..][0 .. 3 * s.n], 3, s.n); + } + try p.in(0); + try p.in(1); + try p.out(2, true); + return true; +} + +fn embedLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + p.bind(e, &.{ 0, 1, 2 }, EmbedArgs{ .T = k.i32_(s.rows), .D = k.i32_(s.n), .NQ = 0, .size = 0 }); + Probe.go(e, .{ k.blocks(s.rows * s.n, 256), 1, 1 }, 256); +} + +fn rowsN(p: *Probe) usize { + return p.shape.rows * p.shape.n; +} + +fn embedEval(p: *Probe, _: k.Mode, out: []R) !void { + const d = p.shape.n; + for (p.get(i32, 1), 0..) |id, t| for (0..d) |i| { + out[t * d + i] = R.bf16(if (id >= 0) p.get(u16, 0)[@as(usize, @intCast(id)) * d + i] else p.get(u16, 2)[t * d + i]); + }; +} + +const embed_def: k.Def = .{ .prep = embedPrep, .launch = embedLaunch, .count = rowsN, .eval = embedEval }; + +/// Codebook tables; stress: a big multiple of 64, an odd integer, then their negations less a little: small sums. +fn codesPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0 or s.rows == 0 or s.k != nq) return false; + const table = try p.alloc(u16, 0, nq * size * s.n); + const codes = try p.alloc(i32, 1, s.rows * nq); + _ = try p.alloc(u16, 2, s.rows * s.n); + for (codes) |*x| x.* = @intCast(rng.uintLessThan(u32, size)); + if (p.fill == .stress) { + @memset(table, 0); + for (0..s.n) |col| { + const top: f32 = @floatFromInt(64 * (64 + rng.uintLessThan(u32, 192))); + const o = odd(rng, 255, true); + const small: f32 = @floatFromInt(rng.uintLessThan(u32, 9)); + for (0..size) |code| { + table[(0 * size + code) * s.n + col] = c.toBf16(top); + table[(1 * size + code) * s.n + col] = c.toBf16(o); + table[(2 * size + code) * s.n + col] = c.toBf16(-top); + table[(3 * size + code) * s.n + col] = c.toBf16(small - o); + } + } + } else { + k.randomBf16(rng, table, -1, 1); + if (p.fill == .nonfinite) { + for (0..nq) |q| codes[q] = 5; + plantRows(table[(2 * size + 5) * s.n ..][0..s.n], 1, s.n); + table[(7 * size + 5) * s.n + 1] = k.inf16; + table[(9 * size + 5) * s.n + 2] = k.ninf16; + table[(4 * size + 5) * s.n + 3] = k.inf16; + table[(8 * size + 5) * s.n + 3] = k.ninf16; + } + } + try p.in(0); + try p.in(1); + try p.out(2, false); + return true; +} + +fn codesLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + p.bind(e, &.{ 0, 1, 2 }, EmbedArgs{ .T = k.i32_(s.rows), .D = k.i32_(s.n), .NQ = nq, .size = size }); + Probe.go(e, .{ k.blocks(s.rows * s.n, 256), 1, 1 }, 256); +} + +fn codesEval(p: *Probe, mode: k.Mode, out: []R) !void { + const d = p.shape.n; + const codes = p.get(i32, 1); + var terms: [nq]R = undefined; + for (0..p.shape.rows) |t| for (0..d) |i| { + for (&terms, 0..) |*x, q| x.* = R.bf16(p.get(u16, 0)[(@as(usize, @intCast(codes[t * nq + q])) + q * size) * d + i]); + out[t * d + i] = k.sum(mode, p.scratch, &terms, nq).bf(); + }; +} + +const codes_def: k.Def = .{ .prep = codesPrep, .launch = codesLaunch, .count = rowsN, .eval = codesEval }; + +fn depthNq(_: c.Shape) usize { + return nq; +} + +fn dinPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.rows == 0 or s.rows > nq or s.n == 0) return false; + const din = try p.alloc(u16, 0, s.rows * s.n); + const mask = try p.alloc(u16, 1, s.n); + const tables = try p.alloc(u16, 2, nq * size * s.n); + const codes = try p.alloc(i32, 3, s.rows); + const masked = try p.alloc(i32, 4, s.rows); + _ = try p.alloc(u16, 5, s.rows * s.n); + k.randomBf16(rng, din, -2, 2); + k.randomBf16(rng, mask, -1, 1); + k.randomBf16(rng, tables, -1, 1); + for (codes, masked) |*x, *m| { + x.* = @intCast(rng.uintLessThan(u32, size)); + m.* = @intFromBool(rng.boolean()); + } + if (p.fill == .nonfinite) plantRows(din, s.rows, s.n); + for (0..5) |i| try p.in(i); + try p.out(5, false); + return true; +} + +fn dinLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1, 2, 3, 4, 5 }, EmbedArgs{ .T = 0, .D = k.i32_(p.shape.n), .NQ = k.i32_(p.shape.rows), .size = size }); + Probe.go(e, .{ k.blocks(p.shape.rows * p.shape.n, 256), 1, 1 }, 256); +} + +fn dinEval(p: *Probe, _: k.Mode, out: []R) !void { + const d = p.shape.n; + for (0..p.shape.rows) |q| for (0..d) |i| { + const e = if (p.get(i32, 4)[q] != 0) p.get(u16, 1)[i] else p.get(u16, 2)[(q * size + @as(usize, @intCast(p.get(i32, 3)[q]))) * d + i]; + out[q * d + i] = R.bf16(p.get(u16, 0)[q * d + i]).add(R.bf16(e)).bf(); + }; +} + +const din_def: k.Def = .{ .prep = dinPrep, .launch = dinLaunch, .count = rowsN, .eval = dinEval }; + +// --- moves: om_copy_rows (shape rows T, n the width, k the source's row stride), om_repeat_rows (k the factor) + +const CopyArgs = extern struct { T: i32, n: i32, lds: i32, ldd: i32 }; + +fn copyPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.k < s.n or s.n == 0) return false; + const src = try p.alloc(u16, 0, s.rows * s.k); + const dst = try p.alloc(u16, 1, s.rows * (s.n + 8)); + k.randomBf16(rng, src, -2, 2); + k.randomBf16(rng, dst, -2, 2); + if (p.fill == .nonfinite) plantRows(src, s.rows, s.k); + try p.in(0); + try p.out(1, true); + return true; +} + +fn copyLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + p.bind(e, &.{ 0, 1 }, CopyArgs{ .T = k.i32_(s.rows), .n = k.i32_(s.n), .lds = k.i32_(s.k), .ldd = k.i32_(s.n + 8) }); + Probe.go(e, .{ k.blocks(s.rows * s.n, 256), 1, 1 }, 256); +} + +fn copyCount(p: *Probe) usize { + return p.shape.rows * (p.shape.n + 8); +} + +fn copyEval(p: *Probe, _: k.Mode, out: []R) !void { + const s = p.shape; + for (out, p.get(u16, 1)) |*o, v| o.* = R.bf16(v); + for (0..s.rows) |t| for (0..s.n) |i| { + out[t * (s.n + 8) + i] = R.bf16(p.get(u16, 0)[t * s.k + i]); + }; +} + +const copy_def: k.Def = .{ .prep = copyPrep, .launch = copyLaunch, .count = copyCount, .eval = copyEval }; + +const RepeatArgs = extern struct { n: i32, D: i32, rep: i32 }; + +fn repPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.k == 0 or s.n == 0) return false; + const x = try p.alloc(u16, 0, s.rows * s.n); + _ = try p.alloc(u16, 1, s.rows * s.k * s.n); + k.randomBf16(rng, x, -2, 2); + if (p.fill == .nonfinite) plantRows(x, s.rows, s.n); + try p.in(0); + try p.out(1, false); + return true; +} + +fn repLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + p.bind(e, &.{ 0, 1 }, RepeatArgs{ .n = k.i32_(s.rows), .D = k.i32_(s.n), .rep = k.i32_(s.k) }); + Probe.go(e, .{ k.blocks(s.rows * s.k * s.n, 256), 1, 1 }, 256); +} + +fn repCount(p: *Probe) usize { + return p.shape.rows * p.shape.k * p.shape.n; +} + +fn repEval(p: *Probe, _: k.Mode, out: []R) !void { + const s = p.shape; + for (out, 0..) |*o, i| o.* = R.bf16(p.get(u16, 0)[(i / s.n / s.k) * s.n + i % s.n]); +} + +const rep_def: k.Def = .{ .prep = repPrep, .launch = repLaunch, .count = repCount, .eval = repEval }; + +fn none(name: []const u8, served: []const c.Shape, in_place: bool) c.Spec { + return .{ .name = name, .accum = .none, .depth = k.noSums, .ops = 4, .out = .bf16, .in_place = in_place, .served = served, .sums = .{ .f16 = 0, .bf16 = 0, .why = "no sums" } }; +} + +pub const entries = [_]k.Entry{ + .{ .def = &conv_def, .spec = .{ .name = "metal lfm3 omni om_short_conv", .accum = .f32, .depth = convDepth, .ops = 2, .out = .bf16, .in_place = true, .tiles = .{ .{ .size = 4, .min = 1 }, .{ .size = 128 }, .{} }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 7 }, .{ .rows = 4, .n = 512, .k = 7 }, .{ .rows = 1, .n = 1024, .k = 7 } } }, .file = .omni, .pipe = "om_short_conv", .control = .{ "const float bx = rb(f(row[d]) * f(row[2 * D + d]));", "const float bx = (f(row[d]) * f(row[2 * D + d]));" }, .what = "drops bf16(B x)" }, + .{ .def = &conv_def, .spec = .{ .name = "metal lfm3 omni om_short_conv_rows", .accum = .f32, .depth = convDepth, .ops = 2, .out = .bf16, .tiles = .{ .{ .size = 64 }, .{}, .{} }, .served = &.{ .{ .rows = 94, .n = 768, .k = 3 }, .{ .rows = 40, .n = 2048, .k = 7 }, .{ .rows = 48, .n = 512, .k = 7 } } }, .file = .omni, .pipe = "om_short_conv_rows", .variant = 1, .control = .{ "Y[(long)t * a.ldy + d] = b(f(BCx[(long)t * a.ld + D + d]) * rb(acc));", "Y[(long)t * a.ldy + d] = b(f(BCx[(long)t * a.ld + D + d]) * (acc));" }, .what = "drops bf16(acc)" }, + .{ .def = &state_def, .spec = none("metal lfm3 omni om_conv_state", &.{ .{ .rows = 94, .n = 768, .k = 3 }, .{ .rows = 40, .n = 2048, .k = 7 }, .{ .rows = 5, .n = 512, .k = 7 } }, true), .file = .omni, .pipe = "om_conv_state", .control = .{ "STATE[d * (K - 1) + j] = STATE[d * (K - 1) + j + T];", "STATE[d * (K - 1) + j] = STATE[d * (K - 1) + j + T - 1];" }, .what = "shifts the kept state by one row less" }, + .{ .def = &c2_def, .spec = .{ .name = "metal lfm3 omni om_conv2d_s2 conv_in, SiLU", .accum = .f32, .depth = depth10, .ops = 8, .mids = 1, .out = .bf16, .served = &.{ .{ .rows = 256, .n = 128, .k = 201 }, .{ .rows = 16, .n = 128, .k = 64 } } }, .file = .omni, .pipe = "om_conv2d_s2", .variant = 0, .control = .{ "float v = rb(acc + f(Bias[c]));", "float v = (acc + f(Bias[c]));" }, .what = "drops bf16(conv + bias) before SiLU" }, + .{ .def = &c2_def, .spec = .{ .name = "metal lfm3 omni om_conv2d_s2 depthwise", .accum = .f32, .depth = depth10, .ops = 4, .out = .bf16, .served = &.{ .{ .rows = 256, .n = 64, .k = 101 }, .{ .rows = 256, .n = 32, .k = 51 } } }, .file = .omni, .pipe = "om_conv2d_s2", .variant = 1, .control = .{ "const int ih = 2 * h - 1 + kh;", "const int ih = 2 * h + kh;" }, .what = "drops the top padding row" }, + .{ .def = &embed_def, .spec = none("metal lfm3 omni om_embed", &.{ .{ .rows = 37, .n = 2048, .k = 1 }, .{ .rows = 1, .n = 2048, .k = 1 } }, true), .file = .omni, .pipe = "om_embed", .control = .{ "if (id >= 0) Y[i] = TABLE[(long)id * a.D + c];", "if (id >= 0) Y[i] = TABLE[(long)id * a.D + (c ^ 1)];" }, .what = "swaps column pairs" }, + .{ .def = &codes_def, .spec = .{ .name = "metal lfm3 omni om_embed_codes", .accum = .f32, .depth = depthNq, .out = .bf16, .served = &.{ .{ .rows = 1, .n = 2048, .k = nq }, .{ .rows = 5, .n = 2048, .k = nq } } }, .file = .omni, .pipe = "om_embed_codes", .control = .{ "for (int q = 0; q < a.NQ; ++q) s += f(TABLE[", "for (int q = 0; q < a.NQ; ++q) s = rb(s) + f(TABLE[" }, .what = "rounds each partial sum to bf16" }, + .{ .def = &din_def, .spec = none("metal lfm3 omni om_depth_in", &.{.{ .rows = nq, .n = 1024, .k = 1 }}, false), .file = .omni, .pipe = "om_depth_in", .control = .{ "X[i] = b(f(DIN[i]) + e);", "X[i] = b(f(DIN[i]) + rb(e * 1.0078125f));" }, .what = "scales the slot's embedding" }, + .{ .def = ©_def, .spec = none("metal lfm3 omni om_copy_rows", &.{ .{ .rows = 12, .n = 1024, .k = 1024 }, .{ .rows = 6, .n = 512, .k = 776 } }, true), .file = .omni, .pipe = "om_copy_rows", .control = .{ "D[(long)t * a.ldd + c] = S[(long)t * a.lds + c];", "D[(long)t * a.ldd + c] = S[(long)t * a.lds + (c ^ 1)];" }, .what = "swaps column pairs" }, + .{ .def = &rep_def, .spec = none("metal lfm3 omni om_repeat_rows", &.{ .{ .rows = 64, .n = 512, .k = 6 }, .{ .rows = 1, .n = 512, .k = 6 } }, false), .file = .omni, .pipe = "om_repeat_rows", .control = .{ "Y[i] = X[(r / a.rep) * a.D + i % a.D];", "Y[i] = X[(r % a.n) * a.D + i % a.D];" }, .what = "repeats the rows in turn" }, +}; diff --git a/zig/tests/contract/metal_omni_elem.zig b/zig/tests/contract/metal_omni_elem.zig new file mode 100644 index 0000000000..f1342bbeff --- /dev/null +++ b/zig/tests/contract/metal_omni_elem.zig @@ -0,0 +1,418 @@ +//! LFM3-Omni's Metal elementwise kernels: activations, casts, the codec's FSQ and spectra, log-mel, word moves. +const std = @import("std"); +const mtl = @import("metal"); +const k = @import("metal_omni_common.zig"); + +const c = k.c; +const R = k.R; +const Probe = k.Probe; + +/// NaN, +Inf, -Inf at three places, and both infinities at two more (indexes past the end are skipped). +fn plant16(x: []u16) void { + const n = x.len; + for ([_]usize{ n / 3, n / 2, n - 1, 1, n / 5 }, [_]u16{ k.nan16, k.inf16, k.ninf16, k.inf16, k.ninf16 }) |i, v| x[i] = v; +} + +fn plant32(x: []f32) void { + const n = x.len; + const nan = std.math.nan(f32); + const inf = std.math.inf(f32); + for ([_]usize{ n / 3, n / 2, n - 1, 1, n / 5 }, [_]f32{ nan, inf, -inf, inf, -inf }) |i, v| x[i] = v; +} + +fn none(name: []const u8, out: c.Elem, served: []const c.Shape) c.Spec { + return .{ .name = name, .accum = .none, .depth = k.noSums, .ops = 8, .out = out, .served = served, .sums = .{ .f16 = 0, .bf16 = 0, .why = "no sums" } }; +} + +fn flat(p: *Probe) usize { + return p.shape.rows * p.shape.n; +} + +// --- om_swiglu: rows of 2F, halves (variant 0) or (a, c) pairs (1) + +const SwiArgs = extern struct { T: i32, F: i32, inter: i32 }; + +fn swiPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0) return false; + const ac = try p.alloc(u16, 0, s.rows * 2 * s.n); + _ = try p.alloc(u16, 1, s.rows * s.n); + k.randomBf16(rng, ac, -6, 6); + if (p.fill == .nonfinite) plant16(ac[0 .. 2 * s.n * @min(s.rows, 4)]); + try p.in(0); + try p.out(1, false); + return true; +} + +fn swiLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1 }, SwiArgs{ .T = k.i32_(p.shape.rows), .F = k.i32_(p.shape.n), .inter = @intCast(p.variant) }); + Probe.go(e, .{ k.blocks(flat(p), 256), 1, 1 }, 256); +} + +fn swiEval(p: *Probe, _: k.Mode, out: []R) !void { + const f = p.shape.n; + const ac = p.get(u16, 0); + for (out, 0..) |*o, i| { + const t = i / f; + const col = i % f; + const ia = if (p.variant == 1) t * 2 * f + 2 * col else t * 2 * f + col; + const ic = if (p.variant == 1) ia + 1 else ia + f; + o.* = k.siluR(R.bf16(ac[ia])).bf().mul(R.bf16(ac[ic])).bf(); + } +} + +const swi_def: k.Def = .{ .prep = swiPrep, .launch = swiLaunch, .count = flat, .eval = swiEval }; + +// --- om_bias_act: act(bf16(A + bias)) from fp32 products; variant 0 column bias, 1 row bias and SiLU, 2 GELU + +const BiasArgs = extern struct { rows: i32, cols: i32, per_row: i32, act: i32, has_bias: i32 }; + +fn biasPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0) return false; + const a = try p.alloc(f32, 0, s.rows * s.n); + const bias = try p.alloc(u16, 1, if (p.variant == 1) s.rows else s.n); + _ = try p.alloc(u16, 2, s.rows * s.n); + for (a) |*x| x.* = (rng.float(f32) * 2 - 1) * 4; + k.randomBf16(rng, bias, -1, 1); + if (p.fill == .nonfinite) plant32(a[0 .. s.n * @min(s.rows, 4)]); + try p.in(0); + try p.in(1); + try p.out(2, false); + return true; +} + +fn biasLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1, 2 }, BiasArgs{ .rows = k.i32_(p.shape.rows), .cols = k.i32_(p.shape.n), .per_row = @intFromBool(p.variant == 1), .act = @intCast(p.variant), .has_bias = 1 }); + Probe.go(e, .{ k.blocks(flat(p), 256), 1, 1 }, 256); +} + +fn biasEval(p: *Probe, _: k.Mode, out: []R) !void { + const n = p.shape.n; + for (out, p.get(f32, 0), 0..) |*o, a, i| { + const b = R.bf16(p.get(u16, 1)[if (p.variant == 1) i / n else i % n]); + const v = R.of(a).add(b).bf(); + o.* = (switch (p.variant) { + 1 => k.siluR(v), + 2 => k.geluR(v), + else => v, + }).bf(); + } +} + +const bias_def: k.Def = .{ .prep = biasPrep, .launch = biasLaunch, .count = flat, .eval = biasEval }; + +// --- casts: om_f32_bf16, om_bf16_f32; shape rows 1, n elements + +fn castPrep(p: *Probe, rng: std.Random) !bool { + const n = flat(p); + if (n == 0) return false; + if (p.variant == 0) { + const x = try p.alloc(f32, 0, n); + for (x, 0..) |*v, i| v.* = (rng.float(f32) * 2 - 1) * (if (i % 7 == 0) @as(f32, 1e5) else 1); + if (p.fill == .nonfinite) plant32(x); + _ = try p.alloc(u16, 1, n); + } else { + const x = try p.alloc(u16, 0, n); + k.randomBf16(rng, x, -3, 3); + if (p.fill == .nonfinite) plant16(x); + _ = try p.alloc(f32, 1, n); + } + try p.in(0); + try p.out(1, false); + return true; +} + +fn castLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1 }, @as(u32, @intCast(flat(p)))); + Probe.go(e, .{ k.blocks(flat(p), 256), 1, 1 }, 256); +} + +fn castEval(p: *Probe, _: k.Mode, out: []R) !void { + for (out, 0..) |*o, i| o.* = if (p.variant == 0) R.of(p.get(f32, 0)[i]).bf() else R.bf16(p.get(u16, 0)[i]); +} + +const cast_def: k.Def = .{ .prep = castPrep, .launch = castLaunch, .count = flat, .eval = castEval }; + +// --- om_fsq_decode: codes [n, 12] to latents [n, 48], levels 8 7 6 6; integer inputs, so a stated rule + +const FsqArgs = extern struct { n: i32, NQ: i32 }; + +fn fsqPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n != 12) return false; + const codes = try p.alloc(i32, 0, s.rows * 12); + _ = try p.alloc(u16, 1, s.rows * 48); + for (codes, 0..) |*x, i| x.* = if (i < 2016) @intCast(i) else @intCast(rng.uintLessThan(u32, 2016)); + try p.in(0); + try p.out(1, false); + return true; +} + +fn fsqLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1 }, FsqArgs{ .n = k.i32_(p.shape.rows), .NQ = 12 }); + Probe.go(e, .{ k.blocks(p.shape.rows * 12, 128), 1, 1 }, 128); +} + +fn fsqCount(p: *Probe) usize { + return p.shape.rows * 48; +} + +fn fsqEval(p: *Probe, _: k.Mode, out: []R) !void { + const levels = [4]i32{ 8, 7, 6, 6 }; + const basis = [4]i32{ 1, 8, 56, 336 }; + for (p.get(i32, 0), 0..) |code, i| for (0..4) |j| { + const hf = @divTrunc(levels[j], 2); + const digit = @mod(@divTrunc(code, basis[j]), levels[j]) - hf; + out[i * 4 + j] = R.of(@floatFromInt(digit)).div(R.of(@floatFromInt(hf))).bf(); + }; +} + +const fsq_def: k.Def = .{ .prep = fsqPrep, .launch = fsqLaunch, .count = fsqCount, .eval = fsqEval, .check = k.checkReplay }; + +// --- the codec's spectra: om_polar (rows of log-magnitude | phase to re | im, fp32) + +const PolarArgs = extern struct { T: i32, NB: i32 }; + +fn polarPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0) return false; + const x = try p.alloc(u16, 0, s.rows * 2 * s.n); + _ = try p.alloc(f32, 1, s.rows * 2 * s.n); + for (0..s.rows) |t| { + k.randomBf16(rng, x[t * 2 * s.n ..][0..s.n], -6, 3); + k.randomBf16(rng, x[t * 2 * s.n + s.n ..][0..s.n], -4, 4); + } + if (p.fill == .nonfinite) plant16(x[0 .. 2 * s.n * @min(s.rows, 4)]); + try p.in(0); + try p.out(1, false); + return true; +} + +fn polarLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1 }, PolarArgs{ .T = k.i32_(p.shape.rows), .NB = k.i32_(p.shape.n) }); + Probe.go(e, .{ k.blocks(flat(p), 256), 1, 1 }, 256); +} + +fn polarCount(p: *Probe) usize { + return 2 * flat(p); +} + +fn cosf(x: f64) f64 { + return @cos(x); +} +fn sinf(x: f64) f64 { + return @sin(x); +} + +fn polarEval(p: *Probe, _: k.Mode, out: []R) !void { + const nb = p.shape.n; + const x = p.get(u16, 0); + for (0..p.shape.rows) |t| for (0..nb) |i| { + const m = R.bf16(x[t * 2 * nb + i]).up(k.exp, k.fn_ulps); + const a = R.bf16(x[t * 2 * nb + nb + i]); + out[t * 2 * nb + i] = m.mul(a.lip(cosf, 1, k.fn_ulps)); + out[t * 2 * nb + nb + i] = m.mul(a.lip(sinf, 1, k.fn_ulps)); + }; +} + +const polar_def: k.Def = .{ .prep = polarPrep, .launch = polarLaunch, .count = polarCount, .eval = polarEval }; + +// --- the log-mel front end: om_mel_frames, om_power, om_logmel (fp32 throughout) + +const MelArgs = extern struct { len: i32, nf: i32, nfft: i32, hop: i32 }; + +/// Shape rows the frames, n the FFT's length, k the hop; the clip ends inside the last frame. +fn framesPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0 or s.k == 0) return false; + const len = (s.rows - 1) * s.k + s.n / 2; + const x = try p.alloc(f32, 0, len + s.n); + const win = try p.alloc(f32, 1, s.n); + _ = try p.alloc(f32, 2, s.rows * s.n); + for (x) |*v| v.* = rng.float(f32) * 2 - 1; + for (win, 0..) |*v, i| v.* = @floatCast(0.5 - 0.5 * @cos(2 * std.math.pi * @as(f64, @floatFromInt(i)) / @as(f64, @floatFromInt(s.n)))); + if (p.fill == .nonfinite) plant32(x[0..len]); + for (0..2) |i| try p.in(i); + try p.out(2, false); + return true; +} + +fn framesLen(p: *const Probe) usize { + return (p.shape.rows - 1) * p.shape.k + p.shape.n / 2; +} + +fn framesLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + p.bind(e, &.{ 0, 1, 2 }, MelArgs{ .len = k.i32_(framesLen(p)), .nf = k.i32_(s.rows), .nfft = k.i32_(s.n), .hop = k.i32_(s.k) }); + Probe.go(e, .{ k.blocks(s.rows * s.n, 256), 1, 1 }, 256); +} + +fn framesEval(p: *Probe, _: k.Mode, out: []R) !void { + const s = p.shape; + for (out, 0..) |*o, i| { + const at = (i / s.n) * s.k + i % s.n; + o.* = if (at < framesLen(p)) R.of(p.get(f32, 0)[at]).mul(R.of(p.get(f32, 1)[i % s.n])) else R.of(0); + } +} + +const frames_def: k.Def = .{ .prep = framesPrep, .launch = framesLaunch, .count = flat, .eval = framesEval }; + +const PowerArgs = extern struct { nf: i32, bins: i32 }; + +/// Spectra rows re | im; stress: integers below 2^10, so both squares and their sum are exact in fp32. +fn powerPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0) return false; + const sp = try p.alloc(f32, 0, s.rows * 2 * s.n); + _ = try p.alloc(f32, 1, s.rows * s.n); + for (sp) |*v| v.* = if (p.fill == .stress) @floatFromInt(@as(i32, @intCast(rng.uintLessThan(u32, 2047))) - 1023) else (rng.float(f32) * 2 - 1) * 30; + if (p.fill == .nonfinite) plant32(sp[0 .. 2 * s.n * @min(s.rows, 4)]); + try p.in(0); + try p.out(1, false); + return true; +} + +fn powerLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1 }, PowerArgs{ .nf = k.i32_(p.shape.rows), .bins = k.i32_(p.shape.n) }); + Probe.go(e, .{ k.blocks(flat(p), 256), 1, 1 }, 256); +} + +fn sqrt(x: f64) f64 { + return @sqrt(x); +} + +fn powerEval(p: *Probe, mode: k.Mode, out: []R) !void { + const nb = p.shape.n; + const sp = p.get(f32, 0); + var t: [2]R = undefined; + for (out, 0..) |*o, i| { + const re = R.of(sp[(i / nb) * 2 * nb + i % nb]); + const im = R.of(sp[(i / nb) * 2 * nb + nb + i % nb]); + t = .{ re.mulExact(re), im.mulExact(im) }; + const m = k.sum(mode, p.scratch, &t, 3).fp(0).up(sqrt, k.sqrt_ulps); + o.* = m.mul(m); + } +} + +fn depth3(_: c.Shape) usize { + return 3; +} + +const power_def: k.Def = .{ .prep = powerPrep, .launch = powerLaunch, .count = flat, .eval = powerEval }; + +const LogmelArgs = extern struct { nf: i32, mels: i32, floor_: f32 }; + +/// Mel rows [nf, mels] with zeros and values under 1e-12 among them (the floor's side). +fn logmelPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0) return false; + const m = try p.alloc(f32, 0, s.rows * s.n); + _ = try p.alloc(f32, 1, s.rows * s.n); + for (m, 0..) |*v, i| v.* = switch (i % 17) { + 3 => 0, + 5 => 1e-14, + else => std.math.exp(rng.float(f32) * 30 - 20), + }; + if (p.fill == .nonfinite) plant32(m); + try p.in(0); + try p.out(1, false); + return true; +} + +fn logmelLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1 }, LogmelArgs{ .nf = k.i32_(p.shape.rows), .mels = k.i32_(p.shape.n), .floor_ = -12 }); + Probe.go(e, .{ k.blocks(flat(p), 256), 1, 1 }, 256); +} + +fn log10(x: f64) f64 { + return std.math.log10(x); +} + +/// The floor as metal::max takes it: a NaN logarithm gives the floor. +fn logmelEval(p: *Probe, _: k.Mode, out: []R) !void { + const s = p.shape; + for (p.get(f32, 0), 0..) |m, i| { + var l = R.of(m).up(log10, k.fn_ulps); + if (std.math.isNan(l.v)) l = R.of(-12) else if (l.v < -12) l = R.of(-12) else if (l.lo < -12) l.lo = -12; + out[(i % s.n) * s.rows + i / s.n] = l.add(R.of(4)).div(R.of(4)); + } +} + +const logmel_def: k.Def = .{ .prep = logmelPrep, .launch = logmelLaunch, .count = flat, .eval = logmelEval, .check = k.checkReplay }; + +// --- word moves: om_set_words (bytes from the arguments), om_zero_words, om_copy_words; raw, checked byte for byte + +fn wordsPrep(p: *Probe, rng: std.Random) !bool { + const n = p.shape.n; + if (n == 0 or (p.variant == 0 and n > 1024)) return false; + const src = try p.alloc(u32, 0, n); + const dst = try p.alloc(u32, 1, n + 8); + for (src) |*x| x.* = rng.int(u32); + for (dst) |*x| x.* = rng.int(u32); + try p.in(0); + try p.out(1, true); + return true; +} + +fn wordsLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const n: u32 = @intCast(p.shape.n); + e.setPipeline(p.pipe); + switch (p.variant) { + 0 => { + e.setBuffer(p.refs[1].buf, p.refs[1].off, 0); + e.setBytes(p.host[0], 1); + e.setBytes(std.mem.asBytes(&n), 2); + }, + 1 => { + e.setBuffer(p.refs[1].buf, p.refs[1].off, 0); + e.setBytes(std.mem.asBytes(&n), 1); + }, + else => { + e.setBuffer(p.refs[0].buf, p.refs[0].off, 0); + e.setBuffer(p.refs[1].buf, p.refs[1].off, 1); + e.setBytes(std.mem.asBytes(&n), 2); + }, + } + Probe.go(e, .{ k.blocks(n, 256), 1, 1 }, 256); +} + +fn wordsCount(p: *Probe) usize { + return (p.shape.n + 8) * 4; +} + +fn wordsEval(_: *Probe, _: k.Mode, out: []R) !void { + @memset(out, R.of(0)); +} + +fn wordsCheck(p: *Probe, _: k.g.Gate, out: []const u8) !?[]const u8 { + const n = p.shape.n; + const want = try p.a.dupe(u32, p.get(u32, 1)); + defer p.a.free(want); + for (want[0..n], p.get(u32, 0)) |*w, s| w.* = if (p.variant == 1) 0 else s; + return k.sameBytes(out, std.mem.sliceAsBytes(want)); +} + +const words_def: k.Def = .{ .prep = wordsPrep, .launch = wordsLaunch, .count = wordsCount, .eval = wordsEval, .check = wordsCheck }; + +fn raw(name: []const u8, served: []const c.Shape) c.Spec { + return .{ .name = name, .accum = .none, .depth = k.noSums, .out = .raw, .in_place = true, .tiles = .{ .{}, .{ .size = 256 }, .{} }, .served = served, .nan = .{ .stated = "moves 32-bit words: every byte as it was" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "no sums" } }; +} + +pub const entries = [_]k.Entry{ + .{ .def = &swi_def, .spec = none("metal lfm3 omni om_swiglu halves", .bf16, &.{ .{ .rows = 94, .n = 1792, .k = 1 }, .{ .rows = 40, .n = 7680, .k = 1 } }), .file = .omni, .pipe = "om_swiglu", .variant = 0, .control = .{ "Y[i] = b(rb(x / (1.0f + precise::exp(-x))) * f(AC[ic]));", "Y[i] = b((x / (1.0f + precise::exp(-x))) * f(AC[ic]));" }, .what = "drops bf16(silu(a))" }, + .{ .def = &swi_def, .spec = none("metal lfm3 omni om_swiglu pairs", .bf16, &.{ .{ .rows = 40, .n = 7680, .k = 1 }, .{ .rows = 48, .n = 2304, .k = 1 } }), .file = .omni, .pipe = "om_swiglu", .variant = 1, .control = .{ "const long ic = a.inter ? ia + 1 : ia + a.F;", "const long ic = a.inter ? ia : ia + a.F;" }, .what = "reads a for c in the pairs layout" }, + .{ .def = &bias_def, .spec = none("metal lfm3 omni om_bias_act bias", .bf16, &.{ .{ .rows = 94, .n = 768, .k = 1 }, .{ .rows = 64, .n = 1282, .k = 1 } }), .file = .omni, .pipe = "om_bias_act", .variant = 0, .control = .{ " v = rb(v);\n if (a.act == 1)", " v = rb(v * 1.0078125f);\n if (a.act == 1)" }, .what = "scales the sum" }, + .{ .def = &bias_def, .spec = none("metal lfm3 omni om_bias_act row bias, SiLU", .bf16, &.{.{ .rows = 256, .n = 1632, .k = 1 }}), .file = .omni, .pipe = "om_bias_act", .variant = 1, .control = .{ "float v = A[i] + (a.has_bias ? f(Bias[a.per_row ? r : c]) : 0.0f);", "float v = A[i] + (a.has_bias ? f(Bias[a.per_row ? c % a.rows : c]) : 0.0f);" }, .what = "takes the row bias by column" }, + .{ .def = &bias_def, .spec = none("metal lfm3 omni om_bias_act GELU", .bf16, &.{.{ .rows = 94, .n = 2048, .k = 1 }}), .file = .omni, .pipe = "om_bias_act", .variant = 2, .control = .{ "else if (a.act == 2) v = 0.5f * v * (1.0f + erf_f(v * 0.70710678118654752f));", "else if (a.act == 2) v = 0.5f * v * (1.0f + rb(erf_f(v * 0.70710678118654752f)));" }, .what = "adds bf16(erf)" }, + .{ .def = &cast_def, .spec = none("metal lfm3 omni om_f32_bf16", .bf16, &.{ .{ .rows = 1, .n = 128 * 3000, .k = 1 }, .{ .rows = 1, .n = 1000, .k = 1 } }), .file = .omni, .pipe = "om_f32_bf16", .variant = 0, .control = .{ "if (gid < n) Y[gid] = b(X[gid]);", "if (gid < n) Y[gid] = bf(as_type(X[gid]) >> 16);" }, .what = "truncates instead of rounding" }, + .{ .def = &cast_def, .spec = none("metal lfm3 omni om_bf16_f32", .f32, &.{ .{ .rows = 1, .n = 12 * 2017, .k = 1 }, .{ .rows = 1, .n = 1000, .k = 1 } }), .file = .omni, .pipe = "om_bf16_f32", .variant = 1, .control = .{ "if (gid < n) Y[gid] = f(X[gid]);", "if (gid < n) Y[gid] = f(X[gid ^ 1]);" }, .what = "swaps neighbours" }, + .{ .def = &fsq_def, .spec = .{ .name = "metal lfm3 omni om_fsq_decode", .accum = .none, .depth = k.noSums, .out = .bf16, .served = &.{ .{ .rows = 200, .n = 12, .k = 1 }, .{ .rows = 1, .n = 12, .k = 1 } }, .nan = .{ .stated = "integer codes in: no non-finite input; every latent inside its replay's range" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "no sums" } }, .file = .omni, .pipe = "om_fsq_decode", .control = .{ "const int levels[4] = {8, 7, 6, 6};", "const int levels[4] = {8, 7, 6, 5};" }, .what = "a wrong level" }, + .{ .def = &polar_def, .spec = none("metal lfm3 omni om_polar", .f32, &.{ .{ .rows = 48, .n = 641, .k = 1 }, .{ .rows = 6, .n = 641, .k = 1 } }), .file = .omni, .pipe = "om_polar", .control = .{ "S[t * 2 * a.NB + k] = m * precise::cos(ph);", "S[t * 2 * a.NB + k] = rb(m) * precise::cos(ph);" }, .what = "adds bf16(magnitude)" }, + .{ .def = &frames_def, .spec = none("metal lfm3 omni om_mel_frames", .f32, &.{.{ .rows = 300, .n = 400, .k = 160 }}), .file = .gemm, .pipe = "om_mel_frames", .control = .{ "FR[i] = s < a.len ? X[s] * WIN[n] : 0.0f;", "FR[i] = s < a.len ? X[s] * WIN[n ^ 1] : 0.0f;" }, .what = "swaps window taps" }, + .{ .def = &power_def, .spec = .{ .name = "metal lfm3 omni om_power", .accum = .f32, .depth = depth3, .ops = 6, .out = .f32, .served = &.{.{ .rows = 300, .n = 201, .k = 2 }} }, .file = .gemm, .pipe = "om_power", .control = .{ "P[i] = m * m;", "P[i] = rb(m) * m;" }, .what = "adds bf16(magnitude)" }, + .{ .def = &logmel_def, .spec = .{ .name = "metal lfm3 omni om_logmel", .accum = .none, .depth = k.noSums, .out = .f32, .served = &.{.{ .rows = 300, .n = 128, .k = 1 }}, .nan = .{ .stated = "metal::max takes the floor for a NaN logarithm; infinities as float64" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "no sums" } }, .file = .gemm, .pipe = "om_logmel", .control = .{ "(max(precise::log10(M[i]), a.floor_) + 4.0f) / 4.0f", "(precise::log10(M[i]) + 4.0f) / 4.0f" }, .what = "drops the floor" }, + .{ .def = &words_def, .spec = raw("metal lfm3 omni om_set_words", &.{ .{ .rows = 1, .n = 36, .k = 1 }, .{ .rows = 1, .n = 1024, .k = 1 } }), .file = .omni, .pipe = "om_set_words", .variant = 0, .control = .{ "if (gid < n) D[gid] = V[gid];", "if (gid < n) D[gid] = V[gid / 2];" }, .what = "reads every other word" }, + .{ .def = &words_def, .spec = raw("metal lfm3 omni om_zero_words", &.{ .{ .rows = 1, .n = 36, .k = 1 }, .{ .rows = 1, .n = 3000, .k = 1 } }), .file = .omni, .pipe = "om_zero_words", .variant = 1, .control = .{ "if (gid < n) D[gid] = 0;", "if (gid + 1 < n) D[gid] = 0;" }, .what = "leaves the last word" }, + .{ .def = &words_def, .spec = raw("metal lfm3 omni om_copy_words", &.{ .{ .rows = 1, .n = 36, .k = 1 }, .{ .rows = 1, .n = 3000, .k = 1 } }), .file = .omni, .pipe = "om_copy_words", .variant = 2, .control = .{ "if (gid < n) D[gid] = S[gid];", "if (gid < n) D[gid] = S[gid] ^ (gid == n - 1 ? 1u : 0u);" }, .what = "flips a bit of the last word" }, +}; diff --git a/zig/tests/contract/metal_omni_gemm.zig b/zig/tests/contract/metal_omni_gemm.zig new file mode 100644 index 0000000000..72f0b4f0c2 --- /dev/null +++ b/zig/tests/contract/metal_omni_gemm.zig @@ -0,0 +1,291 @@ +//! LFM3-Omni's Metal products: the few-row om_rows_R, the simdgroup om_gemm_*, and the tensor-unit om_nax_*. +const std = @import("std"); +const mtl = @import("metal"); +const k = @import("metal_omni_common.zig"); + +const c = k.c; +const R = k.R; +const Probe = k.Probe; + +/// Dot rows for gate 4: integers to 255 times to 31 (partials under 2^23), the second half mirrored with a nudge. +fn stressPair(rng: std.Random, x: []u16, w: []u16) void { + const h = x.len / 2; + const top: u32 = @intCast(@min(31, (1 << 23) / (@max(h, 1) * 255))); + for (0..h) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 255)); + const u: f32 = @floatFromInt(1 + rng.uintLessThan(u32, top)); + x[i] = c.toBf16(v); + x[h + i] = c.toBf16(v + @as(f32, if (i < 8) 1 else 0)); + w[i] = c.toBf16(u); + w[h + i] = c.toBf16(-u); + } + if (x.len % 2 == 1) { + x[x.len - 1] = 0; + w[x.len - 1] = 0; + } +} + +fn plantRows(x: []u16, rows: usize, d: usize) void { + const at = [_][2]usize{ .{ 0, d / 3 }, .{ 1, d / 2 }, .{ 2, d - 1 }, .{ 3, 0 }, .{ 3, d - 1 } }; + const put = [_]u16{ k.nan16, k.inf16, k.ninf16, k.inf16, k.ninf16 }; + for (at, put) |p, v| if (p[0] < rows) { + x[p[0] * d + p[1]] = v; + }; +} + +// --- om_rows_R: variant R | epilogue << 8 (0 bf16 y, 1 fp32 yf, 2 SwiGLU pairs) | batched << 12 (12 products) + +const RowsArgs = extern struct { K: i32, N: i32, T: i32, ldx: i32, ldy: i32, has_y: i32, has_yf: i32, epi: i32, sw: i64, sx: i64, sy: i64 }; + +fn rowsR(p: *const Probe) usize { + return p.variant & 0xff; +} +fn epi(p: *const Probe) u32 { + return (p.variant >> 8) & 0xf; +} +fn batch(p: *const Probe) usize { + return if (p.variant & 0x1000 != 0) 12 else 1; +} +fn yCols(p: *const Probe) usize { + return if (epi(p) == 2) p.shape.n / 2 else p.shape.n; +} + +/// Shape rows T, n the weights' rows, k their width; slots w, x, y. +fn rowsPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.rows == 0 or s.k == 0 or s.k % 8 != 0 or s.n == 0 or (epi(p) == 2 and s.n % 4 != 0)) return false; + const nb = batch(p); + const w = try p.alloc(u16, 0, nb * s.n * s.k); + const x = try p.alloc(u16, 1, nb * s.rows * s.k); + if (epi(p) == 1) _ = try p.alloc(f32, 2, nb * s.rows * s.n) else _ = try p.alloc(u16, 2, nb * s.rows * yCols(p)); + if (p.fill == .stress) { + for (0..nb * s.n) |n| for (0..nb * s.rows) |t| if (t / s.rows == n / s.n) { + var prng = std.Random.DefaultPrng.init(n *% 31 +% t); + stressPair(prng.random(), x[t * s.k ..][0..s.k], w[n * s.k ..][0..s.k]); + }; + } else { + k.randomBf16(rng, w, -1, 1); + k.randomBf16(rng, x, -1, 1); + if (p.fill == .nonfinite) plantRows(x, s.rows, s.k); + } + try p.in(0); + try p.in(1); + try p.out(2, false); + return true; +} + +fn rowsLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const yf = epi(p) == 1; + const ldy = if (yf) s.n else yCols(p); + // more rows than the pipeline takes: a launch per R rows, as successive calls + var r0: usize = 0; + while (r0 < s.rows) : (r0 += rowsR(p)) { + const x = Probe.at(p.refs[1], r0 * s.k * 2); + const y = Probe.at(p.refs[2], r0 * ldy * @as(usize, if (yf) 4 else 2)); + e.setPipeline(p.pipe); + e.setBuffer(p.refs[0].buf, p.refs[0].off, 0); + e.setBuffer(x.buf, x.off, 1); + e.setBuffer(if (yf) p.dummy.buf else y.buf, if (yf) p.dummy.off else y.off, 2); + e.setBuffer(if (yf) y.buf else p.dummy.buf, if (yf) y.off else p.dummy.off, 3); + const a = RowsArgs{ .K = k.i32_(s.k), .N = k.i32_(s.n), .T = k.i32_(@min(rowsR(p), s.rows - r0)), .ldx = k.i32_(s.k), .ldy = k.i32_(yCols(p)), .has_y = @intFromBool(!yf), .has_yf = @intFromBool(yf), .epi = if (epi(p) == 2) 2 else 0, .sw = @intCast(s.n * s.k), .sx = @intCast(s.rows * s.k), .sy = @intCast(s.rows * (if (yf) s.n else yCols(p))) }; + e.setBytes(std.mem.asBytes(&a), 4); + Probe.go(e, .{ k.blocks(s.n, 16), 1, batch(p) }, 128); + } +} + +fn rowsCount(p: *Probe) usize { + return batch(p) * p.shape.rows * yCols(p); +} + +fn dot(mode: k.Mode, p: *Probe, x: []const u16, w: []const u16, t: []R, depth: usize) R { + for (t, x, w) |*q, a, b| q.* = R.bf16(a).mulExact(R.bf16(b)); + return k.sum(mode, p.scratch, t, depth); +} + +fn rowsEval(p: *Probe, mode: k.Mode, out: []R) !void { + const s = p.shape; + const t = try p.a.alloc(R, s.k); + defer p.a.free(t); + const x = p.get(u16, 1); + const w = p.get(u16, 0); + const depth = s.k / 32 + 6; + for (0..batch(p)) |z| for (0..s.rows) |r| { + const xr = x[(z * s.rows + r) * s.k ..][0..s.k]; + const at = (z * s.rows + r) * yCols(p); + if (epi(p) == 2) { + for (0..s.n / 2) |j| { + const a = dot(mode, p, xr, w[(z * s.n + 2 * j) * s.k ..][0..s.k], t, depth).bf(); + const g = dot(mode, p, xr, w[(z * s.n + 2 * j + 1) * s.k ..][0..s.k], t, depth).bf(); + out[at + j] = k.siluR(a).bf().mul(g).bf(); + } + } else for (0..s.n) |n| { + const v = dot(mode, p, xr, w[(z * s.n + n) * s.k ..][0..s.k], t, depth); + out[at + n] = if (epi(p) == 1) v else v.bf(); + } + }; +} + +const rows_def: k.Def = .{ .prep = rowsPrep, .launch = rowsLaunch, .count = rowsCount, .eval = rowsEval }; + +fn rowsDepth(s: c.Shape) usize { + return s.k / 32 + 6; +} + +// --- om_gemm_* and om_nax_*: C [m, n] = A B; variant: ta | tb << 1 | in fp32 << 2 | out fp32 << 3 | tensor units << 4 + +const GemmArgs = extern struct { M: i32, N: i32, K: i32, lda: i32, ldb: i32, ldc: i32, ta: i32, tb: i32, sa: i64, sb: i64, sc: i64 }; +const NaxArgs = extern struct { M: i32, N: i32, K: i32, lda: i32, ldb: i32, ldd: i32, kblocks: i32, pad: i32 = 0, sa: i64, sb: i64, sd: i64 }; + +fn ta(p: *const Probe) bool { + return p.variant & 1 != 0; +} +fn tb(p: *const Probe) bool { + return p.variant & 2 != 0; +} +fn f32in(p: *const Probe) bool { + return p.variant & 4 != 0; +} +fn f32out(p: *const Probe) bool { + return p.variant & 8 != 0; +} +fn nax(p: *const Probe) bool { + return p.variant & 16 != 0; +} + +/// A [m, k] (stored [k, m] when ta) at (i, j), B [k, n] (stored [n, k] when tb) at (i, j): their places. +fn aAt(p: *const Probe, i: usize, j: usize) usize { + return if (ta(p)) j * p.shape.rows + i else i * p.shape.k + j; +} +fn bAt(p: *const Probe, i: usize, j: usize) usize { + return if (tb(p)) j * p.shape.k + i else i * p.shape.n + j; +} + +/// Shape rows m, n n, k the depth; slots a, b, c (bf16 or fp32 each as the variant says). +fn gemmPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.rows == 0 or s.n == 0 or s.k == 0) return false; + const ab = [2]usize{ s.rows * s.k, s.k * s.n }; + if (f32in(p)) { + const a = try p.alloc(f32, 0, ab[0]); + const b = try p.alloc(f32, 1, ab[1]); + for (a) |*v| v.* = rng.float(f32) * 2 - 1; + for (b) |*v| v.* = rng.float(f32) * 2 - 1; + if (p.fill == .stress) { + const xs = try p.a.alloc(u16, s.k); + defer p.a.free(xs); + const ws = try p.a.alloc(u16, s.k); + defer p.a.free(ws); + for (0..s.rows) |i| for (0..s.n) |j| if (i == j % s.rows) { + stressPair(rng, xs, ws); + for (0..s.k) |q| { + a[aAt(p, i, q)] = c.fromBf16(xs[q]); + b[bAt(p, q, j)] = c.fromBf16(ws[q]); + } + }; + } + if (p.fill == .nonfinite) for ([_]usize{ 0, 1, 2 }, [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }) |i, v| if (i < s.rows) { + a[aAt(p, i, s.k / 2)] = v; + }; + } else { + const a = try p.alloc(u16, 0, ab[0]); + const b = try p.alloc(u16, 1, ab[1]); + k.randomBf16(rng, a, -1, 1); + k.randomBf16(rng, b, -1, 1); + if (p.fill == .stress) { + const xs = try p.a.alloc(u16, s.k); + defer p.a.free(xs); + const ws = try p.a.alloc(u16, s.k); + defer p.a.free(ws); + for (0..s.rows) |i| for (0..s.n) |j| if (i == j % s.rows) { + stressPair(rng, xs, ws); + for (0..s.k) |q| { + a[aAt(p, i, q)] = xs[q]; + b[bAt(p, q, j)] = ws[q]; + } + }; + } + if (p.fill == .nonfinite) for ([_]usize{ 0, 1, 2 }, [_]u16{ k.nan16, k.inf16, k.ninf16 }) |i, v| if (i < s.rows) { + a[aAt(p, i, s.k / 2)] = v; + }; + } + if (f32out(p)) _ = try p.alloc(f32, 2, s.rows * s.n) else _ = try p.alloc(u16, 2, s.rows * s.n); + try p.in(0); + try p.in(1); + try p.out(2, false); + return true; +} + +fn gemmLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const lda = if (ta(p)) s.rows else s.k; + const ldb = if (tb(p)) s.k else s.n; + if (nax(p)) { + p.bind(e, &.{ 0, 1, 2 }, NaxArgs{ .M = k.i32_(s.rows), .N = k.i32_(s.n), .K = k.i32_(s.k), .lda = k.i32_(lda), .ldb = k.i32_(ldb), .ldd = k.i32_(s.n), .kblocks = k.i32_(s.k / 256), .sa = 0, .sb = 0, .sd = 0 }); + return Probe.go(e, .{ k.blocks(s.n, 128), k.blocks(s.rows, 64), 1 }, 256); + } + p.bind(e, &.{ 0, 1, 2 }, GemmArgs{ .M = k.i32_(s.rows), .N = k.i32_(s.n), .K = k.i32_(s.k), .lda = k.i32_(lda), .ldb = k.i32_(ldb), .ldc = k.i32_(s.n), .ta = @intFromBool(ta(p)), .tb = @intFromBool(tb(p)), .sa = 0, .sb = 0, .sc = 0 }); + Probe.go(e, .{ k.blocks(s.n, 64), k.blocks(s.rows, 64), 1 }, 128); +} + +fn gemmCount(p: *Probe) usize { + return p.shape.rows * p.shape.n; +} + +fn gemmEval(p: *Probe, mode: k.Mode, out: []R) !void { + const s = p.shape; + const t = try p.a.alloc(R, s.k); + defer p.a.free(t); + for (0..s.rows) |i| for (0..s.n) |j| { + for (t, 0..) |*q, x| { + const a = if (f32in(p)) R.of(p.get(f32, 0)[aAt(p, i, x)]) else R.bf16(p.get(u16, 0)[aAt(p, i, x)]); + const b = if (f32in(p)) R.of(p.get(f32, 1)[bAt(p, x, j)]) else R.bf16(p.get(u16, 1)[bAt(p, x, j)]); + q.* = a.mulExact(b); + } + const v = k.sum(mode, p.scratch, t, s.k + 40); + out[i * s.n + j] = if (f32out(p)) v else v.bf(); + }; +} + +const gemm_def: k.Def = .{ .prep = gemmPrep, .launch = gemmLaunch, .count = gemmCount, .eval = gemmEval }; + +fn gemmDepth(s: c.Shape) usize { + return s.k + 40; +} + +fn rowsSpec(name: []const u8, out: c.Elem, served: []const c.Shape) c.Spec { + return .{ .name = name, .accum = .f32, .depth = rowsDepth, .ops = 8, .out = out, .tiles = .{ .{}, .{ .size = 16 }, .{ .size = 256, .step = 8, .min = 8 } }, .served = served }; +} + +fn gemmSpec(name: []const u8, out: c.Elem, served: []const c.Shape, tk: c.Tile) c.Spec { + return .{ .name = name, .accum = .f32, .depth = gemmDepth, .out = out, .tiles = .{ .{ .size = 64 }, .{ .size = 64 }, tk }, .served = served }; +} + +const loop = "for (int c = lane; c < K / 8; c += 32) {"; +const drop = "for (int c = lane; c < K / 8 - 1; c += 32) {"; +const kloop = "for (int k0 = 0; k0 < a.K; k0 += BK) {"; +const ktrim = "for (int k0 = 0; k0 + BK <= a.K; k0 += BK) {"; +const nrows = "const int rows = min(32, a.M - row), cols = min(32, a.N - col);"; +const nrows_ctl = "const int rows = min(31, a.M - row), cols = min(32, a.N - col);"; +const k16: c.Tile = .{ .size = 16 }; +const k256: c.Tile = .{ .size = 256, .step = 8, .min = 8 }; + +pub const entries = [_]k.Entry{ + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_1", .bf16, &.{ .{ .rows = 1, .n = 2048, .k = 2048 }, .{ .rows = 1, .n = 3072, .k = 2048 } }), .file = .rows, .pipe = "om_rows_1", .variant = 1, .control = .{ loop, drop }, .what = "drops each row's last 8 columns" }, + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_4", .bf16, &.{.{ .rows = 4, .n = 1536, .k = 512 }}), .file = .rows, .pipe = "om_rows_4", .variant = 4, .control = .{ loop, drop }, .what = "drops each row's last 8 columns" }, + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_8", .bf16, &.{.{ .rows = 6, .n = 1536, .k = 512 }}), .file = .rows, .pipe = "om_rows_8", .variant = 8, .control = .{ loop, drop }, .what = "drops each row's last 8 columns" }, + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_12", .bf16, &.{.{ .rows = 12, .n = 1024, .k = 1024 }}), .file = .rows, .pipe = "om_rows_12", .variant = 12, .control = .{ loop, drop }, .what = "drops each row's last 8 columns" }, + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_16", .bf16, &.{.{ .rows = 16, .n = 2048, .k = 2048 }}), .file = .rows, .pipe = "om_rows_16", .variant = 16, .control = .{ loop, drop }, .what = "drops each row's last 8 columns" }, + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_1 fp32 out", .f32, &.{.{ .rows = 1, .n = 12288, .k = 2048 }}), .file = .rows, .pipe = "om_rows_1", .variant = 1 | 1 << 8, .control = .{ "if (a.has_yf) YF[(long)r * a.N + n0 + j] = v[j];", "if (a.has_yf) YF[(long)r * a.N + n0 + j] = rb(v[j]);" }, .what = "adds bf16 to the fp32 output" }, + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_12 batched (depth logits)", .bf16, &.{.{ .rows = 1, .n = 2017, .k = 1024 }}), .file = .rows, .pipe = "om_rows_12", .variant = 12 | 0x1000, .control = .{ "X += (long)tg.z * a.sx;", "X += (long)(tg.z / 2) * a.sx;" }, .what = "reads every other batch member's rows" }, + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_1 SwiGLU pairs", .bf16, &.{.{ .rows = 1, .n = 15360, .k = 2048 }}), .file = .rows, .pipe = "om_rows_1", .variant = 1 | 2 << 8, .control = .{ "const float x = rb(v[2 * p]), c = rb(v[2 * p + 1]);", "const float x = rb(v[2 * p]), c = (v[2 * p + 1]);" }, .what = "drops bf16(c)" }, + .{ .def = &rows_def, .spec = rowsSpec("metal lfm3 omni om_rows_12 SwiGLU pairs", .bf16, &.{.{ .rows = 12, .n = 5632, .k = 1024 }}), .file = .rows, .pipe = "om_rows_12", .variant = 12 | 2 << 8, .control = .{ "Y[(long)r * a.ldy + n0 / 2 + p] = b(rb(x / (1.0f + precise::exp(-x))) * c);", "Y[(long)r * a.ldy + n0 / 2 + p] = b((x / (1.0f + precise::exp(-x))) * c);" }, .what = "drops bf16(silu(a))" }, + .{ .def = &gemm_def, .spec = gemmSpec("metal lfm3 omni om_gemm_bb a b^T", .bf16, &.{ .{ .rows = 40, .n = 768, .k = 768 }, .{ .rows = 24, .n = 1792, .k = 768 } }, k16), .file = .gemm, .pipe = "om_gemm_bb", .variant = 2, .control = .{ kloop, ktrim }, .what = "drops the last partial K tile" }, + .{ .def = &gemm_def, .spec = gemmSpec("metal lfm3 omni om_gemm_bf a b", .f32, &.{.{ .rows = 256, .n = 400, .k = 256 }}, k16), .file = .gemm, .pipe = "om_gemm_bf", .variant = 8, .control = .{ kloop, ktrim }, .what = "drops the last partial K tile" }, + .{ .def = &gemm_def, .spec = gemmSpec("metal lfm3 omni om_gemm_bf a^T b^T", .f32, &.{.{ .rows = 24, .n = 768, .k = 4096 }}, k16), .file = .gemm, .pipe = "om_gemm_bf", .variant = 1 | 2 | 8, .control = .{ "v = ld_in(A, a.ta ? (long)gk * a.lda + gm : (long)gm * a.lda + gk);", "v = ld_in(A, a.ta ? (long)gk * a.lda + (gm ^ 1) : (long)gm * a.lda + gk);" }, .what = "reads A's neighbouring column when transposed" }, + .{ .def = &gemm_def, .spec = gemmSpec("metal lfm3 omni om_gemm_ff a b (log-mel)", .f32, &.{ .{ .rows = 300, .n = 402, .k = 400 }, .{ .rows = 300, .n = 128, .k = 201 } }, k16), .file = .gemm, .pipe = "om_gemm_ff", .variant = 4 | 8, .control = .{ "As[r * BK + k] = v;", "As[r * BK + k] = rb(v);" }, .what = "adds bf16 to the fp32 operand" }, + .{ .def = &gemm_def, .spec = gemmSpec("metal lfm3 omni om_nax_bb_nt", .bf16, &.{ .{ .rows = 64, .n = 512, .k = 2048 }, .{ .rows = 100, .n = 256, .k = 768 } }, k256), .file = .nax, .pipe = "om_nax_bb_nt", .variant = 16 | 2, .control = .{ nrows, nrows_ctl }, .what = "drops a block row" }, + .{ .def = &gemm_def, .spec = gemmSpec("metal lfm3 omni om_nax_bf_nt", .f32, &.{.{ .rows = 100, .n = 512, .k = 768 }}, k256), .file = .nax, .pipe = "om_nax_bf_nt", .variant = 16 | 2 | 8, .control = .{ nrows, nrows_ctl }, .what = "drops a block row" }, + .{ .def = &gemm_def, .spec = gemmSpec("metal lfm3 omni om_nax_bf_nn", .f32, &.{.{ .rows = 256, .n = 400, .k = 256 }}, k256), .file = .nax, .pipe = "om_nax_bf_nn", .variant = 16 | 8, .control = .{ nrows, nrows_ctl }, .what = "drops a block row" }, + .{ .def = &gemm_def, .spec = gemmSpec("metal lfm3 omni om_nax_bf_tt", .f32, &.{ .{ .rows = 24, .n = 768, .k = 4096 }, .{ .rows = 64, .n = 256, .k = 2048 } }, k256), .file = .nax, .pipe = "om_nax_bf_tt", .variant = 16 | 1 | 2 | 8, .control = .{ nrows, nrows_ctl }, .what = "drops a block row" }, +}; diff --git a/zig/tests/contract/metal_omni_norms.zig b/zig/tests/contract/metal_omni_norms.zig new file mode 100644 index 0000000000..90a936770c --- /dev/null +++ b/zig/tests/contract/metal_omni_norms.zig @@ -0,0 +1,432 @@ +//! LFM3-Omni's Metal norms and rotary (om_rmsnorm, om_rmsnorm_1k, om_head_norm, om_layernorm, om_qk_prep, om_rope). +const std = @import("std"); +const mtl = @import("metal"); +const k = @import("metal_omni_common.zig"); + +const c = k.c; +const R = k.R; +const Probe = k.Probe; +const eps: f64 = @as(f32, 1e-5); +const sq_why = "a bf16 output follows the sum of squares as ss^-1/2: a bf16 sum's error is a few output ulps"; + +fn eps32() f32 { + return 1e-5; +} + +/// x as given, then the nonfinite fill's NaN, +Inf, -Inf and both infinities in rows 0 to 3. +fn plant(x: []u16, rows: usize, d: usize) void { + const at = [_][2]usize{ .{ 0, d / 3 }, .{ 1, d / 2 }, .{ 2, d - 1 }, .{ 3, 0 }, .{ 3, d - 1 } }; + const put = [_]u16{ k.nan16, k.inf16, k.ninf16, k.inf16, k.ninf16 }; + for (at, put) |p, v| if (p[0] < rows) { + x[p[0] * d + p[1]] = v; + }; +} + +/// Row values: uniform in [-1, 1] with a larger one every 64, or (stress) integers in [-32, 32]. +fn values(rng: std.Random, x: []u16, fill: k.g.Fill) void { + if (fill == .stress) return k.intsBf16(rng, x, 32); + for (x, 0..) |*v, i| v.* = c.toBf16((rng.float(f32) * 2 - 1) * (if (i % 64 == 7) @as(f32, 20) else 1)); +} + +/// rsqrt(ss / n + eps) as the kernels step it. +fn invNorm(ss: R, n: usize) R { + return ss.div(R.of(@floatFromInt(n))).add(R.of(eps)).rsqrt(); +} + +// --- om_rmsnorm and om_rmsnorm_1k: variant 0 plain, 1 the residual and its post norm, 2 a weight row a row (mode 1) + +const NormArgs = extern struct { D: i32, ld: i32, ldy: i32, eps: f32, mode: i32, ldw: i32, has_r: i32, has_pw: i32 }; + +fn wide(p: *const Probe) bool { + return p.variant & 0x10 != 0; +} + +fn normPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n != s.k or s.k == 0 or (wide(p) and s.k < 1024)) return false; + const v = p.variant & 0xf; + const d = s.k; + const x = try p.alloc(u16, 0, s.rows * d); + const w = try p.alloc(u16, 1, if (v == 2) s.rows * d else d); + const r = try p.alloc(u16, 2, if (v == 1) s.rows * d else 1); + const pw = try p.alloc(u16, 3, if (v == 1) d else 1); + _ = try p.alloc(u16, 4, s.rows * d); + values(rng, x, p.fill); + k.randomBf16(rng, w, 0.5, 1.5); + if (v == 1) { + values(rng, r, p.fill); + if (p.fill == .stress) @memset(pw, 0) else k.randomBf16(rng, pw, 0.5, 1.5); + } + if (p.fill == .nonfinite) plant(x, s.rows, d); + if (v == 1) try p.out(0, true) else try p.in(0); + for (1..4) |i| try p.in(i); + try p.out(4, false); + return true; +} + +fn normLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const v = p.variant & 0xf; + const d = p.shape.k; + p.bind(e, &.{ 0, if (v == 1) 2 else null, 1, 4, if (v == 1) 3 else null }, NormArgs{ .D = k.i32_(d), .ld = k.i32_(d), .ldy = k.i32_(d), .eps = eps32(), .mode = if (v == 2) 1 else 0, .ldw = if (v == 2) k.i32_(d) else 0, .has_r = @intFromBool(v == 1), .has_pw = @intFromBool(v == 1) }); + Probe.go(e, .{ p.shape.rows, 1, 1 }, if (wide(p)) 1024 else 256); +} + +fn normCount(p: *Probe) usize { + return p.shape.rows * p.shape.k * @as(usize, if (p.variant & 0xf == 1) 2 else 1); +} + +fn normEval(p: *Probe, mode: k.Mode, out: []R) !void { + const v = p.variant & 0xf; + const d = p.shape.k; + const nt: usize = if (wide(p)) 1024 else 256; + const depth = (d + nt - 1) / nt + 5 + nt / 32 + 1; + const xs = try p.a.alloc(R, 2 * d); + defer p.a.free(xs); + const t = xs[d..]; + const x = p.get(u16, 0); + const w = p.get(u16, 1); + const r = p.get(u16, 2); + const pw = p.get(u16, 3); + const ys = if (v == 1) out[p.shape.rows * d ..] else out; + for (0..p.shape.rows) |row| { + var inv_r = R.of(0); + if (v == 1) { + for (t, 0..) |*q, i| q.* = R.bf16(r[row * d + i]).mulExact(R.bf16(r[row * d + i])); + inv_r = invNorm(k.sum(mode, p.scratch, t, depth), d); + } + for (xs[0..d], 0..) |*q, i| { + q.* = R.bf16(x[row * d + i]); + if (v == 1) { + const add = R.bf16(pw[i]).mul(R.bf16(r[row * d + i]).mul(inv_r).bf()).bf(); + q.* = q.add(add).bf(); + out[row * d + i] = q.*; + } + t[i] = q.mulExact(q.*); + } + const inv = invNorm(k.sum(mode, p.scratch, t, depth), d); + for (xs[0..d], 0..) |q, i| { + const wi = R.bf16(if (v == 2) w[row * d + i] else w[i]); + ys[row * d + i] = if (v == 2) q.mul(inv).mul(wi).bf() else wi.mul(q.mul(inv).bf()).bf(); + } + } +} + +const norm_def: k.Def = .{ .prep = normPrep, .launch = normLaunch, .count = normCount, .eval = normEval }; + +fn depth256(s: c.Shape) usize { + return (s.k + 255) / 256 + 14; +} +fn depth1k(s: c.Shape) usize { + return (s.k + 1023) / 1024 + 38; +} + +fn normSpec(name: []const u8, is_wide: bool, served: []const c.Shape) c.Spec { + return .{ .name = name, .accum = .f32, .depth = if (is_wide) depth1k else depth256, .ops = 8, .mids = 1, .out = .bf16, .in_place = true, .tiles = .{ .{}, .{}, if (is_wide) .{ .size = 1024, .min = 1024 } else .{ .size = 256 } }, .served = served, .sums = .{ .bf16 = 0, .why = sq_why } }; +} + +// --- om_head_norm: rows of T x H heads of hd, in place; shape rows T, n the row's width (H hd), k hd + +const HeadArgs = extern struct { T: i32, H: i32, hd: i32, ld: i32, eps: f32 }; + +fn headPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.k == 0 or s.n % s.k != 0 or s.k > 128) return false; + const x = try p.alloc(u16, 0, s.rows * s.n); + const w = try p.alloc(u16, 1, s.k); + values(rng, x, p.fill); + k.randomBf16(rng, w, 0.5, 1.5); + if (p.fill == .nonfinite) plant(x, s.rows, s.n); + try p.out(0, true); + try p.in(1); + return true; +} + +fn headLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const h = s.n / s.k; + p.bind(e, &.{ 0, 1 }, HeadArgs{ .T = k.i32_(s.rows), .H = k.i32_(h), .hd = k.i32_(s.k), .ld = k.i32_(s.n), .eps = eps32() }); + Probe.go(e, .{ k.blocks(s.rows * h * 32, 256), 1, 1 }, 256); +} + +fn headCount(p: *Probe) usize { + return p.shape.rows * p.shape.n; +} + +/// One head's RMS norm as om_head_norm and om_qk_prep round it: bf16(w bf16(x inv)). +fn headRow(mode: k.Mode, scratch: []f64, x: []const u16, w: []const u16, t: []R, out: []R) void { + const hd = x.len; + for (t, x) |*q, xv| q.* = R.bf16(xv).mulExact(R.bf16(xv)); + const inv = invNorm(k.sum(mode, scratch, t, (hd + 31) / 32 + 6), hd); + for (out, x, w) |*o, xv, wv| o.* = R.bf16(wv).mul(R.bf16(xv).mul(inv).bf()).bf(); +} + +fn headEval(p: *Probe, mode: k.Mode, out: []R) !void { + const s = p.shape; + const t = try p.a.alloc(R, s.k); + defer p.a.free(t); + const x = p.get(u16, 0); + for (0..s.rows * (s.n / s.k)) |hh| headRow(mode, p.scratch, x[hh * s.k ..][0..s.k], p.get(u16, 1), t, out[hh * s.k ..][0..s.k]); +} + +fn headDepth(s: c.Shape) usize { + return (s.k + 31) / 32 + 6; +} + +const head_def: k.Def = .{ .prep = headPrep, .launch = headLaunch, .count = headCount, .eval = headEval }; + +// --- om_layernorm: the projector's LayerNorm with bias, one rounding + +const LnArgs = extern struct { D: i32, eps: f32 }; + +fn lnPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n != s.k or s.k == 0) return false; + const x = try p.alloc(u16, 0, s.rows * s.k); + const w = try p.alloc(u16, 1, s.k); + const bias = try p.alloc(u16, 2, s.k); + _ = try p.alloc(u16, 3, s.rows * s.k); + if (p.fill == .stress) { + // integers whose row sums are zero: the mean and every deviation exact + const h = s.k / 2; + for (0..s.rows) |row| { + k.intsBf16(rng, x[row * s.k ..][0..h], 64); + for (0..h) |i| x[row * s.k + h + i] = x[row * s.k + i] ^ 0x8000; + if (s.k % 2 == 1) x[row * s.k + s.k - 1] = 0; + } + } else values(rng, x, p.fill); + k.randomBf16(rng, w, 0.5, 1.5); + k.randomBf16(rng, bias, -0.5, 0.5); + if (p.fill == .nonfinite) plant(x, s.rows, s.k); + for (0..3) |i| try p.in(i); + try p.out(3, false); + return true; +} + +fn lnLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1, 2, 3 }, LnArgs{ .D = k.i32_(p.shape.k), .eps = eps32() }); + Probe.go(e, .{ p.shape.rows, 1, 1 }, 256); +} + +fn lnCount(p: *Probe) usize { + return p.shape.rows * p.shape.k; +} + +fn lnEval(p: *Probe, mode: k.Mode, out: []R) !void { + const d = p.shape.k; + const t = try p.a.alloc(R, 2 * d); + defer p.a.free(t); + const dev = t[d..]; + const x = p.get(u16, 0); + const w = p.get(u16, 1); + const bias = p.get(u16, 2); + const depth = depth256(p.shape) - 1; + for (0..p.shape.rows) |row| { + for (t[0..d], 0..) |*q, i| q.* = R.bf16(x[row * d + i]); + const mean = k.sum(mode, p.scratch, t[0..d], depth).div(R.of(@floatFromInt(d))); + for (dev, 0..) |*q, i| { + q.* = R.bf16(x[row * d + i]).sub(mean); + t[i] = q.mulExact(q.*); + } + const rstd = invNorm(k.sum(mode, p.scratch, t[0..d], depth + 1), d); + for (dev, 0..) |q, i| out[row * d + i] = q.mul(rstd).mul(R.bf16(w[i])).add(R.bf16(bias[i])).bf(); + } +} + +const ln_def: k.Def = .{ .prep = lnPrep, .launch = lnLaunch, .count = lnCount, .eval = lnEval }; + +// --- rotary: the stacks' heads, their rotary modes and frequencies + +/// The codec (fp32 cos/sin), the backbone's full layers (bf16 cos/sin, half rotated), the encoder, the depth decoder. +const Rot = struct { h: usize, hkv: usize, hd: usize, theta: f32, fraction: f32, pos0: usize, mode: u32 }; +const rots = [4]Rot{ + .{ .h = 8, .hkv = 2, .hd = 64, .theta = 1e6, .fraction = 1, .pos0 = 90, .mode = 0 }, + .{ .h = 16, .hkv = 4, .hd = 128, .theta = 1e6, .fraction = 0.5, .pos0 = 300, .mode = 1 }, + .{ .h = 8, .hkv = 4, .hd = 96, .theta = 1e6, .fraction = 1, .pos0 = 0, .mode = 2 }, + .{ .h = 32, .hkv = 8, .hd = 32, .theta = 1e6, .fraction = 1, .pos0 = 0, .mode = 1 }, +}; + +/// The engine's inverse frequencies (spec.zig's invFreq), in fp32. +fn invFreq(out: []f32, r: Rot) void { + const rot: usize = @as(usize, @intFromFloat(r.fraction * @as(f32, @floatFromInt(r.hd)))) / 2; + for (out, 0..) |*x, i| x.* = if (i < rot) 1.0 / std.math.pow(f32, r.theta, @as(f32, @floatFromInt(2 * i)) / @as(f32, @floatFromInt(r.hd))) else 0; +} + +fn cos(x: f64) f64 { + return @cos(x); +} +fn sin(x: f64) f64 { + return @sin(x); +} + +/// A head's pair (a, z) at position `pos` and frequency `inv` in rotary mode `m`, as om_rope rounds it. +fn rotate(a: R, z: R, pos: usize, inv: f32, m: u32) [2]R { + const fr = R.of(@floatFromInt(pos)).mul(R.of(inv)); + var cs = fr.lip(cos, 1, k.fn_ulps); + var sn = fr.lip(sin, 1, k.fn_ulps); + if (m != 0) { + cs = cs.bf(); + sn = sn.bf(); + } + if (m == 2) return .{ a.mul(cs).bf().add(z.neg().mul(sn).bf()).bf(), z.mul(cs).bf().add(a.mul(sn).bf()).bf() }; + return .{ a.mul(cs).sub(z.mul(sn)).bf(), z.mul(cs).add(a.mul(sn)).bf() }; +} + +const RopeArgs = extern struct { T: i32, H: i32, hd: i32, ld: i32, pos0: i32, mode: i32 }; + +/// om_rope: shape rows T, n heads, k hd; the variant is the mode (and its stack's frequencies). +fn ropePrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + const r = rots[p.variant]; + if (s.k != r.hd or s.n == 0) return false; + const x = try p.alloc(u16, 0, s.rows * s.n * s.k); + const inv = try p.alloc(f32, 1, s.k / 2); + values(rng, x, if (p.fill == .stress) .random else p.fill); + if (p.fill == .nonfinite) plant(x, s.rows, s.n * s.k); + invFreq(inv, r); + try p.out(0, true); + try p.in(1); + return true; +} + +fn ropeLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const r = rots[p.variant]; + p.bind(e, &.{ 0, 1 }, RopeArgs{ .T = k.i32_(s.rows), .H = k.i32_(s.n), .hd = k.i32_(s.k), .ld = k.i32_(s.n * s.k), .pos0 = k.i32_(r.pos0), .mode = @intCast(r.mode) }); + Probe.go(e, .{ k.blocks(s.rows * s.n * s.k / 2, 256), 1, 1 }, 256); +} + +fn ropeCount(p: *Probe) usize { + return p.shape.rows * p.shape.n * p.shape.k; +} + +fn ropeEval(p: *Probe, _: k.Mode, out: []R) !void { + const s = p.shape; + const x = p.get(u16, 0); + const inv = p.get(f32, 1); + const half = s.k / 2; + for (0..s.rows) |t| for (0..s.n) |h| { + const at = (t * s.n + h) * s.k; + for (0..half) |i| { + const o = rotate(R.bf16(x[at + i]), R.bf16(x[at + i + half]), rots[p.variant].pos0 + t, inv[i], rots[p.variant].mode); + out[at + i] = o[0]; + out[at + i + half] = o[1]; + } + }; +} + +const rope_def: k.Def = .{ .prep = ropePrep, .launch = ropeLaunch, .count = ropeCount, .eval = ropeEval }; + +// --- om_qk_prep: the few-row path's heads, normed and rotated, k and v into the caches; variant as `rots` + +const QkArgs = extern struct { T: i32, H: i32, Hkv: i32, hd: i32, ld: i32, pos0: i32, mode: i32, nk: i32, eps: f32 }; + +fn qkLd(r: Rot) usize { + return r.h * r.hd + 2 * r.hkv * r.hd + r.h; +} + +fn qkPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + const r = rots[p.variant]; + if (s.k != r.hd or s.n != r.h or s.rows == 0) return false; + const ld = qkLd(r); + const nk = r.hkv * r.hd; + const x = try p.alloc(u16, 0, s.rows * ld); + const qn = try p.alloc(u16, 1, r.hd); + const kn = try p.alloc(u16, 2, r.hd); + const inv = try p.alloc(f32, 3, r.hd / 2); + const kc = try p.alloc(u16, 4, (r.pos0 + s.rows) * nk); + const vc = try p.alloc(u16, 5, (r.pos0 + s.rows) * nk); + values(rng, x, p.fill); + k.randomBf16(rng, qn, 0.5, 1.5); + k.randomBf16(rng, kn, 0.5, 1.5); + k.randomBf16(rng, kc, -1, 1); + k.randomBf16(rng, vc, -1, 1); + if (p.fill == .nonfinite) plant(x, s.rows, ld); + invFreq(inv, r); + try p.out(0, true); + for (1..4) |i| try p.in(i); + try p.out(4, true); + try p.out(5, true); + return true; +} + +fn qkLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const r = rots[p.variant]; + p.bind(e, &.{ 0, 1, 2, 3, 4, 5 }, QkArgs{ .T = k.i32_(s.rows), .H = k.i32_(r.h), .Hkv = k.i32_(r.hkv), .hd = k.i32_(r.hd), .ld = k.i32_(qkLd(r)), .pos0 = k.i32_(r.pos0), .mode = @intCast(r.mode), .nk = k.i32_(r.hkv * r.hd), .eps = eps32() }); + Probe.go(e, .{ s.rows * (r.h + r.hkv), 1, 1 }, 32); +} + +fn qkCount(p: *Probe) usize { + const r = rots[p.variant]; + return p.shape.rows * qkLd(r) + 2 * (r.pos0 + p.shape.rows) * r.hkv * r.hd; +} + +fn qkEval(p: *Probe, mode: k.Mode, out: []R) !void { + const s = p.shape; + const r = rots[p.variant]; + const ld = qkLd(r); + const nk = r.hkv * r.hd; + const x = p.get(u16, 0); + const kc_out = out[s.rows * ld ..][0 .. (r.pos0 + s.rows) * nk]; + const vc_out = out[s.rows * ld + kc_out.len ..]; + for (out[0 .. s.rows * ld], x) |*o, v| o.* = R.bf16(v); + for (kc_out, p.get(u16, 4)) |*o, v| o.* = R.bf16(v); + for (vc_out, p.get(u16, 5)) |*o, v| o.* = R.bf16(v); + var t: [128]R = undefined; + var sv: [128]R = undefined; + const half = r.hd / 2; + for (0..s.rows) |row| for (0..r.h + r.hkv) |h| { + const q = h < r.h; + const at = row * ld + (if (q) h * r.hd else r.h * r.hd + (h - r.h) * r.hd); + headRow(mode, p.scratch, x[at..][0..r.hd], p.get(u16, if (q) 1 else 2), t[0..r.hd], sv[0..r.hd]); + const dst = if (q) out[at..][0..r.hd] else kc_out[(r.pos0 + row) * nk + (h - r.h) * r.hd ..][0..r.hd]; + for (0..half) |i| { + const o = rotate(sv[i], sv[i + half], r.pos0 + row, p.get(f32, 3)[i], r.mode); + dst[i] = o[0]; + dst[i + half] = o[1]; + } + if (!q) for (0..r.hd) |i| { + vc_out[(r.pos0 + row) * nk + (h - r.h) * r.hd + i] = R.bf16(x[row * ld + r.h * r.hd + nk + (h - r.h) * r.hd + i]); + }; + }; +} + +const qk_def: k.Def = .{ .prep = qkPrep, .launch = qkLaunch, .count = qkCount, .eval = qkEval }; + +// --- the entries, each with the control that drops (or moves) one rounding point + +const sq = "for (int i = lane; i < a.hd; i += 32) "; +const rope_shapes = [3][]const c.Shape{ + &.{ .{ .rows = 6, .n = 10, .k = 64 }, .{ .rows = 48, .n = 10, .k = 64 } }, + &.{ .{ .rows = 6, .n = 24, .k = 128 }, .{ .rows = 40, .n = 24, .k = 128 } }, + &.{ .{ .rows = 6, .n = 12, .k = 96 }, .{ .rows = 94, .n = 12, .k = 96 } }, +}; + +pub const entries = [_]k.Entry{ + .{ .def = &norm_def, .spec = normSpec("metal lfm3 omni om_rmsnorm", false, &.{ .{ .rows = 94, .n = 768, .k = 768 }, .{ .rows = 48, .n = 512, .k = 512 }, .{ .rows = 64, .n = 2048, .k = 2048 } }), .file = .omni, .pipe = "om_rmsnorm", .variant = 0, .control = .{ "y[i] = a.mode == 0 ? b(f(W[i]) * rb(v))", "y[i] = a.mode == 0 ? b(f(W[i]) * (v))" }, .what = "drops bf16(x rsqrt)" }, + .{ .def = &norm_def, .spec = normSpec("metal lfm3 omni om_rmsnorm residual, post norm", false, &.{ .{ .rows = 48, .n = 512, .k = 512 }, .{ .rows = 64, .n = 2048, .k = 2048 } }), .file = .omni, .pipe = "om_rmsnorm", .variant = 1, .control = .{ "const float add = a.has_pw ? rb(f(PW[i]) * rb(f(r[i]) * inv_r)) : f(r[i]);", "const float add = a.has_pw ? rb(f(PW[i]) * (f(r[i]) * inv_r)) : f(r[i]);" }, .what = "drops bf16(r rsqrt)" }, + .{ .def = &norm_def, .spec = normSpec("metal lfm3 omni om_rmsnorm_1k", true, &.{.{ .rows = 1, .n = 2048, .k = 2048 }}), .file = .omni, .pipe = "om_rmsnorm_1k", .variant = 0x10, .control = .{ "y[i] = a.mode == 0 ? b(f(W[i]) * rb(v))", "y[i] = a.mode == 0 ? b(f(W[i]) * (v))" }, .what = "drops bf16(x rsqrt)" }, + .{ .def = &norm_def, .spec = normSpec("metal lfm3 omni om_rmsnorm_1k residual, post norm", true, &.{ .{ .rows = 1, .n = 2048, .k = 2048 }, .{ .rows = 12, .n = 1024, .k = 1024 } }), .file = .omni, .pipe = "om_rmsnorm_1k", .variant = 0x11, .control = .{ "v = rb(v + add);", "v = (v + add);" }, .what = "drops bf16(x + r)" }, + .{ .def = &norm_def, .spec = normSpec("metal lfm3 omni om_rmsnorm_1k row weights", true, &.{.{ .rows = 12, .n = 1024, .k = 1024 }}), .file = .omni, .pipe = "om_rmsnorm_1k", .variant = 0x12, .control = .{ ": b(v * f(W[i]));", ": b(rb(v) * f(W[i]));" }, .what = "adds bf16(x rsqrt) to mode 1" }, + .{ .def = &head_def, .spec = .{ .name = "metal lfm3 omni om_head_norm", .accum = .f32, .depth = headDepth, .ops = 8, .mids = 1, .out = .bf16, .in_place = true, .tiles = .{ .{}, .{}, .{ .size = 32, .min = 32, .step = 32 } }, .served = &.{ .{ .rows = 40, .n = 2048, .k = 128 }, .{ .rows = 94, .n = 768, .k = 96 }, .{ .rows = 48, .n = 512, .k = 64 }, .{ .rows = 12, .n = 1024, .k = 32 } }, .sums = .{ .bf16 = 0, .why = sq_why } }, .file = .omni, .pipe = "om_head_norm", .control = .{ sq ++ "x[i] = b(f(W[i]) * rb(f(x[i]) * inv));", sq ++ "x[i] = b(f(W[i]) * (f(x[i]) * inv));" }, .what = "drops bf16(x rsqrt)" }, + .{ .def = &ln_def, .spec = .{ .name = "metal lfm3 omni om_layernorm", .accum = .f32, .depth = depth256, .ops = 8, .out = .bf16, .tiles = .{ .{}, .{}, .{ .size = 256 } }, .served = &.{.{ .rows = 94, .n = 768, .k = 768 }}, .sums = .{ .bf16 = 0, .why = "the deviations' squares cannot cancel and the mean's error moves every output alike, under its ulps" } }, .file = .omni, .pipe = "om_layernorm", .control = .{ "b((f(x[i]) - mean) * rstd * f(W[i]) + f(Bb[i]))", "b(rb((f(x[i]) - mean) * rstd) * f(W[i]) + f(Bb[i]))" }, .what = "adds bf16((x - mean) rstd)" }, + .{ .def = &rope_def, .spec = ropeSpec("metal lfm3 omni om_rope fp32 cos sin (codec)", rope_shapes[0]), .file = .omni, .pipe = "om_rope", .variant = 0, .control = .{ "x[p] = b(xa * c - z * s);", "x[p] = b(xa * c - z * rb(s));" }, .what = "adds bf16(sin) to mode 0" }, + .{ .def = &rope_def, .spec = ropeSpec("metal lfm3 omni om_rope bf16 cos sin (backbone)", rope_shapes[1]), .file = .omni, .pipe = "om_rope", .variant = 1, .control = .{ " c = rb(c);\n s = rb(s);", " c = rb(c);\n s = (s);" }, .what = "drops bf16(sin)" }, + .{ .def = &rope_def, .spec = ropeSpec("metal lfm3 omni om_rope bf16 steps (encoder)", rope_shapes[2]), .file = .omni, .pipe = "om_rope", .variant = 2, .control = .{ "x[p] = b(rb(xa * c) + rb(-z * s));", "x[p] = b((xa * c) + rb(-z * s));" }, .what = "drops bf16(a cos)" }, + .{ .def = &qk_def, .spec = qkSpec("metal lfm3 omni om_qk_prep fp32 cos sin (codec)", &.{ .{ .rows = 1, .n = 8, .k = 64 }, .{ .rows = 6, .n = 8, .k = 64 } }), .file = .rows, .pipe = "om_qk_prep", .variant = 0, .control = .{ "out[p] = b(xa * c - z * s);", "out[p] = b(xa * c - rb(z * s));" }, .what = "adds bf16(z sin) to mode 0" }, + .{ .def = &qk_def, .spec = qkSpec("metal lfm3 omni om_qk_prep bf16 cos sin (backbone)", &.{ .{ .rows = 1, .n = 16, .k = 128 }, .{ .rows = 16, .n = 16, .k = 128 } }), .file = .rows, .pipe = "om_qk_prep", .variant = 1, .control = .{ "sv[i] = rb(f(w[i]) * rb(f(x[i]) * inv));", "sv[i] = rb(f(w[i]) * (f(x[i]) * inv));" }, .what = "drops bf16(x rsqrt)" }, + .{ .def = &qk_def, .spec = qkSpec("metal lfm3 omni om_qk_prep bf16 cos sin (depth decoder)", &.{.{ .rows = 12, .n = 32, .k = 32 }}), .file = .rows, .pipe = "om_qk_prep", .variant = 3, .control = .{ " c = rb(c);\n s = rb(s);", " c = (c);\n s = rb(s);" }, .what = "drops bf16(cos)" }, +}; + +fn ropeSpec(name: []const u8, served: []const c.Shape) c.Spec { + return .{ .name = name, .accum = .none, .depth = k.noSums, .ops = 8, .out = .bf16, .in_place = true, .tiles = .{ .{ .size = 16 }, .{}, .{} }, .served = served, .sums = .{ .f16 = 0, .bf16 = 0, .why = "rotations: no sums" } }; +} + +fn qkDepth(s: c.Shape) usize { + return (s.k + 31) / 32 + 6; +} + +fn qkSpec(name: []const u8, served: []const c.Shape) c.Spec { + return .{ .name = name, .accum = .f32, .depth = qkDepth, .ops = 16, .mids = 2, .out = .bf16, .in_place = true, .tiles = .{ .{ .size = 16, .min = 1 }, .{}, .{} }, .served = served, .sums = .{ .bf16 = 0, .why = sq_why } }; +} diff --git a/zig/tests/contract/metal_omni_pick.zig b/zig/tests/contract/metal_omni_pick.zig new file mode 100644 index 0000000000..51464bb77b --- /dev/null +++ b/zig/tests/contract/metal_omni_pick.zig @@ -0,0 +1,307 @@ +//! LFM3-Omni's Metal picks: the greedy text pick's two stages, the depth decoder's code picks and a frame's end. +const std = @import("std"); +const mtl = @import("metal"); +const k = @import("metal_omni_common.zig"); + +const c = k.c; +const R = k.R; +const Probe = k.Probe; +const nq = 12; +const size = 2017; +const eos = 2016; +const hist_cap = 64; +const none_at: i32 = 0x7fffffff; + +/// A logit with few distinct values (ties everywhere), NaN and infinities sprinkled when asked. +fn logit(rng: std.Random, fill: k.g.Fill) u16 { + if (fill == .nonfinite and rng.uintLessThan(u32, 400) < 3) return ([3]u16{ k.nan16, k.inf16, k.ninf16 })[rng.uintLessThan(u32, 3)]; + const v = @as(f32, @floatFromInt(rng.uintLessThan(u32, 61))) / 8 - 4; + return if (v == 0 and rng.boolean()) 0x8000 else c.toBf16(v); +} + +/// Whether `pick` is a valid greedy choice: none surely larger, lowest id on a sure tie, never NaN, else 2^31 - 1. +fn valid(pick: i32, lo: []const f64, hi: []const f64, skip: []const bool) bool { + var any = false; + for (lo, skip) |l, s| any = any or (!s and !std.math.isNan(l)); + if (!any) return pick == none_at; + if (pick < 0 or pick >= lo.len) return false; + const p: usize = @intCast(pick); + if (skip[p] or std.math.isNan(lo[p])) return false; + for (lo, hi, skip, 0..) |l, h, s, j| { + if (j == p or s or std.math.isNan(l)) continue; + if (l > hi[p]) return false; + if (j < p and l == h and lo[p] == hi[p] and l == hi[p]) return false; + } + return true; +} + +// --- om_argmax_part (shape n logits, k a part's width) and om_argmax_final (n parts) + +const PartArgs = extern struct { N: i32, ns: i32, per: i32 }; +const sup = [4]i32{ 125031, 125032, 125030, 125029 }; + +fn partPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n == 0 or s.k == 0) return false; + const parts = k.blocks(s.n, s.k); + const x = try p.alloc(u16, 0, s.n); + const ids = try p.alloc(i32, 1, 4); + _ = try p.alloc(f32, 2, parts); + _ = try p.alloc(i32, 3, parts); + for (x) |*v| v.* = logit(rng, p.fill); + // the largest logits on the excluded ids, and a part left with nothing but excluded and NaN ids + @memcpy(ids, &sup); + for (sup) |id| if (id < s.n) { + x[@intCast(id)] = c.toBf16(9); + }; + if (p.fill == .nonfinite and s.n >= 2 * s.k) @memset(x[s.k .. 2 * s.k], k.nan16); + try p.in(0); + try p.in(1); + try p.out(2, false); + try p.out(3, false); + return true; +} + +fn partLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + p.bind(e, &.{ 0, 1, 2, 3 }, PartArgs{ .N = k.i32_(s.n), .ns = 4, .per = k.i32_(s.k) }); + Probe.go(e, .{ k.blocks(s.n, s.k), 1, 1 }, 256); +} + +fn partCount(p: *Probe) usize { + return 8 * k.blocks(p.shape.n, p.shape.k); +} + +fn zeroEval(_: *Probe, _: k.Mode, out: []R) !void { + @memset(out, R.of(0)); +} + +fn partCheck(p: *Probe, _: k.g.Gate, out: []const u8) !?[]const u8 { + const s = p.shape; + const parts = k.blocks(s.n, s.k); + const x = p.get(u16, 0); + const pv = std.mem.bytesAsSlice(f32, out[0 .. 4 * parts]); + const pi = std.mem.bytesAsSlice(i32, out[4 * parts ..]); + var lo: [4096]f64 = undefined; + var skip: [4096]bool = undefined; + for (0..parts) |g| { + const a = g * s.k; + const b = @min(s.n, a + s.k); + for (a..b) |i| { + lo[i - a] = c.fromBf16(x[i]); + skip[i - a] = std.mem.indexOfScalar(i32, &sup, @intCast(i)) != null; + } + const pick = if (pi[g] == none_at) none_at else pi[g] - @as(i32, @intCast(a)); + if (!valid(pick, lo[0 .. b - a], lo[0 .. b - a], skip[0 .. b - a])) return "a part's pick is not its lowest-id largest"; + const want: f32 = if (pi[g] == none_at) -std.math.inf(f32) else c.fromBf16(x[@intCast(pi[g])]); + if (@as(u32, @bitCast(pv[g])) != @as(u32, @bitCast(want))) return "a part's value is not its pick's logit"; + } + return null; +} + +const part_def: k.Def = .{ .prep = partPrep, .launch = partLaunch, .count = partCount, .eval = zeroEval, .check = partCheck }; + +fn finalPrep(p: *Probe, rng: std.Random) !bool { + const n = p.shape.n; + if (n == 0) return false; + const pv = try p.alloc(f32, 0, n); + const pi = try p.alloc(i32, 1, n); + _ = try p.alloc(i32, 2, 1); + for (pv, pi, 0..) |*v, *i, j| { + v.* = c.fromBf16(logit(rng, if (p.fill == .nonfinite) .nonfinite else .random)); + i.* = @intCast(j * 1024 + rng.uintLessThan(u32, 1024)); + } + // the same best in two parts of one lane, the lower id first, and in a third part of another lane + if (n > 40) { + pv[1] = 9; + pv[33] = 9; + pv[n - 1] = 9; + pi[1] = 5; + pi[33] = 2000; + pi[n - 1] = 7; + } + try p.in(0); + try p.in(1); + try p.out(2, false); + return true; +} + +fn finalLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1, 2 }, @as(i32, k.i32_(p.shape.n))); + Probe.go(e, .{ 1, 1, 1 }, 32); +} + +fn four(_: *Probe) usize { + return 4; +} + +fn finalCheck(p: *Probe, _: k.g.Gate, out: []const u8) !?[]const u8 { + const pv = p.get(f32, 0); + const pi = p.get(i32, 1); + var best = -std.math.inf(f32); + var at: i32 = none_at; + for (pv, pi) |v, i| if (v > best or (v == best and i < at)) { + best = v; + at = i; + }; + return k.sameBytes(out, std.mem.asBytes(&at)); +} + +const final_def: k.Def = .{ .prep = finalPrep, .launch = finalLaunch, .count = four, .eval = zeroEval, .check = finalCheck }; + +// --- om_depth_pick: shape rows the pass's picks, n the codebook's size, k the window's frames; variant 1 sampled + +const PickArgs = extern struct { p0: i32, hn: i32, penalty: f32, temp: f32, seed: u64, frame: i32, NQ: i32, size: i32, pad: i32 = 0 }; +const penalty: f32 = 1.2; +const temp: f32 = 0.4; +const seed: u64 = 1003; +const frame = 37; +const p0 = 4; + +fn mix64(z0: u64) u64 { + var z = z0 +% 0x9E3779B97F4A7C15; + z = (z ^ (z >> 30)) *% 0xBF58476D1CE4E5B9; + z = (z ^ (z >> 27)) *% 0x94D049BB133111EB; + return z ^ (z >> 31); +} + +fn pickPrep(p: *Probe, rng: std.Random) !bool { + const s = p.shape; + if (s.n != size or s.rows == 0 or p0 + s.rows > nq or s.k > hist_cap) return false; + const log = try p.alloc(u16, 0, nq * size); + const frame_ = try p.alloc(i32, 1, 3 * nq); + const hist = try p.alloc(i32, 2, hist_cap * nq); + for (log) |*v| v.* = logit(rng, p.fill); + for (0..nq) |q| frame_[q] = @intCast(q); + rng.shuffle(i32, frame_[0..nq]); + for (frame_[nq .. 2 * nq]) |*x| x.* = 0; + for (frame_[2 * nq ..]) |*x| x.* = 1; + for (hist) |*h| h.* = @intCast(rng.uintLessThan(u32, 64)); + // a repeated code among each codebook's largest logits, and exact ties on unrepeated codes + for (0..nq) |q| { + log[q * size + 3] = c.toBf16(5); + log[q * size + 70] = c.toBf16(5); + log[q * size + 90] = c.toBf16(5); + } + try p.in(0); + try p.in(2); + try p.out(1, true); + return true; +} + +fn pickLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + const s = p.shape; + const codes = Probe.at(p.refs[1], nq * 4); + const masked = Probe.at(p.refs[1], 2 * nq * 4); + e.setPipeline(p.pipe); + e.setBuffer(p.refs[0].buf, p.refs[0].off, 0); + e.setBuffer(p.refs[1].buf, p.refs[1].off, 1); + e.setBuffer(p.refs[2].buf, p.refs[2].off, 2); + e.setBuffer(codes.buf, codes.off, 3); + e.setBuffer(masked.buf, masked.off, 4); + const a = PickArgs{ .p0 = p0, .hn = k.i32_(s.k), .penalty = penalty, .temp = if (p.variant == 1) temp else 0, .seed = seed, .frame = frame, .NQ = nq, .size = size }; + e.setBytes(std.mem.asBytes(&a), 5); + Probe.go(e, .{ s.rows, 1, 1 }, 256); +} + +fn pickCount(_: *Probe) usize { + return 3 * nq * 4; +} + +fn ln(x: f64) f64 { + return @log(x); +} + +fn pickCheck(p: *Probe, _: k.g.Gate, out: []const u8) !?[]const u8 { + const s = p.shape; + const got = std.mem.bytesAsSlice(i32, out); + const init = p.get(i32, 1); + const log = p.get(u16, 0); + const hist = p.get(i32, 2); + var lo: [size]f64 = undefined; + var hi: [size]f64 = undefined; + var skip: [size]bool = @splat(false); + for (0..nq) |j| if (got[j] != init[j]) return "the order was changed"; + for (0..nq) |q| { + const picked = std.mem.indexOfScalar(i32, init[p0 .. p0 + s.rows], @intCast(q)) != null; + if (!picked) { + if (got[nq + q] != init[nq + q] or got[2 * nq + q] != init[2 * nq + q]) return "an unpicked codebook was written"; + continue; + } + if (got[2 * nq + q] != 0) return "a picked codebook stays masked"; + var seen: [size]bool = @splat(false); + for (0..s.k) |j| seen[@intCast(hist[j * nq + q])] = true; + for (0..size) |i| { + var v = R.bf16(log[q * size + i]); + if (seen[i]) v = if (v.v < 0) v.mul(R.of(penalty)) else v.div(R.of(penalty)); + if (p.variant == 1) { + const h = mix64(seed ^ mix64((@as(u64, frame) << 32) ^ (@as(u64, q) << 16) ^ @as(u64, i))); + const u = (@as(f64, @floatFromInt(h >> 40)) + 0.5) * (1.0 / 16777216.0); + const g = R.of(u).fp(0.5).up(ln, k.fn_ulps).neg().lip(ln, 1 / (-@log(u) * (1 - 1e-6)), k.fn_ulps); + v = v.div(R.of(temp)).sub(g); + } + lo[i] = v.lo; + hi[i] = v.hi; + if (std.math.isNan(v.v)) { + lo[i] = std.math.nan(f64); + } + } + if (!valid(got[nq + q], &lo, &hi, &skip)) { + std.debug.print(" codebook {d}: picked {d}\n", .{ q, got[nq + q] }); + return "a code is not the lowest-id largest of its penalized (and noised) logits"; + } + } + return null; +} + +const pick_def: k.Def = .{ .prep = pickPrep, .launch = pickLaunch, .count = pickCount, .eval = zeroEval, .check = pickCheck }; + +// --- om_frame_end: variant 0 a frame without end-of-speech (pushed at a slot), 1 a frame with one + +const EndArgs = extern struct { NQ: i32, eos: i32, slot: i32, push: i32 }; + +fn endPrep(p: *Probe, rng: std.Random) !bool { + const codes = try p.alloc(i32, 0, nq); + const hist = try p.alloc(i32, 1, hist_cap * nq); + for (codes) |*x| x.* = @intCast(rng.uintLessThan(u32, eos)); + for (hist) |*x| x.* = @intCast(rng.uintLessThan(u32, eos)); + if (p.variant == 1) codes[0] = eos; + try p.out(0, true); + try p.out(1, true); + return true; +} + +fn endLaunch(p: *Probe, e: mtl.ComputeEncoder) void { + p.bind(e, &.{ 0, 1 }, EndArgs{ .NQ = nq, .eos = eos, .slot = k.i32_(p.shape.k), .push = 1 }); + Probe.go(e, .{ 1, 1, 1 }, 32); +} + +fn endCount(_: *Probe) usize { + return (nq + hist_cap * nq) * 4; +} + +fn endCheck(p: *Probe, _: k.g.Gate, out: []const u8) !?[]const u8 { + const want = try p.a.alloc(i32, nq + hist_cap * nq); + defer p.a.free(want); + @memcpy(want[0..nq], p.get(i32, 0)); + @memcpy(want[nq..], p.get(i32, 1)); + if (std.mem.indexOfScalar(i32, want[0..nq], eos) != null) @memset(want[0..nq], eos) else @memcpy(want[nq + p.shape.k * nq ..][0..nq], want[0..nq]); + return k.sameBytes(out, std.mem.sliceAsBytes(want)); +} + +const end_def: k.Def = .{ .prep = endPrep, .launch = endLaunch, .count = endCount, .eval = zeroEval, .check = endCheck }; + +fn stated(name: []const u8, rule: []const u8, served: []const c.Shape) c.Spec { + return .{ .name = name, .accum = .none, .depth = k.noSums, .out = .raw, .in_place = true, .served = served, .nan = .{ .stated = rule }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "picks compare values: no sums" } }; +} + +const pick_rule = "a NaN logit is never picked; ties (and -Inf) to the lowest id; none left gives 2^31 - 1"; + +pub const entries = [_]k.Entry{ + .{ .def = &part_def, .spec = stated("metal lfm3 omni om_argmax_part", pick_rule, &.{ .{ .rows = 1, .n = 128000, .k = 1024 }, .{ .rows = 1, .n = 3000, .k = 1024 } }), .file = .rows, .pipe = "om_argmax_part", .control = .{ "if (v > best || (v == best && i < at)) {", "if (v > best || (v == best && i > at)) {" }, .what = "ties to the highest id" }, + .{ .def = &final_def, .spec = stated("metal lfm3 omni om_argmax_final", pick_rule, &.{ .{ .rows = 1, .n = 125, .k = 1 }, .{ .rows = 1, .n = 3, .k = 1 } }), .file = .rows, .pipe = "om_argmax_final", .control = .{ "if (PV[i] > best || (PV[i] == best && PI[i] < at)) {", "if (PV[i] >= best) {" }, .what = "ties to the last part" }, + .{ .def = &pick_def, .spec = stated("metal lfm3 omni om_depth_pick greedy", pick_rule, &.{ .{ .rows = 4, .n = size, .k = hist_cap }, .{ .rows = 4, .n = size, .k = 3 } }), .file = .gemm, .pipe = "om_depth_pick", .variant = 0, .control = .{ "if (v > best || (v == best && i < at)) {", "if (v > best || (v == best && i > at)) {" }, .what = "ties to the highest id" }, + .{ .def = &pick_def, .spec = stated("metal lfm3 omni om_depth_pick sampled", "a NaN logit is never drawn; equal noised values to the lowest id", &.{ .{ .rows = 4, .n = size, .k = hist_cap }, .{ .rows = 4, .n = size, .k = 0 } }), .file = .gemm, .pipe = "om_depth_pick", .variant = 1, .control = .{ "if (seen[i]) v = v < 0.0f ? v * a.penalty : v / a.penalty;", "if (seen[i]) v = v < 0.0f ? v / a.penalty : v * a.penalty;" }, .what = "inverts the repetition penalty" }, + .{ .def = &end_def, .spec = stated("metal lfm3 omni om_frame_end pushed", "integer codes: every word as the rule says", &.{ .{ .rows = 1, .n = nq, .k = 17 }, .{ .rows = 1, .n = nq, .k = 63 } }), .file = .gemm, .pipe = "om_frame_end", .variant = 0, .control = .{ "for (int q = 0; q < a.NQ; ++q) HIST[a.slot * a.NQ + q] = CODES[q];", "for (int q = 1; q < a.NQ; ++q) HIST[a.slot * a.NQ + q] = CODES[q];" }, .what = "pushes all but the first code" }, + .{ .def = &end_def, .spec = stated("metal lfm3 omni om_frame_end end of speech", "integer codes: every word as the rule says", &.{.{ .rows = 1, .n = nq, .k = 5 }}), .file = .gemm, .pipe = "om_frame_end", .variant = 1, .control = .{ "for (int q = 0; q < a.NQ; ++q) any |= CODES[q] == a.eos;", "for (int q = 1; q < a.NQ; ++q) any |= CODES[q] == a.eos;" }, .what = "misses an end-of-speech code in codebook 0" }, +}; From 02fa671a9f00f0704a11a80787935c866f1dffeb Mon Sep 17 00:00:00 2001 From: drowzeys <264866615+drowzeys@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:48:08 -0400 Subject: [PATCH 3/3] docs: LFM3-Omni-2B on Metal, its receipt, the checks' clients and the weights' licence tools/lfm3_omni holds the development clients behind the receipt: the 16 synthetic spoken turns, the reference's stage dumps and spoken replies (transformers with the checkpoint's own code, on CUDA, MPS or the CPU), the stage comparison and its table, a streaming client for /v1/audio/chat, and the UTMOS and speech-vs-text scoring. --- docs/recipes/README.md | 1 + docs/recipes/lfm3-omni-metal.md | 112 ++++++++++++++++++++++++++ tools/lfm3_omni/compare_stages.py | 78 ++++++++++++++++++ tools/lfm3_omni/make_utterances.sh | 17 ++++ tools/lfm3_omni/parity.sh | 20 +++++ tools/lfm3_omni/refcommon.py | 50 ++++++++++++ tools/lfm3_omni/reference_depth.py | 68 ++++++++++++++++ tools/lfm3_omni/reference_dump.py | 95 ++++++++++++++++++++++ tools/lfm3_omni/reference_speech.py | 83 +++++++++++++++++++ tools/lfm3_omni/score_replies.py | 78 ++++++++++++++++++ tools/lfm3_omni/speech_chat_client.py | 88 ++++++++++++++++++++ tools/lfm3_omni/summarize_parity.py | 76 +++++++++++++++++ tools/lfm3_omni/utterances.txt | 16 ++++ 13 files changed, 782 insertions(+) create mode 100644 docs/recipes/lfm3-omni-metal.md create mode 100644 tools/lfm3_omni/compare_stages.py create mode 100644 tools/lfm3_omni/make_utterances.sh create mode 100644 tools/lfm3_omni/parity.sh create mode 100644 tools/lfm3_omni/refcommon.py create mode 100644 tools/lfm3_omni/reference_depth.py create mode 100644 tools/lfm3_omni/reference_dump.py create mode 100644 tools/lfm3_omni/reference_speech.py create mode 100644 tools/lfm3_omni/score_replies.py create mode 100644 tools/lfm3_omni/speech_chat_client.py create mode 100644 tools/lfm3_omni/summarize_parity.py create mode 100644 tools/lfm3_omni/utterances.txt diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 8017124388..79ee07c8c1 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -24,6 +24,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | Kolibri-1 | [Native CUDA, block FP8, research checkpoint](kolibri-1.md) | | LFM2.5-24B-A2B | [Native Metal, MLX 4-bit, shared rounds](lfm2.5-24b-a2b.md) | | Spark-X2.6-4B | [Native CUDA](spark-x2.6-4b-cuda.md), [native Metal, the bf16 release, contexts to 1M](spark-x2.6-4b-metal.md) | +| LFM3-Omni-2B | [Native Metal, speech in, text and speech out](lfm3-omni-metal.md) | ## Capability floor diff --git a/docs/recipes/lfm3-omni-metal.md b/docs/recipes/lfm3-omni-metal.md new file mode 100644 index 0000000000..4c8458b873 --- /dev/null +++ b/docs/recipes/lfm3-omni-metal.md @@ -0,0 +1,112 @@ +# LFM3-Omni-2B speech to speech, native Metal + +`lfm3_omni` runs LiquidAI's [`LiquidAI/LFM3-Omni-2B`](https://huggingface.co/LiquidAI/LFM3-Omni-2B) (revision +`5797d07c125503bd4f4c0e7f915ae2d3d703e8de`) on Apple silicon: a spoken turn in, interleaved text and 24 kHz speech out, +streamed. The weights are the bf16 release as stored, the bundled Liquid codec (`audio_tokenizer/`) included. The +engine is `zig/src/families/lfm3_omni/metal/` over one metallib (`zig/kernels/metal/lfm3_omni/`). Checked on an M5 +Ultra; chips without tensor units (M1 to M4) are untested. The vision tower is not loaded. + +```bash +hf download LiquidAI/LFM3-Omni-2B --revision 5797d07c125503bd4f4c0e7f915ae2d3d703e8de --local-dir "$HOME/models/lfm3-omni-2b" +# the Metal compiler: xcodebuild -downloadComponent MetalToolchain (macOS 26 and later), with Xcode selected +zig build lfm3-omni-generate -Doptimize=ReleaseFast -Dcpu=apple_m1 +# one turn: a WAV in (any rate, resampled to 16 kHz), the reply as a WAV and its text +zig-out/bin/lfm3-omni-generate --model "$HOME/models/lfm3-omni-2b" --wav question.wav --out reply.wav \ + --system "Respond with interleaved text and audio. You are a helpful voice assistant." +# a server: POST /v1/audio/chat, GET /health +zig-out/bin/lfm3-omni-generate --model "$HOME/models/lfm3-omni-2b" --serve 127.0.0.1:8441 --system-file system.txt +``` + +About 4.4 GB of unified memory. + +## The endpoint + +`POST /v1/audio/chat` takes JSON: `audio` (a base64 WAV, required), `system`, `stream` (default true), +`audio_temperature` (0.4), `audio_repetition_penalty` (1.2), `seed`, `max_new` (generated slots, a text token or a +speech frame each). With `stream` the answer is server-sent events as the reply is made: + +```text +data: {"type":"text","delta":"I'm on it,"} +data: {"type":"audio","format":"pcm16","sample_rate":24000,"data":""} +... +data: {"type":"done","text":"...","frames":87,"speech_s":6.96,"first_text_ms":36.0,"first_audio_ms":61.2,"total_ms":640.4} +data: [DONE] +``` + +Without it, one JSON object with the same fields and `audio`, the whole reply as a base64 WAV. The server takes one +request at a time; it is the tool's own (not a route of `tensorfold-native serve`, whose speech route is text in). +Text is greedy, as the reference's `do_sample=False`; speech frames are sampled with the card's settings. + +## What runs + +- Log-mel (16 kHz, 400-point periodic Hann, 128 Slaney bins) as fp32 products on the device; the speech encoder's + three stride-2 convolutions, its 17 LFM2 layers and the projector. +- The 24-layer LFM3 backbone (gated short conv, full and sliding attention), the prompt in chunks of 512 rows. +- The depth decoder's three mask passes over a frame's 12 codebook slots, each codebook's pick on the GPU (repetition + penalty, greedy or a Gumbel-max draw, exactly a softmax sample). +- The Liquid codec's decoder and its spectrum head, the inverse STFT as an fp32 product; overlap-add on the host. +- Products of many rows (prompt chunks, the encoder, codec chunks) on the tensor unit through the prefill GEMM's + blocks (`zig/kernels/metal/prefill/gemm_nax.metal`), fp32 sums; a GPU without tensor units takes a + simdgroup-matrix product (`LFM3_OMNI_NAX=0` forces it). +- Up to 16 rows (a backbone row, the depth decoder's 12 slots): a weight-streaming product, the heads' norm, rotary + and cache write in one dispatch, one-row attention with LFM3's output gate inside, SwiGLU as the product's + epilogue; the greedy text pick in two stages. +- One serial command buffer from one copy back to the next: a backbone row runs while the host does the overlap-add; + the first frame is decoded alone, then every 8. + +Every op rounds to bf16 where the reference's op rounds. `LFM3_OMNI_PROFILE=1` logs each phase's GPU time per reply. + +## Checks + +`tools/lfm3_omni/` holds the checks' clients (Python dev tools, not part of the engine): `make_utterances.sh` makes +the 16 synthetic spoken turns of `utterances.txt` with macOS voices; `reference_dump.py` and `reference_depth.py` dump +the reference's stages (transformers with the checkpoint's own code); `parity.sh` runs `lfm3-omni-generate --dump`, +`--force` and `--depth-check` against them and `summarize_parity.py` prints the table below; +`speech_chat_client.py` streams turns through the server and times them; `reference_speech.py` does the same with the +reference runtime; `score_replies.py` scores replies (UTMOS22 strong through SpeechMOS, and the checkpoint's own ASR +against each reply's text). + +## Receipt + +M5 Ultra (256 GB), macOS 27.0.1, Xcode 27.1; the reference is `transformers` 5.21.0 with the checkpoint's own code on +PyTorch 2.14.1 (MPS, bf16), on the same Mac. The 16 turns of `tools/lfm3_omni/utterances.txt` (macOS voices), a +voice-assistant system text, identity codebook order, greedy text and codes: + +| stage | check | Metal vs MPS bf16 | +| --- | --- | --- | +| log-mel | cosine (min), max abs | 1.000000, 3.1e-4 | +| speech encoder | cosine (min of 16) | 0.99999 | +| projector | cosine (min) | 0.99990 | +| prompt | token ids | 16/16 identical | +| backbone prefill | row cosine (median of medians, min of medians) | 0.99982, 0.99977 | +| text | teacher-forced greedy top-1 | 520/524 (99.2%) | +| depth decoder | on the reference's own states: greedy agreement, logits cosine median / min | 95.5%, 0.999994 / 0.99989 | +| codes | teacher-forced through the whole loop | 89.2% | +| codec | the reference's codes to audio: SNR median (min); spectrum rows cosine (min) | 36.6 dB (25.1 dB); 0.99994 | + +Text and code misses are near ties (within 2 bf16 ulps of the reference's pick); 15% of the reference's own greedy +codebook picks are exact bf16 ties, so free-running codes part after a frame or two. Against MPS fp32 the port sits +where the reference's own bf16 run sits (encoder 0.999995, prefill rows 0.99988 median, for both). + +Replies with the card's audio settings (temperature 0.4, penalty 1.2, seeds 1000 + i, the server and +`speech_chat_client.py`; the reference through `reference_speech.py`, its first audio being the first sampled frame +plus that frame's codec decode). Two rounds, the servers and the reference alternated; other services were resident +on the Mac but idle (runs paused while its interactive voice service was in use): + +| | TensorFold Metal | reference runtime (MPS bf16) | +| --- | --- | --- | +| first text, median | 0.034 s | | +| first audio, median | 0.063 s | 0.172 to 0.191 s | +| whole reply, median (6.5 s of speech) | 0.67 to 0.68 s | 2.41 to 2.54 s | +| UTMOS22 strong, median (min) | 4.43 (4.35) | 4.44 (4.15) | +| speech vs its own text, WER | 1.5% | 1.2% | + +WER is the checkpoint's own ASR on each reply's audio against the reply's streamed text. + + +## Licence + +The weights are under Liquid AI's LFM Open License v1.0 (`license: other`, `lfm1.0` on the model card; the text is the +`LICENSE` file in the checkpoint). Use, changes and redistribution are allowed; commercial use is licensed only to an +organisation under US$10 million in annual revenue, and a qualified non-profit may use it for research or other +non-commercial purposes regardless. TensorFold's code stays Apache-2.0. diff --git a/tools/lfm3_omni/compare_stages.py b/tools/lfm3_omni/compare_stages.py new file mode 100644 index 0000000000..e8a75d614f --- /dev/null +++ b/tools/lfm3_omni/compare_stages.py @@ -0,0 +1,78 @@ +"""Stage parity of lfm3-omni-generate's dumps against reference_dump.py's: python compare_stages.py REF_DIR TF_DIR.""" +import json +import sys +from pathlib import Path + +import numpy as np + +MIXED_START = 125015 + + +def stat(a, b): + """Cosine, relative L2 and largest absolute difference in float64.""" + a, b = a.astype(np.float64).ravel(), b.astype(np.float64).ravel() + cos = float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30)) + rel = float(np.linalg.norm(a - b) / (np.linalg.norm(a) + 1e-30)) + return {"cos": round(cos, 6), "rel_l2": round(rel, 5), "max_abs": round(float(np.abs(a - b).max()), 5)} + + +def first_split(a, b): + """The length of the common prefix of two sequences.""" + m = min(len(a), len(b)) + return next((i for i in range(m) if np.any(a[i] != b[i])), m) + + +def compare(r, t): + """Every stage the two dumps share.""" + def raw(name, dtype=np.float32): + p = t / name + return np.fromfile(p, dtype=dtype) if p.exists() else None + + out = {} + mel = np.load(r / "mel.npy") + tm = raw("mel.f32").reshape(128, -1) + out["mel"] = stat(mel, tm[:, : mel.shape[1]]) | {"frames": [mel.shape[1], tm.shape[1]]} + enc = np.load(r / "enc.npy") + te = raw("enc.f32").reshape(-1, enc.shape[1]) + out["encoder"] = stat(enc, te) | {"tokens": [enc.shape[0], te.shape[0]]} + emb = np.load(r / "audio_emb.npy") + out["projector"] = stat(emb, raw("audio_emb.f32").reshape(emb.shape)) + ids, ti = np.load(r / "input_ids.npy"), raw("input_ids.i32", np.int32) + out["prompt_ids_equal"] = bool(len(ti) == len(ids) + 1 and (ti[: len(ids)] == ids).all() and ti[-1] == MIXED_START) + ph = np.load(r / "prefill_hidden.npy") + tph = raw("prefill_hidden.f32").reshape(-1, ph.shape[1]) + n = ph.shape[0] + out["prefill_hidden"] = stat(ph, tph[:n]) | {"last_row": stat(ph[-1], tph[n - 1])} + out["prefill_hidden"]["min_row_cos"] = min(stat(ph[i], tph[i])["cos"] for i in range(n)) + dl = np.load(r / "depth_logits0.npy") + tdl = raw("depth_logits0.f32").reshape(dl.shape) + out["depth_logits_frame0"] = stat(dl, tdl) | {"argmax_equal": int((dl.argmax(-1) == tdl.argmax(-1)).sum())} + seq = np.load(r / "gen_seq.npy") + ts = np.concatenate([[MIXED_START], raw("gen_seq.i32", np.int32)]) + div = first_split(seq, ts) + out["free_run_slots"] = {"ref": int(len(seq)), "tf": int(len(ts)), "identical_prefix": int(div)} + codes, tc = np.load(r / "gen_codes.npy").T, raw("gen_codes.i32", np.int32).reshape(-1, 12) + same = int(first_split(codes, tc)) + out["free_run_frames"] = {"ref": int(len(codes)), "tf": int(len(tc)), "identical_prefix": same} + sh, tsh = np.load(r / "step_hidden.npy"), raw("step_hidden.f32").reshape(-1, ph.shape[1]) + k = min(len(sh), len(tsh), div) + min_cos = round(min(stat(sh[i], tsh[i])["cos"] for i in range(k)), 6) if k else None + out["step_hidden_before_divergence"] = {"steps": int(k), "min_cos": min_cos, "first": stat(sh[0], tsh[0])} + cl = raw("codec_lin.f32") + if cl is not None and (r / "codec_lin.npy").exists(): + rl = np.load(r / "codec_lin.npy") + cl = cl.reshape(-1, rl.shape[1]) + k = min(len(rl), len(cl)) + out["codec_lin"] = stat(rl[:k], cl[:k]) | {"rows": [len(rl), len(cl)]} + rw, tw = np.load(r / "reply.npy"), raw("reply.f32") + k = min(len(rw), len(tw)) + snr = 10 * np.log10(np.sum(rw[:k] ** 2) / (np.sum((rw[:k] - tw[:k]) ** 2) + 1e-30)) + out["codec_wave"] = {"samples": [len(rw), len(tw)], "snr_db": round(float(snr), 2)} | stat(rw[:k], tw[:k]) + return out + + +if __name__ == "__main__": + R, T = Path(sys.argv[1]), Path(sys.argv[2]) + res = compare(R, T) + print(json.dumps(res, indent=1)) + json.dump(res, open(T / "parity.json", "w"), indent=1) diff --git a/tools/lfm3_omni/make_utterances.sh b/tools/lfm3_omni/make_utterances.sh new file mode 100644 index 0000000000..9c1c376a98 --- /dev/null +++ b/tools/lfm3_omni/make_utterances.sh @@ -0,0 +1,17 @@ +#!/bin/sh +# The 16 synthetic spoken turns (macOS `say`, eight system voices in turn) as 16 kHz WAVs: make_utterances.sh OUT_DIR +set -eu +here=$(cd "$(dirname "$0")" && pwd) +out=$1 +mkdir -p "$out" +: > "$out/refs.tsv" +set -- Samantha Daniel Karen Moira Rishi Tessa Fred Albert +i=0 +while IFS= read -r line; do + [ -z "$line" ] && continue + v=$(eval echo "\${$((i % 8 + 1))}") + f=$(printf "u%02d.wav" "$i") + say -v "$v" -o "$out/$f" --data-format=LEI16@16000 "$line" + printf "%s\t%s\t%s\n" "$f" "$v" "$line" >> "$out/refs.tsv" + i=$((i + 1)) +done < "$here/utterances.txt" diff --git a/tools/lfm3_omni/parity.sh b/tools/lfm3_omni/parity.sh new file mode 100644 index 0000000000..9ebc20cb23 --- /dev/null +++ b/tools/lfm3_omni/parity.sh @@ -0,0 +1,20 @@ +#!/bin/sh +# Stage parity over WAV_DIR/u*.wav: parity.sh BIN MODEL REF_DIR OUT_DIR SYSTEM_FILE WAV_DIR [DEV] [DTYPE] ($PYTHON) +set -u +here=$(cd "$(dirname "$0")" && pwd) +B=$1; M=$2; REF=$3; OUT=$4; SYS=$5; WAVS=$6; DEV=${7:-cuda}; DT=${8:-bf16}; P=${PYTHON:-python3} +busy() { [ -n "${LFM3_OMNI_BUSY:-}" ] && sh -c "$LFM3_OMNI_BUSY"; } +for w in "$WAVS"/u*.wav; do + u=$(basename "$w" .wav); mkdir -p "$OUT/$u" + while busy; do sleep 5; done + "$B" --model "$M" --wav "$w" --system-file "$SYS" --audio-temperature 0 --identity-order 1 --dump "$OUT/$u" \ + --codec-ref "$REF/$u" > "$OUT/$u/run.log" 2>&1 + "$P" "$here/compare_stages.py" "$REF/$u" "$OUT/$u" > /dev/null + "$B" --model "$M" --wav "$w" --system-file "$SYS" --audio-temperature 0 --identity-order 1 --force "$REF/$u" \ + > "$OUT/$u/forced.log" 2>&1 + [ -f "$REF/$u/ref_depth.npy" ] || "$P" "$here/reference_depth.py" "$M" "$REF/$u" prep "$DEV" "$DT" > /dev/null + "$B" --model "$M" --depth-check "$REF/$u" > "$OUT/$u/depth.log" 2>&1 + "$P" "$here/reference_depth.py" "$M" "$REF/$u" score > "$OUT/$u/depth_parity.json" + echo "$u $(grep FORCED "$OUT/$u/forced.log") $(cat "$OUT/$u/depth_parity.json")" +done +"$P" "$here/summarize_parity.py" "$REF" "$OUT" diff --git a/tools/lfm3_omni/refcommon.py b/tools/lfm3_omni/refcommon.py new file mode 100644 index 0000000000..c6b38e50f0 --- /dev/null +++ b/tools/lfm3_omni/refcommon.py @@ -0,0 +1,50 @@ +"""The reference runtime for LFM3-Omni-2B's checks: transformers with the checkpoint's own code, on cuda, mps or cpu.""" +import os +import subprocess +import time +from pathlib import Path + +import torch + + +def dtype_of(name): + """bf16 (the release) or f32 (the reference's own rounding noise).""" + return {"bf16": torch.bfloat16, "f32": torch.float32}[name] + + +def sync(dev): + """Waits for the device so host timers cover its work.""" + if dev == "cuda": + torch.cuda.synchronize() + elif dev == "mps": + torch.mps.synchronize() + + +def load(model_dir, dev, dtype): + """The model and processor, loaded on the CPU then moved (an MPS fp32 load straight to the device can fail).""" + from transformers import AutoModelForMultimodalLM, AutoProcessor + model = AutoModelForMultimodalLM.from_pretrained(Path(model_dir), dtype=dtype, device_map="cpu", + trust_remote_code=True) + model = model.eval().to(dev) + proc = AutoProcessor.from_pretrained(Path(model_dir), trust_remote_code=True) + return model, proc + + +def chat(proc, system, wav, dev): + """The processor's inputs for a system text and one spoken turn.""" + conv = [{"role": "system", "content": system}, {"role": "user", "content": [{"type": "audio", "audio": str(wav)}]}] + out = proc.apply_chat_template(conv, add_generation_prompt=True, tokenize=True, return_dict=True, + return_tensors="pt") + return out.to(dev) + + +def wait_idle(): + """Waits while $LFM3_OMNI_BUSY (a command) exits 0, so a shared machine's interactive work goes first.""" + cmd = os.environ.get("LFM3_OMNI_BUSY") + while cmd and subprocess.run(cmd, shell=True).returncode == 0: + time.sleep(5) + + +def identity_order(): + """Greedy codes in identity codebook order: the reveal order is the only randomness left under greedy audio.""" + torch.randperm = lambda n, *a, device=None, **k: torch.arange(n, device=device) diff --git a/tools/lfm3_omni/reference_depth.py b/tools/lfm3_omni/reference_depth.py new file mode 100644 index 0000000000..eb3d62b647 --- /dev/null +++ b/tools/lfm3_omni/reference_depth.py @@ -0,0 +1,68 @@ +"""The depth decoder on the reference reply's states: reference_depth.py MODEL DUMP prep|score [DEV] [DTYPE].""" +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +import refcommon as rc + +AUDIO_OUT = 125032 + + +def inputs(d): + """The state before each reference frame and its codes, as lfm3-omni-generate --depth-check reads them.""" + seq, sh, codes = np.load(d / "gen_seq.npy"), np.load(d / "step_hidden.npy"), np.load(d / "gen_codes.npy").T + slots = [j for j in range(1, len(seq)) if seq[j] == AUDIO_OUT][: len(codes)] + return np.stack([sh[j - 1] for j in slots]).astype(np.float32), codes[: len(slots)].astype(np.int32) + + +def prep(model_dir, d, dev, dtype): + """Writes depth_h.f32 and depth_codes.i32, and the reference's three passes of logits per frame.""" + h_all, c_all = inputs(d) + h_all.tofile(d / "depth_h.f32") + c_all.tofile(d / "depth_codes.i32") + model, _ = rc.load(model_dir, dev, dtype) + dd, out = model.model.depth_decoder, [] + with torch.no_grad(): + for h, c in zip(h_all, c_all): + ht = torch.from_numpy(h)[None].to(dev, dtype) + ct = torch.zeros((1, 12), dtype=torch.long, device=dev) + masked = torch.ones((1, 12), dtype=torch.bool, device=dev) + rows = [] + for p in range(3): + rows.append(dd(ht, ct, masked)[0].float().cpu().numpy()) + for q in range(4 * p, 4 * p + 4): + ct[0, q] = int(c[q]) + masked[0, q] = False + out.append(np.stack(rows)) + np.save(d / "ref_depth.npy", np.stack(out)) + + +def score(d): + """Greedy agreement and logits cosine of depth_check.f32 against the reference's passes.""" + ref = np.load(d / "ref_depth.npy") + tf = np.fromfile(d / "depth_check.f32", dtype=np.float32).reshape(ref.shape) + agree = total = 0 + cos = [] + for p in range(3): + for q in range(4 * p, 4 * p + 4): + a, b = ref[:, p, q], tf[:, p, q] + agree += int((a.argmax(-1) == b.argmax(-1)).sum()) + total += len(a) + for x, y in zip(a.astype(np.float64), b.astype(np.float64)): + cos.append(float(x @ y / (np.linalg.norm(x) * np.linalg.norm(y)))) + r = {"frames": int(ref.shape[0]), "argmax_agree": agree, "scored": total, "rate": round(agree / total, 4), + "logit_cos_min": round(min(cos), 6), "logit_cos_median": round(float(np.median(cos)), 6)} + print(json.dumps(r)) + json.dump(r, open(d / "depth_parity.json", "w")) + + +if __name__ == "__main__": + D = Path(sys.argv[2]) + if sys.argv[3] == "prep": + dev = sys.argv[4] if len(sys.argv) > 4 else "cuda" + prep(sys.argv[1], D, dev, rc.dtype_of(sys.argv[5] if len(sys.argv) > 5 else "bf16")) + else: + score(D) diff --git a/tools/lfm3_omni/reference_dump.py b/tools/lfm3_omni/reference_dump.py new file mode 100644 index 0000000000..122485a955 --- /dev/null +++ b/tools/lfm3_omni/reference_dump.py @@ -0,0 +1,95 @@ +"""The reference's stage dumps for the parity: reference_dump.py MODEL OUT DEV DTYPE SYSTEM_FILE WAV...""" +import json +import sys +import time +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch + +import refcommon as rc + + +def f32(t): + """A tensor as a float32 array on the host.""" + return t.detach().float().cpu().numpy() + + +def dump(model, proc, codec, system, wav, d, dev, dtype): + """One spoken turn: front end, encoder, prompt, prefill rows, a greedy reply with its states, depth and codec.""" + d.mkdir(parents=True, exist_ok=True) + inputs = rc.chat(proc, system, wav, dev) + ids = inputs["input_ids"][0] + np.save(d / "input_ids.npy", ids.cpu().numpy().astype(np.int32)) + feats = inputs["input_features"] + np.save(d / "mel.npy", f32(feats[0])) + with torch.no_grad(): + enc, _ = model.model.audio_encoder(feats.to(dtype), inputs["feature_attention_mask"].sum(-1)) + np.save(d / "enc.npy", f32(enc[0])) + np.save(d / "audio_emb.npy", f32(model.model.get_audio_features(feats, inputs["feature_attention_mask"]))) + out = model.model(**dict(inputs.items()), use_cache=False) + np.save(d / "prefill_hidden.npy", f32(out.last_hidden_state[0])) + np.save(d / "first_logits.npy", f32(model.lm_head(out.last_hidden_state[0, -1]))) + lm, steps = model.model.language_model, [] + orig = lm.forward + + def record(*a, **k): + o = orig(*a, **k) + if len(steps) < 400: + steps.append(o.last_hidden_state[0, -1].detach().float().cpu()) + return o + + lm.forward = record + rc.sync(dev) + t = time.perf_counter() + with torch.no_grad(): + g = model.generate(**inputs, generation_mode="interleaved", max_new_tokens=512, do_sample=False, + audio_temperature=0.0, audio_repetition_penalty=1.2) + rc.sync(dev) + total = time.perf_counter() - t + lm.forward = orig + seq = g.sequences[0, ids.shape[0]:g.sequence_lengths[0]].cpu().numpy().astype(np.int32) + n = int(g.audio_lengths[0]) + codes = g.audio_codes[0, :, :n] + np.save(d / "gen_seq.npy", seq) + np.save(d / "gen_codes.npy", codes.cpu().numpy().astype(np.int32)) + np.save(d / "step_hidden.npy", torch.stack(steps).numpy()) + seq.tofile(d / "gen_seq.i32") + codes.cpu().numpy().astype(np.int32).tofile(d / "gen_codes.i32") + keep = [x for x in seq.tolist() if x != model.config.audio_out_token_id] + text = proc.tokenizer.decode(keep, skip_special_tokens=True) + with torch.no_grad(): + z = torch.zeros((1, 12), dtype=torch.long, device=dev) + masked = torch.ones((1, 12), dtype=torch.bool, device=dev) + np.save(d / "depth_logits0.npy", f32(model.model.depth_decoder(out.last_hidden_state[0, -1:], z, masked)[0])) + if n: + q = codec.quantizer.decode(codes[None]).mT.to(codec.from_quantizer.weight.dtype) + e = codec.from_quantizer(q) + dec = codec.decoder + up = torch.nn.functional.interpolate(e.mT, e.shape[1] * dec.frame_stack, mode="nearest-exact").mT + np.save(d / "codec_lin.npy", f32(dec.lin(dec.lfm(inputs_embeds=up, use_cache=False).last_hidden_state)[0])) + wave = codec.decode(codes[None]).audio_values[0].squeeze(0).float().cpu().numpy() + np.save(d / "reply.npy", wave) + sf.write(d / "reply.wav", wave, 24000) + meta = {"wav": Path(wav).name, "text": text, "frames": n, "slots": len(seq), "gen_s": round(total, 3)} + json.dump(meta, open(d / "meta.json", "w")) + print(Path(wav).stem, n, len(seq), round(total, 2), text[:100], flush=True) + + +def main(): + """Every WAV's dumps into OUT/.""" + out, dev, dtype = Path(sys.argv[2]), sys.argv[3], rc.dtype_of(sys.argv[4]) + system = Path(sys.argv[5]).read_text().strip() + rc.identity_order() + model, proc = rc.load(sys.argv[1], dev, dtype) + codec = proc.audio_tokenizer.to(device=dev).eval() + out.mkdir(parents=True, exist_ok=True) + json.dump({"system": system, "device": dev, "dtype": sys.argv[4]}, open(out / "system.json", "w")) + for wav in sys.argv[6:]: + rc.wait_idle() + dump(model, proc, codec, system, wav, out / Path(wav).stem, dev, dtype) + + +if __name__ == "__main__": + main() diff --git a/tools/lfm3_omni/reference_speech.py b/tools/lfm3_omni/reference_speech.py new file mode 100644 index 0000000000..e1f34f2748 --- /dev/null +++ b/tools/lfm3_omni/reference_speech.py @@ -0,0 +1,83 @@ +"""The reference runtime's spoken replies and their timing: reference_speech.py MODEL DEV SYSTEM_FILE OUT_DIR WAV...""" +import json +import statistics +import sys +import time +from pathlib import Path + +import soundfile as sf +import torch +from transformers.generation.streamers import BaseStreamer + +import refcommon as rc + + +class First(BaseStreamer): + """Notes when the first speech frame is sampled.""" + + def __init__(self, dev): + self.dev, self.t = dev, None + + def put(self, value): + if self.t is None: + rc.sync(self.dev) + self.t = time.perf_counter() + + def end(self): + pass + + +def turn(model, proc, codec, system, wav, seed, dev): + """One reply with the card's audio settings; first audio = the first frame sampled plus its codec decode alone.""" + torch.manual_seed(seed) + first = First(dev) + rc.sync(dev) + t = time.perf_counter() + inputs = rc.chat(proc, system, wav, dev) + with torch.no_grad(): + out = model.generate(**inputs, generation_mode="interleaved", max_new_tokens=512, do_sample=False, + audio_temperature=0.4, audio_repetition_penalty=1.2, audio_streamer=first) + rc.sync(dev) + total = time.perf_counter() - t + n = int(out.audio_lengths[0]) + codes = out.audio_codes[:1, :, :n] + rc.sync(dev) + tc = time.perf_counter() + with torch.no_grad(): + codec.decode(codes[:, :, :1]) + rc.sync(dev) + one = time.perf_counter() - tc + seq = out.sequences[0, inputs["input_ids"].shape[1]:out.sequence_lengths[0]].tolist() + text = proc.tokenizer.decode([x for x in seq if x != model.config.audio_out_token_id], skip_special_tokens=True) + with torch.no_grad(): + audio = codec.decode(codes).audio_values[0].squeeze(0).float().cpu().numpy() if n else None + row = {"wav": Path(wav).name, "reply": text.strip(), "frames": n, "first_frame_s": round((first.t or t) - t, 4), + "first_audio_s": round((first.t or t) - t + one, 4), "total_s": round(total, 3), + "speech_s": round(n * 1920 / 24000, 2)} + return row, audio + + +def main(): + """A warm-up turn, then every WAV (seed 1000 + its index), replies saved as OUT_DIR/rNN.wav.""" + dev, system, out = sys.argv[2], Path(sys.argv[3]).read_text().strip(), Path(sys.argv[4]) + out.mkdir(parents=True, exist_ok=True) + model, proc = rc.load(sys.argv[1], dev, torch.bfloat16) + codec = proc.audio_tokenizer.to(device=dev).eval() + wavs = sys.argv[5:] + turn(model, proc, codec, system, wavs[0], 0, dev) + rows = [] + for i, w in enumerate(wavs): + rc.wait_idle() + row, audio = turn(model, proc, codec, system, w, 1000 + i, dev) + if audio is not None: + sf.write(out / f"r{i:02d}.wav", audio, 24000, subtype="PCM_16") + rows.append(row) + print(json.dumps(row), flush=True) + keys = ("first_frame_s", "first_audio_s", "total_s", "speech_s") + median = {k: round(statistics.median(r[k] for r in rows), 4) for k in keys} + print("MEDIAN", json.dumps(median)) + json.dump({"system": system, "median": median, "rows": rows}, open(out / "replies.json", "w"), indent=1) + + +if __name__ == "__main__": + main() diff --git a/tools/lfm3_omni/score_replies.py b/tools/lfm3_omni/score_replies.py new file mode 100644 index 0000000000..0cfcbcd9e9 --- /dev/null +++ b/tools/lfm3_omni/score_replies.py @@ -0,0 +1,78 @@ +"""UTMOS22 and speech-vs-text WER of saved replies: score_replies.py mos|asr MODEL DEV OUT.json NAME:REPLY_DIR ...""" +import glob +import json +import re +import statistics +import sys +import wave +from pathlib import Path + +import numpy as np +import torch + + +def load(path): + """A WAV as mono float32 and its rate.""" + with wave.open(str(path)) as w: + sr, ch, width, raw = w.getframerate(), w.getnchannels(), w.getsampwidth(), w.readframes(w.getnframes()) + a = np.frombuffer(raw, dtype={2: np.int16, 4: np.int32}[width]).astype(np.float32) / float(1 << (8 * width - 1)) + return a.reshape(-1, ch).mean(1), sr + + +def norm(t): + """Text for WER: lower case, numbers as words, punctuation dropped.""" + from num2words import num2words + t = t.lower().replace("%", " percent").replace("’", "'") + t = re.sub(r"(\d+):(\d+)", r"\1 by \2", t) + t = re.sub(r"\d+", lambda m: " " + num2words(int(m.group())) + " ", t) + t = re.sub(r"[^a-z' ]", " ", t.replace("-", " ").replace("‑", " ")) + return re.sub(r"\s+", " ", t).strip() + + +def mos_pass(out, specs): + """UTMOS22 strong (SpeechMOS, needs torchaudio) of every rNN.wav, into OUT.mos.json.""" + model = torch.hub.load("tarepan/SpeechMOS:v1.2.0", "utmos22_strong", trust_repo=True) + res = {} + for spec in specs: + name, d = spec.split(":", 1) + scores = [] + for w in sorted(glob.glob(str(Path(d) / "r*.wav"))): + a, sr = load(w) + with torch.no_grad(): + scores.append(round(model(torch.from_numpy(np.ascontiguousarray(a))[None], sr).item(), 4)) + res[name] = scores + print(name, "utmos median", statistics.median(scores), flush=True) + json.dump(res, open(out.with_suffix(".mos.json"), "w")) + + +def asr_pass(model_dir, dev, out, specs): + """The checkpoint's own ASR of every reply against the reply's text (replies.json), merged with the UTMOS pass.""" + import jiwer + sys.path.insert(0, str(Path(__file__).parent)) + import refcommon as rc + model, proc = rc.load(model_dir, dev, torch.bfloat16) + mos, res = json.load(open(out.with_suffix(".mos.json"))), {} + for spec in specs: + name, d = spec.split(":", 1) + rows = json.load(open(Path(d) / "replies.json"))["rows"] + hyps = [] + for w in sorted(glob.glob(str(Path(d) / "r*.wav"))): + inputs = rc.chat(proc, "Perform ASR in English.", w, dev) + with torch.no_grad(): + o = model.generate(**inputs, max_new_tokens=160, generation_mode="text", do_sample=False) + ids = o.sequences[0, inputs["input_ids"].shape[1]:o.sequence_lengths[0]] + hyps.append(proc.decode(ids, skip_special_tokens=True).strip()) + refs, scores = [r["reply"] for r in rows], mos[name] + wer = jiwer.wer([norm(r) for r in refs], [norm(h) for h in hyps]) + res[name] = {"n": len(hyps), "utmos_median": round(statistics.median(scores), 3), + "utmos_min": round(min(scores), 3), "speech_vs_text_wer": round(wer, 4), + "per_reply": [{"text": r, "heard": h, "utmos": s} for r, h, s in zip(refs, hyps, scores)]} + print(name, json.dumps({k: v for k, v in res[name].items() if k != "per_reply"}), flush=True) + json.dump(res, open(out, "w"), indent=1) + + +if __name__ == "__main__": + if sys.argv[1] == "mos": + mos_pass(Path(sys.argv[4]), sys.argv[5:]) + else: + asr_pass(sys.argv[2], sys.argv[3], Path(sys.argv[4]), sys.argv[5:]) diff --git a/tools/lfm3_omni/speech_chat_client.py b/tools/lfm3_omni/speech_chat_client.py new file mode 100644 index 0000000000..3175e8fcab --- /dev/null +++ b/tools/lfm3_omni/speech_chat_client.py @@ -0,0 +1,88 @@ +"""Streams spoken turns through POST /v1/audio/chat and times them: speech_chat_client.py HOST:PORT OUT_DIR WAV...""" +import argparse +import base64 +import http.client +import json +import os +import statistics +import subprocess +import time +import wave +from pathlib import Path + + +def busy(): + """True while $LFM3_OMNI_BUSY (a command) exits 0.""" + cmd = os.environ.get("LFM3_OMNI_BUSY") + return bool(cmd) and subprocess.run(cmd, shell=True).returncode == 0 + + +def turn(host, port, body): + """One request read as it arrives: the reply's text, its 16-bit 24 kHz PCM, client-side times and the done event.""" + c = http.client.HTTPConnection(host, port, timeout=600) + t0 = time.perf_counter() + c.request("POST", "/v1/audio/chat", body=json.dumps(body), headers={"Content-Type": "application/json"}) + r = c.getresponse() + if r.status != 200: + raise RuntimeError(f"HTTP {r.status}: {r.read()[:200]!r}") + first_text = first_audio = None + pcm, text, done = bytearray(), "", {} + while line := r.readline(): + line = line.strip() + if not line.startswith(b"data: ") or line == b"data: [DONE]": + continue + ev = json.loads(line[6:]) + if ev["type"] == "text": + first_text = first_text or time.perf_counter() + text += ev["delta"] + elif ev["type"] == "audio": + first_audio = first_audio or time.perf_counter() + pcm += base64.b64decode(ev["data"]) + elif ev["type"] == "done": + done = ev + total = time.perf_counter() - t0 + return text.strip(), bytes(pcm), { + "first_text_s": round((first_text or t0) - t0, 3), "first_audio_s": round((first_audio or t0) - t0, 3), + "total_s": round(total, 3), "speech_s": round(len(pcm) / 48000, 2), + "server_first_audio_ms": done.get("first_audio_ms"), "frames": done.get("frames")} + + +def main(): + """Every WAV in order (seed + its index), each reply saved as OUT_DIR/rNN.wav, medians in OUT_DIR/replies.json.""" + ap = argparse.ArgumentParser() + ap.add_argument("server") + ap.add_argument("out") + ap.add_argument("wavs", nargs="+") + ap.add_argument("--system-file") + ap.add_argument("--audio-temperature", type=float, default=0.4) + ap.add_argument("--audio-repetition-penalty", type=float, default=1.2) + ap.add_argument("--seed", type=int, default=1000) + a = ap.parse_args() + host, port = a.server.rsplit(":", 1) + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + system = Path(a.system_file).read_text().strip() if a.system_file else None + rows = [] + for i, w in enumerate(a.wavs): + while busy(): + time.sleep(5) + body = {"audio": base64.b64encode(Path(w).read_bytes()).decode(), "stream": True, "seed": a.seed + i, + "audio_temperature": a.audio_temperature, "audio_repetition_penalty": a.audio_repetition_penalty} + if system is not None: + body["system"] = system + text, pcm, row = turn(host, int(port), body) + with wave.open(str(out / f"r{i:02d}.wav"), "wb") as f: + f.setnchannels(1) + f.setsampwidth(2) + f.setframerate(24000) + f.writeframes(pcm) + rows.append({"wav": Path(w).name, "reply": text} | row) + print(json.dumps(rows[-1]), flush=True) + keys = ("first_text_s", "first_audio_s", "total_s", "speech_s") + median = {k: round(statistics.median(r[k] for r in rows), 3) for k in keys} + print("MEDIAN", json.dumps(median)) + json.dump({"system": system, "median": median, "rows": rows}, open(out / "replies.json", "w"), indent=1) + + +if __name__ == "__main__": + main() diff --git a/tools/lfm3_omni/summarize_parity.py b/tools/lfm3_omni/summarize_parity.py new file mode 100644 index 0000000000..002316ff9b --- /dev/null +++ b/tools/lfm3_omni/summarize_parity.py @@ -0,0 +1,76 @@ +"""The stage table over parity.sh's output: python summarize_parity.py REF_DIR OUT_DIR [NAME] (writes summary.json).""" +import json +import re +import statistics +import sys +from pathlib import Path + +import numpy as np + +FORCED = re.compile(r"FORCED text (\d+)/(\d+) codes (\d+)/(\d+) frames (\d+)/(\d+)") + + +def cos(a, b): + """Cosine in float64.""" + a, b = a.astype(np.float64).ravel(), b.astype(np.float64).ravel() + return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30)) + + +def pct(a, b): + """a/b with its percentage.""" + return f"{a}/{b} ({100 * a / max(1, b):.1f}%)" + + +def summarize(r, t, name): + """Per-stage minima and medians over every utterance with a parity.json.""" + keys = ("mel_cos", "mel_max", "enc", "proj", "ids", "pre_med", "pre_min", "lin", "snr") + rows = {k: [] for k in keys} + forced, depth, dmed, dmin = [0] * 6, [0, 0], [], [] + us = sorted(p.name for p in t.iterdir() if (p / "parity.json").exists()) + for u in us: + p = json.load(open(t / u / "parity.json")) + rows["mel_cos"].append(p["mel"]["cos"]) + rows["mel_max"].append(p["mel"]["max_abs"]) + rows["enc"].append(p["encoder"]["cos"]) + rows["proj"].append(p["projector"]["cos"]) + rows["ids"].append(p["prompt_ids_equal"]) + ph = np.load(r / u / "prefill_hidden.npy") + tph = np.fromfile(t / u / "prefill_hidden.f32", dtype=np.float32).reshape(-1, ph.shape[1])[: len(ph)] + rcos = [cos(ph[i], tph[i]) for i in range(len(ph))] + rows["pre_med"].append(statistics.median(rcos)) + rows["pre_min"].append(min(rcos)) + if "codec_lin" in p: + rows["lin"].append(p["codec_lin"]["cos"]) + rows["snr"].append(p["codec_wave"]["snr_db"]) + m = FORCED.search(open(t / u / "forced.log").read()) + if m: + forced = [x + int(m.group(i + 1)) for i, x in enumerate(forced)] + dp = t / u / "depth_parity.json" + if dp.exists() and dp.stat().st_size: + d = json.load(open(dp)) + depth = [depth[0] + d["argmax_agree"], depth[1] + d["scored"]] + dmed.append(d["logit_cos_median"]) + dmin.append(d["logit_cos_min"]) + med = statistics.median + return { + "name": name, "utterances": len(us), + "log-mel cos (min), max abs (max)": [round(min(rows["mel_cos"]), 6), max(rows["mel_max"])], + "speech encoder cos (min)": round(min(rows["enc"]), 6), + "projector cos (min)": round(min(rows["proj"]), 6), + "prompt ids identical": f'{sum(rows["ids"])}/{len(us)}', + "prefill row cos (median of medians, min of medians, min row)": + [round(med(rows["pre_med"]), 6), round(min(rows["pre_med"]), 6), round(min(rows["pre_min"]), 6)], + "teacher-forced text top-1": pct(forced[0], forced[1]), + "teacher-forced codes": pct(forced[2], forced[3]), + "depth on the reference's states: greedy agreement": pct(depth[0], depth[1]), + "depth logits cos (median, min)": [round(med(dmed), 6), round(min(dmin), 6)] if dmed else None, + "codec lin cos (min)": round(min(rows["lin"]), 6) if rows["lin"] else None, + "codec SNR dB (median, min)": [round(med(rows["snr"]), 2), round(min(rows["snr"]), 2)] if rows["snr"] else None, + } + + +if __name__ == "__main__": + R, T = Path(sys.argv[1]), Path(sys.argv[2]) + out = summarize(R, T, sys.argv[3] if len(sys.argv) > 3 else T.name) + print(json.dumps(out, indent=1)) + json.dump(out, open(T / "summary.json", "w"), indent=1) diff --git a/tools/lfm3_omni/utterances.txt b/tools/lfm3_omni/utterances.txt new file mode 100644 index 0000000000..8c1f8bf186 --- /dev/null +++ b/tools/lfm3_omni/utterances.txt @@ -0,0 +1,16 @@ +Hey, can you give me an idea for a short horror film set in a lighthouse? +How long should a teaser trailer be? +Make a five second clip of a red fox running through snow at dusk. +Change the aspect ratio to nine by sixteen and render it again. +Cast Tanaka as the lead and give her a softer voice for the night scenes. +What is the difference between a dolly zoom and a crash zoom? +Add a slow push in on the lighthouse keeper as the storm hits. +Please export the timeline as a ProRes file at twenty four frames per second. +Can you queue three variations with seeds forty two, seven and nineteen? +Remind me what the second act break is in the Last Tram script. +Turn the music down by six decibels under the dialogue. +I want the color grade to feel like an eighties horror movie, warm and grainy. +Stop. Actually, never mind, go back to the previous shot. +Give me a one sentence logline for a heist movie on a moving train. +Lower the temperature of the key light and add some haze. +How many frames are in a ten second clip at twenty four frames per second?