diff --git a/zig/build/dsv41_metal.zig b/zig/build/dsv41_metal.zig index 1ee5d6bd14..54ed7bceff 100644 --- a/zig/build/dsv41_metal.zig +++ b/zig/build/dsv41_metal.zig @@ -4,7 +4,7 @@ const std = @import("std"); /// The package over the family's neutral host module; `metal` and `sources` are the engine's own modules. pub fn package(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module, optimize: std.builtin.OptimizeMode) struct { pkg: *std.Build.Module, host: *std.Build.Module } { const host = @import("dsv41_shared.zig").moduleAt(b, b.graph.host, optimize, "neutral.zig"); - // the lanes and the engine API the backend's modules import + // the lanes, the engine API and the served engine (serve_lanes) the backend's target and pass plug into @import("dsv41_draft.zig").addImportsAt(b, host, b.graph.host, optimize); const pkg = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/metal/metal.zig"), @@ -13,9 +13,24 @@ pub fn package(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Modu .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "dsv41", .module = host } }, }); + pkg.addImport("core", host.import_table.get("core").?); // the picture decoder image input reads return .{ .pkg = pkg, .host = host }; } +/// The native server's engine over the package: metal/entry.zig, the family's own serving modules beside it. +pub fn engine(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const p = package(b, metal, sources, optimize); + const m = b.createModule(.{ + .root_source_file = b.path("zig/src/families/deepseek_v41/metal/entry.zig"), + .target = b.graph.host, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "dsv41_metal", .module = p.pkg }, .{ .name = "dsv41", .module = p.host }, .{ .name = "metal", .module = metal } }, + }); + for ([_][]const u8{ "engine_api", "core", "tokenizer", "dsv41_serve", "lanes", "json" }) |name| m.addImport(name, p.host.import_table.get(name).?); + return m; +} + /// `test-dsv41-metal`: package, layer and forward checks on this GPU; `test-dsv41-metal-real`: a TF_DSV41_REAL layer. pub fn steps(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module, test_step: *std.Build.Step) void { const p = package(b, metal, sources, .Debug); @@ -26,13 +41,39 @@ pub fn steps(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module .link_libc = true, .imports = &.{ .{ .name = "dsv41_metal", .module = p.pkg }, .{ .name = "metal", .module = metal }, .{ .name = "dsv41", .module = p.host } }, }); + // the package's own checks run with the default test step; its device suites run from their named steps + test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = p.pkg })).step); const step = b.step("test-dsv41-metal", "DeepSeek-V4.1 Metal layer and forward: float64 stages, rows at once == one at a time"); - const pkg_run = &b.addRunArtifact(b.addTest(.{ .root_module = p.pkg })).step; - step.dependOn(pkg_run); + // serve_lanes over the CPU twin, checked here so no CUDA step has to depend on the neutral served engine + const lanes_test = b.createModule(.{ + .root_source_file = b.path("zig/src/families/deepseek_v41/serve_lanes_test.zig"), + .target = b.graph.host, + .optimize = .Debug, + .link_libc = true, + .imports = &.{ .{ .name = "core", .module = p.host.import_table.get("core").? }, .{ .name = "engine_api", .module = p.host.import_table.get("engine_api").? }, .{ .name = "lanes", .module = p.host.import_table.get("lanes").? }, .{ .name = "tokenizer", .module = p.host.import_table.get("tokenizer").? } }, + }); + test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = lanes_test })).step); const run = b.addRunArtifact(b.addTest(.{ .root_module = tests })); run.has_side_effects = true; // a device run: never cached step.dependOn(&run.step); - test_step.dependOn(pkg_run); // the package's host checks only; the device run stays in its own step + // the served engine over the tiny model: LaneHost replies, drafts, the prompt cache, seeds, cancels + const serve = b.createModule(.{ + .root_source_file = b.path("zig/tests/dsv41_metal/serve_test.zig"), + .target = b.graph.host, + .optimize = .Debug, + .link_libc = true, + .imports = &.{ .{ .name = "dsv41_metal", .module = p.pkg }, .{ .name = "metal", .module = metal }, .{ .name = "dsv41", .module = p.host } }, + }); + serve.addImport("dsv41_serve_fixtures", p.host.import_table.get("dsv41_serve").?.import_table.get("dsv41_serve_fixtures").?); + const serve_run = b.addRunArtifact(b.addTest(.{ .root_module = serve })); + serve_run.has_side_effects = true; + step.dependOn(&serve_run.step); + const serve_step = b.step("test-dsv41-metal-serve", "DeepSeek-V4.1 Metal served through LaneHost: concurrent, drafted, cached, seeded, cancelled"); + serve_step.dependOn(&serve_run.step); + // the release-checkpoint steps (TF_DSV41_REAL) stay opt-in, so the umbrella holds the default device suites + const all = b.step("test-dsv41-metal-all", "DeepSeek-V4.1 Metal: every device test step on this GPU (layer, served)"); + all.dependOn(step); + all.dependOn(serve_step); const fast = package(b, metal, sources, .ReleaseSafe); const real = b.createModule(.{ .root_source_file = b.path("zig/tests/dsv41_metal/real_test.zig"), @@ -66,4 +107,7 @@ pub fn steps(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module const ptime_run = b.addRunArtifact(b.addTest(.{ .root_module = ptime })); ptime_run.has_side_effects = true; b.step("test-dsv41-metal-prompt-time", "DeepSeek-V4.1 Metal: ms a stage of one release layer's prompt chunks (TF_DSV41_REAL)").dependOn(&ptime_run.step); + // tensorfold-dsv41 over Metal: the native server with this family's lane engine (the CUDA pair's macOS twin) + const exe = @import("dsv41_draft.zig").serverOn(b, b.graph.host, .ReleaseSafe, engine(b, metal, sources, .ReleaseFast)); + b.step("dsv41-metal-serve", "Install tensorfold-dsv41 over Metal (DeepSeek-V4.1 release checkpoints on this Mac)").dependOn(&b.addInstallArtifact(exe, .{}).step); } diff --git a/zig/kernels/metal/deepseek_v41/sources.zig b/zig/kernels/metal/deepseek_v41/sources.zig index d69ea64ddd..f8275519cb 100644 --- a/zig/kernels/metal/deepseek_v41/sources.zig +++ b/zig/kernels/metal/deepseek_v41/sources.zig @@ -21,6 +21,8 @@ pub const fold = @embedFile("fold.metal"); pub const compress = @embedFile("compress.metal"); /// DeepSeek-V4.1 router, indexer, compressor and head matvecs: lane chains over K, one butterfly, any row count. pub const mv = @embedFile("mv.metal"); +/// DeepSeek V4.1's vision tower: epilogues, RoPE, attention, SwiGLU, the aligner's unshuffle and the stream merge. +pub const vision = @embedFile("vision.metal"); /// DeepSeek-V4.1 candidate blocks: block maxima, the newest block pinned, an exact top-2048 as a byte mask. pub const cand = @embedFile("cand.metal"); /// DeepSeek V4.1's indexer scores: FP4 query and key dots in one stated order, causal and block masks. diff --git a/zig/kernels/metal/deepseek_v41/vision.metal b/zig/kernels/metal/deepseek_v41/vision.metal new file mode 100644 index 0000000000..1323795f61 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/vision.metal @@ -0,0 +1,181 @@ +// DeepSeek V4.1's vision tower: epilogues, 2D RoPE, full attention, SwiGLU, the aligner's unshuffle, the merge. +#include +using namespace metal; +#pragma clang fp contract(off) + +#ifndef TF_CTL +#define TF_CTL 0 // planted faults: 1 bias dropped, 2 sin/cos swapped, 3 max not subtracted, 4 gate/up swapped, 5 a stride off, 6 a delimiter late +#endif + +inline float tf_vis_bf(float v) { + return float(bfloat(v)); +} + +// erf to fp32 precision (W. J. Cody's rational approximations, as libm's erff). +inline float tf_vis_erf(float x) { + const float ax = fabs(x); + if (ax < 0.84375f) { + const float z = x * x; + const float r = 1.28379167e-01f + z * (-3.25042108e-01f + z * (-2.84817495e-02f + z * (-5.77027029e-03f + z * -2.37630166e-05f))); + const float s = 1.0f + z * (3.97917233e-01f + z * (6.50222499e-02f + z * (5.08130628e-03f + z * (1.32494738e-04f + z * -3.96022827e-06f)))); + return x + x * (r / s); + } + if (ax < 1.25f) { + const float s = ax - 1.0f; + const float P = -2.36211857e-03f + s * (4.14856106e-01f + s * (-3.72207481e-01f + s * (3.18346620e-01f + s * (-1.10894694e-01f + s * (3.54783031e-02f + s * -2.16637559e-03f))))); + const float Q = 1.0f + s * (1.06420882e-01f + s * (5.40397942e-01f + s * (7.18286547e-02f + s * (1.26171216e-01f + s * (1.36370836e-02f + s * 1.19844998e-02f))))); + const float e = 8.45062911e-01f + P / Q; + return x >= 0 ? e : -e; + } + if (ax >= 6.0f) return x >= 0 ? 1.0f : -1.0f; + const float s = 1.0f / (ax * ax); + float R, S; + if (ax < 1.0f / 0.35f) { + R = -9.86494403e-03f + s * (-6.93858570e-01f + s * (-1.05586262e+01f + s * (-6.23753357e+01f + s * (-1.62396660e+02f + s * (-1.84605087e+02f + s * (-8.12874374e+01f + s * -9.81432934e+00f)))))); + S = 1.0f + s * (1.96512714e+01f + s * (1.37657754e+02f + s * (4.34565887e+02f + s * (6.45387268e+02f + s * (4.29008148e+02f + s * (1.08635002e+02f + s * (6.57024977e+00f + s * -6.04244141e-02f))))))); + } else { + R = -9.86494310e-03f + s * (-7.99283242e-01f + s * (-1.77579556e+01f + s * (-1.60636383e+02f + s * (-6.37566467e+02f + s * (-1.02509509e+03f + s * -4.83519196e+02f))))); + S = 1.0f + s * (3.03380604e+01f + s * (3.25792511e+02f + s * (1.53672961e+03f + s * (3.19985820e+03f + s * (2.55305029e+03f + s * (4.74528534e+02f + s * -2.24409515e+01f)))))); + } + const float zh = as_type(as_type(ax) & 0xfffff000u); + const float r = metal::precise::exp(-zh * zh - 0.5625f) * metal::precise::exp((zh - ax) * (zh + ax) + R / S); + const float e = 1.0f - r / ax; + return x >= 0 ? e : -e; +} + +struct TfVisEpiArgs { + uint rows, n, gelu, has_bias, has_res; +}; + +// A linear's fp32 sums to bf16: + bias, one rounding; then GELU (rounded) and the residual add (rounded) when asked. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_vis_epi(const device float* ACC [[buffer(0)]], const device bfloat* BIAS [[buffer(1)]], + const device bfloat* RES [[buffer(2)]], device bfloat* Y [[buffer(3)]], + constant TfVisEpiArgs& a [[buffer(4)]], uint2 g [[thread_position_in_grid]]) { + if (g.x >= a.n || g.y >= a.rows) return; + const ulong at = ulong(g.y) * a.n + g.x; + float v = ACC[at]; + if (a.has_bias != 0 && TF_CTL != 1) v = v + float(BIAS[g.x]); // control 1: the bias dropped + float z = tf_vis_bf(v); + if (a.gelu != 0) z = tf_vis_bf(0.5f * z * (1.0f + tf_vis_erf(z * 0.70710678118654752f))); + if (a.has_res != 0) z = tf_vis_bf(float(RES[at]) + z); + Y[at] = bfloat(z); +} + +struct TfVisQkvArgs { + uint rows, heads, hd; +}; + +// The qkv linear's sums [rows][3][heads][hd] to bf16 with the bias, q and k rotated by their patch's (cos, sin) pairs. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_vis_qkv(const device float* ACC [[buffer(0)]], const device bfloat* BIAS [[buffer(1)]], + const device float2* CS [[buffer(2)]], device bfloat* QKV [[buffer(3)]], + constant TfVisQkvArgs& a [[buffer(4)]], uint3 g [[thread_position_in_grid]]) { + const uint half_hd = a.hd / 2; + if (g.x >= half_hd || g.y >= a.rows || g.z >= 3 * a.heads) return; + const uint part = g.z / a.heads, head = g.z % a.heads; + const uint D = a.heads * a.hd; + const uint c1 = part * D + head * a.hd + g.x; + const ulong at = ulong(g.y) * 3 * D + c1; + const float x1 = tf_vis_bf(ACC[at] + float(BIAS[c1])); + const float x2 = tf_vis_bf(ACC[at + half_hd] + float(BIAS[c1 + half_hd])); + if (part == 2) { + QKV[at] = bfloat(x1); + QKV[at + half_hd] = bfloat(x2); + return; + } + const float2 raw = CS[ulong(g.y) * half_hd + g.x]; + const float2 cs = TF_CTL == 2 ? float2(raw.y, raw.x) : raw; // control 2: sin and cos swapped + const float p = x1 * cs.x, q = x2 * cs.y, r = x2 * cs.x, s = x1 * cs.y; + QKV[at] = bfloat(p - q); + QKV[at + half_hd] = bfloat(r + s); +} + +struct TfVisAttnArgs { + uint n, heads; + float scale; +}; + +// One simdgroup a (query, head): 32 keys a step, online softmax in fp32, values by the shuffled weights. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_vis_attn( + const device bfloat* QKV [[buffer(0)]], device bfloat* O [[buffer(1)]], constant TfVisAttnArgs& a [[buffer(2)]], + uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { + const uint D = a.heads * 64u; + const ulong stride = 3ul * D; + const device bfloat* qp = QKV + ulong(tg.x) * stride + tg.y * 64u; + float q[64]; + for (uint d = 0; d < 64u; ++d) q[d] = float(qp[d]); + float m = -3.0e38f, l = 0.0f, acc0 = 0.0f, acc1 = 0.0f; + for (uint j0 = 0; j0 < a.n; j0 += 32u) { + const uint j = j0 + lane; + float s = -3.0e38f; + if (j < a.n) { + const device bfloat* kp = QKV + ulong(j) * stride + D + tg.y * 64u; + float t = 0.0f; + for (uint d = 0; d < 64u; ++d) t = t + q[d] * float(kp[d]); + s = t * a.scale; + } + const float mn = max(m, simd_max(s)); + const float alpha = metal::precise::exp(m - mn); + const float p = j < a.n ? metal::precise::exp(s - (TF_CTL == 3 ? 0.0f : mn)) : 0.0f; // control 3: the max not subtracted + l = l * alpha + simd_sum(p); + acc0 = acc0 * alpha; + acc1 = acc1 * alpha; + const uint cnt = min(32u, a.n - j0); + for (uint jj = 0; jj < cnt; ++jj) { + const float pj = simd_shuffle(p, ushort(jj)); + const device bfloat* vp = QKV + ulong(j0 + jj) * stride + 2u * D + tg.y * 64u + 2u * lane; + acc0 = acc0 + pj * float(vp[0]); + acc1 = acc1 + pj * float(vp[1]); + } + m = mn; + } + device bfloat* op = O + ulong(tg.x) * D + tg.y * 64u + 2u * lane; + op[0] = bfloat(acc0 / l); + op[1] = bfloat(acc1 / l); +} + +struct TfVisSwiArgs { + uint rows, inter; +}; + +// The MLP's w1 sums [rows][2][inter] (gate, then up), each rounded to bf16: silu(gate) rounded, times up, rounded. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_vis_swiglu(const device float* ACC [[buffer(0)]], device bfloat* Y [[buffer(1)]], + constant TfVisSwiArgs& a [[buffer(2)]], uint2 g [[thread_position_in_grid]]) { + if (g.x >= a.inter || g.y >= a.rows) return; + const ulong at = ulong(g.y) * 2ul * a.inter + g.x; + const float gt = tf_vis_bf(ACC[at + (TF_CTL == 4 ? a.inter : 0u)]); // control 4: gate and up swapped + const float up = tf_vis_bf(ACC[at + (TF_CTL == 4 ? 0u : a.inter)]); + const float s = tf_vis_bf(gt / (1.0f + metal::precise::exp(-gt))); + Y[ulong(g.y) * a.inter + g.x] = bfloat(s * up); +} + +struct TfVisUnshufArgs { + uint nh, nw, lh, lw, dim, r; +}; + +// The aligner's F.unfold of the zero-padded patch grid: out [lh * lw][dim * r * r], column c * r * r + ky * r + kx. +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_vis_unshuffle(const device bfloat* X [[buffer(0)]], device bfloat* Y [[buffer(1)]], + constant TfVisUnshufArgs& a [[buffer(2)]], uint2 g [[thread_position_in_grid]]) { + const uint rr = a.r * a.r; + if (g.x >= a.dim * rr || g.y >= a.lh * a.lw) return; + const uint c = g.x / rr, k = g.x % rr; + const uint rd = TF_CTL == 5 ? a.r + 1u : a.r; // control 5: the unfold's stride off by one + const uint py = (g.y / a.lw) * a.r + k / rd, px = (g.y % a.lw) * a.r + k % rd; + const bool in = py < a.nh && px < a.nw; + Y[ulong(g.y) * a.dim * rr + g.x] = in ? X[(ulong(py) * a.nw + px) * a.dim + c] : bfloat(0.0f); +} + +struct TfVisMergeArgs { + uint dim, rows; +}; + +// Stream rows [rows][4][dim] take an aligner row (SRC >= 0) or a delimiter (-1 start, -2 end, -3 newline). +[[kernel, max_total_threads_per_threadgroup(32)]] void tf_dsv41_vis_merge(const device bfloat* FEAT [[buffer(0)]], const device bfloat* DELIM [[buffer(1)]], + const device int* SRC [[buffer(2)]], device bfloat* Y [[buffer(3)]], + constant TfVisMergeArgs& a [[buffer(4)]], uint2 g [[thread_position_in_grid]]) { + if (g.x >= a.dim || g.y >= a.rows) return; + const int s = SRC[g.y]; + if (s <= -4) return; + const uint di = TF_CTL == 6 ? (uint(-1 - s) + 1u) % 3u : uint(-1 - s); // control 6: one delimiter a row late + const bfloat v = s >= 0 ? FEAT[ulong(s) * a.dim + g.x] : DELIM[ulong(di) * a.dim + g.x]; + for (uint k = 0; k < 4u; ++k) Y[(ulong(g.y) * 4ul + k) * a.dim + g.x] = v; +} diff --git a/zig/src/core/contract/dsv41_metal_records.zig b/zig/src/core/contract/dsv41_metal_records.zig index b24ae23ad2..d4f6610714 100644 --- a/zig/src/core/contract/dsv41_metal_records.zig +++ b/zig/src/core/contract/dsv41_metal_records.zig @@ -7,7 +7,7 @@ const head = @import("dsv41_fp8_records.zig").records ++ @import("dsv41_fp8_mma_ @import("dsv41_engram_records.zig").records ++ @import("dsv41_glue_records.zig").records ++ @import("dsv41_compress_records.zig").records ++ @import("dsv41_cand_records.zig").records ++ @import("dsv41_xmma_records.zig").records ++ @import("dsv41_mv_records.zig").records ++ - @import("dsv41_hc_records.zig").records; + @import("dsv41_vision_records.zig").records ++ @import("dsv41_hc_records.zig").records; pub const records = head ++ @import("dsv41_fold_records.zig").records ++ @import("dsv41_glue_tmp_records.zig").records ++ @import("dsv41_fused_records.zig").records ++ @import("dsv41_route_records.zig").records; diff --git a/zig/src/core/contract/dsv41_vision_records.zig b/zig/src/core/contract/dsv41_vision_records.zig new file mode 100644 index 0000000000..a6eab020a5 --- /dev/null +++ b/zig/src/core/contract/dsv41_vision_records.zig @@ -0,0 +1,15 @@ +//! DeepSeek V4.1's vision launches on Metal: each reserves its thread count, so none depends on M1/M2. +const std = @import("std"); +const M = @import("limits.zig").MetalLaunch; +pub const records = [_]M{ + .{ .name = "dsv41 vision epilogue tf_dsv41_vis_epi", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 vision qkv rope tf_dsv41_vis_qkv", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 vision attention tf_dsv41_vis_attn", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 vision swiglu tf_dsv41_vis_swiglu", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 vision unshuffle tf_dsv41_vis_unshuffle", .threads = 32, .tg_memory = 0, .reserved = 32 }, + .{ .name = "dsv41 vision merge tf_dsv41_vis_merge", .threads = 32, .tg_memory = 0, .reserved = 32 }, +}; + +test "the DeepSeek V4.1 vision 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/families/deepseek_v41/draft/iface.zig b/zig/src/families/deepseek_v41/draft/iface.zig index 39aa01f471..f8bfae834e 100644 --- a/zig/src/families/deepseek_v41/draft/iface.zig +++ b/zig/src/families/deepseek_v41/draft/iface.zig @@ -46,8 +46,15 @@ pub const Target = struct { replays: ?*const fn (ptr: *anyopaque, list: []const Final) anyerror!void = null, /// Engram reads ahead of the next window (TF_DSV41_ENGRAM_WARM); reads only, window bits are unaffected. warm: ?*const fn (ptr: *anyopaque, asks: []const Warm) void = null, + /// After `prefill`: the prompt rows the slot took from kept state (its first row computed is this one). + resumed: ?*const fn (ptr: *anyopaque, slot: u32) u64 = null, }; + /// The prompt rows the slot's last `prefill` took from kept state (0: none, or a target that keeps none). + pub fn resumedRows(t: Target, slot: u32) u64 { + return if (t.vtable.resumed) |f| f(t.ptr, slot) else 0; + } + /// A slot's whole prompt whose rows but the last are in (its `finish` / `tail` follows). pub const Final = struct { slot: u32, ids: []const u32 }; diff --git a/zig/src/families/deepseek_v41/draft/lanes.zig b/zig/src/families/deepseek_v41/draft/lanes.zig index 35dc43ebbb..124c1bed92 100644 --- a/zig/src/families/deepseek_v41/draft/lanes.zig +++ b/zig/src/families/deepseek_v41/draft/lanes.zig @@ -189,6 +189,7 @@ pub const Lanes = struct { const ids = s.prompt(); x.target.admit(slot, s.max_new); st.first = try x.target.prefill(slot, ids, s.sampling, ids.len); + s.cached = @intCast(x.target.resumedRows(slot)); st.len = ids.len; _ = x.scratch.reset(.retain_capacity); // the taps hold the prompt's last rows (after a verify-tail prompt only that row), at their positions diff --git a/zig/src/families/deepseek_v41/metal/boot.zig b/zig/src/families/deepseek_v41/metal/boot.zig new file mode 100644 index 0000000000..88e3d53e2b --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/boot.zig @@ -0,0 +1,159 @@ +//! A release checkpoint folder on this Mac: config, kernels, mapped weights, DSpark, Engram rows; the served model. +const std = @import("std"); +const mtl = @import("metal"); +const host = @import("dsv41"); +const pipelines = @import("pipelines.zig"); +const model_weights = @import("model_weights.zig"); +const dspark = @import("dspark.zig"); +const rope = @import("rope.zig"); +const scratch = @import("scratch.zig"); +const forward = @import("forward.zig"); +const serve = @import("serve.zig"); +const target_ = @import("target.zig"); + +const Config = host.config.Config; +const eh = host.engram_host; +const er = host.engram_rows; +const Allocator = std.mem.Allocator; + +/// The files a boot reads beside the shards: the Engram token map (here or under TF_DSV41_ASSETS). +pub const token_map = "engram-token-map.bin"; + +pub const Options = struct { + slots: u32 = 4, + /// positions a slot holds (null: 262,144, at most the checkpoint's) + context: ?u64 = null, + drafts: bool = true, + cache_bytes: u64 = 4 << 30, + /// a checkpoint of the release's format at other shapes (TF_DSV41_SMALL=1: the smoke test's tiny model) + small: bool = false, + /// Engram row shards (null: TF_DSV41_ENGRAM_DIR, else the config's engram_table_dir beside the checkpoint) + engram_dir: ?[]const u8 = null, +}; + +/// Everything a served model borrows, owned here and closed after the model's target and pass. +pub const Boot = struct { + a: Allocator, + c: Config, + device: mtl.Device, + queue: mtl.Queue, + p: pipelines.Pipelines, + w: *model_weights.Model, + ds: ?dspark.Weights = null, + tables: rope.Tables, + s: scratch.Scratch, + map_bytes: []align(4) u8 = &.{}, + lbuf: [eh.max_layers]u32 = undefined, + eng: eh.Host = undefined, + rows: ?*er.Rows = null, + + /// The checkpoint in `dir` loaded and served by a `serve.Metal` (the caller's `serve_lanes.Served` closes it). + pub fn open(a: Allocator, io: std.Io, dir: []const u8, o: Options) !*serve.Metal { + const b = try a.create(Boot); + errdefer a.destroy(b); + b.a = a; + b.ds = null; + b.map_bytes = &.{}; + b.rows = null; + b.c = try readConfig(a, io, dir, o.small); + const c = &b.c; + const limit: u32 = @intCast(@min(o.context orelse 262144, c.rope.max_positions)); + b.device = try mtl.Device.init(); + errdefer b.device.deinit(); + b.queue = try b.device.queue(); + errdefer b.queue.deinit(); + b.p = try pipelines.Pipelines.load(a, b.device, c); + errdefer b.p.deinit(); + var pack = try host.pack.Pack.open(a, io, dir); + defer pack.deinit(); + b.w = try model_weights.Model.load(a, io, b.device, b.queue, &pack, c, null); + errdefer b.w.deinit(); + if (o.drafts and c.mtp_layers > 0) b.ds = try dspark.loadWeights(a, io, b.device, &pack, c); + errdefer if (b.ds) |*w| w.deinit(); + b.tables = try rope.Tables.init(b.device, c, limit); + errdefer b.tables.deinit(); + b.s = try scratch.Scratch.init(b.device, c, limit); + errdefer b.s.deinit(); + const lookup = try b.engram(io, dir, o); + errdefer b.closeEngram(); + const r = b.w.report(); + std.log.info("dsv41 metal: {d} layers, {d} shards ({d} MiB mapped, {d} MiB copied), {d} slots of {d} positions, DSpark {s}", .{ c.layers, r.shards, r.shard_bytes >> 20, r.arena_bytes >> 20, o.slots, limit, if (b.ds != null) "on" else "off" }); + const env: target_.Env = .{ .device = b.device, .queue = b.queue, .c = c, .p = &b.p, .tables = &b.tables, .s = &b.s, .model = b.w.view(), .lookup = lookup }; + return serve.Metal.open(a, .{ .env = env, .top = &b.w.top, .dspark = if (b.ds) |*w| w else null, .owner = .{ .ctx = b, .close = close } }, .{ .slots = o.slots, .limit = limit, .drafts = o.drafts, .cache_bytes = o.cache_bytes }); + } + + fn readConfig(a: Allocator, io: std.Io, dir: []const u8, small: bool) !Config { + if (!small) return Config.read(a, io, dir); + const path = try std.fs.path.join(a, &.{ dir, "config.json" }); + defer a.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 24)); + defer a.free(bytes); + return Config.parseShapes(a, bytes); + } + + /// The Engram hashes (the canonical token map) and the packed row shards; null without Engram layers. + fn engram(b: *Boot, io: std.Io, dir: []const u8, o: Options) !?forward.Lookup { + const c = &b.c; + if (c.engram_layers.len == 0) return null; + const assets = if (std.c.getenv("TF_DSV41_ASSETS")) |v| std.mem.span(v) else dir; + const path = try std.fs.path.join(b.a, &.{ assets, token_map }); + defer b.a.free(path); + b.map_bytes = std.Io.Dir.cwd().readFileAllocOptions(io, path, b.a, .limited(1 << 26), .@"4", null) catch |e| { + std.log.err("dsv41 metal: the Engram token map {s} ({t}): put it beside the checkpoint or under TF_DSV41_ASSETS", .{ path, e }); + return error.EngramTokenMap; + }; + errdefer b.a.free(b.map_bytes); + b.eng = try eh.Host.load(b.map_bytes, c.engram_heads, c.engram_pad, c.engram_vocab, &b.lbuf); + const rows_dir = try b.rowsDir(dir, o); + defer b.a.free(rows_dir); + // a boot that serves Engram layers proves the tables match the release: the manifest, headers, sampled rows + const full = std.mem.eql(u8, std.mem.span(std.c.getenv("TF_DSV41_ENGRAM_VERIFY") orelse ""), "full"); + const em = host.engram_manifest; + const ck = em.check(b.a, io, dir, rows_dir, full) catch |e| { + switch (e) { + error.EngramManifestMissing => std.log.err("dsv41 metal: no {s} beside the Engram tables {s}: run tf-dsv41-engram-pack pack", .{ em.filename, rows_dir }), + error.EngramManifestBad => std.log.err("dsv41 metal: the Engram manifest {s}/{s} is missing or unreadable", .{ rows_dir, em.filename }), + error.EngramSourceMismatch => std.log.err("dsv41 metal: the release no longer matches the Engram tables in {s}: the shards moved or changed", .{rows_dir}), + error.EngramShardCorrupt => std.log.err("dsv41 metal: an Engram table in {s} is corrupt or does not match its manifest", .{rows_dir}), + error.EngramSampleMismatch => std.log.err("dsv41 metal: an Engram table in {s} differs from the release shards on a sampled row", .{rows_dir}), + else => std.log.err("dsv41 metal: the Engram tables in {s} cannot be verified ({s})", .{ rows_dir, @errorName(e) }), + } + return e; + }; + std.log.info("dsv41 metal: Engram tables checked against the release ({d} sampled rows{s})", .{ ck.samples, if (full) " and every byte" else "" }); + const cap: u32 = if (std.c.getenv("TF_DSV41_ENGRAM_CACHE")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else 65536; + b.rows = try er.Rows.open(b.a, rows_dir, 0, 1, b.eng.layer_ids, cap); + std.log.info("dsv41 metal: Engram rows from {s} ({d} records a layer cached)", .{ rows_dir, cap }); + return .{ .a = b.a, .host = &b.eng, .source = b.rows.?.source() }; + } + + /// The row shards' folder: the option, TF_DSV41_ENGRAM_DIR, else engram_table_dir (relative to the checkpoint). + fn rowsDir(b: *Boot, dir: []const u8, o: Options) ![]u8 { + if (o.engram_dir) |d| return b.a.dupe(u8, d); + if (std.c.getenv("TF_DSV41_ENGRAM_DIR")) |v| return b.a.dupe(u8, std.mem.span(v)); + const t = b.c.engram_table_dir.items(); + if (t.len == 0) return error.EngramRowsUnset; + if (std.fs.path.isAbsolute(t)) return b.a.dupe(u8, t); + return std.fs.path.join(b.a, &.{ dir, t }); + } + + fn closeEngram(b: *Boot) void { + if (b.rows) |r| r.close(); + b.rows = null; + if (b.map_bytes.len > 0) b.a.free(b.map_bytes); + b.map_bytes = &.{}; + } + + fn close(ctx: *anyopaque) void { + const b: *Boot = @ptrCast(@alignCast(ctx)); + b.closeEngram(); + b.s.deinit(); + b.tables.deinit(); + if (b.ds) |*w| w.deinit(); + b.w.deinit(); + b.p.deinit(); + b.queue.deinit(); + b.device.deinit(); + b.a.destroy(b); + } +}; diff --git a/zig/src/families/deepseek_v41/metal/contract_exports.zig b/zig/src/families/deepseek_v41/metal/contract_exports.zig index 79b356064d..ddb5dd306b 100644 --- a/zig/src/families/deepseek_v41/metal/contract_exports.zig +++ b/zig/src/families/deepseek_v41/metal/contract_exports.zig @@ -7,6 +7,7 @@ pub const attn = @import("csa2_attn.zig"); pub const experts = @import("experts.zig"); pub const mv = @import("mv.zig"); pub const glue = @import("glue_k.zig"); +pub const vision = @import("vision.zig"); pub const compress = @import("compress.zig"); pub const xmma = @import("experts_mma.zig"); pub const hc = @import("hc.zig"); diff --git a/zig/src/families/deepseek_v41/metal/entry.zig b/zig/src/families/deepseek_v41/metal/entry.zig new file mode 100644 index 0000000000..7eb864254b --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/entry.zig @@ -0,0 +1,44 @@ +//! DeepSeek-V4.1's Metal entry point for the native server: a release checkpoint served through the lane core. +const std = @import("std"); +const api = @import("engine_api"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); + +const serve_lanes = host.serve_lanes; +const Allocator = std.mem.Allocator; + +pub const backends: []const []const u8 = &.{"metal"}; +pub const families: []const api.Family = &.{.{ .model_type = "deepseek_v41", .formats = &.{"fp8"}, .raw_text = false }}; + +pub fn chip(_: Allocator) ?[]const u8 { + return null; +} + +/// The engine for `o.dir`, or null with `problem` set (the native_engines contract). +pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + if (@import("sys").experimental.refusal(o.model_type, @import("sys").getenv("TF_EXPERIMENTAL")) != null) { + problem.* = "DeepSeek-V4.1 on Metal is experimental: its served lanes have no whole-model gate yet; set TF_EXPERIMENTAL=deepseek to opt in"; + return null; + } + if (!std.mem.eql(u8, o.model_type, "deepseek_v41")) { + problem.* = try std.fmt.allocPrint(a, "this build serves deepseek_v41 checkpoints only, not {s}", .{o.model_type}); + return null; + } + const opts: pkg.boot.Options = .{ + .slots = @max(1, @min(o.lanes, 16)), + .context = if (o.context) |c| @intCast(@max(c, 1)) else null, + .drafts = o.drafts, + .cache_bytes = if (o.prompt_cache_gib) |g| @intFromFloat(@max(g, 0) * (1 << 30)) else 4 << 30, + .small = if (std.c.getenv("TF_DSV41_SMALL")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false, + }; + const m = pkg.boot.Boot.open(gpa, io, o.dir, opts) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the DeepSeek-V4.1 Metal engine cannot load {s} ({s})", .{ o.dir, @errorName(e) }); + return null; + }; + const s = serve_lanes.Served.open(gpa, io, m.model(), o.dir, o.context) catch |e| { + m.model().close(m); + problem.* = try std.fmt.allocPrint(a, "the DeepSeek-V4.1 Metal engine cannot start its lanes ({s})", .{@errorName(e)}); + return null; + }; + return .{ .engine = s.engine(), .close = serve_lanes.Served.close, .ctx = s, .halt = serve_lanes.Served.halt }; +} diff --git a/zig/src/families/deepseek_v41/metal/forward.zig b/zig/src/families/deepseek_v41/metal/forward.zig index c6022aa34c..622ec616b4 100644 --- a/zig/src/families/deepseek_v41/metal/forward.zig +++ b/zig/src/families/deepseek_v41/metal/forward.zig @@ -245,11 +245,11 @@ const max_buffers = 16; /// Layers lo .. hi of the forward, the embedding before layer 0 and the head after the last. fn encodeRange(x: *const Ctx, e: mtl.ComputeEncoder, m: Model, b: *const Bufs, lo: usize, hi: usize) !void { if (lo == 0) { - encodeEmbed(x, e, m, b); + try encodeEmbed(x, e, m, b); e.barrier(); } for (lo..hi) |i| try encodeLayer(x, e, m, b, i); - if (hi == m.layers.len) encodeHead(x, e, m, b); + if (hi == m.layers.len) try encodeHead(x, e, m, b); } test "a round's command buffers: one layer, then four even parts of the rest" { diff --git a/zig/src/families/deepseek_v41/metal/metal.zig b/zig/src/families/deepseek_v41/metal/metal.zig index 55f4e53333..bec87c2930 100644 --- a/zig/src/families/deepseek_v41/metal/metal.zig +++ b/zig/src/families/deepseek_v41/metal/metal.zig @@ -1,4 +1,4 @@ -//! DeepSeek V4.1 on Metal, the backend package: the decode forward (1-16 rows), the prompt path and DSpark. +//! DeepSeek V4.1 on Metal, the backend package: the decode forward (1-16 rows), the prompt path, the served lanes. const std = @import("std"); pub const host = @import("dsv41"); @@ -25,6 +25,11 @@ pub const dspark_plan = @import("dspark_plan.zig"); pub const dspark_stage = @import("dspark_stage.zig"); pub const dspark_weights = @import("dspark_weights.zig"); pub const engram = pipelines.engram; +pub const snapshots = @import("snapshots.zig"); +pub const target = @import("target.zig"); +pub const pass = @import("pass.zig"); +pub const serve = @import("serve.zig"); +pub const boot = @import("boot.zig"); test { std.testing.refAllDecls(@This()); @@ -34,4 +39,5 @@ test { _ = prompt_scratch; _ = dspark; _ = dspark_plan; + _ = pass; } diff --git a/zig/src/families/deepseek_v41/metal/pass.zig b/zig/src/families/deepseek_v41/metal/pass.zig new file mode 100644 index 0000000000..520be10a3e --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/pass.zig @@ -0,0 +1,57 @@ +//! DSpark's pass for the lanes (iface.Pass) over the Metal drafter: the target's taps in, blocks of drafts out. +const std = @import("std"); +const host = @import("dsv41"); +const dspark = @import("dspark.zig"); +const target_ = @import("target.zig"); + +const iface = host.iface; + +pub const Pass = struct { + x: *dspark.Pass, + /// the target whose taps the ingests name (by address) + t: *const target_.Target, + passes: u64 = 0, + ingests: u64 = 0, + + pub fn pass(p: *Pass) iface.Pass { + return .{ .ptr = p, .vtable = &.{ .ingest = ingest, .propose = propose, .reset = reset, .markov = markov } }; + } + + fn self(ptr: *anyopaque) *Pass { + return @ptrCast(@alignCast(ptr)); + } + + fn ingest(ptr: *anyopaque, slot: u32, start: u64, taps: iface.Taps, rows: []const u32) anyerror!void { + const p = self(ptr); + if (rows.len == 0) return; + const at = p.t.bufferOf(taps.address) orelse return error.UnknownTaps; + p.ingests += 1; + try p.x.ingest(slot, start, .{ .buf = at.buf, .off = at.off, .rows = taps.rows }, rows); + } + + fn propose(ptr: *anyopaque, asks: []const iface.Ask, out: []iface.Proposal) anyerror!void { + const p = self(ptr); + var xs: [16]dspark.Ask = undefined; + var ys: [16]dspark.Proposal = undefined; + if (asks.len > xs.len or out.len != asks.len) return error.TooManySlots; + for (asks, out, xs[0..asks.len], ys[0..asks.len]) |a, o, *x, *y| { + x.* = .{ .slot = a.slot, .anchor = a.anchor, .start = a.start, .params = a.params }; + y.* = .{ .drafts = o.drafts, .conf = o.conf, .cand = o.cand, .base = o.base }; + } + p.passes += 1; + try p.x.propose(xs[0..asks.len], ys[0..asks.len]); + } + + fn reset(ptr: *anyopaque, slot: u32) void { + self(ptr).x.reset(slot); + } + + fn markov(ptr: *anyopaque) ?host.dspark.Markov { + return self(ptr).x.markov(); + } +}; + +test "the pass's asks and proposals are the drafter's, field for field" { + try std.testing.expectEqual(@sizeOf(iface.Ask), @sizeOf(dspark.Ask)); + try std.testing.expectEqual(@sizeOf(iface.Proposal), @sizeOf(dspark.Proposal)); +} diff --git a/zig/src/families/deepseek_v41/metal/serve.zig b/zig/src/families/deepseek_v41/metal/serve.zig new file mode 100644 index 0000000000..b0d03fc2f1 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/serve.zig @@ -0,0 +1,121 @@ +//! DeepSeek-V4.1 on Metal as the lanes serve it (serve_lanes.Model): the target, DSpark's pass, costs and hooks. +const std = @import("std"); +const mtl = @import("metal"); +const host = @import("dsv41"); +const weights = @import("weights.zig"); +const dspark = @import("dspark.zig"); +const prompt = @import("prompt.zig"); +const target_ = @import("target.zig"); +const pass_ = @import("pass.zig"); +const pipelines = @import("pipelines.zig"); + +const serve_lanes = host.serve_lanes; +const iface = host.iface; +const api = host.api; +const Allocator = std.mem.Allocator; + +/// The parts a served model borrows; `owner` (boot.zig's model, or a test's rig) is closed after them. +pub const Parts = struct { + env: target_.Env, + top: *const weights.Top, + /// the drafter's weights (null: serial windows, no drafts) + dspark: ?*const dspark.Weights, + owner: ?Owner = null, +}; + +pub const Owner = struct { ctx: *anyopaque, close: *const fn (ctx: *anyopaque) void }; + +pub const Options = struct { + slots: u32 = 4, + limit: u32, + /// DSpark's pass when the checkpoint has one (false: serial windows through the same lanes) + drafts: bool = true, + cache_bytes: u64 = 0, + prompt: prompt.Opts = .{}, +}; + +/// A position's draft candidates the pass keeps (the shape the lanes size their buffers by). +const candidates = 64; + +pub const Metal = struct { + a: Allocator, + parts: Parts, + opt: Options, + t: *target_.Target, + dp: ?*dspark.Pass = null, + ad: pass_.Pass = undefined, + + pub fn open(a: Allocator, parts: Parts, opt: Options) !*Metal { + const m = try a.create(Metal); + errdefer a.destroy(m); + m.* = .{ .a = a, .parts = parts, .opt = opt, .t = undefined }; + m.t = try target_.Target.init(a, parts.env, .{ .slots = opt.slots, .limit = opt.limit, .prompt = opt.prompt, .cache_bytes = opt.cache_bytes }); + errdefer m.t.deinit(); + const e = parts.env; + if (opt.drafts) if (parts.dspark) |w| { + m.dp = try dspark.Pass.init(a, e.device, e.queue, e.c, e.p, e.s, e.tables, parts.top, w, .{ .slots = opt.slots, .candidates = @min(candidates, e.c.vocab) }); + m.ad = .{ .x = m.dp.?, .t = m.t }; + }; + return m; + } + + /// The lanes' view: one slot a stream, rows of 16 a forward, the release's DSpark shape on one device. + pub fn model(m: *Metal) serve_lanes.Model { + const c = m.parts.env.c; + const shape: host.dspark.Shape = .{ .block = c.dspark_block, .noise = c.dspark_noise_token, .window = c.window, .rank = c.dspark_markov_rank, .hidden = c.hidden, .candidates = @min(candidates, c.vocab) }; + return .{ + .ptr = m, + .name = "dsv41-metal", + .slots = m.opt.slots, + .limit = m.opt.limit, + .world = 1, + .prefill_step = m.opt.prompt.chunk, + .forward_rows = pipelines.max_rows, + .shape = shape, + .target = targetFn, + .pass = passFn, + .calibrate = calibrate, + .memory = .{ .ctx = m, .read = memory }, + .health = .{ .ctx = m, .read = healthy }, + .close = close, + // the default cost table: TF_DSV41_CALIB has nothing to time on Metal + .calib = .{ .mode = "default" }, + }; + } + + fn self(ctx: *anyopaque) *Metal { + return @ptrCast(@alignCast(ctx)); + } + + fn targetFn(ctx: *anyopaque) iface.Target { + return self(ctx).t.target(); + } + + fn passFn(ctx: *anyopaque) ?iface.Pass { + const m = self(ctx); + return if (m.dp != null) m.ad.pass() else null; + } + + fn calibrate(_: *anyopaque, _: Allocator, _: std.Io, _: []const u8, _: host.calib_env.Knobs, _: host.calib_env.Shape) anyerror!host.calib_lanes.Boot { + return error.Unmeasured; // the mode is "default": never asked to measure + } + + /// The device's allocated bytes now (the weights, slots, prompt buffers and kept-state host copies aside). + fn memory(ctx: ?*anyopaque, _: bool) ?api.Memory { + const m = self(ctx.?); + const used = m.parts.env.device.allocated(); + return .{ .active = used, .cache = m.t.cache.held, .peak = used }; + } + + fn healthy(ctx: *anyopaque) bool { + return !self(ctx).t.failed; + } + + fn close(ctx: *anyopaque) void { + const m = self(ctx); + if (m.dp) |p| p.deinit(); + m.t.deinit(); + if (m.parts.owner) |o| o.close(o.ctx); + m.a.destroy(m); + } +}; diff --git a/zig/src/families/deepseek_v41/metal/snapshots.zig b/zig/src/families/deepseek_v41/metal/snapshots.zig new file mode 100644 index 0000000000..dc14dd0b0d --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/snapshots.zig @@ -0,0 +1,154 @@ +//! Kept prompt states on the host: a slot's regions up to its rows, its ids and DSpark taps, restored into any slot. +const std = @import("std"); +const mtl = @import("metal"); +const Config = @import("dsv41").config.Config; +const kv = @import("kv.zig"); + +/// One kept state: the slot after `ids` (every row run), and the choice's logits when it ended a prompt. +pub const Entry = struct { + ids: []u32, + /// each region's window, pool rows, index key rows and carry, in layer order + bytes: []u8, + /// the taps ring: position p's row at p % ring rows + taps: []u16, + logits: ?[]f32, + tick: u64 = 0, + + fn size(e: *const Entry) u64 { + return e.ids.len * 4 + e.bytes.len + e.taps.len * 2 + if (e.logits) |l| l.len * 4 else 0; + } + + fn free(e: *Entry, a: std.mem.Allocator) void { + a.free(e.ids); + a.free(e.bytes); + a.free(e.taps); + if (e.logits) |l| a.free(l); + a.destroy(e); + } +}; + +/// The pool and key rows a state of `len` positions has written at a KV source (and the row it may be filling). +pub fn poolRowsOf(c: *const Config, slot: *const kv.Slot, layer: u32, len: u64) usize { + const all = kv.poolRows(c, layer, slot.limit); + return @min(all, @as(usize, @intCast(len / c.compressRatio(layer))) + 1); +} + +/// Every region's bytes a state of `len` positions reads, as (buffer, bytes) runs in order. +fn runs(c: *const Config, slot: *const kv.Slot, len: u64, out: *[4 * 64]Run) []Run { + var n: usize = 0; + for (slot.regions, 0..) |r, l| { + out[n] = .{ .buf = r.window, .len = r.window.length() }; + n += 1; + if (r.pool) |p| { + const rows = poolRowsOf(c, slot, @intCast(l), len); + out[n] = .{ .buf = p, .len = rows * c.head_dim * 2 }; + out[n + 1] = .{ .buf = r.keys.?, .len = rows * kv.index_key_bytes }; + n += 2; + } + if (r.carry) |k| { + out[n] = .{ .buf = k, .len = k.length() }; + n += 1; + } + } + return out[0..n]; +} + +const Run = struct { buf: mtl.Buffer, len: usize }; + +/// Zeroes what a state of up to `high` positions may have written: a reused slot reads as a new one. +pub fn clear(c: *const Config, slot: *const kv.Slot, high: u64) void { + var buf: [4 * 64]Run = undefined; + for (runs(c, slot, high, &buf)) |r| @memset(r.buf.contents()[0..r.len], 0); + for (slot.regions) |*r| kv.resetCarry(r, c); +} + +pub const Store = struct { + a: std.mem.Allocator, + /// bytes the entries may hold (0: no prompt cache) + cap: u64, + held: u64 = 0, + tick: u64 = 0, + entries: std.ArrayList(*Entry) = .empty, + + pub fn deinit(st: *Store) void { + for (st.entries.items) |e| e.free(st.a); + st.entries.deinit(st.a); + } + + pub fn on(st: *const Store) bool { + return st.cap > 0; + } + + /// The longest entry that starts `ids`: shorter than them, or as long with the choice's logits. + pub fn best(st: *Store, ids: []const u32) ?*Entry { + var out: ?*Entry = null; + for (st.entries.items) |e| { + if (e.ids.len > ids.len or !std.mem.eql(u32, e.ids, ids[0..e.ids.len])) continue; + if (e.ids.len == ids.len and e.logits == null) continue; + if (out == null or e.ids.len > out.?.ids.len) out = e; + } + if (out) |e| { + st.tick += 1; + e.tick = st.tick; + } + return out; + } + + /// The slot's state after `ids` kept (the oldest entries make room); one entry a distinct `ids`. + pub fn save(st: *Store, c: *const Config, slot: *const kv.Slot, ids: []const u32, ring: []const u16, logits: ?[]const f32) !void { + if (!st.on() or ids.len == 0) return; + st.tick += 1; + for (st.entries.items) |e| if (std.mem.eql(u32, e.ids, ids)) { + e.tick = st.tick; + if (e.logits != null or logits == null) return; + e.logits = try st.a.dupe(f32, logits.?); + st.held += logits.?.len * 4; + return; + }; + var buf: [4 * 64]Run = undefined; + const rs = runs(c, slot, ids.len, &buf); + var total: usize = 0; + for (rs) |r| total += r.len; + const size = ids.len * 4 + total + ring.len * 2 + if (logits) |l| l.len * 4 else 0; + if (size > st.cap) return; + while (st.held + size > st.cap) st.evict(); + const e = try st.a.create(Entry); + errdefer st.a.destroy(e); + e.* = .{ .ids = try st.a.dupe(u32, ids), .bytes = undefined, .taps = undefined, .logits = null, .tick = st.tick }; + errdefer st.a.free(e.ids); + e.bytes = try st.a.alloc(u8, total); + errdefer st.a.free(e.bytes); + var at: usize = 0; + for (rs) |r| { + @memcpy(e.bytes[at..][0..r.len], r.buf.contents()[0..r.len]); + at += r.len; + } + e.taps = try st.a.dupe(u16, ring); + errdefer st.a.free(e.taps); + if (logits) |l| e.logits = try st.a.dupe(f32, l); + try st.entries.append(st.a, e); + st.held += e.size(); + } + + /// The entry's state into `slot` (cleared first up to `high`, its old rows) and its taps into `ring`. + pub fn restore(_: *const Store, c: *const Config, e: *const Entry, slot: *const kv.Slot, high: u64, ring: []u16) void { + clear(c, slot, high); + var buf: [4 * 64]Run = undefined; + var at: usize = 0; + for (runs(c, slot, e.ids.len, &buf)) |r| { + @memcpy(r.buf.contents()[0..r.len], e.bytes[at..][0..r.len]); + at += r.len; + } + @memcpy(ring, e.taps); + } + + fn evict(st: *Store) void { + var old: usize = 0; + for (st.entries.items, 0..) |e, i| if (e.tick < st.entries.items[old].tick) { + old = i; + }; + const e = st.entries.swapRemove(old); + st.held -= e.size(); + e.free(st.a); + } +}; diff --git a/zig/src/families/deepseek_v41/metal/target.zig b/zig/src/families/deepseek_v41/metal/target.zig new file mode 100644 index 0000000000..2e83b64b23 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/target.zig @@ -0,0 +1,418 @@ +//! The lanes' verify target on Metal: a slot's prompt (prompt.zig), its windows of rows (forward.zig), keeps and reuse. +const std = @import("std"); +const mtl = @import("metal"); +const host = @import("dsv41"); +const kv = @import("kv.zig"); +const rope = @import("rope.zig"); +const ctx_ = @import("ctx.zig"); +const scratch = @import("scratch.zig"); +const forward = @import("forward.zig"); +const prompt = @import("prompt.zig"); +const pipelines = @import("pipelines.zig"); +const snapshots = @import("snapshots.zig"); + +const iface = host.iface; +const lanes = host.lanes; +const Config = host.config.Config; +const Allocator = std.mem.Allocator; +const R = pipelines.max_rows; + +/// What the target reads and borrows: the device, kernels, RoPE tables, layer scratch, weights and Engram rows. +pub const Env = struct { + device: mtl.Device, + queue: mtl.Queue, + c: *const Config, + p: *const pipelines.Pipelines, + tables: *const rope.Tables, + s: *const scratch.Scratch, + model: forward.Model, + lookup: ?forward.Lookup, +}; + +pub const Options = struct { + slots: u32, + /// positions a slot holds (the context limit) + limit: u32, + prompt: prompt.Opts = .{}, + /// bytes of kept prompt states (0: no prompt cache) + cache_bytes: u64 = 0, +}; + +/// What the target did, for the tests and the logs. +pub const Stats = struct { prompts: u64 = 0, hits: u64 = 0, cached_rows: u64 = 0, windows: u64 = 0, rows: u64 = 0, rewinds: u64 = 0, launches: u64 = 0 }; + +/// A slot: its device state, prompt tail, ids (rows run, the last window's included until its keep) and taps. +const Slot = struct { + kv: kv.Slot, + tail: prompt.Tail, + hist: std.ArrayList(u32) = .empty, + /// the last window: its first position and rows; `open` until its keep + base: u64 = 0, + rows: u32 = 0, + open: bool = false, + /// the taps of the last `ring rows` positions, position p at row p % ring rows + ring: []u16, + /// the taps `taps` names after a prompt: its last rows in order + ptaps: mtl.Buffer, + ptaps_rows: u32 = 0, + after_prompt: bool = false, + /// the most positions any request wrote here since the slot was last cleared + high: u64 = 0, + live: bool = false, + /// every launch since the request began completed (a release keeps only a whole state) + whole: bool = false, + /// the rows the last prompt took from kept state + resumed: u32 = 0, +}; + +pub const Target = struct { + a: Allocator, + env: Env, + opt: Options, + b: forward.Bufs, + pr: prompt.Prompt, + slots: []Slot, + cache: snapshots.Store, + stats: Stats = .{}, + /// a command buffer failed: the device state is no longer trusted (the health hook) + failed: bool = false, + + pub fn init(a: Allocator, env: Env, opt: Options) !*Target { + const c = env.c; + if (opt.slots == 0 or opt.limit == 0 or opt.limit > env.tables.positions) return error.BadOptions; + _ = try prompt.check(c); + const t = try a.create(Target); + errdefer a.destroy(t); + t.* = .{ .a = a, .env = env, .opt = opt, .b = undefined, .pr = undefined, .slots = &.{}, .cache = .{ .a = a, .cap = opt.cache_bytes } }; + t.b = try forward.Bufs.init(env.device, c); + errdefer t.b.deinit(); + t.pr = try prompt.Prompt.init(env.device, c, opt.prompt, opt.limit); + errdefer t.pr.deinit(); + const slots = try a.alloc(Slot, opt.slots); + errdefer a.free(slots); + var made: usize = 0; + errdefer for (slots[0..made]) |*s| freeSlot(a, s); + const cap = try prompt.tailRows(c, opt.prompt); + const ptb = @max(4, @as(usize, ringRows(c)) * tapWidth(c) * 2); + for (slots) |*s| { + var st = try kv.Slot.init(a, env.device, c, opt.limit); + errdefer st.deinit(); + var tail = try prompt.Tail.init(env.device, c, cap); + errdefer tail.deinit(); + const ptaps = try env.device.buffer(ptb, mtl.ResourceOptions.shared); + errdefer ptaps.deinit(); + const ring = try a.alloc(u16, @as(usize, ringRows(c)) * tapWidth(c)); + @memset(ring, 0); + s.* = .{ .kv = st, .tail = tail, .ring = ring, .ptaps = ptaps }; + made += 1; + } + t.slots = slots; + return t; + } + + fn freeSlot(a: Allocator, s: *Slot) void { + s.kv.deinit(); + s.tail.deinit(); + s.ptaps.deinit(); + a.free(s.ring); + s.hist.deinit(a); + } + + pub fn deinit(t: *Target) void { + for (t.slots) |*s| freeSlot(t.a, s); + t.a.free(t.slots); + t.cache.deinit(); + t.pr.deinit(); + t.b.deinit(); + t.a.destroy(t); + } + + /// The taps ring's rows: the drafter's context window. + fn ringRows(c: *const Config) u32 { + return c.window; + } + + /// Values in a taps row: every target layer's hidden. + fn tapWidth(c: *const Config) usize { + return c.dspark_targets.len * c.hidden; + } + + pub fn target(t: *Target) iface.Target { + return .{ .ptr = t, .vtable = &.{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .resumed = resumedFn } }; + } + + fn self(ptr: *anyopaque) *Target { + return @ptrCast(@alignCast(ptr)); + } + + /// The device buffer and offset of an address a `taps` named (the pass's ingest reads it there). + pub fn bufferOf(t: *const Target, address: u64) ?struct { buf: mtl.Buffer, off: usize } { + for (t.slots) |s| { + const bufs = [_]?mtl.Buffer{ s.ptaps, s.kv.taps }; + for (bufs) |ob| if (ob) |b| { + const lo = @intFromPtr(b.contents()); + if (address >= lo and address < lo + b.length()) return .{ .buf = b, .off = @intCast(address - lo) }; + }; + } + return null; + } + + // prompts ------------------------------------------------------------------------------------------------------ + + fn resumedFn(ptr: *anyopaque, slot: u32) u64 { + const t = self(ptr); + return if (slot < t.slots.len) t.slots[slot].resumed else 0; + } + + fn prefill(ptr: *anyopaque, slot: u32, ids: []const u32, smp: ?lanes.Sampling, draw: u64) anyerror!u32 { + const t = self(ptr); + if (slot >= t.slots.len) return error.NoSlot; + const s = &t.slots[slot]; + const c = t.env.c; + if (s.live) return error.SlotLive; + if (ids.len == 0) return error.EmptyPrompt; + if (ids.len >= t.opt.limit) return error.PastLimit; + for (ids) |id| if (id >= c.vocab) return error.BadToken; + s.whole = false; + s.open = false; + s.rows = 0; + t.stats.prompts += 1; + // a kept state that starts this prompt, else a cleared slot + const hit = if (t.cache.on()) t.cache.best(ids) else null; + const at: u32 = if (hit) |e| @intCast(e.ids.len) else 0; + if (hit) |e| { + t.cache.restore(c, e, &s.kv, s.high, s.ring); + t.stats.hits += 1; + t.stats.cached_rows += at; + } else snapshots.clear(c, &s.kv, s.high); + s.resumed = at; + s.tail.end = at; + s.tail.done = at; + s.hist.clearRetainingCapacity(); + try s.hist.appendSlice(t.a, ids); + s.high = @max(at, ids.len); + const n: u32 = @intCast(ids.len); + const V: usize = c.vocab; + var logits: []const f32 = undefined; + if (at == n) logits = hit.?.logits.? else { + t.pr.run(t.promptEnv(), &s.kv, &s.tail, ids) catch |e| return t.fail(e); + const all = t.pr.logits(c); + logits = all[(t.pr.head - 1) * V ..][0..V]; + const tr = t.pr.tapRows(c); + for (0..t.pr.taps) |i| t.ringPut(s, n - t.pr.taps + i, tr[i * tapWidth(c) ..][0..tapWidth(c)]); + } + t.promptTaps(s, n); + const choice = try t.choose(logits, draw, smp); + s.base = n; + s.live = true; + s.whole = true; + if (t.cache.on() and at < n) t.cache.save(c, &s.kv, ids, s.ring, logits) catch {}; + return choice; + } + + fn promptEnv(t: *const Target) prompt.Env { + const e = t.env; + return .{ .c = e.c, .p = e.p, .tables = e.tables, .model = e.model, .lookup = e.lookup, .queue = e.queue }; + } + + fn ringPut(t: *const Target, s: *Slot, pos: u64, row: []const u16) void { + const w = tapWidth(t.env.c); + if (w == 0) return; + const at: usize = @intCast(pos % ringRows(t.env.c)); + @memcpy(s.ring[at * w ..][0..w], row); + } + + /// The taps of the prompt's last rows (the ring's, up to `ring rows`) in order, for `taps` after a prefill. + fn promptTaps(t: *const Target, s: *Slot, n: u32) void { + const c = t.env.c; + const w = tapWidth(c); + s.after_prompt = true; + s.ptaps_rows = if (w == 0) 0 else @min(n, ringRows(c)); + const out = s.ptaps.slice(u16, @as(usize, s.ptaps_rows) * w); + for (0..s.ptaps_rows) |i| { + const pos = n - s.ptaps_rows + i; + @memcpy(out[i * w ..][0..w], s.ring[(pos % ringRows(c)) * w ..][0..w]); + } + } + + /// The keyed choice of a logits row over the whole vocabulary (greedy: the first maximum, ties to the lower id). + fn choose(t: *const Target, row: []const f32, draw: u64, smp: ?lanes.Sampling) !u32 { + const s = smp orelse return argmax(row); + if (s.temperature <= 0) return argmax(row); + const vals = try t.a.alloc(f64, row.len); + defer t.a.free(vals); + const ids = try t.a.alloc(u64, row.len); + defer t.a.free(ids); + for (vals, ids, row, 0..) |*v, *id, x, i| { + v.* = x; + id.* = i; + } + return @intCast(try lanes.sampling.choose(t.a, vals, ids, draw, s)); + } + + fn argmax(row: []const f32) u32 { + var best: usize = 0; + for (row, 0..) |x, i| if (x > row[best]) { + best = i; + }; + return @intCast(best); + } + + fn fail(t: *Target, e: anyerror) anyerror { + if (e == error.CommandBufferFailed) t.failed = true; + return e; + } + + // windows ------------------------------------------------------------------------------------------------------ + + /// A round's windows: slots' rows share a launch (up to R rows, a slot once), each launch one forward. + fn window(ptr: *anyopaque, segments: []const iface.Segment, choices: [][]u32) anyerror!void { + const t = self(ptr); + var from: usize = 0; + while (from < segments.len) { + var rows: usize = 0; + var to = from; + while (to < segments.len) : (to += 1) { + const g = segments[to]; + if (g.slot >= t.slots.len) return error.NoSlot; + if (to > from and (rows + g.tokens.len > R or dup(segments[from..to], g.slot))) break; + rows += g.tokens.len; + } + try t.windowGroup(segments[from..to], choices[from..to]); + from = to; + } + } + + fn dup(done: []const iface.Segment, slot: u32) bool { + for (done) |o| if (o.slot == slot) return true; + return false; + } + + /// One launch over the rows of these segments (distinct slots, at most R rows in all). + fn windowGroup(t: *Target, group: []const iface.Segment, choices: [][]u32) !void { + const c = t.env.c; + const V: usize = c.vocab; + var total: u32 = 0; + for (group, choices) |g, out| { + const s = &t.slots[g.slot]; + const n = g.tokens.len; + if (!s.live) return error.SlotNotLive; + if (g.start != s.hist.items.len) return error.PositionMismatch; + if (n == 0 or total + n > R or g.draws.len != n or out.len < n) return error.BadRows; + if (g.start + n > t.opt.limit) return error.PastLimit; + // the decode forward runs consecutive positions: a tree's rows must be one chain + if (g.parents) |ps| for (ps, 0..) |q, r| if (q != @as(i32, @intCast(r)) - 1) return error.TreeWindow; + for (g.tokens) |id| if (id >= c.vocab) return error.BadToken; + total += @intCast(n); + } + for (group) |g| { + const s = &t.slots[g.slot]; + s.whole = false; + try s.hist.appendSlice(t.a, g.tokens); + s.high = @max(s.high, s.hist.items.len); + } + t.launch(group) catch |e| return t.fail(e); + const lg = t.b.logitsOf(c, total); + const w = tapWidth(c); + var r0: usize = 0; + for (group, choices) |g, out| { + const s = &t.slots[g.slot]; + const n = g.tokens.len; + for (0..n) |r| out[r] = try t.choose(lg[(r0 + r) * V ..][0..V], g.draws[r], g.sampling); + if (s.kv.taps) |tb| for (0..n) |r| t.ringPut(s, g.start + r, tb.slice(u16, n * w)[r * w ..][0..w]); + s.base = g.start; + s.rows = @intCast(n); + s.open = true; + s.after_prompt = false; + s.tail.end = @intCast(s.hist.items.len); + s.tail.done = s.tail.end; + s.whole = true; + t.stats.windows += 1; + t.stats.rows += n; + r0 += n; + } + t.stats.launches += 1; + } + + /// One launch of the decode forward over the segments' rows, each slot's ids at its positions from `start`. + fn launch(t: *Target, group: []const iface.Segment) !void { + const e = t.env; + const c = e.c; + const look = c.engram_max_ngram - 1; + var pieces: [R]forward.Piece = undefined; + var spans: [R]scratch.Span = undefined; + var segs: [R]ctx_.Seg = undefined; + var r0: u32 = 0; + for (group, 0..) |g, i| { + const s = &t.slots[g.slot]; + const n: u32 = @intCast(g.tokens.len); + const start: u32 = @intCast(g.start); + pieces[i] = .{ .ids = g.tokens, .lookback = s.hist.items[start -| look..start] }; + spans[i] = .{ .p0 = start, .n = n }; + segs[i] = .{ .slot = &s.kv, .p0 = start, .r0 = r0, .n = n }; + r0 += n; + } + try forward.prepareSegs(&t.b, c, pieces[0..group.len], e.lookup); + try e.s.writeSpans(c, spans[0..group.len]); + const x: ctx_.Ctx = .{ .c = c, .p = e.p, .s = e.s, .slot = segs[0].slot, .tables = e.tables, .p0 = segs[0].p0, .rows = r0, .segs = segs[0..group.len] }; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + try forward.run(&x, e.queue, e.model, &t.b); + } + + /// Keep the window's first rows (a chain from row 0): later rows' ids go, the compressor's carry rewinds. + fn keep(ptr: *anyopaque, slot: u32, path: []const u32) anyerror!void { + const t = self(ptr); + if (slot >= t.slots.len) return error.NoSlot; + const s = &t.slots[slot]; + if (!s.open) return error.NoWindow; + if (path.len == 0 or path.len > s.rows) return error.NotAPath; + for (path, 0..) |r, i| if (r != i) return error.NotAPath; + s.open = false; + const m: u32 = @intCast(path.len); + if (m == s.rows) return; + const len = s.base + m; + s.hist.shrinkRetainingCapacity(@intCast(len)); + for (s.kv.regions) |*r| kv.rewindCarry(r, t.env.c, s.base, m); + s.tail.end = @intCast(len); + s.tail.done = s.tail.end; + t.stats.rewinds += 1; + } + + fn taps(ptr: *anyopaque, slot: u32) iface.Taps { + const t = self(ptr); + const s = &t.slots[slot]; + const w: u32 = @intCast(tapWidth(t.env.c)); + if (s.after_prompt) return .{ .address = @intFromPtr(s.ptaps.contents()), .rows = s.ptaps_rows, .stride = w }; + const tb = s.kv.taps orelse return .{ .address = 0, .rows = 0, .stride = w }; + return .{ .address = @intFromPtr(tb.contents()), .rows = s.rows, .stride = w }; + } + + /// The request ended: its state is kept (the prompt cache) when every launch of it completed. + fn release(ptr: *anyopaque, slot: u32) void { + const t = self(ptr); + if (slot >= t.slots.len) return; + const s = &t.slots[slot]; + if (!s.live) return; + s.live = false; + s.open = false; + if (s.whole and t.cache.on()) t.cache.save(t.env.c, &s.kv, s.hist.items, s.ring, null) catch {}; + } + + /// The serial reference without lanes: the prompt, then one row a step, each drawn at its own position. + pub fn serial(t: *Target, slot: u32, ids: []const u32, n: usize, smp: ?lanes.Sampling) ![]u32 { + var out: std.ArrayList(u32) = .empty; + errdefer out.deinit(t.a); + var tok = try prefill(t, slot, ids, smp, ids.len); + while (out.items.len < n) { + try out.append(t.a, tok); + const at: u64 = ids.len + out.items.len - 1; + var ch: [1]u32 = undefined; + var cs = [_][]u32{&ch}; + try window(t, &.{.{ .slot = slot, .start = at, .tokens = &.{tok}, .parents = null, .draws = &.{at + 1}, .sampling = smp }}, &cs); + tok = ch[0]; + } + release(t, slot); + return out.toOwnedSlice(t.a); + } +}; diff --git a/zig/src/families/deepseek_v41/metal/vision.zig b/zig/src/families/deepseek_v41/metal/vision.zig new file mode 100644 index 0000000000..a577a614d9 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/vision.zig @@ -0,0 +1,227 @@ +//! DeepSeek V4.1's vision tower on Metal: the release's ViT (2D RoPE, full attention) and 3x3 aligner. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const glue = @import("glue.zig"); + +pub const Ref = glue.Ref; + +/// The release's `vision_config` (and the text hidden width the aligner maps to); defaults are the shipped values. +pub const Config = struct { + layers: u32 = 32, + dim: u32 = 1024, + heads: u32 = 16, + inter: u32 = 2816, + patch: u32 = 14, + theta: f64 = 10000, + ratio: u32 = 3, + max_tokens: u32 = 1024, + min_pixels: u32 = 295936, + max_wh_ratio: ?u32 = null, + out_dim: u32 = 5120, + image_token_id: u32 = 129264, + eps: f32 = 1e-6, + + pub fn headDim(c: Config) u32 { + return c.dim / c.heads; + } + + /// Values in a flattened patch: three channels of patch x patch. + pub fn patchValues(c: Config) u32 { + return 3 * c.patch * c.patch; + } + + /// Values in an aligner input row: ratio x ratio patches of the tower's width. + pub fn alignIn(c: Config) u32 { + return c.dim * c.ratio * c.ratio; + } + + /// The most patches an image can have: every token of the budget a ratio x ratio block. + pub fn maxPatches(c: Config) u32 { + return c.max_tokens * c.ratio * c.ratio; + } + + /// The `vision_config` and `image_token_id` of a config.json (null: no tower); `out_dim` is the text hidden width. + pub fn parse(a: std.mem.Allocator, text: []const u8, out_dim: u32) !?Config { + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + const root = try std.json.parseFromSliceLeaky(std.json.Value, arena.allocator(), text, .{}); + if (root != .object) return error.BadConfig; + const v = (root.object.get("vision_config") orelse return null); + if (v != .object) return error.BadConfig; + var c: Config = .{ .out_dim = out_dim }; + const U = struct { + fn num(o: std.json.ObjectMap, key: []const u8) ?f64 { + const x = o.get(key) orelse return null; + return switch (x) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => null, + }; + } + }; + if (U.num(v.object, "num_hidden_layers")) |x| c.layers = @intFromFloat(x); + if (U.num(v.object, "hidden_size")) |x| c.dim = @intFromFloat(x); + if (U.num(v.object, "num_attention_heads")) |x| c.heads = @intFromFloat(x); + if (U.num(v.object, "intermediate_size")) |x| c.inter = @intFromFloat(x); + if (U.num(v.object, "patch_size")) |x| c.patch = @intFromFloat(x); + if (U.num(v.object, "rope_theta")) |x| c.theta = x; + if (U.num(v.object, "downsample_ratio")) |x| c.ratio = @intFromFloat(x); + if (U.num(v.object, "max_image_tokens")) |x| c.max_tokens = @intFromFloat(x); + if (U.num(v.object, "min_pixels")) |x| c.min_pixels = @intFromFloat(x); + if (U.num(v.object, "max_wh_ratio")) |x| c.max_wh_ratio = @intFromFloat(x); + if (U.num(root.object, "image_token_id")) |x| c.image_token_id = @intFromFloat(x); + try c.check(); + return c; + } + + pub fn check(c: Config) !void { + if (c.headDim() != 64 or c.dim % c.heads != 0) return error.UnsupportedVision; + if (c.layers == 0 or c.ratio == 0 or c.patch == 0 or c.inter == 0) return error.UnsupportedVision; + } +}; + +const names = [_][:0]const u8{ "tf_dsv41_vis_epi", "tf_dsv41_vis_qkv", "tf_dsv41_vis_attn", "tf_dsv41_vis_swiglu", "tf_dsv41_vis_unshuffle", "tf_dsv41_vis_merge" }; + +/// The tower's kernels (deepseek_v41/vision.metal). +pub const Pipelines = struct { + p: [names.len]mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + return loadSource(device, sources.dsv41.vision); + } + + /// The same kernels from other text: the contract's planted controls prefix a TF_CTL define. + pub fn loadSource(device: mtl.Device, text: []const u8) !Pipelines { + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var out: Pipelines = undefined; + for (names, 0..) |n, i| out.p[i] = try mtl.Pipeline.init(device, lib, n, false); + return out; + } + + pub fn deinit(p: *Pipelines) void { + for (p.p) |x| x.deinit(); + } +}; + +fn pipe(p: *const Pipelines, k: usize) mtl.Pipeline { + return p.p[k]; +} + +const k_epi = 0; +const k_qkv = 1; +const k_attn = 2; +const k_swiglu = 3; +const k_unshuffle = 4; +const k_merge = 5; + +fn bind(e: mtl.ComputeEncoder, refs: []const Ref) void { + for (refs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); +} + +fn grid(e: mtl.ComputeEncoder, x: usize, y: usize, z: usize) void { + e.dispatchThreads(mtl.Size.of(x, y, z), mtl.Size.of(@min(x, 32), 1, 1)); +} + +/// A linear's fp32 sums acc [rows][n] to bf16 y: + bias (bf16 [n]) when given, GELU when asked, then + res [rows][n]. +pub fn epilogue(p: *const Pipelines, e: mtl.ComputeEncoder, acc: Ref, bias: ?Ref, res: ?Ref, y: Ref, rows: u32, n: u32, gelu: bool) void { + e.setPipeline(pipe(p, k_epi)); + bind(e, &.{ acc, bias orelse acc, res orelse acc, y }); + e.setValue(extern struct { rows: u32, n: u32, gelu: u32, has_bias: u32, has_res: u32 }{ .rows = rows, .n = n, .gelu = @intFromBool(gelu), .has_bias = @intFromBool(bias != null), .has_res = @intFromBool(res != null) }, 4); + grid(e, n, rows, 1); +} + +/// The qkv sums acc [rows][3 * dim] to bf16 qkv with the bias; q and k rotated by cs [rows][head_dim / 2] (cos, sin). +pub fn qkv(p: *const Pipelines, e: mtl.ComputeEncoder, acc: Ref, bias: Ref, cs: Ref, out: Ref, rows: u32, heads: u32, hd: u32) void { + e.setPipeline(pipe(p, k_qkv)); + bind(e, &.{ acc, bias, cs, out }); + e.setValue(extern struct { rows: u32, heads: u32, hd: u32 }{ .rows = rows, .heads = heads, .hd = hd }, 4); + grid(e, hd / 2, rows, 3 * heads); +} + +/// Full attention of every row over all rows, one simdgroup a (row, head): out [rows][dim] from qkv [rows][3 * dim]. +pub fn attention(p: *const Pipelines, e: mtl.ComputeEncoder, qkv_: Ref, out: Ref, rows: u32, heads: u32) void { + e.setPipeline(pipe(p, k_attn)); + bind(e, &.{ qkv_, out }); + e.setValue(extern struct { n: u32, heads: u32, scale: f32 }{ .n = rows, .heads = heads, .scale = 0.125 }, 2); + e.dispatchGroups(mtl.Size.of(rows, heads, 1), mtl.Size.of(32, 1, 1)); +} + +/// The MLP's w1 sums acc [rows][2 * inter] to y [rows][inter] bf16: silu(gate) * up. +pub fn swiglu(p: *const Pipelines, e: mtl.ComputeEncoder, acc: Ref, y: Ref, rows: u32, inter: u32) void { + e.setPipeline(pipe(p, k_swiglu)); + bind(e, &.{ acc, y }); + e.setValue(extern struct { rows: u32, inter: u32 }{ .rows = rows, .inter = inter }, 2); + grid(e, inter, rows, 1); +} + +/// The aligner's unfold: x [nh * nw][dim] to y [lh * lw][dim * r * r], the grid zero-padded to whole blocks. +pub fn unshuffle(p: *const Pipelines, e: mtl.ComputeEncoder, x: Ref, y: Ref, nh: u32, nw: u32, dim: u32, r: u32) void { + const lh = (nh + r - 1) / r; + const lw = (nw + r - 1) / r; + e.setPipeline(pipe(p, k_unshuffle)); + bind(e, &.{ x, y }); + e.setValue(extern struct { nh: u32, nw: u32, lh: u32, lw: u32, dim: u32, r: u32 }{ .nh = nh, .nw = nw, .lh = lh, .lw = lw, .dim = dim, .r = r }, 2); + grid(e, dim * r * r, lh * lw, 1); +} + +/// Stream rows [rows][4][dim] take feat's rows or the delimiters as `src` [rows] (int32) names them; others stay. +pub fn merge(p: *const Pipelines, e: mtl.ComputeEncoder, feat: Ref, delims: Ref, src: Ref, y: Ref, rows: u32, dim: u32) void { + e.setPipeline(pipe(p, k_merge)); + bind(e, &.{ feat, delims, src, y }); + e.setValue(extern struct { dim: u32, rows: u32 }{ .dim = dim, .rows = rows }, 4); + grid(e, dim, rows, 1); +} + +/// The tower's kernels and the glue ones it shares (RMS norm, bf16 matvecs), loaded together. +pub const Kit = struct { + v: Pipelines, + g: glue.Pipelines, + + pub fn load(device: mtl.Device) !Kit { + var v = try Pipelines.load(device); + errdefer v.deinit(); + return .{ .v = v, .g = try glue.Pipelines.load(device, sources.dsv41.glue_tmp) }; + } + + pub fn deinit(k: *Kit) void { + k.v.deinit(); + k.g.deinit(); + } +}; + +/// The 2D RoPE's (cos, sin) pairs [nh * nw][head_dim / 2]: the first half rotates by row, the second by column. +pub fn ropeTable(out: [][2]f32, nh: u32, nw: u32, hd: u32, theta: f64) void { + const half = hd / 2; + const m = half / 2; + std.debug.assert(out.len == @as(usize, nh) * nw * half); + var inv: [32]f32 = undefined; + for (0..m) |i| inv[i] = @floatCast(1.0 / std.math.pow(f64, theta, @as(f64, @floatFromInt(2 * i)) / @as(f64, @floatFromInt(half)))); + for (0..nh) |h| for (0..nw) |w| for (0..half) |j| { + const pos: f32 = @floatFromInt(if (j < m) h else w); + const angle: f64 = pos * inv[j % m]; + out[(h * nw + w) * half + j] = .{ @floatCast(@cos(angle)), @floatCast(@sin(angle)) }; + }; +} + +test "the release's vision_config parses to the defaults" { + const text = + \\{"image_token_id": 129264, "vision_config": {"num_hidden_layers": 32, "hidden_size": 1024, "num_attention_heads": 16, + \\ "intermediate_size": 2816, "patch_size": 14, "rope_theta": 10000, "downsample_ratio": 3, "max_image_tokens": 1024, + \\ "min_pixels": 295936, "max_wh_ratio": null}} + ; + const c = (try Config.parse(std.testing.allocator, text, 5120)).?; + try std.testing.expectEqual(Config{}, c); + try std.testing.expect((try Config.parse(std.testing.allocator, "{}", 5120)) == null); +} + +test "the 2D RoPE table: row angles first, column angles after, position 0 is the identity" { + var t: [6 * 32][2]f32 = undefined; + ropeTable(&t, 2, 3, 64, 10000); + for (t[0..32]) |cs| try std.testing.expectEqual([2]f32{ 1, 0 }, cs); + // patch (1, 2): the first 16 pairs rotate by 1 * inv_freq, the next 16 by 2 * inv_freq + const at = (1 * 3 + 2) * 32; + try std.testing.expectApproxEqAbs(@as(f32, @floatCast(@cos(@as(f64, 1.0)))), t[at][0], 1e-7); + try std.testing.expectApproxEqAbs(@as(f32, @floatCast(@cos(@as(f64, 2.0)))), t[at + 16][0], 1e-7); +} diff --git a/zig/src/families/deepseek_v41/phases_timed.zig b/zig/src/families/deepseek_v41/phases_timed.zig index 5c195fbdbd..b2c38fadae 100644 --- a/zig/src/families/deepseek_v41/phases_timed.zig +++ b/zig/src/families/deepseek_v41/phases_timed.zig @@ -17,7 +17,11 @@ pub const TimedTarget = struct { return .{ .ptr = t, .vtable = &vtable }; } - const vtable: iface.Target.VTable = .{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .admit = admit, .begin = begin, .piece = piece, .finish = finish, .tail = tail, .pieces = pieces, .replays = replays, .warm = warm }; + const vtable: iface.Target.VTable = .{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .admit = admit, .begin = begin, .piece = piece, .finish = finish, .tail = tail, .pieces = pieces, .replays = replays, .warm = warm, .resumed = resumed }; + + fn resumed(p: *anyopaque, slot: u32) u64 { + return self(p).inner.resumedRows(slot); + } fn admit(p: *anyopaque, slot: u32, max_new: u64) void { self(p).inner.admit(slot, max_new); diff --git a/zig/tests/contract/dsv41_vision_ref.zig b/zig/tests/contract/dsv41_vision_ref.zig new file mode 100644 index 0000000000..23af0002c2 --- /dev/null +++ b/zig/tests/contract/dsv41_vision_ref.zig @@ -0,0 +1,73 @@ +//! Float64 oracles of DeepSeek V4.1's vision kernels from the math alone: epilogues, attention, SwiGLU. +const std = @import("std"); + +/// erf by its series below |x| = 4 (about 1e-9 absolute) and the complementary asymptotic series past it; erf(inf) = 1. +pub fn erf(x: f64) f64 { + if (std.math.isNan(x)) return x; + const ax = @abs(x); + if (ax >= 6) return std.math.sign(x); + if (ax < 4) { + var term = ax; + var sum = ax; + var n: f64 = 0; + while (@abs(term) > 1e-18 * @abs(sum) and ax > 0) { + n += 1; + term *= -ax * ax / n; + sum += term / (2 * n + 1); + } + return std.math.sign(x) * sum * 2 / @sqrt(std.math.pi); + } + var t: f64 = 1; + var s: f64 = 1; + for (1..14) |i| { + t *= -@as(f64, @floatFromInt(2 * i - 1)) / (2 * ax * ax); + s += t; + } + return std.math.sign(x) * (1 - @exp(-ax * ax) / (ax * @sqrt(std.math.pi)) * s); +} + +pub fn gelu(x: f64) f64 { + return 0.5 * x * (1 + erf(x * 0.7071067811865476)); +} + +pub fn silu(x: f64) f64 { + return x / (1 + @exp(-x)); +} + +/// One query row of one head against `n` keys: the softmax-weighted values, and the weights' magnitude of the values. +pub fn attend(out: []f64, mag: []f64, q: []const f64, keys: []const f64, vals: []const f64, n: usize, hd: usize, stride: usize) void { + var scores: [512]f64 = undefined; + var mx: f64 = -std.math.inf(f64); + for (0..n) |j| { + var s: f64 = 0; + for (0..hd) |d| s += q[d] * keys[j * stride + d]; + scores[j] = s / 8.0; + if (std.math.isNan(scores[j])) mx = std.math.nan(f64) else if (!std.math.isNan(mx)) mx = @max(mx, scores[j]); + } + var sum: f64 = 0; + for (0..n) |j| { + scores[j] = @exp(scores[j] - mx); + sum += scores[j]; + } + for (0..hd) |d| { + var acc: f64 = 0; + var m: f64 = 0; + for (0..n) |j| { + acc += scores[j] * vals[j * stride + d]; + m += @abs(scores[j] * vals[j * stride + d]); + } + out[d] = acc / sum; + mag[d] = m / sum; + } +} + +test "the oracle's erf and attention" { + try std.testing.expectApproxEqAbs(@as(f64, 0.8427007929497149), erf(1), 1e-12); + try std.testing.expectEqual(@as(f64, 1), erf(std.math.inf(f64))); + var out: [2]f64 = undefined; + var mag: [2]f64 = undefined; + // two equal keys: the mean of the values + attend(&out, &mag, &.{ 1, 0 }, &.{ 1, 0, 1, 0 }, &.{ 2, 4, 6, 8 }, 2, 2, 2); + try std.testing.expectApproxEqAbs(@as(f64, 4), out[0], 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 6), out[1], 1e-12); +} diff --git a/zig/tests/contract/metal_dsv41_limits.zig b/zig/tests/contract/metal_dsv41_limits.zig index 6c36a77634..a4972dd0f4 100644 --- a/zig/tests/contract/metal_dsv41_limits.zig +++ b/zig/tests/contract/metal_dsv41_limits.zig @@ -34,6 +34,7 @@ pub fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch, if (std.mem.startsWith(u8, d.name, "dsv41 cand ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.cand, .pre_m5 = sources.dsv41.cand }; if (std.mem.startsWith(u8, d.name, "dsv41 compress ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.compress, .pre_m5 = sources.dsv41.compress }; if (std.mem.startsWith(u8, d.name, "dsv41 mv ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.mv, .pre_m5 = sources.dsv41.mv }; + if (std.mem.startsWith(u8, d.name, "dsv41 vision ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.vision, .pre_m5 = sources.dsv41.vision }; if (std.mem.startsWith(u8, d.name, "dsv41 fold ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.fold, .pre_m5 = sources.dsv41.fold }; if (std.mem.startsWith(u8, d.name, "dsv41 hc ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.hc, .pre_m5 = sources.dsv41.hc }; if (std.mem.startsWith(u8, d.name, "dsv41 fp8 rows ")) return .{ .decl = d, .function = fun, .here = sources.dsv41.fp8_rows, .pre_m5 = sources.dsv41.fp8_rows }; diff --git a/zig/tests/contract/metal_dsv41_main.zig b/zig/tests/contract/metal_dsv41_main.zig index dae57e7be3..bc00f234de 100644 --- a/zig/tests/contract/metal_dsv41_main.zig +++ b/zig/tests/contract/metal_dsv41_main.zig @@ -17,6 +17,7 @@ const dsv41_engram = @import("metal_dsv41_engram.zig"); const dsv41_glue = @import("metal_dsv41_glue.zig"); const dsv41_compress = @import("metal_dsv41_compress.zig"); const dsv41_mv = @import("metal_dsv41_mv.zig"); +const dsv41_vision = @import("metal_dsv41_vision.zig"); const dsv41_hc = @import("metal_dsv41_hc.zig"); const dsv41_fold = @import("metal_dsv41_fold.zig"); const dsv41_glue_tmp = @import("metal_dsv41_glue_tmp.zig"); @@ -85,6 +86,7 @@ pub const Suites = struct { glue: *dsv41_glue.Suite, compress: *dsv41_compress.Suite, mv: *dsv41_mv.Suite, + vision: *dsv41_vision.Suite, hc: *dsv41_hc.Suite, fold: *dsv41_fold.Suite, glue_tmp: *dsv41_glue_tmp.Suite, @@ -120,6 +122,8 @@ pub const Suites = struct { errdefer compress.destroy(); const mv = try dsv41_mv.Suite.create(a, device, queue, kernels, only); errdefer mv.destroy(); + const vision = try dsv41_vision.Suite.create(a, device, queue, kernels, only); + errdefer vision.destroy(); const hc = try dsv41_hc.Suite.create(a, device, queue, kernels, only); errdefer hc.destroy(); const fold = try dsv41_fold.Suite.create(a, device, queue, kernels, only); @@ -138,7 +142,7 @@ pub const Suites = struct { errdefer fused5.destroy(); const glue_tmp2 = try dsv41_glue_tmp2.Suite.create(a, device, queue, kernels, only); errdefer glue_tmp2.destroy(); - return .{ .fp8 = fp8, .fp8_mma = fp8_mma, .topk = topk, .cand = cand, .attn = attn, .experts = experts, .xmma = xmma, .ix = ix, .engram = engram, .glue = glue, .compress = compress, .mv = mv, .hc = hc, .fold = fold, .glue_tmp = glue_tmp, .fused = fused, .fused2 = fused2, .fused3 = fused3, .fused4 = fused4, .fused5 = fused5, .glue_tmp2 = glue_tmp2 }; + return .{ .fp8 = fp8, .fp8_mma = fp8_mma, .topk = topk, .cand = cand, .attn = attn, .experts = experts, .xmma = xmma, .ix = ix, .engram = engram, .glue = glue, .compress = compress, .mv = mv, .vision = vision, .hc = hc, .fold = fold, .glue_tmp = glue_tmp, .fused = fused, .fused2 = fused2, .fused3 = fused3, .fused4 = fused4, .fused5 = fused5, .glue_tmp2 = glue_tmp2 }; } pub fn destroy(s: *Suites) void { @@ -151,6 +155,7 @@ pub const Suites = struct { s.glue_tmp.destroy(); s.fold.destroy(); s.hc.destroy(); + s.vision.destroy(); s.mv.destroy(); s.compress.destroy(); s.glue.destroy(); @@ -192,6 +197,7 @@ pub const Suites = struct { t.failed += try s.mv.head(io, seed, std.c.getenv("TF_DSV41_TIME") != null); } t.failed += try s.glue.tables(); + t.failed += try s.vision.invariance(seed); t.failed += try s.compress.invariance(seed); t.failed += try s.compress.equality(seed); t.failed += try s.hc.invariance(seed); diff --git a/zig/tests/contract/metal_dsv41_vision.zig b/zig/tests/contract/metal_dsv41_vision.zig new file mode 100644 index 0000000000..15ee6cae90 --- /dev/null +++ b/zig/tests/contract/metal_dsv41_vision.zig @@ -0,0 +1,409 @@ +//! DeepSeek V4.1's vision kernels under the contract on Metal: epilogues, RoPE, attention, SwiGLU, layout. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const ref = @import("dsv41_vision_ref.zig"); + +const c = contract.spec; +const g = contract.gates; +const vis = tf.dsv41_metal.vision; +const base = "dsv41 vision "; +const u: f64 = 0x1p-8; // half a bf16 ulp, relative + +pub const Fam = enum { epi_bias, epi_gelu, epi_res, epi_res0, qkv, swiglu, unshuffle, merge, attn }; + +fn noSums(_: c.Shape) usize { + return 0; +} + +pub fn spec(f: Fam, name: []const u8) c.Spec { + const t32: c.Tile = .{ .size = 32 }; + return switch (f) { + .epi_bias, .epi_gelu, .epi_res, .epi_res0 => .{ .name = name, .accum = .none, .depth = noSums, .out = .bf16, .tiles = .{ .{}, t32, .{} }, .served = &.{ + .{ .rows = 300, .n = 1024, .k = 1 }, .{ .rows = 64, .n = 5120, .k = 1 }, .{ .rows = 33, .n = 3072, .k = 1 }, .{ .rows = 5, .n = 2816, .k = 1 }, .{ .rows = 1, .n = 1024, .k = 1 }, + } }, + .qkv => .{ .name = name, .accum = .none, .depth = noSums, .out = .bf16, .tiles = .{ .{}, .{}, .{ .size = 64, .step = 2, .min = 2 } }, .served = &.{ + .{ .rows = 96, .n = 16, .k = 64 }, .{ .rows = 7, .n = 2, .k = 64 }, .{ .rows = 1, .n = 16, .k = 64 }, .{ .rows = 33, .n = 3, .k = 64 }, + } }, + .swiglu => .{ .name = name, .accum = .none, .depth = noSums, .out = .bf16, .tiles = .{ .{}, t32, .{} }, .served = &.{ + .{ .rows = 128, .n = 2816, .k = 1 }, .{ .rows = 5, .n = 192, .k = 1 }, .{ .rows = 1, .n = 2816, .k = 1 }, .{ .rows = 40, .n = 1000, .k = 1 }, + } }, + .unshuffle => .{ .name = name, .accum = .none, .depth = noSums, .out = .bf16, .tiles = .{ .{ .size = 3 }, .{ .size = 3 }, .{ .size = 32 } }, .served = &.{ + .{ .rows = 36, .n = 43, .k = 1024 }, .{ .rows = 5, .n = 7, .k = 128 }, .{ .rows = 1, .n = 1, .k = 64 }, .{ .rows = 9, .n = 4, .k = 100 }, + } }, + .merge => .{ .name = name, .accum = .none, .depth = noSums, .out = .bf16, .tiles = .{ .{}, t32, .{} }, .served = &.{ + .{ .rows = 200, .n = 5120, .k = 40 }, .{ .rows = 16, .n = 256, .k = 9 }, .{ .rows = 1, .n = 5120, .k = 1 }, .{ .rows = 70, .n = 100, .k = 20 }, + } }, + .attn => .{ .name = name, .accum = .none, .depth = noSums, .out = .bf16, .tiles = .{ .{ .size = 32 }, .{}, .{} }, .served = &.{ + .{ .rows = 96, .n = 16, .k = 64 }, .{ .rows = 37, .n = 4, .k = 64 }, .{ .rows = 1, .n = 16, .k = 64 }, .{ .rows = 130, .n = 2, .k = 64 }, + } }, + }; +} + +fn pow2(e: i32) f32 { + return std.math.pow(f32, 2, @floatFromInt(e)); +} + +fn spread(rng: std.Random, lo: i32, hi: i32) f32 { + return (rng.float(f32) * 2 - 1) * pow2(rng.intRangeAtMost(i32, lo, hi)); +} + +fn at(r: rig_.Ref, extra: usize) vis.Ref { + return .{ .buf = r.buf, .off = r.off + extra }; +} + +/// The five planted values of the rows family (a NaN, both infinities, two at one row's ends) into fp32 rows. +fn plant32(x: []f32, rows: usize, width: usize) usize { + const place = [_][2]usize{ .{ 0, width / 3 }, .{ 1, width / 2 }, .{ 2, width - 1 }, .{ 3, 0 }, .{ 3, width - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + var n: usize = 0; + for (place, put) |p, v| if (p[0] < rows) { + x[p[0] * width + p[1]] = v; + n += 1; + }; + return n; +} + +fn bf16s(a: std.mem.Allocator, n: usize, rng: std.Random, lo: i32, hi: i32) ![]u16 { + const v = try a.alloc(u16, n); + for (v) |*o| o.* = c.toBf16(spread(rng, lo, hi)); + return v; +} + +fn f32s(a: std.mem.Allocator, n: usize, rng: std.Random, lo: i32, hi: i32) ![]f32 { + const v = try a.alloc(f32, n); + for (v) |*o| o.* = spread(rng, lo, hi); + return v; +} + +const vt = g.hooks(Proj); + +pub const Proj = struct { + a: std.mem.Allocator, + arena: std.heap.ArenaAllocator, + rig: rig_.Rig, + fam: Fam, + name: []const u8, + /// the planted fault this kernel's text carries (0: the kernel as served) + ctl: u8 = 0, + pipes: vis.Pipelines, + shape: c.Shape = undefined, + acc: []f32 = &.{}, + x: []u16 = &.{}, + bias: []u16 = &.{}, + res: []u16 = &.{}, + cs: [][2]f32 = &.{}, + src: []i32 = &.{}, + in: [4]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, f: Fam, name: []const u8, ctl: u8) !Proj { + const sources = @import("kernel_sources"); + const text = try std.fmt.allocPrint(a, "#define TF_CTL {d}\n{s}", .{ ctl, sources.dsv41.vision }); + defer a.free(text); + return .{ .a = a, .arena = std.heap.ArenaAllocator.init(a), .rig = .{ .a = a, .device = device, .queue = queue }, .fam = f, .name = name, .ctl = ctl, .pipes = try vis.Pipelines.loadSource(device, text) }; + } + + pub fn deinit(p: *Proj) void { + p.arena.deinit(); + p.pipes.deinit(); + p.rig.deinit(); + } + + pub fn kernel(p: *Proj) g.Kernel { + const control: ?g.Gate = if (p.ctl != 0) .shapes else null; + return .{ .spec = spec(p.fam, p.name), .ctx = p, .vt = &vt, .control = control }; + } + + fn isEpi(p: *const Proj) bool { + return p.fam == .epi_bias or p.fam == .epi_gelu or p.fam == .epi_res or p.fam == .epi_res0; + } + + /// qkv: the three parts of every head; the others: a row's outputs. + pub fn count(p: *Proj) usize { + const s = p.shape; + return switch (p.fam) { + .qkv => s.rows * 3 * s.n * s.k, + .swiglu => s.rows * s.n, + .unshuffle => ((s.rows + 2) / 3) * ((s.n + 2) / 3) * s.k * 9, + .merge => s.rows * 4 * s.n, + .attn => s.rows * s.n * 64, + else => s.rows * s.n, + }; + } + + pub fn prepare(p: *Proj, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (p.fam == .qkv and shape.k % 2 != 0) return false; + if (p.fam == .attn and (shape.k != 64 or shape.rows > 500)) return false; + p.rig.reset(); + _ = p.arena.reset(.retain_capacity); + const al = p.arena.allocator(); + p.shape = shape; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const nf = fill == .nonfinite; + const s = shape; + switch (p.fam) { + .epi_bias, .epi_gelu, .epi_res, .epi_res0 => { + p.acc = try f32s(al, s.rows * s.n, rng, -5, 3); + p.bias = try bf16s(al, s.n, rng, -4, 1); + p.res = try bf16s(al, s.rows * s.n, rng, -4, 3); + if (nf) _ = plant32(p.acc, s.rows, s.n); + if (nf and p.fam == .epi_res) p.res[s.n + 1] = c.toBf16(-std.math.inf(f32)); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.acc)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.bias)); + p.in[2] = try p.rig.input(std.mem.sliceAsBytes(p.res)); + }, + .qkv => { + const w = 3 * s.n * s.k; + p.acc = try f32s(al, s.rows * w, rng, -4, 2); + p.bias = try bf16s(al, w, rng, -4, 1); + p.cs = try al.alloc([2]f32, s.rows * s.k / 2); + for (p.cs) |*o| { + const angle = (rng.float(f64) * 2 - 1) * std.math.pi; + o.* = .{ @floatCast(@cos(angle)), @floatCast(@sin(angle)) }; + } + if (nf) _ = plant32(p.acc, s.rows, w); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.acc)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.bias)); + p.in[2] = try p.rig.input(std.mem.sliceAsBytes(p.cs)); + }, + .swiglu => { + p.acc = try f32s(al, s.rows * 2 * s.n, rng, -5, 3); + if (nf) _ = plant32(p.acc, s.rows, 2 * s.n); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.acc)); + }, + .unshuffle => { + p.x = try bf16s(al, s.rows * s.n * s.k, rng, -3, 3); + if (nf) p.x[0] = c.toBf16(std.math.nan(f32)); + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + }, + .merge => { + p.x = try bf16s(al, s.k * s.n, rng, -3, 3); + p.bias = try bf16s(al, 3 * s.n, rng, -3, 3); + p.src = try al.alloc(i32, s.rows); + for (p.src, 0..) |*o, i| o.* = if (i % 5 == 4) -1 - @as(i32, @intCast(i % 3)) else @intCast(rng.uintLessThan(usize, s.k)); + if (nf) { + // a NaN feature row and an infinite delimiter, each read by a row + p.x[0] = c.toBf16(std.math.nan(f32)); + p.bias[s.n + 1] = c.toBf16(std.math.inf(f32)); + p.src[0] = 0; + p.src[4] = -2; + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + p.in[1] = try p.rig.input(std.mem.sliceAsBytes(p.bias)); + p.in[2] = try p.rig.input(std.mem.sliceAsBytes(p.src)); + }, + .attn => { + const D = s.n * 64; + p.x = try bf16s(al, s.rows * 3 * D, rng, -2, 1); + if (nf) { + // NaN and infinities in four queries of the first head only: those rows' head goes non-finite + const qs = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32) }; + for (qs, 0..) |v, r| if (r < s.rows) { + p.x[r * 3 * D + 3 + r] = c.toBf16(v); + }; + } + p.in[0] = try p.rig.input(std.mem.sliceAsBytes(p.x)); + }, + } + p.out = try p.rig.output(p.count() * 2); + return true; + } + + /// Rows r0 to r0 + n of the prepared call as one launch (the attention and unshuffle take whole calls). + pub fn encode(p: *Proj, e: anytype, r0: usize, n: usize) !void { + const s = p.shape; + const pl = &p.pipes; + const n32: u32 = @intCast(n); + switch (p.fam) { + .epi_bias, .epi_gelu, .epi_res, .epi_res0 => { + const w = s.n; + const bias: ?vis.Ref = if (p.fam == .epi_res0) null else at(p.in[1], 0); + const res: ?vis.Ref = if (p.fam == .epi_res or p.fam == .epi_res0) at(p.in[2], r0 * w * 2) else null; + vis.epilogue(pl, e, at(p.in[0], r0 * w * 4), bias, res, at(p.out, r0 * w * 2), n32, @intCast(w), p.fam == .epi_gelu); + }, + .qkv => vis.qkv(pl, e, at(p.in[0], r0 * 3 * s.n * s.k * 4), at(p.in[1], 0), at(p.in[2], r0 * s.k / 2 * 8), at(p.out, r0 * 3 * s.n * s.k * 2), n32, @intCast(s.n), @intCast(s.k)), + .swiglu => vis.swiglu(pl, e, at(p.in[0], r0 * 2 * s.n * 4), at(p.out, r0 * s.n * 2), n32, @intCast(s.n)), + .unshuffle => vis.unshuffle(pl, e, at(p.in[0], 0), at(p.out, 0), @intCast(s.rows), @intCast(s.n), @intCast(s.k), 3), + .merge => vis.merge(pl, e, at(p.in[0], 0), at(p.in[1], 0), at(p.in[2], r0 * 4), at(p.out, r0 * 4 * s.n * 2), n32, @intCast(s.n)), + .attn => vis.attention(pl, e, at(p.in[0], 0), at(p.out, 0), @intCast(s.rows), @intCast(s.n)), + } + } + + pub fn run(p: *Proj) !void { + const pass = p.rig.begin(); + try p.encode(pass.e, 0, p.shape.rows); + try p.rig.finish(pass); + } + + pub fn reference(p: *Proj, want: []f64, mag: []f64) !void { + const s = p.shape; + @memset(mag, 0); + switch (p.fam) { + .epi_bias, .epi_gelu, .epi_res, .epi_res0 => for (want, mag, 0..) |*w, *m, i| { + const z: f64 = @as(f64, p.acc[i]) + (if (p.fam == .epi_res0) 0 else c.fromBf16(p.bias[i % s.n])); + const y = if (p.fam == .epi_gelu) ref.gelu(z) else z; + const r: f64 = if (p.fam == .epi_res or p.fam == .epi_res0) c.fromBf16(p.res[i]) else 0; + w.* = y + r; + m.* = @abs(z) + @abs(r); + }, + .qkv => for (0..s.rows) |r| for (0..3 * s.n) |ph| for (0..s.k / 2) |j| { + const part = ph / s.n; + const o = (r * 3 * s.n + ph) * s.k + j; + const col = ph * s.k + j; + const x1: f64 = @as(f64, p.acc[r * 3 * s.n * s.k + col]) + c.fromBf16(p.bias[col]); + const x2: f64 = @as(f64, p.acc[r * 3 * s.n * s.k + col + s.k / 2]) + c.fromBf16(p.bias[col + s.k / 2]); + if (part == 2) { + want[o] = x1; + want[o + s.k / 2] = x2; + mag[o] = @abs(x1); + mag[o + s.k / 2] = @abs(x2); + continue; + } + const cs = p.cs[r * s.k / 2 + j]; + const co: f64 = cs[0]; + const si: f64 = cs[1]; + want[o] = x1 * co - x2 * si; + want[o + s.k / 2] = x2 * co + x1 * si; + mag[o] = @abs(x1 * co) + @abs(x2 * si); + mag[o + s.k / 2] = @abs(x2 * co) + @abs(x1 * si); + }, + .swiglu => for (0..s.rows) |r| for (0..s.n) |j| { + const gt: f64 = p.acc[r * 2 * s.n + j]; + const up: f64 = p.acc[r * 2 * s.n + s.n + j]; + want[r * s.n + j] = ref.silu(gt) * up; + mag[r * s.n + j] = @abs(gt); // the gate's own size sets how far its rounding moves the output + }, + .unshuffle => { + const lw = (s.n + 2) / 3; + for (0..((s.rows + 2) / 3) * lw) |l| for (0..s.k) |ch| for (0..9) |kk| { + const py = (l / lw) * 3 + kk / 3; + const px = (l % lw) * 3 + kk % 3; + want[(l * s.k + ch) * 9 + kk] = if (py < s.rows and px < s.n) c.fromBf16(p.x[(py * s.n + px) * s.k + ch]) else 0; + }; + }, + .merge => for (0..s.rows) |r| for (0..4) |st| for (0..s.n) |d| { + const sr = p.src[r]; + want[(r * 4 + st) * s.n + d] = if (sr >= 0) c.fromBf16(p.x[@as(usize, @intCast(sr)) * s.n + d]) else c.fromBf16(p.bias[@as(usize, @intCast(-1 - sr)) * s.n + d]); + }, + .attn => { + const D = s.n * 64; + const q = try p.a.alloc(f64, 64); + defer p.a.free(q); + const kk = try p.a.alloc(f64, s.rows * 3 * D); + defer p.a.free(kk); + for (kk, p.x) |*o, v| o.* = c.fromBf16(v); + for (0..s.rows) |r| for (0..s.n) |h| { + for (q, 0..) |*o, d| o.* = kk[r * 3 * D + h * 64 + d]; + ref.attend(want[(r * s.n + h) * 64 ..][0..64], mag[(r * s.n + h) * 64 ..][0..64], q, kk[D + h * 64 ..], kk[2 * D + h * 64 ..], s.rows, 64, 3 * D); + }; + }, + } + } + + /// The allowed |got - want|: the intermediate bf16 roundings of each formula, fp32's own, then the output's. + pub fn bounds(p: *Proj, want: []const f64, mag: []const f64, allow: []f64) !void { + const e32: f64 = 0x1p-24; + for (allow, want, mag) |*l, w, m| { + const work: f64 = switch (p.fam) { + .epi_bias => e32 * m, + .epi_gelu => 1.2 * u * m + 0x1p-20 * m, + .epi_res, .epi_res0 => u * m + e32 * m, + .qkv => (u + 4 * e32) * m, + .swiglu => u * (3 + m) * @abs(w) + 0x1p-20 * @abs(w), + .attn => 0x1p-14 * m, + else => 0, + }; + l.* = work + (if (p.fam == .unshuffle or p.fam == .merge) 0 else u * (@abs(w) + work)) + 0x1p-133; + } + } +}; + +/// Every registered variant and the launch-split probes after the gates. +pub const Suite = struct { + a: std.mem.Allocator, + probes: std.ArrayList(*Proj) = .empty, + + const E = struct { f: Fam, name: []const u8, ctl: u8 = 0 }; + const list = [_]E{ + .{ .f = .epi_bias, .name = base ++ "epilogue bias tf_dsv41_vis_epi" }, + .{ .f = .epi_gelu, .name = base ++ "epilogue bias gelu tf_dsv41_vis_epi" }, + .{ .f = .epi_res, .name = base ++ "epilogue bias residual tf_dsv41_vis_epi" }, + .{ .f = .epi_res0, .name = base ++ "epilogue residual tf_dsv41_vis_epi" }, + .{ .f = .qkv, .name = base ++ "qkv bias rope tf_dsv41_vis_qkv" }, + .{ .f = .swiglu, .name = base ++ "swiglu tf_dsv41_vis_swiglu" }, + .{ .f = .unshuffle, .name = base ++ "unshuffle tf_dsv41_vis_unshuffle" }, + .{ .f = .merge, .name = base ++ "merge tf_dsv41_vis_merge" }, + .{ .f = .attn, .name = base ++ "attention tf_dsv41_vis_attn" }, + .{ .f = .epi_bias, .name = base ++ "epilogue bias control: the bias dropped", .ctl = 1 }, + .{ .f = .qkv, .name = base ++ "qkv control: sin and cos swapped", .ctl = 2 }, + .{ .f = .attn, .name = base ++ "attention control: the max not subtracted", .ctl = 3 }, + .{ .f = .swiglu, .name = base ++ "swiglu control: gate and up swapped", .ctl = 4 }, + .{ .f = .unshuffle, .name = base ++ "unshuffle control: a stride off by one", .ctl = 5 }, + .{ .f = .merge, .name = base ++ "merge control: a delimiter one row late", .ctl = 6 }, + }; + + pub fn create(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !*Suite { + const s = try a.create(Suite); + s.* = .{ .a = a }; + for (list) |entry| { + if (only) |text| if (std.mem.indexOf(u8, entry.name, text) == null) continue; + const p = try a.create(Proj); + p.* = try Proj.init(a, device, queue, entry.f, entry.name, entry.ctl); + try s.probes.append(a, p); + try kernels.append(a, p.kernel()); + } + return s; + } + + pub fn destroy(s: *Suite) void { + for (s.probes.items) |p| { + p.deinit(); + s.a.destroy(p); + } + s.probes.deinit(s.a); + s.a.destroy(s); + } + + fn splits(f: Fam) bool { + return f != .attn and f != .unshuffle; + } + + /// Launches of 1, 7 and rows/3 rows give the whole launch's bytes. + pub fn invariance(s: *Suite, seed: u64) !usize { + var bad: usize = 0; + var launches: usize = 0; + var kernels: usize = 0; + for (s.probes.items) |p| { + if (!splits(p.fam)) continue; + kernels += 1; + const shape = spec(p.fam, p.name).served[1]; + if (!try p.prepare(.{ .rows = shape.rows + 20, .n = shape.n, .k = shape.k }, .random, seed)) return error.Refused; + try p.run(); + const len = p.count() * 2; + const whole = try s.a.dupe(u8, p.rig.bytes(p.out, len)); + defer s.a.free(whole); + const rows = p.shape.rows; + for ([_]usize{ 1, 7, @max(1, rows / 3) }) |step| { + @memset(p.rig.bytes(p.out, len), 0x55); + const pass = p.rig.begin(); + var r0: usize = 0; + while (r0 < rows) : (r0 += step) { + try p.encode(pass.e, r0, @min(step, rows - r0)); + launches += 1; + } + try p.rig.finish(pass); + if (!std.mem.eql(u8, whole, p.rig.bytes(p.out, len))) { + bad += 1; + std.debug.print("FAIL row invariance {s}: {d} rows a launch differs from one launch\n", .{ p.name, step }); + } + } + } + if (kernels == 0) return 0; + std.debug.print("{s} row invariance {s}: {d} kernels at 1, 7 and rows/3 rows a launch ({d} launches) byte-equal to one launch, {d} differ\n", .{ if (bad == 0) "PASS" else "FAIL", base, kernels, launches, bad }); + return bad; + } +}; diff --git a/zig/tests/dsv41_metal/boot_test.zig b/zig/tests/dsv41_metal/boot_test.zig new file mode 100644 index 0000000000..62f45a6ca4 --- /dev/null +++ b/zig/tests/dsv41_metal/boot_test.zig @@ -0,0 +1,171 @@ +//! A small release-format checkpoint folder booted the way the server boots one (boot.zig), and served a reply. +const std = @import("std"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); +const fixture = @import("fixture.zig"); +const fixtures = @import("dsv41_serve_fixtures"); + +const api = host.api; +const Served = host.serve_lanes.Served; +const Config = host.config.Config; +const testing = std.testing; + +/// The mini tokenizer's special tokens move from ids 1000 .. 1010 to right after its 266 byte-level ones. +const special0 = 266; + +/// The tiny drafted model without Engram, its vocabulary the mini tokenizer's (277 ids in 384; BOS 266, EOS 267). +pub fn config() Config { + var c = fixture.dsparkConfig(); + c.vocab = 384; + c.engram_layers = host.config.Ids.of(&.{}); + c.engram_rows = host.config.List(u64, 8).of(&.{}); + c.bos = special0; + c.eos = special0 + 1; + return c; +} + +/// The mini tokenizer with its special tokens at special0 .. +fn tokenizer(a: std.mem.Allocator) ![]u8 { + var text = try a.dupe(u8, fixtures.mini_tokenizer); + for (0..11) |i| { + var from: [16]u8 = undefined; + var to: [16]u8 = undefined; + const next = try std.mem.replaceOwned(u8, a, text, try std.fmt.bufPrint(&from, "\"id\": {d},", .{1000 + i}), try std.fmt.bufPrint(&to, "\"id\": {d},", .{special0 + i})); + a.free(text); + text = next; + } + return text; +} + +fn list(w: *std.Io.Writer, xs: anytype) !void { + try w.writeByte('['); + for (xs, 0..) |x, i| try w.print("{s}{d}", .{ if (i == 0) "" else ",", x }); + try w.writeByte(']'); +} + +/// config.json as a release checkpoint writes it, at `c`'s shapes. +fn configJson(a: std.mem.Allocator, c: *const Config) ![]u8 { + var out: std.Io.Writer.Allocating = .init(a); + defer out.deinit(); + const w = &out.writer; + try w.print("{{\"model_type\":\"deepseek_v41\",\"bos_token_id\":{d},\"eos_token_id\":{d},", .{ c.bos, c.eos }); + try w.writeAll("\"quantization_config\":{\"quant_method\":\"fp8\",\"weight_block_size\":[32,32],\"scale_fmt\":\"ue8m0\",\"expert_dtype\":\"fp4\"},"); + const ints = .{ + .{ "vocab_size", c.vocab }, .{ "hidden_size", c.hidden }, .{ "num_hidden_layers", c.layers }, + .{ "num_attention_heads", c.heads }, .{ "head_dim", c.head_dim }, .{ "qk_rope_head_dim", c.rope_dim }, + .{ "q_lora_rank", c.q_lora }, .{ "o_lora_rank", c.o_lora }, .{ "o_groups", c.o_groups }, + .{ "sliding_window", c.window }, .{ "n_routed_experts", c.experts }, .{ "n_shared_experts", c.shared_experts }, + .{ "num_experts_per_tok", c.topk }, .{ "moe_intermediate_size", c.expert_width }, .{ "index_n_heads", c.index_heads }, + .{ "index_head_dim", c.index_dim }, .{ "index_topk", c.index_topk }, .{ "candidate_source_layer_id", c.candidate_source }, + .{ "candidate_topk_blocks", c.candidate_blocks }, .{ "candidate_block_size", c.candidate_block_size }, .{ "hc_mult", c.hc_mult }, + .{ "hc_sinkhorn_iters", c.hc_sinkhorn_iters }, .{ "engram_max_ngram_size", c.engram_max_ngram }, .{ "engram_vocab_size", c.engram_vocab }, + .{ "engram_n_heads", c.engram_heads }, .{ "engram_head_dim", c.engram_head_dim }, .{ "engram_pad_token_id", c.engram_pad }, + .{ "engram_compressed_vocab_size", c.engram_compressed_vocab }, .{ "num_nextn_predict_layers", c.mtp_layers }, .{ "dspark_block_size", c.dspark_block }, + .{ "dspark_noise_token_id", c.dspark_noise_token }, .{ "dspark_markov_rank", c.dspark_markov_rank }, .{ "dspark_n_routed_experts", c.dspark_experts }, + .{ "dspark_num_experts_per_tok", c.dspark_topk }, .{ "num_key_value_heads", 1 }, .{ "max_position_embeddings", c.rope.max_positions }, + }; + inline for (ints) |p| try w.print("\"{s}\":{d},", .{ p[0], p[1] }); + const floats = .{ + .{ "rms_norm_eps", c.eps }, .{ "routed_scaling_factor", c.routed_scale }, .{ "swiglu_limit", c.swiglu_limit }, + .{ "hc_eps", c.hc_eps }, .{ "hc_post_alpha", c.hc_post_alpha }, .{ "engram_gate_clamp", c.engram_gate_clamp }, + .{ "rope_theta", c.rope.theta }, .{ "compress_rope_theta", c.rope.compress_theta }, + }; + inline for (floats) |p| try w.print("\"{s}\":{e},", .{ p[0], p[1] }); + const lists = .{ + .{ "compress_ratios", c.compress_ratios.items() }, .{ "kv_source_layer_ids", c.kv_sources.items() }, + .{ "index_source_layer_ids", c.index_sources.items() }, .{ "engram_layer_ids", c.engram_layers.items() }, + .{ "dspark_target_layer_ids", c.dspark_targets.items() }, .{ "engram_num_embeddings", c.engram_rows.items() }, + }; + inline for (lists) |p| { + try w.print("\"{s}\":", .{p[0]}); + try list(w, p[1]); + try w.writeByte(','); + } + try w.print("\"norm_topk_prob\":{},\"scoring_func\":\"sqrtsoftplus\",\"topk_method\":\"noaux_tc\",", .{c.norm_topk}); + try w.print("\"rope_scaling\":{{\"rope_type\":\"yarn\",\"factor\":{e},\"beta_fast\":{e},\"beta_slow\":{e},\"original_max_position_embeddings\":{d}}}}}", .{ c.rope.factor, c.rope.beta_fast, c.rope.beta_slow, c.rope.original_max_positions }); + return out.toOwnedSlice(); +} + +/// The folder: config.json, the weights (3 shards and their index), the mini tokenizer. +pub fn write(a: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, seed: u64) !void { + var c = config(); + try c.validate(); + const json = try configJson(a, &c); + defer a.free(json); + try dir.writeFile(io, .{ .sub_path = "config.json", .data = json }); + const img = try fixture.image(a, &c, seed); + defer a.free(img); + try @import("shards.zig").write(a, io, dir, img, 3, 0); + const tok = try tokenizer(a); + defer a.free(tok); + try dir.writeFile(io, .{ .sub_path = "tokenizer.json", .data = tok }); +} + +test "dsv41 metal boot: a small checkpoint folder parses back to its config, boots and replies, drafted == draft:false" { + const a = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try write(a, io, tmp.dir, 81); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + // TF_DSV41_SMOKE_DIR: the same folder written there too, for a server's smoke boot + if (std.c.getenv("TF_DSV41_SMOKE_DIR")) |d| { + var out = try std.Io.Dir.cwd().openDir(io, std.mem.span(d), .{}); + defer out.close(io); + try write(a, io, out, 81); + } + const want = config(); + const m = try pkg.boot.Boot.open(a, io, dir, .{ .slots = 2, .context = 1024, .small = true, .cache_bytes = 0 }); + const got = m.parts.env.c; + try testing.expectEqual(want.vocab, got.vocab); + try testing.expectEqualSlices(u32, want.compress_ratios.items(), got.compress_ratios.items()); + try testing.expectEqual(want.eos, got.eos); + const s = try Served.open(a, io, m.model(), dir, null); + defer Served.close(s); + const e = s.engine(); + const prompt = [_]u32{ special0, special0 + 5, 70, 71, 72, 73, special0 + 6 }; + var replies: [2][24]u32 = undefined; + for ([_]bool{ true, false }, &replies, 0..) |drafts, *r, i| { + var box: Box = .{}; + const req: api.Request = .{ .prompt = &prompt, .max_tokens = 24, .drafts = drafts }; + try e.submit(@intCast(900 + i), &req, .{ .ctx = &box, .event = Box.event }); + try box.wait(io); + try testing.expectEqual(@as(usize, 24), box.n); + r.* = box.out; + } + std.debug.print("{s} dsv41 metal boot: {d} layers from 3 shards, a reply of 24 tokens, drafted == draft:false\n", .{ if (std.mem.eql(u32, &replies[0], &replies[1])) "PASS" else "FAIL", got.layers }); + try testing.expectEqualSlices(u32, &replies[1], &replies[0]); +} + +const Box = struct { + mutex: std.Io.Mutex = .init, + out: [24]u32 = undefined, + n: usize = 0, + done: bool = false, + + fn event(ctx: *anyopaque, _: api.Id, e: *const api.Event) void { + const b: *Box = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(testing.io); + defer b.mutex.unlock(testing.io); + switch (e.*) { + .tokens => |ts| for (ts) |t| if (b.n < b.out.len) { + b.out[b.n] = t; + b.n += 1; + }, + .finished => b.done = true, + else => {}, + } + } + + fn wait(b: *Box, io: std.Io) !void { + for (0..300_000) |_| { + b.mutex.lockUncancelable(io); + const done = b.done; + b.mutex.unlock(io); + if (done) return; + std.Io.sleep(io, .fromMilliseconds(1), .awake) catch {}; + } + return error.Timeout; + } +}; diff --git a/zig/tests/dsv41_metal/fixture_serve.zig b/zig/tests/dsv41_metal/fixture_serve.zig new file mode 100644 index 0000000000..8f276cc7c5 --- /dev/null +++ b/zig/tests/dsv41_metal/fixture_serve.zig @@ -0,0 +1,69 @@ +//! The tiny model served on Metal: prompt split, Engram before it, a DSpark drafter; its serve.Metal over the rig. +const std = @import("std"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); +const fixture = @import("fixture.zig"); +const ff = @import("fixture_fwd.zig"); + +pub const limit = 1280; + +/// The prompt model (ced's split at layer 4, Engram at 1 and 3) with fixture.dsparkConfig's drafter. +pub fn config() host.config.Config { + var c = fixture.promptConfig(); + const d = fixture.dsparkConfig(); + c.mtp_layers = d.mtp_layers; + c.dspark_block = d.dspark_block; + c.dspark_noise_token = d.dspark_noise_token; + c.dspark_markov_rank = d.dspark_markov_rank; + c.dspark_experts = d.dspark_experts; + c.dspark_topk = d.dspark_topk; + return c; +} + +/// The rig (weights, Engram rows and hashes) and the drafter's weights the served model borrows. +const Held = struct { + a: std.mem.Allocator, + f: *ff.Fwd, + w: pkg.dspark.Weights, + + fn close(ctx: *anyopaque) void { + const h: *Held = @ptrCast(@alignCast(ctx)); + h.w.deinit(); + h.f.deinit(); + h.a.destroy(h); + } +}; + +pub const Opts = struct { slots: u32 = 4, drafts: bool = true, cache_bytes: u64 = 0, chunk: u32 = 64 }; + +/// The tiny model from `seed` served on this GPU; the returned model owns the rig (closed with it). +pub fn open(a: std.mem.Allocator, io: std.Io, seed: u64, o: Opts) !*pkg.serve.Metal { + const h = try a.create(Held); + errdefer a.destroy(h); + h.a = a; + h.f = try ff.Fwd.initWith(a, io, config(), limit, seed); + errdefer h.f.deinit(); + const r = h.f.r; + { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const img = try fixture.image(a, &r.c, seed); + defer a.free(img); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = img }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var pack = try host.pack.Pack.open(a, io, dir); + defer pack.deinit(); + h.w = try pkg.dspark.loadWeights(a, io, r.device, &pack, &r.c); + } + errdefer h.w.deinit(); + const env: pkg.target.Env = .{ .device = r.device, .queue = r.queue, .c = &r.c, .p = &r.p, .tables = &r.tables, .s = &r.s, .model = h.f.model(), .lookup = h.f.lookup() }; + const parts: pkg.serve.Parts = .{ .env = env, .top = &r.top.?, .dspark = &h.w, .owner = .{ .ctx = h, .close = Held.close } }; + return pkg.serve.Metal.open(a, parts, .{ .slots = o.slots, .limit = limit, .drafts = o.drafts, .cache_bytes = o.cache_bytes, .prompt = .{ .chunk = o.chunk } }); +} + +/// Closes a model `open` made (the served engine closes its own through serve_lanes). +pub fn close(m: *pkg.serve.Metal) void { + const md = m.model(); + md.close(md.ptr); +} diff --git a/zig/tests/dsv41_metal/layer_test.zig b/zig/tests/dsv41_metal/layer_test.zig index 0cb448c372..83ed237f9b 100644 --- a/zig/tests/dsv41_metal/layer_test.zig +++ b/zig/tests/dsv41_metal/layer_test.zig @@ -17,6 +17,7 @@ test { _ = @import("zerocopy_test.zig"); _ = @import("prompt_test.zig"); _ = @import("dspark_test.zig"); + _ = @import("multi_test.zig"); } test "dsv41 metal: every stage of every layer kind within its float64 bound" { diff --git a/zig/tests/dsv41_metal/multi_test.zig b/zig/tests/dsv41_metal/multi_test.zig new file mode 100644 index 0000000000..4819c786c6 --- /dev/null +++ b/zig/tests/dsv41_metal/multi_test.zig @@ -0,0 +1,219 @@ +//! Rows of several slots in one decode launch equal each slot's rows launched alone, bit for bit. +const std = @import("std"); +const mtl = @import("metal"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); +const fixture = @import("fixture.zig"); +const ff = @import("fixture_fwd.zig"); +const fs = @import("fixture_serve.zig"); + +const testing = std.testing; +const fwd = pkg.forward; +const limit = 1280; +const look = 3; + +/// One slot's share of a launch: where it stands and the ids its rows carry (after `look` ids of lookback). +const Part = struct { p0: u32, hist: []u32, n: usize }; + +fn finish(r: *fixture.Rig, f: *ff.Fwd, x: *const pkg.ctx.Ctx, rows: u32) !void { + const pass = fixture.Pass.begin(r.queue); + try fwd.encode(x, pass.e, f.model(), &f.b); + try pass.finish(); + _ = rows; +} + +/// Each part's rows alone on its slot (the single-stream path): logits [rows][vocab] and taps per part, concatenated. +fn alone(a: std.mem.Allocator, f: *ff.Fwd, slots: []pkg.kv.Slot, parts: []const Part, logits: []f32, taps: []u16) !void { + const r = f.r; + const c = &r.c; + const V: usize = c.vocab; + const TH = c.dspark_targets.len * c.hidden; + _ = a; + var r0: usize = 0; + for (parts, slots) |pt, *sl| { + try fwd.prepare(&f.b, c, pt.hist[look..][0..pt.n], pt.hist[0..look], f.lookup()); + try r.s.writeRows(c, pt.p0, @intCast(pt.n)); + const x = r.ctx(sl, pt.p0, @intCast(pt.n)); + try finish(r, f, &x, @intCast(pt.n)); + @memcpy(logits[r0 * V ..][0 .. pt.n * V], f.b.logitsOf(c, pt.n)); + @memcpy(taps[r0 * TH ..][0 .. pt.n * TH], sl.taps.?.slice(u16, pt.n * TH)); + r0 += pt.n; + } +} + +/// All parts' rows in one launch; the taps are read back from each slot. +fn together(f: *ff.Fwd, slots: []pkg.kv.Slot, parts: []const Part, logits: []f32, taps: []u16) !void { + const r = f.r; + const c = &r.c; + const V: usize = c.vocab; + const TH = c.dspark_targets.len * c.hidden; + var pieces: [16]fwd.Piece = undefined; + var spans: [16]pkg.scratch.Span = undefined; + var segs: [16]pkg.ctx.Seg = undefined; + var r0: u32 = 0; + for (parts, slots, 0..) |pt, *sl, i| { + pieces[i] = .{ .ids = pt.hist[look..][0..pt.n], .lookback = pt.hist[0..look] }; + spans[i] = .{ .p0 = pt.p0, .n = @intCast(pt.n) }; + segs[i] = .{ .slot = sl, .p0 = pt.p0, .r0 = r0, .n = @intCast(pt.n) }; + r0 += @intCast(pt.n); + } + const k = parts.len; + try fwd.prepareSegs(&f.b, c, pieces[0..k], f.lookup()); + try r.s.writeSpans(c, spans[0..k]); + var x = r.ctx(&slots[0], parts[0].p0, r0); + x.segs = segs[0..k]; + try finish(r, f, &x, r0); + @memcpy(logits[0 .. r0 * V], f.b.logitsOf(c, r0)); + var at: usize = 0; + for (parts, slots) |pt, *sl| { + @memcpy(taps[at * TH ..][0 .. pt.n * TH], sl.taps.?.slice(u16, pt.n * TH)); + at += pt.n; + } +} + +fn reset(r: *fixture.Rig, to: *const pkg.kv.Slot, from: *const pkg.kv.Slot) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = r.queue.commandBuffer(); + to.copyFrom(cb, from); + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.CommandBufferFailed; +} + +/// Bytes that differ between two slots' every region: window, pool, keys, carry and the carry's inputs. +fn stateDiff(x: *const pkg.kv.Slot, y: *const pkg.kv.Slot) usize { + var n: usize = 0; + for (x.regions, y.regions) |p, q| { + const pairs = [_][2]?mtl.Buffer{ .{ p.window, q.window }, .{ p.pool, q.pool }, .{ p.keys, q.keys }, .{ p.carry, q.carry }, .{ p.inputs, q.inputs } }; + for (pairs) |pr| if (pr[0]) |bx| { + const len = bx.length(); + for (bx.contents()[0..len], pr[1].?.contents()[0..len]) |s, t| n += @intFromBool(s != t); + }; + } + return n; +} + +/// Slots at `starts` with `rows` rows each: one launch of all vs one launch each, bit for bit. +fn compare(a: std.mem.Allocator, f: *ff.Fwd, starts: []const u32, rows: []const usize, seed: u64) !usize { + const r = f.r; + const c = &r.c; + const k = starts.len; + var pristine: [4]pkg.kv.Slot = undefined; + var one: [4]pkg.kv.Slot = undefined; + var many: [4]pkg.kv.Slot = undefined; + var parts: [4]Part = undefined; + var made: usize = 0; + defer for (0..made) |i| { + pristine[i].deinit(); + one[i].deinit(); + many[i].deinit(); + a.free(parts[i].hist); + }; + var total: usize = 0; + for (0..k) |i| { + pristine[i] = try pkg.kv.Slot.init(a, r.device, c, limit); + one[i] = try pkg.kv.Slot.init(a, r.device, c, limit); + many[i] = try pkg.kv.Slot.init(a, r.device, c, limit); + parts[i] = .{ .p0 = starts[i], .hist = try ff.ids(a, c, look + rows[i], seed + 10 * i), .n = rows[i] }; + made = i + 1; + fixture.fillPast(c, &pristine[i], starts[i], seed + 100 * i); + try reset(r, &one[i], &pristine[i]); + try reset(r, &many[i], &pristine[i]); + total += rows[i]; + } + const TH = c.dspark_targets.len * c.hidden; + const l1 = try a.alloc(f32, total * c.vocab); + defer a.free(l1); + const ln = try a.alloc(f32, total * c.vocab); + defer a.free(ln); + const t1 = try a.alloc(u16, total * TH); + defer a.free(t1); + const tn = try a.alloc(u16, total * TH); + defer a.free(tn); + try alone(a, f, one[0..k], parts[0..k], l1, t1); + try together(f, many[0..k], parts[0..k], ln, tn); + var bad: usize = 0; + for (l1, ln) |s, t| bad += @intFromBool(@as(u32, @bitCast(s)) != @as(u32, @bitCast(t))); + var tb: usize = 0; + for (t1, tn) |s, t| tb += @intFromBool(s != t); + var st: usize = 0; + for (0..k) |i| st += stateDiff(&one[i], &many[i]); + const word = if (bad + tb + st == 0) "PASS" else "FAIL"; + std.debug.print("{s} dsv41 metal multi-slot forward: {d} slots, {d} rows in one launch vs one launch a slot: {d} logits, {d} tap values, {d} state bytes differ\n", .{ word, k, total, bad, tb, st }); + return bad + tb + st; +} + +test "dsv41 metal multi-slot forward: rows of 2, 3 and 4 slots in one launch equal each slot's launch alone" { + const a = testing.allocator; + const f = try ff.Fwd.init(a, testing.io, limit, 61); + defer f.deinit(); + var bad: usize = 0; + // positions of both parities, a slot near the start and one past the window; ratio 2 pools complete on odd rows + bad += try compare(a, f, &.{ 1101, 1104 }, &.{ 3, 5 }, 300); + bad += try compare(a, f, &.{ 1103, 640, 1100 }, &.{ 1, 4, 2 }, 310); + bad += try compare(a, f, &.{ 1101, 1104, 1099, 97 }, &.{ 4, 4, 4, 4 }, 320); + bad += try compare(a, f, &.{ 1101, 1102, 1103, 1104 }, &.{ 1, 1, 1, 1 }, 330); + bad += try compare(a, f, &.{ 7, 128, 1101, 1279 - 6 }, &.{ 2, 6, 1, 7 }, 340); + try testing.expectEqual(@as(usize, 0), bad); +} + +/// A model with the tiny drafter's config and prompts of different lengths on its slots. +fn rounds(m: *pkg.serve.Metal, packed_: bool, out: *std.ArrayList(u32)) !void { + const a = testing.allocator; + const tg = m.t.target(); + const lens = [4]usize{ 150, 211, 64, 333 }; + var next: [4]u32 = undefined; + var prompts: [4][]u32 = undefined; + for (&prompts, 0..) |*p, i| p.* = try ff.ids(a, &fs.config(), lens[i], 400 + i); + defer for (prompts) |p| a.free(p); + for (0..4) |i| next[i] = try tg.prefill(@intCast(i), prompts[i], null, lens[i]); + var pos = lens; + // verify rounds of one anchor and three drafts a slot; the keep cuts back a different number each time + for (0..4) |round| { + var toks: [4][4]u32 = undefined; + var draws: [4][4]u64 = undefined; + var chs: [4][4]u32 = undefined; + var segs: [4]host.iface.Segment = undefined; + var cs: [4][]u32 = undefined; + for (0..4) |i| { + toks[i][0] = next[i]; + for (1..4) |j| toks[i][j] = @intCast((next[i] * 7 + j * 13 + round * 5 + i) % 200); + for (&draws[i], 0..) |*d, j| d.* = pos[i] + 1 + j; + segs[i] = .{ .slot = @intCast(i), .start = pos[i], .tokens = &toks[i], .parents = null, .draws = &draws[i], .sampling = null }; + cs[i] = &chs[i]; + } + if (packed_) try tg.window(&segs, &cs) else for (0..4) |i| try tg.window(segs[i .. i + 1], cs[i .. i + 1]); + for (0..4) |i| { + const kept: u32 = @intCast(1 + (i + round) % 4); + const path = [_]u32{ 0, 1, 2, 3 }; + try tg.keep(@intCast(i), path[0..kept]); + try out.appendSlice(a, &chs[i]); + pos[i] += kept; + next[i] = chs[i][kept - 1]; + } + } + for (0..4) |i| tg.release(@intCast(i)); +} + +test "dsv41 metal multi-slot forward: a 4-slot DSpark verify round in one launch equals per-slot launches" { + const a = testing.allocator; + const m1 = try fs.open(a, testing.io, 71, .{}); + defer fs.close(m1); + const m2 = try fs.open(a, testing.io, 71, .{}); + defer fs.close(m2); + var o1: std.ArrayList(u32) = .empty; + defer o1.deinit(a); + var o2: std.ArrayList(u32) = .empty; + defer o2.deinit(a); + try rounds(m1, true, &o1); + try rounds(m2, false, &o2); + var st: usize = 0; + for (m1.t.slots, m2.t.slots) |*x, *y| st += stateDiff(&x.kv, &y.kv); + const same = std.mem.eql(u32, o1.items, o2.items); + const word = if (same and st == 0) "PASS" else "FAIL"; + std.debug.print("{s} dsv41 metal multi-slot verify rounds: {d} choices, {d} launches in one vs {d} apart, {d} state bytes differ\n", .{ word, o1.items.len, m1.t.stats.launches, m2.t.stats.launches, st }); + try testing.expect(same); + try testing.expectEqual(@as(usize, 0), st); + try testing.expect(m1.t.stats.launches < m2.t.stats.launches); +} diff --git a/zig/tests/dsv41_metal/serve_test.zig b/zig/tests/dsv41_metal/serve_test.zig new file mode 100644 index 0000000000..9bf82fd51a --- /dev/null +++ b/zig/tests/dsv41_metal/serve_test.zig @@ -0,0 +1,267 @@ +//! The tiny model served on Metal through serve_lanes and LaneHost: replies equal the target's own serial decode. +const std = @import("std"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); +const fs = @import("fixture_serve.zig"); +const ff = @import("fixture_fwd.zig"); + +const api = host.api; +const Served = host.serve_lanes.Served; +const testing = std.testing; +const gpa = testing.allocator; +const io = testing.io; +const seed = 71; + +test { + _ = @import("boot_test.zig"); +} + +/// A reply's tokens and finish, written on LaneHost's thread and read on the test's. +const Box = struct { + mutex: std.Io.Mutex = .init, + out: std.ArrayList(u32) = .empty, + done: bool = false, + reason: ?api.Reason = null, + drafted: u64 = 0, + cached: u32 = 0, + + fn event(ctx: *anyopaque, _: api.Id, e: *const api.Event) void { + const b: *Box = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(io); + defer b.mutex.unlock(io); + switch (e.*) { + .tokens => |ts| b.out.appendSlice(gpa, ts) catch {}, + .finished => |fin| { + b.done = true; + b.reason = fin.reason; + b.drafted = fin.stats.drafted; + }, + .prefilled => |n| b.cached = n, + else => {}, + } + } + + fn wait(b: *Box) !void { + for (0..600_000) |_| { + b.mutex.lockUncancelable(io); + const done = b.done; + b.mutex.unlock(io); + if (done) return; + std.Io.sleep(io, .fromMilliseconds(1), .awake) catch {}; + } + return error.Timeout; + } + + fn deinit(b: *Box) void { + b.out.deinit(gpa); + } +}; + +/// One request to send: its prompt, length, drafts and sampling. +const Ask = struct { prompt: []const u32, max: u32 = 24, drafts: bool = true, sampling: ?api.Sampling = null }; + +var next_id: api.Id = 1; +/// The prompt rows each reply of the last `serve` call took from kept state (its `prefilled` event), in ask order. +var cached_rows: [8]u32 = undefined; + +/// The asks sent together (or each after the last finished); each reply returned gpa-owned, the drafted rows summed. +fn serve(e: api.Engine, asks: []const Ask, together: bool, out: [][]u32, drafted: *u64) !void { + const boxes = try gpa.alloc(Box, asks.len); + defer gpa.free(boxes); + @memset(boxes, .{}); + defer for (boxes) |*b| b.deinit(); + const reqs = try gpa.alloc(api.Request, asks.len); + defer gpa.free(reqs); + for (reqs, asks) |*r, k| r.* = .{ .prompt = k.prompt, .max_tokens = k.max, .drafts = k.drafts, .sampling = k.sampling }; + for (reqs, boxes) |*r, *b| { + next_id += 1; + try e.submit(next_id, r, .{ .ctx = b, .event = Box.event }); + if (!together) try b.wait(); + } + for (boxes, out, asks, 0..) |*b, *o, k, i| { + try b.wait(); + try testing.expectEqual(api.Reason.length, b.reason.?); + try testing.expectEqual(@as(usize, k.max), b.out.items.len); + drafted.* += b.drafted; + cached_rows[i] = b.cached; + o.* = try b.out.toOwnedSlice(gpa); + } +} + +fn freeAll(xs: [][]u32) void { + for (xs) |x| gpa.free(x); +} + +/// The reference: the target's serial decode of each ask on a model of its own (no lanes, no cache, no drafts). +fn reference(asks: []const Ask, out: [][]u32) !void { + const m = try fs.open(gpa, io, seed, .{ .slots = 1, .drafts = false }); + defer fs.close(m); + for (asks, out) |k, *o| o.* = try m.t.serial(0, k.prompt, k.max, k.sampling); +} + +fn ids(n: usize, s: u64) ![]u32 { + return ff.ids(gpa, &fs.config(), n, s); +} + +fn say(ok: bool, comptime fmt: []const u8, args: anytype) void { + std.debug.print("{s} dsv41 metal serve: ", .{if (ok) "PASS" else "FAIL"}); + std.debug.print(fmt ++ "\n", args); +} + +fn same(want: []const []u32, got: []const []u32) usize { + var bad: usize = 0; + for (want, got) |w, g| bad += @intFromBool(!std.mem.eql(u32, w, g)); + return bad; +} + +test "dsv41 metal serve: (a) concurrent == solo == serial, (b) drafted == draft:false, through LaneHost" { + const p0 = try ids(150, 1); + defer gpa.free(p0); + const p1 = try ids(620, 2); + defer gpa.free(p1); + const asks = [_]Ask{ .{ .prompt = p0 }, .{ .prompt = p1 } }; + var want: [2][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const m = try fs.open(gpa, io, seed, .{}); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + const e = s.engine(); + var drafted: u64 = 0; + var solo: [2][]u32 = undefined; + try serve(e, &asks, false, &solo, &drafted); + defer freeAll(&solo); + var both: [2][]u32 = undefined; + try serve(e, &asks, true, &both, &drafted); + defer freeAll(&both); + const plain_asks = [_]Ask{ .{ .prompt = p0, .drafts = false }, .{ .prompt = p1, .drafts = false } }; + var plain_drafted: u64 = 0; + var plain: [2][]u32 = undefined; + try serve(e, &plain_asks, true, &plain, &plain_drafted); + defer freeAll(&plain); + const bad = [3]usize{ same(&want, &solo), same(&want, &both), same(&want, &plain) }; + say(bad[0] == 0, "(a) solo vs serial: {d} of 2 replies differ", .{bad[0]}); + say(bad[1] == 0, "(a) together vs serial: {d} of 2 replies differ ({d} windows of {d} rows in all)", .{ bad[1], m.t.stats.windows, m.t.stats.rows }); + say(bad[2] == 0 and drafted > 0 and plain_drafted == 0, "(b) draft:false vs serial: {d} differ; drafted rows {d} (drafts on), {d} (off), {d} pass launches, {d} windows cut back", .{ bad[2], drafted, plain_drafted, m.ad.passes, m.t.stats.rewinds }); + try testing.expectEqual([3]usize{ 0, 0, 0 }, bad); + try testing.expect(drafted > 0 and m.t.stats.rewinds > 0); + try testing.expectEqual(@as(u64, 0), plain_drafted); +} + +test "dsv41 metal serve: (c) prompt-cache hits (the same prompt, a later turn) == fresh runs" { + const p = try ids(300, 3); + defer gpa.free(p); + const m = try fs.open(gpa, io, seed, .{ .cache_bytes = 1 << 30 }); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + const e = s.engine(); + var drafted: u64 = 0; + var first: [1][]u32 = undefined; + try serve(e, &.{.{ .prompt = p, .max = 40 }}, false, &first, &drafted); + defer freeAll(&first); + const first_cached = cached_rows[0]; + // a later turn: the prompt, the reply but its last token, new ids; then the first prompt again + const extra = try ids(50, 4); + defer gpa.free(extra); + const turn = try std.mem.concat(gpa, u32, &.{ p, first[0][0..39], extra }); + defer gpa.free(turn); + const branch = try std.mem.concat(gpa, u32, &.{ p, extra }); + defer gpa.free(branch); + const asks = [_]Ask{ .{ .prompt = turn, .max = 30 }, .{ .prompt = p, .max = 40 }, .{ .prompt = branch, .max = 30 } }; + const hits0 = m.t.stats.hits; + var got: [3][]u32 = undefined; + try serve(e, &asks, false, &got, &drafted); + defer freeAll(&got); + const hits = m.t.stats.hits - hits0; + const reported = cached_rows[0..3].*; + var want: [3][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const bad = same(&want, &got); + say(bad == 0 and hits == 3, "(c) 3 cached prompts ({d} hits, {d} rows restored in all) vs fresh: {d} replies differ", .{ hits, m.t.stats.cached_rows, bad }); + try testing.expectEqual(@as(usize, 0), bad); + try testing.expectEqual(@as(u64, 3), hits); + // each reply reports the rows it resumed (prompt_tokens_details.cached_tokens) + say(reported[0] > p.len and reported[1] >= p.len and reported[2] >= p.len, "(c) cached rows reported {any} for prompts of {d}, {d} and {d} tokens", .{ reported, turn.len, p.len, branch.len }); + for (reported, [3]usize{ turn.len, p.len, branch.len }) |r, n| try testing.expect(r > 0 and r <= n); + try testing.expect(reported[0] > p.len and reported[1] >= p.len and reported[2] >= p.len); + try testing.expectEqual(@as(u32, 0), first_cached); + // the later turn resumed past the first prompt: its reply's rows came from the cache too + try testing.expect(m.t.stats.cached_rows >= 3 * p.len + 30); +} + +test "dsv41 metal serve: (d) sampled replies are their seed's: again, drafted or not, together, serial" { + const p0 = try ids(200, 5); + defer gpa.free(p0); + const p1 = try ids(90, 6); + defer gpa.free(p1); + const hot: api.Sampling = .{ .seed = 1234, .temperature = 0.9, .top_k = 0, .top_p = 1.0 }; + const nucleus: api.Sampling = .{ .seed = 99, .temperature = 0.7, .top_k = 40, .top_p = 0.9 }; + const asks = [_]Ask{ .{ .prompt = p0, .sampling = hot }, .{ .prompt = p1, .sampling = nucleus } }; + var want: [2][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const m = try fs.open(gpa, io, seed, .{}); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + const e = s.engine(); + var drafted: u64 = 0; + var a1: [2][]u32 = undefined; + try serve(e, &asks, true, &a1, &drafted); + defer freeAll(&a1); + var a2: [2][]u32 = undefined; + try serve(e, &asks, false, &a2, &drafted); + defer freeAll(&a2); + const plain_asks = [_]Ask{ .{ .prompt = p0, .sampling = hot, .drafts = false }, .{ .prompt = p1, .sampling = nucleus, .drafts = false } }; + var a3: [2][]u32 = undefined; + try serve(e, &plain_asks, true, &a3, &drafted); + defer freeAll(&a3); + var other_seed = hot; + other_seed.seed = 1235; + var a4: [1][]u32 = undefined; + try serve(e, &.{.{ .prompt = p0, .sampling = other_seed }}, false, &a4, &drafted); + defer freeAll(&a4); + const bad = [3]usize{ same(&want, &a1), same(&want, &a2), same(&want, &a3) }; + const moved = !std.mem.eql(u32, a4[0], want[0]); + say(bad[0] + bad[1] + bad[2] == 0 and moved, "(d) together / again / draft:false vs serial: {d} / {d} / {d} differ; another seed's reply {s}", .{ bad[0], bad[1], bad[2], if (moved) "differs" else "is the same" }); + try testing.expectEqual([3]usize{ 0, 0, 0 }, bad); + try testing.expect(moved); +} + +test "dsv41 metal serve: (e) a request cancelled during its prompt leaves the slot and the next request valid" { + const long = try ids(1100, 7); + defer gpa.free(long); + const p0 = try ids(120, 8); + defer gpa.free(p0); + const p1 = try ids(140, 9); + defer gpa.free(p1); + const asks = [_]Ask{ .{ .prompt = p0, .max = 40 }, .{ .prompt = p1 } }; + var want: [2][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const m = try fs.open(gpa, io, seed, .{ .slots = 2, .cache_bytes = 1 << 30 }); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + const e = s.engine(); + // a decoding stream, then the long prompt (cancelled as soon as it is submitted), then a new request + var bg: Box = .{}; + defer bg.deinit(); + const r0: api.Request = .{ .prompt = p0, .max_tokens = 40 }; + try e.submit(500, &r0, .{ .ctx = &bg, .event = Box.event }); + var gone: Box = .{}; + defer gone.deinit(); + const rl: api.Request = .{ .prompt = long, .max_tokens = 64 }; + try e.submit(501, &rl, .{ .ctx = &gone, .event = Box.event }); + e.cancel(501); + try gone.wait(); + var drafted: u64 = 0; + var after: [1][]u32 = undefined; + try serve(e, asks[1..], false, &after, &drafted); + defer freeAll(&after); + try bg.wait(); + const ok0 = std.mem.eql(u32, want[0], bg.out.items); + const ok1 = std.mem.eql(u32, want[1], after[0]); + say(ok0 and ok1, "(e) cancelled ({t}, {d} tokens out): the decoding reply {s}, the next request's {s} the serial one", .{ gone.reason.?, gone.out.items.len, if (ok0) "equals" else "differs from", if (ok1) "equals" else "differs from" }); + try testing.expect(gone.reason.? == .cancelled); + try testing.expect(ok0 and ok1); +}