dsv41 metal: the vision tower's kernels under the kernel contract

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

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.

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 <metal_stdlib>
+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<float>(as_type<uint>(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;
+}

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;

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

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

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