From 73d0de6874452d5a2e39b4173605fa0312770c65 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:09:43 -0700 Subject: [PATCH 01/15] DeepSeek V4.1 Flash: configuration, Engram tables and resident weight staging Config parsing, the Engram hash layout, table plans over the original shards' embed weight and scale, and the EXL3 resident weight staging behind paced reads. Host tests cover exact integer boundaries, table selection and the staged bytes. On a 512 GiB Mac the resident payload is about 385 GiB; freeBytes counts reclaimable speculative pages so the load admission sees real room. No kernels; the served engine comes with the lane backend. --- zig/build/deepseek_v41.zig | 22 ++ zig/deepseek_v41_weights_checks.zig | 8 + .../families/deepseek_v41/resident_load.zig | 249 ++++++++++++++++ zig/src/families/deepseek_v41/view.zig | 25 ++ zig/src/families/deepseek_v41/weights.zig | 273 ++++++++++++++++++ zig/src/tensorfold.zig | 5 + 6 files changed, 582 insertions(+) create mode 100644 zig/build/deepseek_v41.zig create mode 100644 zig/deepseek_v41_weights_checks.zig create mode 100644 zig/src/families/deepseek_v41/resident_load.zig create mode 100644 zig/src/families/deepseek_v41/view.zig create mode 100644 zig/src/families/deepseek_v41/weights.zig diff --git a/zig/build/deepseek_v41.zig b/zig/build/deepseek_v41.zig new file mode 100644 index 0000000000..6b196fedff --- /dev/null +++ b/zig/build/deepseek_v41.zig @@ -0,0 +1,22 @@ +//! Experimental DS4.1 weight-loading qualification targets. +const std = @import("std"); + +pub fn weights( + b: *std.Build, + target: std.Build.ResolvedTarget, + optimize: std.builtin.OptimizeMode, + metal: *std.Build.Module, +) *std.Build.Step { + const host = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/deepseek_v41_weights_checks.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "metal", .module = metal }, + }, + }) }); + const step = b.step("test-deepseek-v41-weights", "V4.1 configuration, Engram table plans and resident weight staging"); + step.dependOn(&b.addRunArtifact(host).step); + return step; +} diff --git a/zig/deepseek_v41_weights_checks.zig b/zig/deepseek_v41_weights_checks.zig new file mode 100644 index 0000000000..b3d4fedcac --- /dev/null +++ b/zig/deepseek_v41_weights_checks.zig @@ -0,0 +1,8 @@ +//! Host checks for V4.1 configuration, Engram table plans and resident weight staging; no kernels compiled. +test { + _ = @import("src/families/deepseek_v41/config.zig"); + _ = @import("src/families/deepseek_v41/view.zig"); + _ = @import("src/families/deepseek_v41/engram_hash.zig"); + _ = @import("src/families/deepseek_v41/load.zig"); + _ = @import("src/families/deepseek_v41/weights.zig"); +} diff --git a/zig/src/families/deepseek_v41/resident_load.zig b/zig/src/families/deepseek_v41/resident_load.zig new file mode 100644 index 0000000000..422af3555c --- /dev/null +++ b/zig/src/families/deepseek_v41/resident_load.zig @@ -0,0 +1,249 @@ +//! Header-first resident loading: no mapped weights, deferred disk lookup or duplicate Engram projections. +const std = @import("std"); +const mtl = @import("metal"); +const st = @import("../../core/safetensors.zig"); +const cp = @import("../../core/checkpoint_metal.zig"); +const paced = @import("../../core/paced_read.zig"); + +pub const Reference = struct { shard: usize, entry: st.Entry }; +pub const Shard = struct { + path: [:0]const u8, + file_bytes: usize, + data_start: usize, + begin: usize, + end: usize, + names: st.Header, +}; + +pub const Plan = struct { + arena: std.heap.ArenaAllocator, + shards: std.ArrayList(Shard) = .empty, + tensors: std.StringHashMapUnmanaged(Reference) = .empty, + bytes: usize = 0, + + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, table_dir: []const u8) !Plan { + var self = Plan{ .arena = std.heap.ArenaAllocator.init(gpa) }; + errdefer self.deinit(); + const a = self.arena.allocator(); + const index_path = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + const index_bytes = try std.Io.Dir.cwd().readFileAlloc(io, index_path, a, .limited(64 << 20)); + const index = try std.json.parseFromSliceLeaky(struct { weight_map: std.json.Value }, a, index_bytes, .{ + .ignore_unknown_fields = true, + }); + if (index.weight_map != .object or index.weight_map.object.count() == 0) return error.BadCheckpointIndex; + var files: std.StringHashMapUnmanaged(void) = .empty; + var entries = index.weight_map.object.iterator(); + while (entries.next()) |item| { + if (item.value_ptr.* != .string) return error.BadCheckpointIndex; + const file = item.value_ptr.string; + if (!safeBasename(file)) return error.BadCheckpointIndex; + try files.put(a, file, {}); + } + const ordered = try a.alloc([]const u8, files.count()); + var file_it = files.keyIterator(); + var n: usize = 0; + while (file_it.next()) |file| : (n += 1) ordered[n] = file.*; + std.mem.sort([]const u8, ordered, {}, less); + for (ordered) |file| try self.add(try std.fs.path.join(a, &.{ dir, file }), null); + if (self.tensors.count() != index.weight_map.object.count()) return error.CheckpointIndexMismatch; + entries = index.weight_map.object.iterator(); + while (entries.next()) |item| { + const ref = self.tensors.get(item.key_ptr.*) orelse return error.CheckpointIndexMismatch; + if (!std.mem.eql(u8, std.fs.path.basename(self.shards.items[ref.shard].path), item.value_ptr.string)) + return error.CheckpointIndexMismatch; + } + for ([_]struct { layer: u32, file: []const u8, rows: usize }{ + .{ .layer = 1, .file = "model-00047-of-00048.safetensors", .rows = 384006168 }, + .{ .layer = 14, .file = "model-00048-of-00048.safetensors", .rows = 384016682 }, + }) |table| { + const prefix = try std.fmt.allocPrint(a, "layers.{d}.engram.embed.", .{table.layer}); + try self.add(try std.fs.path.join(a, &.{ table_dir, table.file }), prefix); + const weight = try std.fmt.allocPrint(a, "{s}weight", .{prefix}); + const scale = try std.fmt.allocPrint(a, "{s}scale", .{prefix}); + try self.require(weight, .f8_e4m3, &.{ table.rows, 256 }); + try self.require(scale, .f8_e8m0, &.{ table.rows, 8 }); + if (self.shards.items[self.shards.items.len - 1].names.count() != 2) return error.UnexpectedEngramTensor; + } + return self; + } + + pub fn deinit(self: *Plan) void { + self.arena.deinit(); + self.* = undefined; + } + + pub fn get(self: *const Plan, name: []const u8) !st.Entry { + return (self.tensors.get(name) orelse return error.MissingTensor).entry; + } + + pub fn has(self: *const Plan, name: []const u8) bool { + return self.tensors.contains(name); + } + + pub fn require(self: *const Plan, name: []const u8, dtype: st.DType, shape: []const usize) !void { + const e = try self.get(name); + if (e.dtype != dtype or e.rank != shape.len or !std.mem.eql(usize, e.shape[0..e.rank], shape)) { + std.log.err("unsupported tensor {s}: dtype {s}, shape {any}", .{ name, @tagName(e.dtype), e.shape[0..e.rank] }); + return error.UnsupportedTensor; + } + } + + fn add(self: *Plan, path: []const u8, prefix: ?[]const u8) !void { + const a = self.arena.allocator(); + const own_path = try a.dupeSentinel(u8, path, 0); + const fd = std.c.open(own_path, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + defer _ = std.c.close(fd); + const length = std.c.lseek(fd, 0, std.c.SEEK.END); + if (length < 8) return error.BadSafetensors; + const file_bytes: usize = @intCast(length); + var head: [8]u8 = undefined; + try readAll(fd, &head, 0); + const header_bytes = std.mem.readInt(u64, &head, .little); + if (header_bytes > @min(file_bytes - 8, 64 << 20)) return error.BadSafetensors; + const json = try a.alloc(u8, @intCast(header_bytes)); + try readAll(fd, json, 8); + const data_start: usize = 8 + @as(usize, @intCast(header_bytes)); + const names = try st.parseHeaderPrefix(a, json, file_bytes - data_start, prefix); + if (names.count() == 0) return error.MissingTensor; + var begin: usize = if (prefix == null) 0 else file_bytes - data_start; + var end: usize = if (prefix == null) file_bytes - data_start else 0; + for (names.keys(), names.values()) |name, entry| { + begin = @min(begin, entry.begin); + end = @max(end, entry.end); + const slot = try self.tensors.getOrPut(a, name); + if (slot.found_existing) return error.DuplicateTensor; + slot.value_ptr.* = .{ .shard = self.shards.items.len, .entry = entry }; + } + if (prefix != null) { + var selected_bytes: usize = 0; + for (names.values()) |entry| selected_bytes = try std.math.add(usize, selected_bytes, entry.end - entry.begin); + if (selected_bytes != end - begin) return error.NoncontiguousEngramTables; + } + self.bytes = try std.math.add(usize, self.bytes, end - begin); + try self.shards.append(a, .{ .path = own_path, .file_bytes = file_bytes, .data_start = data_start, .begin = begin, .end = end, .names = names }); + } + + /// Only called after architecture validation and the complete memory admission calculation. + pub fn materialize(self: *const Plan, gpa: std.mem.Allocator, device: mtl.Device, residency: mtl.ResidencySet, queue: mtl.Queue) !cp.Checkpoint { + var out = cp.Checkpoint.init(gpa); + errdefer out.deinit(); + const stage = try paced.stage(gpa, 32 << 20); + defer gpa.free(stage); + var warm = mtl.keepalive.Target{ .queue = queue, .sets = @as(*const [1]mtl.ResidencySet, &residency) }; + var last_tick: f64 = -std.math.inf(f64); + var short: paced.Memory = undefined; + for (self.shards.items, 0..) |shard, i| { + const fd = std.c.open(shard.path, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + defer _ = std.c.close(fd); + if (std.c.fcntl(fd, 48, @as(c_int, 1)) != 0) return error.UncachedReadUnavailable; + if (std.c.lseek(fd, 0, std.c.SEEK.END) != shard.file_bytes) return error.CheckpointChanged; + const bytes = shard.end - shard.begin; + const buffer = try device.buffer(@max(bytes, 16), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + var owned = false; + errdefer if (!owned) buffer.deinit(); + try out.own(buffer, bytes); + owned = true; + residency.add(buffer); + residency.commit(); + residency.requestResidency(); + var copied: usize = 0; + while (copied < bytes) { + const count = @min(bytes - copied, stage.len - paced.alignment); + paced.read(fd, buffer.contents()[copied..][0..count], shard.data_start + shard.begin + copied, stage, &short) catch |err| { + if (err == error.MemoryShort) { + var msg: [512]u8 = undefined; + std.log.err("deepseek-v41: {s}", .{paced.refusal(&msg, short, paced.floor())}); + } + return err; + }; + copied += count; + const now = mtl.clock.seconds(); + if (now - last_tick >= 0.4) { + mtl.keepalive.Target.tick(&warm); + last_tick = now; + } + } + for (shard.names.keys(), shard.names.values()) |name, e| { + try out.set(name, .{ .buffer = buffer, .offset = e.begin - shard.begin, .bytes = e.end - e.begin, .dtype = e.dtype, .shape = e.shape, .rank = e.rank }); + } + std.log.info("resident checkpoint shard {d}/{d}: {d} bytes, residency set {d} bytes", .{ i + 1, self.shards.items.len, bytes, residency.allocatedSize() }); + } + residency.commit(); + residency.requestResidency(); + return out; + } +}; + +fn readAll(fd: std.c.fd_t, dest: []u8, at: usize) !void { + var done: usize = 0; + while (done < dest.len) { + const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + if (n <= 0) return error.ShortRead; + done += @intCast(n); + } +} + +fn safeBasename(s: []const u8) bool { + return s.len > 0 and !std.mem.eql(u8, s, ".") and !std.mem.eql(u8, s, "..") and + std.mem.indexOfAny(u8, s, "/\\") == null and std.mem.indexOfScalar(u8, s, 0) == null and + std.mem.endsWith(u8, s, ".safetensors"); +} + +fn less(_: void, a: []const u8, b: []const u8) bool { + return std.mem.lessThan(u8, a, b); +} + +test "selected table plan excludes projections and preserves tensor byte offsets" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = + \\{"padding":{"dtype":"U8","shape":[3],"data_offsets":[0,3]}, + \\ "engram.embed.weight":{"dtype":"U8","shape":[4],"data_offsets":[3,7]}, + \\ "engram.embed.scale":{"dtype":"U8","shape":[2],"data_offsets":[7,9]}, + \\ "engram.wkv":{"dtype":"U8","shape":[1],"data_offsets":[9,10]}} + ; + var data: [8 + header.len + 10]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + @memset(data[8 + header.len ..], 0); + try tmp.dir.writeFile(io, .{ .sub_path = "table.safetensors", .data = &data }); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}/table.safetensors", .{tmp.sub_path}); + defer a.free(path); + var plan = Plan{ .arena = std.heap.ArenaAllocator.init(a) }; + defer plan.deinit(); + try plan.add(path, "engram.embed."); + try std.testing.expectEqual(@as(usize, 6), plan.bytes); + try std.testing.expectEqual(@as(usize, 3), plan.shards.items[0].begin); + try std.testing.expectEqual(@as(usize, 9), plan.shards.items[0].end); + try std.testing.expectEqual(@as(usize, 7), (try plan.get("engram.embed.scale")).begin); + try std.testing.expectError(error.MissingTensor, plan.get("engram.wkv")); + try std.testing.expectError(error.DuplicateTensor, plan.add(path, "engram.embed.")); +} + +test "selected table payload must be contiguous and indexed files cannot escape checkpoint directory" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = + \\{"engram.embed.weight":{"dtype":"U8","shape":[4],"data_offsets":[0,4]}, + \\ "projection":{"dtype":"U8","shape":[3],"data_offsets":[4,7]}, + \\ "engram.embed.scale":{"dtype":"U8","shape":[2],"data_offsets":[7,9]}} + ; + var data: [8 + header.len + 9]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + @memset(data[8 + header.len ..], 0); + try tmp.dir.writeFile(io, .{ .sub_path = "gap.safetensors", .data = &data }); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}/gap.safetensors", .{tmp.sub_path}); + defer a.free(path); + var plan = Plan{ .arena = std.heap.ArenaAllocator.init(a) }; + defer plan.deinit(); + try std.testing.expectError(error.NoncontiguousEngramTables, plan.add(path, "engram.embed.")); + for ([_][]const u8{ "../x.safetensors", "/x.safetensors", "a\\x.safetensors", ".", "..", "x.safetensors\x00", "x.json", "" }) |name| try std.testing.expect(!safeBasename(name)); + try std.testing.expect(safeBasename("model-00001-of-00039.safetensors")); +} diff --git a/zig/src/families/deepseek_v41/view.zig b/zig/src/families/deepseek_v41/view.zig new file mode 100644 index 0000000000..fbaa08ce39 --- /dev/null +++ b/zig/src/families/deepseek_v41/view.zig @@ -0,0 +1,25 @@ +//! A borrowed Metal buffer region; element precision is carried by each operator's explicit descriptor. +const std = @import("std"); +const mtl = @import("metal"); + +pub const View = struct { + buffer: mtl.Buffer, + offset: usize = 0, + + pub fn shifted(self: View, bytes: usize) View { + return .{ .buffer = self.buffer, .offset = self.offset + bytes }; + } + + pub fn require(self: View, bytes: usize, alignment: usize) !void { + if (self.offset % alignment != 0 or self.offset > self.buffer.length() or + bytes > self.buffer.length() - self.offset) return error.InvalidTensorView; + } + + /// Rows are strided elements; only the final row's used width is required. + pub fn requireRows(self: View, rows: usize, stride: usize, width: usize, element_bytes: usize) !void { + if (rows == 0 or width == 0 or stride < width) return error.InvalidTensorView; + const preceding = try std.math.mul(usize, rows - 1, stride); + const elements = try std.math.add(usize, preceding, width); + try self.require(try std.math.mul(usize, elements, element_bytes), element_bytes); + } +}; diff --git a/zig/src/families/deepseek_v41/weights.zig b/zig/src/families/deepseek_v41/weights.zig new file mode 100644 index 0000000000..1377466c9c --- /dev/null +++ b/zig/src/families/deepseek_v41/weights.zig @@ -0,0 +1,273 @@ +//! Strict pinned backbone schema and borrowed resident bindings; only expert descriptor tables are owned. +const std = @import("std"); +const mtl = @import("metal"); +const cp = @import("../../core/checkpoint_metal.zig"); +const st = @import("../../core/safetensors.zig"); +const exl3 = @import("../../core/exl3_projection.zig"); +const load = @import("resident_load.zig"); + +pub const Tensor = cp.Tensor; +pub const Matrix = exl3.Matrix; +pub const Hc = struct { projection: Tensor, scale: Tensor, base: Tensor }; +pub const Compressor = struct { wkv: Matrix, wgate: ?Matrix, norm: Tensor }; +pub const Indexer = struct { wq_b: Matrix, wk: ?Matrix, weights_proj: Tensor, k_norm: ?Tensor }; +pub const Attention = struct { + wq_a: Matrix, + wq_b: Matrix, + wkv: Matrix, + wo_b: Matrix, + wo_a: [8]Matrix, + sink: Tensor, + q_norm: Tensor, + kv_norm: Tensor, + compressor: ?Compressor, + indexer: ?Indexer, +}; +/// Projection order is w1 (gate), w3 (up), w2 (down), for shared and routed experts alike. +pub const Moe = struct { + gate: Tensor, + bias: Tensor, + shared: [3]Matrix, + descriptors: [3]mtl.Buffer, +}; +pub const Engram = struct { table: Tensor, scales: Tensor, q_weight: Tensor, k_weight: Tensor, wkv: Matrix }; +pub const Layer = struct { + attn_norm: Tensor, + ffn_norm: Tensor, + hc_attn: Hc, + hc_ffn: Hc, + attn: Attention, + moe: Moe, + engram: ?Engram, +}; +pub const Model = struct { + embed: Tensor, + head: Matrix, + norm: Tensor, + layers: [40]Layer, + + /// Checkpoint buffers must outlive this borrowed model and every submitted command using it. + pub fn deinit(self: *Model) void { + for (&self.layers) |*layer| for (layer.moe.descriptors) |buffer| buffer.deinit(); + self.* = undefined; + } + pub fn registerResidency(self: *const Model, set: mtl.ResidencySet) void { + for (&self.layers) |*layer| for (layer.moe.descriptors) |buffer| set.add(buffer); + } + pub fn residentBytes(_: *const Model) usize { + return sizeofDescriptors(); + } +}; + +/// Admission adds these bytes to load.Plan.bytes; no dequantized weights or duplicate projections are allocated. +pub fn sizeofDescriptors() usize { + return 40 * 3 * 384 * @sizeOf(exl3.ExpertDescriptor); +} + +/// Every executed tensor is checked using headers only, before checkpoint materialization or model allocation. +pub fn validate(plan: *const load.Plan) !void { + try populate(false, plan, null); +} + +/// Bind resident tensors and marker payloads; route descriptors replace redundant host matrix metadata. +pub fn bind(device: mtl.Device, checkpoint: *const cp.Checkpoint) !Model { + var model: Model = undefined; + try populate(true, checkpoint, &model); + var owned: usize = 0; + errdefer for (0..owned) |i| model.layers[i / 3].moe.descriptors[i % 3].deinit(); + var prefix: [128]u8 = undefined; + for (&model.layers, 0..) |*layer, layer_id| { + inline for (.{ "w1", "w3", "w2" }, 0..) |part, projection| { + const bytes = 384 * @sizeOf(exl3.ExpertDescriptor); + const buffer = try device.buffer(bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + layer.moe.descriptors[projection] = buffer; + owned += 1; + const out: [*]exl3.ExpertDescriptor = @ptrCast(@alignCast(buffer.contents())); + for (0..384) |expert| { + var m: Matrix = undefined; + const k: usize = if (projection == 2) 2304 else 5120; + const n: usize = if (projection == 2) 5120 else 2304; + try matrix(true, checkpoint, try named(&prefix, "layers.{d}.ffn.experts.{d}.{s}", .{ layer_id, expert, part }), k, n, &m); + out[expert] = try exl3.descriptor(m); + } + } + } + return model; +} + +/// A borrowed CPU/GPU matrix view for checkpoint qualification, using the same strict storage contract. +pub fn lookup(checkpoint: *const cp.Checkpoint, prefix: []const u8, k: usize, n: usize) !Matrix { + var out: Matrix = undefined; + try matrix(true, checkpoint, prefix, k, n, &out); + return out; +} + +fn named(buffer: []u8, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(buffer, fmt, args); +} + +fn plain(comptime binding: bool, source: anytype, name: []const u8, dtype: st.DType, shape: []const usize, destination: ?*Tensor) !void { + const tensor = try source.get(name); + if (tensor.dtype != dtype or tensor.rank != shape.len or + !std.mem.eql(usize, tensor.shape[0..tensor.rank], shape)) + { + std.log.err("unsupported backbone tensor {s}: {s} {any}", .{ name, @tagName(tensor.dtype), tensor.shape[0..tensor.rank] }); + return error.UnsupportedTensor; + } + const offset = if (binding) tensor.offset else tensor.begin; + if (offset % dtype.size() != 0) return error.InvalidTensorView; + if (binding) { + var count: usize = 1; + for (shape) |dim| count = try std.math.mul(usize, count, dim); + if (tensor.bytes != try std.math.mul(usize, count, dtype.size())) return error.UnsupportedTensor; + if (tensor.offset % dtype.size() != 0 or tensor.offset > tensor.buffer.length() or + tensor.bytes > tensor.buffer.length() - tensor.offset) return error.InvalidTensorView; + destination.?.* = tensor; + } +} + +fn has(source: anytype, name: []const u8) bool { + return source.has(name); +} + +const Scale = struct { name: []const u8, format: exl3.ScaleFormat }; +fn scale(comptime binding: bool, source: anytype, prefix: []const u8, size: usize, comptime fp: []const u8, comptime signs: []const u8, storage: []u8) !Scale { + var alternative: [160]u8 = undefined; + const f = try named(storage, "{s}.{s}", .{ prefix, fp }); + const s = try named(&alternative, "{s}.{s}", .{ prefix, signs }); + const have_fp = has(source, f); + if (have_fp == has(source, s)) return error.MissingOrAmbiguousExl3Scale; + const name = if (have_fp) f else try named(storage, "{s}.{s}", .{ prefix, signs }); + const t = try source.get(name); + if (t.rank != 1 or t.shape[0] != (if (have_fp) size else size / 16) or + t.dtype != (if (have_fp) st.DType.f16 else st.DType.i16)) return error.UnsupportedExl3Scale; + const offset = if (binding) t.offset else t.begin; + if (offset % 2 != 0) return error.InvalidTensorView; + if (binding) { + const bytes = (if (have_fp) size else size / 16) * 2; + if (t.bytes != bytes or t.offset % 2 != 0 or t.offset > t.buffer.length() or + bytes > t.buffer.length() - t.offset) return error.InvalidTensorView; + } + return .{ .name = name, .format = if (have_fp) .f16 else .packed_signs }; +} + +fn matrix(comptime binding: bool, source: anytype, prefix: []const u8, k: usize, n: usize, destination: ?*Matrix) !void { + var name: [160]u8 = undefined; + const trellis = try source.get(try named(&name, "{s}.trellis", .{prefix})); + if (trellis.dtype != .i16 or trellis.rank != 3 or trellis.shape[0] != k / 16 or + trellis.shape[1] != n / 16 or trellis.shape[2] % 8 != 0) return error.UnsupportedExl3Shape; + if ((if (binding) trellis.offset else trellis.begin) % 2 != 0) return error.InvalidTensorView; + const halves = trellis.shape[2] / 8; + if (halves < 2 or halves > 16 or (halves % 2 != 0 and halves > 7)) return error.UnsupportedExl3Width; + // Pinned checkpoint uses mul1 throughout; missing markers cannot silently select the legacy codebook. + const marker = try source.get(try named(&name, "{s}.mul1", .{prefix})); + if (marker.dtype != .i32 or marker.rank != 0) return error.UnsupportedExl3Marker; + if ((if (binding) marker.offset else marker.begin) % 4 != 0) return error.InvalidTensorView; + if (has(source, try named(&name, "{s}.mcg", .{prefix})) or + has(source, try named(&name, "{s}.3inst", .{prefix})) or + has(source, try named(&name, "{s}.bias", .{prefix})) or + has(source, try named(&name, "{s}.weight", .{prefix}))) return error.UnsupportedExl3Parts; + var su_name: [160]u8 = undefined; + var sv_name: [160]u8 = undefined; + const su = try scale(binding, source, prefix, k, "suh", "su", &su_name); + const sv = try scale(binding, source, prefix, n, "svh", "sv", &sv_name); + if (binding) { + const expected = k / 16 * (n / 16) * halves * 16; + if (trellis.bytes != expected or trellis.offset % 2 != 0 or + trellis.offset > trellis.buffer.length() or expected > trellis.buffer.length() - trellis.offset) + return error.InvalidTensorView; + if (marker.bytes != 4 or marker.offset % 4 != 0 or marker.offset > marker.buffer.length() or + 4 > marker.buffer.length() - marker.offset) return error.InvalidTensorView; + const marker_bytes = marker.buffer.contents()[marker.offset..][0..4]; + if (std.mem.readInt(u32, marker_bytes, .little) != exl3.format.Mul1Mul) return error.UnsupportedExl3Marker; + const suh = try source.get(su.name); + const svh = try source.get(sv.name); + destination.?.* = .{ + .trellis = .{ .buffer = trellis.buffer, .offset = trellis.offset }, + .suh = .{ .buffer = suh.buffer, .offset = suh.offset }, + .svh = .{ .buffer = svh.buffer, .offset = svh.offset }, + .k = k, + .n = n, + .width = .{ .halves = @intCast(halves) }, + .codebook = .mul1, + .suh_format = su.format, + .svh_format = sv.format, + }; + try destination.?.validate(); + } +} + +fn populate(comptime binding: bool, source: anytype, model: ?*Model) !void { + try plain(binding, source, "embed.weight", .bf16, &.{ 129280, 5120 }, if (binding) &model.?.embed else null); + try plain(binding, source, "norm.weight", .bf16, &.{5120}, if (binding) &model.?.norm else null); + try matrix(binding, source, "head", 5120, 129280, if (binding) &model.?.head else null); + var name: [160]u8 = undefined; + var prefix: [128]u8 = undefined; + for (0..40) |layer| { + const out: ?*Layer = if (binding) &model.?.layers[layer] else null; + try plain(binding, source, try named(&name, "layers.{d}.attn_norm.weight", .{layer}), .bf16, &.{5120}, if (binding) &out.?.attn_norm else null); + try plain(binding, source, try named(&name, "layers.{d}.ffn_norm.weight", .{layer}), .bf16, &.{5120}, if (binding) &out.?.ffn_norm else null); + inline for (.{ "attn", "ffn" }) |kind| { + const hc: ?*Hc = if (binding) (if (std.mem.eql(u8, kind, "attn")) &out.?.hc_attn else &out.?.hc_ffn) else null; + try plain(binding, source, try named(&name, "layers.{d}.hc_{s}_fn", .{ layer, kind }), .f32, &.{ 24, 20480 }, if (binding) &hc.?.projection else null); + try plain(binding, source, try named(&name, "layers.{d}.hc_{s}_scale", .{ layer, kind }), .f32, &.{3}, if (binding) &hc.?.scale else null); + try plain(binding, source, try named(&name, "layers.{d}.hc_{s}_base", .{ layer, kind }), .f32, &.{24}, if (binding) &hc.?.base else null); + } + const attn: ?*Attention = if (binding) &out.?.attn else null; + inline for (.{ "wq_a", "wq_b", "wkv", "wo_b" }, .{ 5120, 1280, 5120, 8192 }, .{ 1280, 32768, 512, 5120 }) |part, k, n| { + try matrix(binding, source, try named(&prefix, "layers.{d}.attn.{s}", .{ layer, part }), k, n, if (binding) &@field(attn.?, part) else null); + } + for (0..8) |slice| try matrix(binding, source, try named(&prefix, "layers.{d}.attn.wo_a.slice.{d}", .{ layer, slice }), 4096, 1024, if (binding) &attn.?.wo_a[slice] else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.attn_sink", .{layer}), .f32, &.{64}, if (binding) &attn.?.sink else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.q_norm.weight", .{layer}), .bf16, &.{1280}, if (binding) &attn.?.q_norm else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.kv_norm.weight", .{layer}), .bf16, &.{512}, if (binding) &attn.?.kv_norm else null); + const owner = layer == 2 or layer == 8 or layer == 14 or layer == 20; + if (binding) { + attn.?.compressor = if (owner) @as(Compressor, undefined) else null; + attn.?.indexer = if (owner or layer == 24 or layer == 28 or layer == 32 or layer == 36) + @as(Indexer, undefined) + else + null; + } + if (owner) { + const compressor: ?*Compressor = if (binding) &attn.?.compressor.? else null; + try matrix(binding, source, try named(&prefix, "layers.{d}.attn.compressor.wkv", .{layer}), 5120, 512, if (binding) &compressor.?.wkv else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.compressor.norm.weight", .{layer}), .bf16, &.{512}, if (binding) &compressor.?.norm else null); + if (binding) compressor.?.wgate = if (layer < 20) @as(Matrix, undefined) else null; + if (layer < 20) try matrix(binding, source, try named(&prefix, "layers.{d}.attn.compressor.wgate", .{layer}), 5120, 512, if (binding) &compressor.?.wgate.? else null); + } + if (owner or layer == 24 or layer == 28 or layer == 32 or layer == 36) { + const indexer: ?*Indexer = if (binding) &attn.?.indexer.? else null; + try matrix(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wq_b", .{layer}), 1280, 4096, if (binding) &indexer.?.wq_b else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.indexer.weights_proj.weight", .{layer}), .f16, &.{ 32, 5120 }, if (binding) &indexer.?.weights_proj else null); + if (binding) { + indexer.?.wk = if (owner) @as(Matrix, undefined) else null; + indexer.?.k_norm = if (owner) @as(Tensor, undefined) else null; + } + if (owner) { + try matrix(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wk", .{layer}), 512, 128, if (binding) &indexer.?.wk.? else null); + try plain(binding, source, try named(&name, "layers.{d}.attn.indexer.k_norm.weight", .{layer}), .bf16, &.{128}, if (binding) &indexer.?.k_norm.? else null); + } + } + const moe: ?*Moe = if (binding) &out.?.moe else null; + try plain(binding, source, try named(&name, "layers.{d}.ffn.gate.weight", .{layer}), .f16, &.{ 384, 5120 }, if (binding) &moe.?.gate else null); + try plain(binding, source, try named(&name, "layers.{d}.ffn.gate.bias", .{layer}), .f16, &.{384}, if (binding) &moe.?.bias else null); + inline for (.{ "w1", "w3", "w2" }, 0..) |part, projection| { + const k: usize = if (projection == 2) 2304 else 5120; + const n: usize = if (projection == 2) 5120 else 2304; + try matrix(binding, source, try named(&prefix, "layers.{d}.ffn.shared_experts.{s}", .{ layer, part }), k, n, if (binding) &moe.?.shared[projection] else null); + if (!binding) for (0..384) |expert| try matrix(false, source, try named(&prefix, "layers.{d}.ffn.experts.{d}.{s}", .{ layer, expert, part }), k, n, null); + } + const engram = layer == 1 or layer == 14; + if (binding) out.?.engram = if (engram) @as(Engram, undefined) else null; + if (engram) { + const e: ?*Engram = if (binding) &out.?.engram.? else null; + const rows: usize = if (layer == 1) 384006168 else 384016682; + try plain(binding, source, try named(&name, "layers.{d}.engram.embed.weight", .{layer}), .f8_e4m3, &.{ rows, 256 }, if (binding) &e.?.table else null); + try plain(binding, source, try named(&name, "layers.{d}.engram.embed.scale", .{layer}), .f8_e8m0, &.{ rows, 8 }, if (binding) &e.?.scales else null); + try plain(binding, source, try named(&name, "layers.{d}.engram.q_weight", .{layer}), .f32, &.{ 4, 5120 }, if (binding) &e.?.q_weight else null); + try plain(binding, source, try named(&name, "layers.{d}.engram.k_weight", .{layer}), .f32, &.{ 4, 5120 }, if (binding) &e.?.k_weight else null); + try matrix(binding, source, try named(&prefix, "layers.{d}.engram.wkv", .{layer}), 6144, 25600, if (binding) &e.?.wkv else null); + } + } +} diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 26b68d7925..0cf0c9f5a4 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -24,6 +24,11 @@ pub const frags = @import("core/frags.zig"); pub const affine_mm = @import("core/affine_mm.zig"); /// The checkpoint formats (MLX affine and its fp64 reader for the device checks, exact widen). pub const quant = @import("core/quant/quant.zig"); +pub const deepseek_v41_config = @import("families/deepseek_v41/config.zig"); +pub const deepseek_v41_load = @import("families/deepseek_v41/load.zig"); +pub const deepseek_v41_weights = @import("families/deepseek_v41/weights.zig"); +pub const deepseek_v41_view = @import("families/deepseek_v41/view.zig"); +pub const deepseek_v41_engram_hash = @import("families/deepseek_v41/engram_hash.zig"); pub const exl3_format = @import("core/exl3_format.zig"); pub const exl3_projection = @import("core/exl3_projection.zig"); pub const exlr8 = @import("core/exlr8.zig"); From 7a7d55091795ee81ec35fc25b69e0f1b69c5637a Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:12:42 -0700 Subject: [PATCH 02/15] DeepSeek V4.1 Flash: Metal operators with fp64-dot references and the qualification runner The Engram hash/commit, sparse and compressed attention (CSA2), the hyperconnection mix, routed and shared experts, the cache and rotary glue as source Metal kernels, with zig/tests/deepseek_v41_reference as the independent CPU oracle: plain reductions in fp64, no native statistics, no injected coefficients. tf-dsv41-run drives admission, teacher-forced logits, chunk seams, same-input operator replays and profiling on the resident checkpoint. The retained whole-model gate stays failed at 4.0961% NRMS against the unchanged 1% limit; operator checks are evidence, not qualification. --- build.zig | 3 +- zig/build/deepseek_v41.zig | 38 ++ zig/build/metal_targets.zig | 6 + zig/deepseek_v41_checks.zig | 13 + zig/deepseek_v41_reference_checks.zig | 7 + zig/deepseek_v41_run.zig | 306 +++++++++++++ zig/kernels/metal/core/expert_groups.metal | 36 ++ zig/kernels/metal/deepseek_v41/csa2.metal | 102 +++++ zig/kernels/metal/deepseek_v41/engram.metal | 150 +++++++ zig/kernels/metal/deepseek_v41/index.metal | 103 +++++ zig/kernels/metal/deepseek_v41/kv.metal | 141 ++++++ zig/kernels/metal/deepseek_v41/mhc.metal | 120 +++++ zig/kernels/metal/deepseek_v41/ops.metal | 124 ++++++ zig/kernels/metal/deepseek_v41/rope.metal | 17 + zig/kernels/metal/deepseek_v41/route.metal | 53 +++ zig/kernels/metal/sources.zig | 11 + zig/src/core/expert_groups.zig | 101 +++++ zig/src/core/gpu_counters.zig | 78 +++- zig/src/families/deepseek_v41/attention.zig | 258 +++++++++++ zig/src/families/deepseek_v41/csa2.zig | 77 ++++ zig/src/families/deepseek_v41/csa2_index.zig | 139 ++++++ zig/src/families/deepseek_v41/csa2_test.zig | 309 +++++++++++++ zig/src/families/deepseek_v41/engram.zig | 203 +++++++++ zig/src/families/deepseek_v41/engram_test.zig | 247 +++++++++++ .../fixtures/hc_attention_rms.json | 25 ++ .../deepseek_v41/index_precision_test.zig | 57 +++ zig/src/families/deepseek_v41/kv.zig | 162 +++++++ zig/src/families/deepseek_v41/kv_test.zig | 234 ++++++++++ zig/src/families/deepseek_v41/metal_kv.zig | 111 +++++ zig/src/families/deepseek_v41/metal_model.zig | 363 ++++++++++++++++ zig/src/families/deepseek_v41/mhc.zig | 95 ++++ zig/src/families/deepseek_v41/mhc_test.zig | 209 +++++++++ zig/src/families/deepseek_v41/ops.zig | 110 +++++ zig/src/families/deepseek_v41/ops_check.zig | 191 ++++++++ zig/src/families/deepseek_v41/profiler.zig | 93 ++++ zig/src/families/deepseek_v41/rope.zig | 134 ++++++ zig/src/families/deepseek_v41/route.zig | 39 ++ zig/src/families/deepseek_v41/route_test.zig | 152 +++++++ .../families/deepseek_v41/test_support.zig | 40 ++ zig/src/families/deepseek_v41/workspace.zig | 107 +++++ zig/src/metal/counters.zig | 4 +- zig/src/metal/metal.zig | 1 + zig/src/tensorfold.zig | 4 + .../deepseek_v41_reference/attention_math.zig | 298 +++++++++++++ .../deepseek_v41_reference/engram_math.zig | 106 +++++ zig/tests/deepseek_v41_reference/math.zig | 289 +++++++++++++ .../operator_attention.zig | 155 +++++++ .../operator_diagnostic.zig | 124 ++++++ .../deepseek_v41_reference/operator_hc.zig | 98 +++++ .../operator_history.zig | 298 +++++++++++++ .../deepseek_v41_reference/operator_moe.zig | 409 ++++++++++++++++++ .../deepseek_v41_reference/operators.zig | 92 ++++ zig/tests/deepseek_v41_reference/profile.zig | 52 +++ .../deepseek_v41_reference/projection.zig | 261 +++++++++++ zig/tests/deepseek_v41_reference/qualify.zig | 97 +++++ .../deepseek_v41_reference/reference.zig | 176 ++++++++ .../rope_frequencies.json | 14 + zig/tests/deepseek_v41_reference/trace.zig | 123 ++++++ 58 files changed, 7340 insertions(+), 25 deletions(-) create mode 100644 zig/deepseek_v41_checks.zig create mode 100644 zig/deepseek_v41_reference_checks.zig create mode 100644 zig/deepseek_v41_run.zig create mode 100644 zig/kernels/metal/core/expert_groups.metal create mode 100644 zig/kernels/metal/deepseek_v41/csa2.metal create mode 100644 zig/kernels/metal/deepseek_v41/engram.metal create mode 100644 zig/kernels/metal/deepseek_v41/index.metal create mode 100644 zig/kernels/metal/deepseek_v41/kv.metal create mode 100644 zig/kernels/metal/deepseek_v41/mhc.metal create mode 100644 zig/kernels/metal/deepseek_v41/ops.metal create mode 100644 zig/kernels/metal/deepseek_v41/rope.metal create mode 100644 zig/kernels/metal/deepseek_v41/route.metal create mode 100644 zig/src/core/expert_groups.zig create mode 100644 zig/src/families/deepseek_v41/attention.zig create mode 100644 zig/src/families/deepseek_v41/csa2.zig create mode 100644 zig/src/families/deepseek_v41/csa2_index.zig create mode 100644 zig/src/families/deepseek_v41/csa2_test.zig create mode 100644 zig/src/families/deepseek_v41/engram.zig create mode 100644 zig/src/families/deepseek_v41/engram_test.zig create mode 100644 zig/src/families/deepseek_v41/fixtures/hc_attention_rms.json create mode 100644 zig/src/families/deepseek_v41/index_precision_test.zig create mode 100644 zig/src/families/deepseek_v41/kv.zig create mode 100644 zig/src/families/deepseek_v41/kv_test.zig create mode 100644 zig/src/families/deepseek_v41/metal_kv.zig create mode 100644 zig/src/families/deepseek_v41/metal_model.zig create mode 100644 zig/src/families/deepseek_v41/mhc.zig create mode 100644 zig/src/families/deepseek_v41/mhc_test.zig create mode 100644 zig/src/families/deepseek_v41/ops.zig create mode 100644 zig/src/families/deepseek_v41/ops_check.zig create mode 100644 zig/src/families/deepseek_v41/profiler.zig create mode 100644 zig/src/families/deepseek_v41/rope.zig create mode 100644 zig/src/families/deepseek_v41/route.zig create mode 100644 zig/src/families/deepseek_v41/route_test.zig create mode 100644 zig/src/families/deepseek_v41/test_support.zig create mode 100644 zig/src/families/deepseek_v41/workspace.zig create mode 100644 zig/tests/deepseek_v41_reference/attention_math.zig create mode 100644 zig/tests/deepseek_v41_reference/engram_math.zig create mode 100644 zig/tests/deepseek_v41_reference/math.zig create mode 100644 zig/tests/deepseek_v41_reference/operator_attention.zig create mode 100644 zig/tests/deepseek_v41_reference/operator_diagnostic.zig create mode 100644 zig/tests/deepseek_v41_reference/operator_hc.zig create mode 100644 zig/tests/deepseek_v41_reference/operator_history.zig create mode 100644 zig/tests/deepseek_v41_reference/operator_moe.zig create mode 100644 zig/tests/deepseek_v41_reference/operators.zig create mode 100644 zig/tests/deepseek_v41_reference/profile.zig create mode 100644 zig/tests/deepseek_v41_reference/projection.zig create mode 100644 zig/tests/deepseek_v41_reference/qualify.zig create mode 100644 zig/tests/deepseek_v41_reference/reference.zig create mode 100644 zig/tests/deepseek_v41_reference/rope_frequencies.json create mode 100644 zig/tests/deepseek_v41_reference/trace.zig diff --git a/build.zig b/build.zig index b12fb3c40f..b8884fe5f0 100644 --- a/build.zig +++ b/build.zig @@ -208,11 +208,10 @@ pub fn build(b: *std.Build) void { } /// The native Metal server build is owned by the shared Metal target module. +const metalTargets = metal_build.metalTargets; const nativeServer = metal_build.nativeServer; /// macOS: the metallib, the Metal runtime's test programs and the engine over Metal. -const metalTargets = metal_build.metalTargets; - /// The server and distribution use the same embedded-source Metal engine. const metalEngineModules = metal_build.metalEngineModules; diff --git a/zig/build/deepseek_v41.zig b/zig/build/deepseek_v41.zig index 6b196fedff..2e419e9755 100644 --- a/zig/build/deepseek_v41.zig +++ b/zig/build/deepseek_v41.zig @@ -20,3 +20,41 @@ pub fn weights( step.dependOn(&b.addRunArtifact(host).step); return step; } + +pub fn checks( + b: *std.Build, + target: std.Build.ResolvedTarget, + optimize: std.builtin.OptimizeMode, + metal: *std.Build.Module, + sources: *std.Build.Module, + engine: *std.Build.Module, +) struct { device: *std.Build.Step, reference: *std.Build.Step } { + const device = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/deepseek_v41_checks.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "metal", .module = metal }, + .{ .name = "kernel_sources", .module = sources }, + }, + }) }); + const device_step = b.step("test-deepseek-v41", "Experimental DS4.1 EXL3, Engram, mHC and sparse attention on Metal"); + device_step.dependOn(&b.addRunArtifact(device).step); + const reference = b.addTest(.{ + .root_module = b.createModule(.{ + .root_source_file = b.path("zig/deepseek_v41_reference_checks.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "tensorfold", .module = engine }, + .{ .name = "metal", .module = metal }, + }, + }), + .filters = &.{"tests.deepseek_v41_reference."}, + }); + const reference_step = b.step("test-deepseek-v41-reference", "Independent CPU teacher-oracle arithmetic boundaries"); + reference_step.dependOn(&b.addRunArtifact(reference).step); + return .{ .device = device_step, .reference = reference_step }; +} diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index d5ab369dd7..646305427e 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -1,5 +1,6 @@ //! Metal build targets and their shared source/library/module construction. const std = @import("std"); +const dsv41_build = @import("deepseek_v41.zig"); const builtin = @import("builtin"); const qwen27_build = @import("qwen27.zig"); const core_checks = @import("core_checks.zig"); @@ -170,6 +171,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s const programs = [_]struct { name: []const u8, path: []const u8, about: []const u8 }{ .{ .name = "tf-metal-tests", .path = "zig/tests/gpu_tests.zig", .about = "GPU runtime tests" }, .{ .name = "tf-dispatch-bench", .path = "zig/tests/dispatch_bench.zig", .about = "Dispatch boundary benchmark" }, + .{ .name = "tf-dsv41-run", .path = "zig/deepseek_v41_run.zig", .about = "Full-resident DeepSeek V4.1 admission, teacher-forced logits, chunk seams and decode timing" }, .{ .name = "tf-kernel-exact", .path = "zig/tests/kernel_exact.zig", .about = "Bit-exact check of a kernel against an MLX oracle case" }, }; for (programs) |p| { @@ -276,6 +278,10 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s const load_admission_run = b.addRunArtifact(load_admission_cases); b.step("test-flashnext-load-admission", "CPU-only Flash Next runtime quantization admission").dependOn(&load_admission_run.step); test_step.dependOn(&load_admission_run.step); + test_step.dependOn(dsv41_build.weights(b, target, optimize, metal)); + const dsv41 = dsv41_build.checks(b, target, optimize, metal, mods.sources, mods.engine); + test_step.dependOn(dsv41.device); + test_step.dependOn(dsv41.reference); const pack_cache_cases = b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/flashnext_pack_cache_cases.zig"), .target = target, diff --git a/zig/deepseek_v41_checks.zig b/zig/deepseek_v41_checks.zig new file mode 100644 index 0000000000..755a35c618 --- /dev/null +++ b/zig/deepseek_v41_checks.zig @@ -0,0 +1,13 @@ +//! Source-built V4.1 operator qualification; real Metal execution, no checkpoint needed. +test { + _ = @import("src/core/exl3_projection_gpu_test.zig"); + _ = @import("src/families/deepseek_v41/engram.zig"); + _ = @import("src/families/deepseek_v41/csa2.zig"); + _ = @import("src/families/deepseek_v41/ops_check.zig"); + _ = @import("src/families/deepseek_v41/load.zig"); + _ = @import("src/families/deepseek_v41/index_precision_test.zig"); + _ = @import("src/families/deepseek_v41/config.zig"); + _ = @import("src/families/deepseek_v41/profiler.zig"); + _ = @import("src/core/gpu_counters.zig"); + _ = @import("src/core/expert_groups.zig"); +} diff --git a/zig/deepseek_v41_reference_checks.zig b/zig/deepseek_v41_reference_checks.zig new file mode 100644 index 0000000000..61aa30f788 --- /dev/null +++ b/zig/deepseek_v41_reference_checks.zig @@ -0,0 +1,7 @@ +//! Independent CPU arithmetic boundaries used by the resident-checkpoint teacher oracle. +test { + _ = @import("tests/deepseek_v41_reference/projection.zig"); + _ = @import("tests/deepseek_v41_reference/math.zig"); + _ = @import("tests/deepseek_v41_reference/attention_math.zig"); + _ = @import("tests/deepseek_v41_reference/qualify.zig"); +} diff --git a/zig/deepseek_v41_run.zig b/zig/deepseek_v41_run.zig new file mode 100644 index 0000000000..e6ca81e783 --- /dev/null +++ b/zig/deepseek_v41_run.zig @@ -0,0 +1,306 @@ +//! Full-checkpoint admission and execution receipts; no synthetic weights or replay substitute. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const ds = tf.deepseek_v41; +const config = tf.deepseek_v41_config; +const load = tf.deepseek_v41_load; +const weights = tf.deepseek_v41_weights; +const process_memory = @import("src/server/process_memory.zig"); +const oracle = @import("tests/deepseek_v41_reference/qualify.zig"); +const profile = @import("tests/deepseek_v41_reference/profile.zig"); +const trace = @import("tests/deepseek_v41_reference/trace.zig"); +const operators = @import("tests/deepseek_v41_reference/operators.zig"); +const attention_operators = @import("tests/deepseek_v41_reference/operator_attention.zig"); +const history_operators = @import("tests/deepseek_v41_reference/operator_history.zig"); +const hc_operators = @import("tests/deepseek_v41_reference/operator_hc.zig"); +const moe_operators = @import("tests/deepseek_v41_reference/operator_moe.zig"); + +fn usage(code: u8) noreturn { + std.debug.print("usage: tf-dsv41-run plan MODEL TABLES\n" ++ + " tf-dsv41-run teacher MODEL TABLES IDS_JSON OUT_DIR [CHUNK]\n" ++ + " tf-dsv41-run seams MODEL TABLES IDS_JSON OUT_DIR\n" ++ + " tf-dsv41-run oracle MODEL TABLES IDS_JSON OUT_DIR\n" ++ + " tf-dsv41-run profile MODEL TABLES IDS_JSON OUT_DIR\n" ++ + " tf-dsv41-run trace MODEL TABLES IDS_JSON OUT_DIR\n" ++ + " tf-dsv41-run operators MODEL TABLES OUT_DIR\n" ++ + " tf-dsv41-run attention-sameinput MODEL TABLES TRACE_CPU_BIN OUT_DIR [LAYER]\n" ++ + " tf-dsv41-run attention-history-sameinput MODEL TABLES TRACE_DIR OUT_DIR TOKEN LAYER\n" ++ + " tf-dsv41-run hc-sameinput MODEL TABLES TRACE_CPU_BIN OUT_DIR [LAYER]\n" ++ + " tf-dsv41-run moe-sameinput MODEL TABLES TRACE_GPU_BIN OUT_DIR [LAYER]\n" ++ + " tf-dsv41-run generate MODEL TABLES IDS_JSON OUT_DIR [TOKENS]\n" ++ + "IDS_JSON: nonempty JSON array of token IDs below 129280.\n" ++ + "Limits: oracle/trace 16 IDs; profile 256 IDs; teacher/seams 4096 IDs.\n" ++ + "History: TOKEN below 4096; retained normalized prefix at most 256 MiB.\n" ++ + "Local diagnostics do not qualify the model. Oracle retains the independent 1% NRMS gate.\n", .{}); + std.process.exit(code); +} +fn number(s: []const u8) !u32 { + return std.fmt.parseInt(u32, s, 10); +} +fn validArguments(args: []const [:0]const u8) bool { + if (args.len < 2) return false; + const mode = args[1]; + for ([_][]const u8{ "plan", "operators" }, [_]usize{ 4, 5 }) |name, count| + if (std.mem.eql(u8, mode, name)) return args.len == count; + if (std.mem.eql(u8, mode, "attention-history-sameinput")) return args.len == 8; + for ([_][]const u8{ "teacher", "generate", "attention-sameinput", "hc-sameinput", "moe-sameinput" }) |name| + if (std.mem.eql(u8, mode, name)) return args.len == 6 or args.len == 7; + for ([_][]const u8{ "seams", "oracle", "profile", "trace" }) |name| + if (std.mem.eql(u8, mode, name)) return args.len == 6; + return false; +} +fn argmax(row: []const f32) !u32 { + var best: u32 = 0; + if (row.len == 0) return error.EmptyLogits; + var score = -std.math.inf(f32); + for (row, 0..) |v, i| { + if (!std.math.isFinite(v)) return error.NonfiniteLogits; + if (v > score) { + best = @intCast(i); + score = v; + } + } + return best; +} +fn write(io: std.Io, a: std.mem.Allocator, dir: []const u8, name: []const u8, data: []const u8) !void { + const path = try std.fs.path.join(a, &.{ dir, name }); + defer a.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); +} +const Counters = struct { disk_read: u64, pageins: u64, footprint: u64 }; +/// Darwin's rusage_info_v4 fields the runner reads around execution (page-ins, disk reads, footprint). +const RUsage = extern struct { + uuid: [16]u8, + user_time: u64, + system_time: u64, + pkg_idle_wkups: u64, + interrupt_wkups: u64, + pageins: u64, + wired_size: u64, + resident_size: u64, + phys_footprint: u64, + proc_start_abstime: u64, + proc_exit_abstime: u64, + child_user_time: u64, + child_system_time: u64, + child_pkg_idle_wkups: u64, + child_interrupt_wkups: u64, + child_pageins: u64, + child_elapsed_abstime: u64, + diskio_bytesread: u64, +}; +fn counters() !Counters { + const proc_pid_rusage = struct { + extern "c" fn proc_pid_rusage(pid: c_int, flavor: c_int, buffer: *anyopaque) c_int; + }.proc_pid_rusage; + const getpid = struct { + extern "c" fn getpid() c_int; + }.getpid; + var r: RUsage = undefined; + if (proc_pid_rusage(getpid(), 4, &r) != 0) return error.MemoryCountersUnavailable; + return .{ .disk_read = r.diskio_bytesread, .pageins = r.pageins, .footprint = r.phys_footprint }; +} +fn teacher(m: *ds.Model, ids: []const u32, chunk: u32, out: []f32) !f64 { + if (chunk == 0 or chunk > m.workspace.plan.rows or ids.len > m.cache.plan.capacity) + return error.InvalidTeacherWindow; + if (out.len != ids.len * 129280) return error.InvalidTeacherOutputLength; + try m.reset(); + var elapsed: f64 = 0; + var start: usize = 0; + while (start < ids.len) { + const end = @min(start + chunk, ids.len); + const row = try m.forward(ids[start..end], .all); + @memcpy(out[start * 129280 .. end * 129280], row); + for (row) |value| if (!std.math.isFinite(value)) return error.NonfiniteLogits; + elapsed += m.last_gpu_seconds; + start = end; + } + return elapsed; +} + +pub fn main(init: std.process.Init) !void { + const a = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len == 2 and (std.mem.eql(u8, args[1], "--help") or std.mem.eql(u8, args[1], "-h"))) usage(0); + if (!validArguments(args)) usage(2); + const arena = init.arena.allocator(); + const mode = args[1]; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + if (std.mem.eql(u8, mode, "plan")) { + const cfg = try config.Config.load(arena, io, args[2]); + var plan = try load.Plan.open(a, io, args[2], args[3]); + defer plan.deinit(); + try weights.validate(&plan); + const device = try mtl.Device.init(); + defer device.deinit(); + std.debug.print("{{\"mode\":\"plan\",\"weight_bytes\":{d},\"shards\":{d},\"tensors\":{d},\"context\":{d},\"metal_working_set_bytes\":{d}}}\n", .{ plan.bytes, plan.shards.items.len, plan.tensors.count(), cfg.text_config.max_position_embeddings, device.maxWorkingSet() }); + return; + } + const is_operators = std.mem.eql(u8, mode, "operators"); + const is_attention = std.mem.eql(u8, mode, "attention-sameinput"); + const is_attention_history = std.mem.eql(u8, mode, "attention-history-sameinput"); + const is_hc = std.mem.eql(u8, mode, "hc-sameinput"); + const is_moe = std.mem.eql(u8, mode, "moe-sameinput"); + const max_ids: usize = if (std.mem.eql(u8, mode, "oracle") or std.mem.eql(u8, mode, "trace")) 16 else if (std.mem.eql(u8, mode, "profile")) ds.prefill_step else if (std.mem.eql(u8, mode, "teacher") or std.mem.eql(u8, mode, "seams")) 4096 else 1048576; + const input = if (is_operators or is_attention or is_attention_history or is_hc or is_moe) &.{} else try std.Io.Dir.cwd().readFileAlloc(io, args[4], arena, .limited(max_ids * 12 + 2)); + const ids: []const u32 = if (is_operators or is_attention or is_attention_history or is_hc or is_moe) &.{} else try std.json.parseFromSliceLeaky([]const u32, arena, input, .{}); + if (!is_operators and !is_attention and !is_attention_history and !is_hc and !is_moe and (ids.len == 0 or ids.len > 1048576)) return error.InvalidPrompt; + if (ids.len > max_ids) return error.InvalidPrompt; + for (ids) |id| if (id >= 129280) return error.InvalidToken; + const is_oracle = std.mem.eql(u8, mode, "oracle"); + const is_trace = std.mem.eql(u8, mode, "trace"); + const dir = args[if (is_operators) @as(usize, 4) else 5]; + try std.Io.Dir.cwd().createDirPath(io, dir); + const map = try std.fs.path.join(arena, &.{ args[2], "engram-token-map.bin" }); + const map_blob = if (is_oracle or is_trace) + try std.Io.Dir.cwd().readFileAlloc(io, map, arena, .limited(2 << 20)) + else + &.{}; + const model = try a.create(ds.Model); + defer a.destroy(model); + const begin = mtl.clock.seconds(); + model.* = try ds.Model.open(a, io, .{ .dir = args[2], .table_dir = args[3], .token_map = map }); + defer model.deinit(); + const ready = mtl.clock.seconds(); + const before = try counters(); + var execution_after: ?Counters = null; + var execution_error: ?anyerror = null; + execute(a, io, model, args, ids, map_blob, dir, before, &execution_after) catch |err| { + execution_error = err; + }; + const after = execution_after orelse try counters(); + if (after.disk_read < before.disk_read or after.pageins < before.pageins) return error.NonmonotonicCounters; + const resident_gate_passed = after.disk_read == before.disk_read; + const receipt = try std.json.Stringify.valueAlloc(a, .{ + .mode = mode, + .evidence_origin = "current_invocation_not_historical_qualification", + .context_capacity = model.cache.plan.capacity, + .consumed_tokens = model.cache.position, + .weight_bytes = model.checkpoint.residentBytes(), + .resident_allocation_bytes = model.residency.allocatedSize(), + .ready_footprint_bytes = before.footprint, + .final_footprint_bytes = after.footprint, + .load_seconds = ready - begin, + .measurement_scope = if (execution_after != null) "model_execution" else "execution_and_diagnostic_io", + .execution_disk_read_bytes = after.disk_read - before.disk_read, + .execution_pageins = after.pageins - before.pageins, + .resident_gate_passed = resident_gate_passed, + .execution_error = if (execution_error) |err| @as(?[]const u8, @errorName(err)) else null, + }, .{}); + defer a.free(receipt); + std.debug.print("{s}\n", .{receipt}); + try write(io, a, dir, "residency.json", receipt); + if (!resident_gate_passed) std.debug.print("residency gate failed: ResidentExecutionReadFromDisk\n", .{}); + if (execution_error) |err| return err; + if (!resident_gate_passed) return error.ResidentExecutionReadFromDisk; +} + +fn execute(a: std.mem.Allocator, io: std.Io, model: *ds.Model, args: []const [:0]const u8, ids: []const u32, map_blob: []const u8, dir: []const u8, before: Counters, execution_after: *?Counters) !void { + const mode = args[1]; + const is_operators = std.mem.eql(u8, mode, "operators"); + const is_attention = std.mem.eql(u8, mode, "attention-sameinput"); + const is_attention_history = std.mem.eql(u8, mode, "attention-history-sameinput"); + const is_hc = std.mem.eql(u8, mode, "hc-sameinput"); + const is_moe = std.mem.eql(u8, mode, "moe-sameinput"); + const is_oracle = std.mem.eql(u8, mode, "oracle"); + const is_profile = std.mem.eql(u8, mode, "profile"); + const is_trace = std.mem.eql(u8, mode, "trace"); + const is_teacher = std.mem.eql(u8, mode, "teacher"); + const is_seams = std.mem.eql(u8, mode, "seams"); + if (is_operators) { + try operators.run(a, io, model, dir); + } else if (is_attention) { + try attention_operators.runAttention(a, io, model, args[4], dir, if (args.len > 6) try number(args[6]) else 0); + } else if (is_attention_history) { + try history_operators.runAttentionHistory(a, io, model, args[4], dir, try number(args[6]), try number(args[7])); + } else if (is_hc) { + try hc_operators.runHyperconnections(a, io, model, args[4], dir, if (args.len > 6) try number(args[6]) else 0); + } else if (is_moe) { + try moe_operators.runMoe(a, io, model, args[4], dir, if (args.len > 6) try number(args[6]) else 21); + } else if (is_oracle) { + try operators.run(a, io, model, dir); + _ = try oracle.run(a, io, model, ids, map_blob, dir); + } else if (is_profile) { + try profile.run(a, io, model, ids, dir); + } else if (is_trace) { + try trace.run(a, io, model, ids, map_blob, dir); + } else if (is_teacher or is_seams) { + const chunk = if (args.len > 6) try number(args[6]) else model.workspace.plan.rows; + if (chunk == 0 or chunk > model.workspace.plan.rows) return error.InvalidPrefillChunk; + const outputs = try a.alloc(f32, ids.len * 129280); + defer a.free(outputs); + const elapsed = try teacher(model, ids, chunk, outputs); + // Capture execution before writing the receipt, not filesystem metadata reads from the receipt itself. + if (is_seams) { + const alternative = try a.alloc(f32, outputs.len); + defer a.free(alternative); + for ([_]u32{ 1, 2, 3, 7, 31, 127, 128, 129, 255, 256 }) |split| { + if (split > model.workspace.plan.rows) continue; + _ = try teacher(model, ids, split, alternative); + var error_max: f64 = 0; + var disagreements: usize = 0; + var bit_disagreements: usize = 0; + for (outputs, alternative) |x, y| { + error_max = @max(error_max, @abs(@as(f64, x) - y)); + bit_disagreements += @intFromBool(@as(u32, @bitCast(x)) != @as(u32, @bitCast(y))); + } + for (0..ids.len) |i| { + if (try argmax(outputs[i * 129280 ..][0..129280]) != try argmax(alternative[i * 129280 ..][0..129280])) + disagreements += 1; + } + std.debug.print("{{\"mode\":\"seam\",\"chunk\":{d},\"max_logit_error\":{d},\"bit_disagreements\":{d},\"greedy_disagreements\":{d}}}\n", .{ split, error_max, bit_disagreements, disagreements }); + if (disagreements != 0 or bit_disagreements != 0) { + var name: [48]u8 = undefined; + const file = try std.fmt.bufPrint(&name, "seam-{d}-logits.f32", .{split}); + execution_after.* = try counters(); + try write(io, a, dir, file, std.mem.sliceAsBytes(alternative)); + return error.ChunkSeamMismatch; + } + } + } + std.debug.print("{{\"mode\":\"teacher\",\"tokens\":{d},\"gpu_seconds\":{d}}}\n", .{ ids.len, elapsed }); + execution_after.* = try counters(); + try write(io, a, dir, "logits.f32", std.mem.sliceAsBytes(outputs)); + } else if (std.mem.eql(u8, mode, "generate")) { + const n = if (args.len > 6) try number(args[6]) else 64; + if (ids.len > model.cache.plan.capacity or n == 0 or n > model.cache.plan.capacity - ids.len) + return error.ContextFull; + try model.reset(); + const prefill_start = mtl.clock.seconds(); + var row: []const f32 = &.{}; + var start: usize = 0; + var progress_at = prefill_start; + while (start < ids.len) { + const end = @min(start + model.workspace.plan.rows, ids.len); + row = try model.forward(ids[start..end], if (end == ids.len) .last else .none); + start = end; + const now = mtl.clock.seconds(); + if (now - progress_at >= 60 or start == ids.len) { + const memory = try counters(); + std.debug.print("{{\"mode\":\"prefill_progress\",\"completed_tokens\":{d},\"prompt_tokens\":{d},\"elapsed_seconds\":{d},\"footprint_bytes\":{d},\"disk_read_bytes\":{d},\"pageins\":{d}}}\n", .{ start, ids.len, now - prefill_start, memory.footprint, memory.disk_read - before.disk_read, memory.pageins - before.pageins }); + progress_at = now; + } + } + const prefill_end = mtl.clock.seconds(); + const generated = try a.alloc(u32, n); + defer a.free(generated); + var gpu_seconds: f64 = 0; + const decode_start = mtl.clock.seconds(); + for (generated, 0..) |*token, i| { + token.* = try argmax(row); + if (i + 1 < generated.len) { + row = try model.forward(&.{token.*}, .last); + gpu_seconds += model.last_gpu_seconds; + } + } + const decode_end = mtl.clock.seconds(); + execution_after.* = try counters(); + const json = try std.json.Stringify.valueAlloc(a, generated, .{}); + defer a.free(json); + try write(io, a, dir, "tokens.json", json); + std.debug.print("{{\"mode\":\"generate\",\"prompt_tokens\":{d},\"completion_tokens\":{d},\"prefill_seconds\":{d},\"decode_seconds\":{d},\"decode_gpu_seconds\":{d}}}\n", .{ ids.len, n, prefill_end - prefill_start, decode_end - decode_start, gpu_seconds }); + } else usage(2); +} diff --git a/zig/kernels/metal/core/expert_groups.metal b/zig/kernels/metal/core/expert_groups.metal new file mode 100644 index 0000000000..be543d5780 --- /dev/null +++ b/zig/kernels/metal/core/expert_groups.metal @@ -0,0 +1,36 @@ +// Group routed slots into padded eight-row expert tiles; scatter preserves original token and output IDs. +#include +using namespace metal; +struct TfExpertGroups { uint rows, experts, repeat, padded; }; +#define TF_GROUP_ARGS \ + const device uint* ids [[buffer(0)]], device atomic_uint* counts [[buffer(1)]], \ + device uint* offsets [[buffer(2)]], device uint* ranks [[buffer(3)]], \ + device uint* expert_ids [[buffer(4)]], device uint* token_ids [[buffer(5)]], \ + device uint* slot_ids [[buffer(6)]], constant TfExpertGroups& d [[buffer(7)]], \ + uint i [[thread_position_in_grid]] + +kernel void tf_expert_group_clear(TF_GROUP_ARGS) { + if (i < d.experts) atomic_store_explicit(counts + i, 0u, memory_order_relaxed); + if (i < d.padded) expert_ids[i] = 0xffffffffu; +} +kernel void tf_expert_group_count(TF_GROUP_ARGS) { + if (i >= d.rows || ids[i] >= d.experts) return; + ranks[i] = atomic_fetch_add_explicit(counts + ids[i], 1u, memory_order_relaxed); +} +kernel void tf_expert_group_prefix(TF_GROUP_ARGS) { + if (i != 0) return; + uint offset = 0; + for (uint expert = 0; expert < d.experts; ++expert) { + offsets[expert] = offset; + const uint count = atomic_load_explicit(counts + expert, memory_order_relaxed); + offset += (count + 7u) & ~7u; + } +} +kernel void tf_expert_group_scatter(TF_GROUP_ARGS) { + if (i >= d.rows || ids[i] >= d.experts) return; + const uint destination = offsets[ids[i]] + ranks[i]; + if (destination >= d.padded) return; + expert_ids[destination] = ids[i]; + token_ids[destination] = i / d.repeat; + slot_ids[destination] = i; +} diff --git a/zig/kernels/metal/deepseek_v41/csa2.metal b/zig/kernels/metal/deepseek_v41/csa2.metal new file mode 100644 index 0000000000..9d16cfcce5 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/csa2.metal @@ -0,0 +1,102 @@ +// Source is compiled after kv.metal, which provides exact packed-cache decoders. +struct Dsv41Attention { + uint rows, heads, width, q_stride, position, chunk_rows, window, ring_stride; + uint chunk_stride, main_stride, main_block, ratio, compressed_rows, topk, index_stride; + float scale; +}; +// BF16 query/key products are exact in FP32; compensate accumulation and lane reduction. +inline void dsv41_dot_add(thread float &sum, thread float &error, float value) { + float next = sum + value; + float recovered = next - sum; + error += (sum - (next - recovered)) + (value - recovered); + sum = next; +} +inline float dsv41_dot_reduce(float sum, float error) { + for (uint offset = 16; offset != 0; offset >>= 1) { + float other = simd_shuffle_xor(sum, offset); + float other_error = simd_shuffle_xor(error, offset); + dsv41_dot_add(sum, error, other); + dsv41_dot_add(sum, error, other_error); + } + return sum + error; +} +// BF16 probability rounding uses fixed query-based tiles, never a chunk-dependent physical key layout. +inline uint dsv41_window_columns(constant Dsv41Attention &p, uint row) { + return p.position + row == 0 ? 1 : p.window; +} +inline device const uchar *dsv41_attn_key(uint slot, uint row, device const uchar *ring, + device const uchar *chunk, device const uchar *main, device const int *indices, + constant Dsv41Attention &p, thread uint &kind) { + uint query = p.position + row; + uint columns = dsv41_window_columns(p, row); + if (slot < columns) { + uint pos; + if (query == 0) { + pos = 0; + } else if (query + 1 < p.window) { + pos = (query + 1 + slot) % p.window; + } else { + uint begin = query + 1 > p.window ? query + 1 - p.window : 0; + pos = begin + slot; + } + if (pos > query) return nullptr; + kind = 0; + return pos >= p.position ? chunk + (pos - p.position) * p.chunk_stride + : ring + (pos % p.window) * p.ring_stride; + } + if (!p.ratio) return nullptr; + int pos = indices[row * p.index_stride + slot - columns]; + if (pos < 0 || uint(pos) >= min(p.compressed_rows, (query + 1) / p.ratio)) return nullptr; + kind = 1; + return main + uint(pos) * p.main_stride; +} +kernel void dsv41_sparse_attention(device const ushort *q [[buffer(0)]], device const uchar *ring [[buffer(1)]], + device const uchar *chunk [[buffer(2)]], device const uchar *main [[buffer(3)]], + device const int *indices [[buffer(4)]], device const float *sink [[buffer(5)]], + device ushort *out [[buffer(6)]], constant Dsv41Attention &p [[buffer(7)]], + uint2 g [[threadgroup_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + uint head = g.x, row = g.y; + if (head >= p.heads || row >= p.rows) return; + threadgroup float scores[64]; + float query[16], acc[16]; + for (uint j = 0; j < p.width / 32; ++j) { + query[j] = dsv41_bf(q[row * p.q_stride + head * p.width + lane + j * 32]); acc[j] = 0; + } + float mx = -1e30f, sum = 0; + uint query_position = p.position + row; + uint compressed_columns = p.ratio ? min(p.topk, min(p.compressed_rows, (query_position + 1) / p.ratio)) : 0; + uint count = dsv41_window_columns(p, row) + compressed_columns; + for (uint base = 0; base < count; base += 64) { + float next_max = mx; + for (uint j = 0; j < 64; ++j) { + uint kind = 0; + device const uchar *key = base + j < count ? dsv41_attn_key(base + j, row, ring, chunk, main, indices, p, kind) : nullptr; + float dot = 0, dot_error = 0; + if (key) for (uint d = 0; d < p.width / 32; ++d) + dsv41_dot_add(dot, dot_error, query[d] * dsv41_unpack(key, lane + d * 32, p.width, kind ? p.main_block : 32, kind)); + dot = key ? dsv41_dot_reduce(dot, dot_error) * p.scale : -INFINITY; + next_max = max(next_max, dot); + if (lane == 0) scores[j] = dot; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float correction = exp(mx - next_max); + float tile_sum = 0; + for (uint d = 0; d < p.width / 32; ++d) acc[d] *= correction; + for (uint j = 0; j < 64; ++j) { + float probability = exp(scores[j] - next_max); + tile_sum += probability; + float rounded = dsv41_bf(dsv41_round(probability)); + uint kind = 0; + device const uchar *key = base + j < count ? dsv41_attn_key(base + j, row, ring, chunk, main, indices, p, kind) : nullptr; + if (key) for (uint d = 0; d < p.width / 32; ++d) + acc[d] += rounded * dsv41_unpack(key, lane + d * 32, p.width, kind ? p.main_block : 32, kind); + } + sum = sum * correction + tile_sum; + mx = next_max; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + // The source adds sink mass to the completed online denominator before the FP32 division. + float denom = sum + exp(sink[head] - mx); + for (uint d = 0; d < p.width / 32; ++d) + out[row * p.q_stride + head * p.width + lane + d * 32] = dsv41_round(acc[d] / denom); +} diff --git a/zig/kernels/metal/deepseek_v41/engram.metal b/zig/kernels/metal/deepseek_v41/engram.metal new file mode 100644 index 0000000000..532bbd01ee --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/engram.metal @@ -0,0 +1,150 @@ +// Math/layout attribution: DeepSeek AI DeepSeek-V4.1-Flash inference/{engram,model}.py (MIT). +#include +using namespace metal; + +struct EngramLayer { + uint layer_id, rows; + ulong multipliers[4]; + uint primes[24], offsets[24]; +}; +struct EngramHistory { int tail[3]; uint available; }; +struct EngramHashArgs { uint tokens, pad, masked; }; + +inline int engram_token(const device uint* ids, const device uint* map, + const device uchar* mask, uint pos, uint masked) { + return masked && mask[pos] == 0 ? -1 : int(map[ids[pos]]); +} + +// One lane per token; history is immutable until all layers have hashed this chunk. +kernel void tf_v41_engram_hash(const device uint* ids [[buffer(0)]], + const device uint* map [[buffer(1)]], + const device uchar* mask [[buffer(2)]], + const device EngramHistory& history [[buffer(3)]], + device uint* out [[buffer(4)]], + constant EngramLayer& layer [[buffer(5)]], + constant EngramHashArgs& args [[buffer(6)]], + uint pos [[thread_position_in_grid]]) { + if (pos >= args.tokens) return; + bool blocked = false; + ulong rolling = 0; + for (uint shift = 0; shift < 4; shift++) { + int token; + if (shift <= pos) token = engram_token(ids, map, mask, pos - shift, args.masked); + else { + const uint back = shift - pos - 1; + token = back < history.available ? history.tail[back] : -1; + } + blocked = blocked || token == -1; + const ulong compressed = blocked ? args.pad : uint(token); + rolling ^= compressed * layer.multipliers[shift]; + if (shift != 0) for (uint head = 0; head < 8; head++) { + const uint col = (shift - 1) * 8 + head; + out[size_t(pos) * 24 + col] = uint(rolling % layer.primes[col]) + layer.offsets[col]; + } + } +} + +// Called exactly once after all Engram layers hash a chunk, never once per layer. +kernel void tf_v41_engram_commit(const device uint* ids [[buffer(0)]], + const device uint* map [[buffer(1)]], + const device uchar* mask [[buffer(2)]], + device EngramHistory& history [[buffer(3)]], + constant EngramHashArgs& args [[buffer(4)]]) { + int tail[3]; + for (uint i = 0; i < 3; i++) { + tail[i] = i < args.tokens ? engram_token(ids, map, mask, args.tokens - i - 1, args.masked) + : i - args.tokens < history.available ? history.tail[i - args.tokens] : -1; + } + for (uint i = 0; i < 3; i++) history.tail[i] = tail[i]; + history.available = min(3u, history.available + args.tokens); +} + +kernel void tf_v41_engram_reset(device EngramHistory& history [[buffer(0)]]) { + for (uint i = 0; i < 3; i++) history.tail[i] = -1; + history.available = 0; +} + +inline float engram_e4m3(uchar bits) { + const uint magnitude = bits & 127, exponent = magnitude >> 3, mantissa = magnitude & 7; + float value = magnitude == 127 ? NAN : exponent == 0 ? float(mantissa) * 0x1p-9f + : ldexp(1.0f + float(mantissa) * 0.125f, int(exponent) - 7); + return bits & 128 ? -value : value; +} +inline float engram_ue8m0(uchar bits) { + return bits == 255 ? NAN : as_type(bits == 0 ? 0x00400000u : uint(bits) << 23); +} + +struct EngramGatherArgs { + ulong weight_stride, scale_stride, output_stride; + uint tokens, columns, head_dim, scale_block; +}; + +// Read only the selected E4M3FN rows; one UE8M0 scale per 32 values of each row. +kernel void tf_v41_engram_gather(const device uint* ids [[buffer(0)]], + const device uchar* table [[buffer(1)]], + const device uchar* scales [[buffer(2)]], + device bfloat* out [[buffer(3)]], + constant EngramGatherArgs& args [[buffer(4)]], + uint2 pos [[thread_position_in_grid]]) { + if (pos.y >= args.tokens || pos.x >= args.columns * args.head_dim) return; + const uint col = pos.x / args.head_dim, d = pos.x % args.head_dim; + const size_t row = ids[size_t(pos.y) * args.columns + col]; + const float value = engram_e4m3(table[row * args.weight_stride + d]); + const float scale = engram_ue8m0(scales[row * args.scale_stride + d / args.scale_block]); + out[size_t(pos.y) * args.output_stride + pos.x] = bfloat(value * scale); +} + +struct EngramGateArgs { + ulong hidden_stride, kv_stride, output_stride; + uint dim, copies, masked; + float eps; +}; + +// Each copy gets its own RMS dot gate; all copies share the value following the flattened keys. +kernel void tf_v41_engram_gate(const device bfloat* hidden [[buffer(0)]], + const device bfloat* kv [[buffer(1)]], + const device float* q_weight [[buffer(2)]], + const device float* k_weight [[buffer(3)]], + const device uchar* mask [[buffer(4)]], + device bfloat* out [[buffer(5)]], + constant EngramGateArgs& args [[buffer(6)]], + uint2 group [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float3 sums[8]; + threadgroup float gate; + const size_t hbase = size_t(group.x) * args.hidden_stride + size_t(group.y) * args.dim; + const size_t kbase = size_t(group.x) * args.kv_stride + size_t(group.y) * args.dim; + const size_t wbase = size_t(group.y) * args.dim; + float3 acc = 0; + for (uint d = t; d < args.dim; d += 256) { + const float h = float(hidden[hbase + d]), k = float(kv[kbase + d]); + const float weight = q_weight[wbase + d] * k_weight[wbase + d]; + acc += float3(h * h, k * k, (h * weight) * k); + } + acc = float3(simd_sum(acc.x), simd_sum(acc.y), simd_sum(acc.z)); + if (lane == 0) sums[sg] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + acc = lane < 8 ? sums[lane] : float3(0); + acc = float3(simd_sum(acc.x), simd_sum(acc.y), simd_sum(acc.z)); + if (lane == 0) { + float rstd; + { +#pragma clang fp contract(off) + const float factor = 1.0f / float(args.dim); + rstd = metal::precise::rsqrt(acc.x * factor + args.eps) + * metal::precise::rsqrt(acc.y * factor + args.eps); + } + const float dot = (acc.z * rstd) * metal::precise::rsqrt(float(args.dim)); + const float signed_root = copysign(metal::precise::sqrt(max(abs(dot), 1e-6f)), dot); + gate = args.masked && mask[group.x] == 0 ? 0.0f : 1.0f / (1.0f + metal::precise::exp(-signed_root)); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const size_t vbase = size_t(group.x) * args.kv_stride + size_t(args.copies) * args.dim; + const size_t obase = size_t(group.x) * args.output_stride + size_t(group.y) * args.dim; + for (uint d = t; d < args.dim; d += 256) + out[obase + d] = bfloat(float(hidden[hbase + d]) + gate * float(kv[vbase + d])); +} diff --git a/zig/kernels/metal/deepseek_v41/index.metal b/zig/kernels/metal/deepseek_v41/index.metal new file mode 100644 index 0000000000..8d62337624 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/index.metal @@ -0,0 +1,103 @@ +// Source is compiled after kv.metal; no query-by-key-by-head score tensor is materialized. +struct Dsv41Index { + uint rows, heads, width, q_stride, weight_stride, position, ratio, keys, key_stride; + uint tile_start, tile_width, topk, candidate_topk, block_size, use_candidates, make_candidates; + float scale; +}; +struct Dsv41Pick { float score; int position; }; +inline bool dsv41_less(Dsv41Pick a, Dsv41Pick b) { + return a.score < b.score || (a.score == b.score && a.position > b.position); +} +inline void dsv41_heap(device Dsv41Pick *h, uint size, Dsv41Pick v) { + if (!dsv41_less(h[0], v)) return; + uint at = 0; + while (2 * at + 1 < size) { + uint child = 2 * at + 1; + if (child + 1 < size && dsv41_less(h[child + 1], h[child])) ++child; + if (!dsv41_less(h[child], v)) break; + h[at] = h[child]; at = child; + } + h[at] = v; +} +inline int dsv41_index_position(uint local, uint row, device const Dsv41Pick *candidates, + constant Dsv41Index &p) { + uint slot = p.tile_start + local; + if (!p.use_candidates) return int(slot); + uint block = slot / p.block_size; + if (block >= p.candidate_topk) return -1; + int id = candidates[row * p.candidate_topk + block].position; + return id < 0 ? -1 : id * int(p.block_size) + int(slot % p.block_size); +} +kernel void dsv41_index_score(device const ushort *q [[buffer(0)]], device const ushort *weight [[buffer(1)]], + device const uchar *key [[buffer(2)]], device const Dsv41Pick *candidates [[buffer(3)]], + device float *scores [[buffer(4)]], constant Dsv41Index &p [[buffer(5)]], + uint2 g [[threadgroup_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + uint local = g.x, row = g.y; + if (local >= p.tile_width || row >= p.rows) return; + int pos = dsv41_index_position(local, row, candidates, p); + uint visible = min(p.keys, (p.position + row + 1) / p.ratio); + if (pos < 0 || uint(pos) >= visible) { + if (lane == 0) scores[row * p.tile_width + local] = -INFINITY; + return; + } + device const uchar *k = key + uint(pos) * p.key_stride; + float key_values[4]; + for (uint d = 0; d < p.width / 32; ++d) key_values[d] = dsv41_unpack(k, lane + 32 * d, p.width, 32, 2); + float total = 0; + for (uint head = 0; head < p.heads; ++head) { + float dot = 0; + for (uint d = 0; d < p.width / 32; ++d) + dot += dsv41_bf(q[row * p.q_stride + head * p.width + lane + 32 * d]) * key_values[d]; + dot = dsv41_bf(dsv41_round(simd_sum(dot))); + const float scaled_weight = dsv41_bf(dsv41_round(dsv41_bf(weight[row * p.weight_stride + head]) * p.scale)); + total += dsv41_bf(dsv41_round(max(dot, 0.0f) * scaled_weight)); + } + if (lane == 0) scores[row * p.tile_width + local] = dsv41_bf(dsv41_round(total)); +} +kernel void dsv41_index_clear(device Dsv41Pick *top [[buffer(0)]], device Dsv41Pick *candidate [[buffer(1)]], + constant Dsv41Index &p [[buffer(2)]], uint2 i [[thread_position_in_grid]]) { + if (i.y >= p.rows) return; + if (i.x < p.topk) top[i.y * p.topk + i.x] = {-INFINITY, -1}; + if (p.make_candidates && i.x < p.candidate_topk) candidate[i.y * p.candidate_topk + i.x] = {-INFINITY, -1}; +} +kernel void dsv41_index_merge(device const float *scores [[buffer(0)]], device Dsv41Pick *top [[buffer(1)]], + device Dsv41Pick *candidate [[buffer(2)]], constant Dsv41Index &p [[buffer(3)]], + uint row [[thread_position_in_grid]]) { + if (row >= p.rows) return; + device Dsv41Pick *heap = top + row * p.topk; + for (uint i = 0; i < p.tile_width; ++i) { + float score = scores[row * p.tile_width + i]; + int pos = dsv41_index_position(i, row, candidate, p); + if (isfinite(score)) dsv41_heap(heap, p.topk, {score, pos}); + } + if (!p.make_candidates) return; + uint visible = min(p.keys, (p.position + row + 1) / p.ratio); + device Dsv41Pick *blocks = candidate + row * p.candidate_topk; + for (uint base = 0; base < p.tile_width; base += p.block_size) { + float score = -INFINITY; + for (uint j = 0; j < p.block_size; ++j) score = max(score, scores[row * p.tile_width + base + j]); + uint block = (p.tile_start + base) / p.block_size; + if (visible && block == (visible - 1) / p.block_size) score = INFINITY; + if (score > -INFINITY) dsv41_heap(blocks, p.candidate_topk, {score, int(block)}); + } +} +inline void dsv41_sort_positions(device Dsv41Pick *p, uint size) { + // Selection is score-ordered; attention sees the reference's ascending position order. + for (uint i = 1; i < size; ++i) { + Dsv41Pick value = p[i]; uint j = i; + while (j > 0 && (p[j - 1].position < 0 || + (value.position >= 0 && p[j - 1].position > value.position))) { + if (value.position < 0) break; + p[j] = p[j - 1]; --j; + } + p[j] = value; + } +} +kernel void dsv41_index_finish(device Dsv41Pick *top [[buffer(0)]], device Dsv41Pick *candidate [[buffer(1)]], + device int *positions [[buffer(2)]], constant Dsv41Index &p [[buffer(3)]], uint row [[thread_position_in_grid]]) { + if (row >= p.rows) return; + device Dsv41Pick *heap = top + row * p.topk; + dsv41_sort_positions(heap, p.topk); + for (uint i = 0; i < p.topk; ++i) positions[row * p.topk + i] = heap[i].position; + if (p.make_candidates) dsv41_sort_positions(candidate + row * p.candidate_topk, p.candidate_topk); +} diff --git a/zig/kernels/metal/deepseek_v41/kv.metal b/zig/kernels/metal/deepseek_v41/kv.metal new file mode 100644 index 0000000000..6ac2017efb --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/kv.metal @@ -0,0 +1,141 @@ +// DeepSeek V4.1 cache quantization and pooling, following the official inference kernels. +#include +using namespace metal; +inline float dsv41_bf(ushort x) { return as_type(uint(x) << 16); } +inline ushort dsv41_round(float x) { + uint u = as_type(x); + return ushort((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} +inline float dsv41_e4(uchar c) { + uint e = (c >> 3) & 15, m = c & 7; + float v = e == 0 ? float(m) * 0x1p-9f : ldexp(1.0f + float(m) / 8.0f, int(e) - 7); + return (c & 128) ? -v : v; +} +inline uchar dsv41_e4code(float x) { + float a = min(abs(x), 448.0f); + if (a == 0) return signbit(x) ? 128 : 0; + int e = max(-6, int((as_type(a) >> 23) & 255) - 127); + int q = int(rint(ldexp(a, 3 - e))); + uint c = e == -6 && q < 8 ? uint(q) : uint((e + 7) * 8 + q - 8); + return uchar(min(c, 126u) | (signbit(x) ? 128u : 0u)); +} +inline float dsv41_e2(uchar c) { + const float v[8] = {0, .5f, 1, 1.5f, 2, 3, 4, 6}; + return (c & 8) ? -v[c & 7] : v[c & 7]; +} +inline uchar dsv41_e2code(float x) { + float a = min(abs(x), 6.0f); + uint pick = a <= .25f ? 0 : a < .75f ? 1 : a <= 1.25f ? 2 : a < 1.75f ? 3 + : a <= 2.5f ? 4 : a < 3.5f ? 5 : a <= 5.0f ? 6 : 7; + return uchar(pick | (signbit(x) ? 8 : 0)); +} +inline float dsv41_unpack(device const uchar *p, uint d, uint width, uint block, uint kind) { + if (kind == 0) { + float s = as_type(uint(p[width + d / 32]) << 23); + return dsv41_bf(dsv41_round(dsv41_e4(p[d]) * s)); + } + uchar c = p[d / 2]; + float s = kind == 1 ? dsv41_e4(p[width / 2 + d / block]) + : as_type(uint(p[width / 2 + d / block]) << 23); + return dsv41_bf(dsv41_round(dsv41_e2((d & 1) ? c >> 4 : c & 15) * s)); +} +struct Dsv41Pack { uint rows, width, stride, packed_stride, block, kind, first, ring, inplace; }; +kernel void dsv41_pack(device ushort *x [[buffer(0)]], device uchar *out [[buffer(1)]], + constant Dsv41Pack &p [[buffer(2)]], uint2 g [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_threadgroup]]) { + uint row = g.y, group = g.x, d = group * p.block + lane; + if (row >= p.rows) return; + float v = lane < p.block ? dsv41_bf(x[row * p.stride + d]) : 0; + float a = simd_max(abs(v)); + float scale; uchar sc; + if (p.kind == 1) { + sc = dsv41_e4code(max(a, 6.0f * 0x1p-9f) / 6.0f); scale = dsv41_e4(sc); + } else { + float minimum = p.kind == 0 ? 1e-4f : 6.0f * 0x1p-126f; + float z = max(a, minimum) / (p.kind == 0 ? 448.0f : 6.0f); + uint bits = as_type(z); + int exp = int((bits >> 23) & 255) - 127 + ((bits & 0x7fffff) != 0); + sc = uchar(exp + 127); scale = ldexp(1.0f, exp); + } + uint dst = p.ring ? (p.first + row) % p.ring : p.first + row; + device uchar *packed = out + dst * p.packed_stride; + if (lane == 0) packed[(p.kind == 0 ? p.width : p.width / 2) + group] = sc; + if (lane >= p.block) return; + if (p.kind == 0) { + uchar c = dsv41_e4code(v / scale); packed[d] = c; + if (p.inplace) x[row * p.stride + d] = dsv41_round(dsv41_e4(c) * scale); + } else { + uchar c = dsv41_e2code(v / scale); + uchar other = uchar(simd_shuffle_xor(uint(c), 1)); + if (!(lane & 1)) packed[d / 2] = c | (other << 4); + if (p.inplace) x[row * p.stride + d] = dsv41_round(dsv41_e2(c) * scale); + } +} +struct Dsv41Pool { uint rows, width, ratio, position, stride, input_f32; float eps; }; +kernel void dsv41_pool(device const uchar *values [[buffer(0)]], device const float *gate [[buffer(1)]], + device float *partial [[buffer(2)]], device const ushort *weight [[buffer(3)]], + device ushort *latent [[buffer(4)]], constant Dsv41Pool &p [[buffer(5)]], + uint group [[threadgroup_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + // A group reads the prior chunk's tail before the independent tail-update dispatch. + uint start = (p.position / p.ratio + group) * p.ratio; + threadgroup float pooled[512]; + float square = 0; + for (uint d = lane; d < p.width; d += 128) { + float mx = -INFINITY, sum = 0, total = 0; + for (uint j = 0; j < p.ratio; ++j) { + uint pos = start + j; + float score = p.ratio == 1 ? 0 : (pos < p.position ? partial[p.ratio * p.width + (pos % p.ratio) * p.width + d] + : gate[(pos - p.position) * p.stride + d]); + mx = max(mx, score); + } + float slots[2], exponent[2]; + for (uint j = 0; j < p.ratio; ++j) { + uint pos = start + j; + slots[j] = pos < p.position ? partial[(pos % p.ratio) * p.width + d] + : (p.input_f32 ? reinterpret_cast(values)[(pos - p.position) * p.stride + d] + : dsv41_bf(reinterpret_cast(values)[(pos - p.position) * p.stride + d])); + float score = p.ratio == 1 ? 0 : (pos < p.position ? partial[p.ratio * p.width + (pos % p.ratio) * p.width + d] + : gate[(pos - p.position) * p.stride + d]); + exponent[j] = exp(score - mx); sum += exponent[j]; + } + for (uint j = 0; j < p.ratio; ++j) total += slots[j] * (exponent[j] / sum); + float v = dsv41_bf(dsv41_round(total)); pooled[d] = v; square += v * v; + } + threadgroup float sums[4]; + float ss = simd_sum(square); + if (!(lane % 32)) sums[lane / 32] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float inv = rsqrt((sums[0] + sums[1] + sums[2] + sums[3]) / p.width + p.eps); + for (uint d = lane; d < p.width; d += 128) + latent[group * p.width + d] = dsv41_round(pooled[d] * inv * dsv41_bf(weight[d])); +} +kernel void dsv41_pool_tail(device const float *values [[buffer(0)]], device const float *gate [[buffer(1)]], + device float *partial [[buffer(2)]], constant Dsv41Pool &p [[buffer(3)]], + uint d [[thread_position_in_grid]]) { + if (d >= p.width || p.ratio == 1) return; + uint end = p.position + p.rows, tail = end % p.ratio; + for (uint j = 0; j < tail; ++j) { + uint pos = end - tail + j; + if (pos < p.position) continue; + partial[j * p.width + d] = values[(pos - p.position) * p.stride + d]; + partial[p.ratio * p.width + j * p.width + d] = gate[(pos - p.position) * p.stride + d]; + } +} +kernel void dsv41_copy(device const uchar *src [[buffer(0)]], device uchar *dst [[buffer(1)]], + constant uint &bytes [[buffer(2)]], uint i [[thread_position_in_grid]]) { + if (i < bytes) dst[i] = src[i]; +} +kernel void dsv41_publish_window(device const uchar *chunk [[buffer(0)]], device uchar *ring [[buffer(1)]], + constant uint4 &p [[buffer(2)]], uint2 i [[thread_position_in_grid]]) { + uint count = min(p.x, 128u), first = p.x - count; + if (i.x >= p.z || i.y >= count) return; + uint row = first + i.y, slot = (p.y + row) % 128; + ring[slot * p.w + i.x] = chunk[row * p.z + i.x]; +} +kernel void dsv41_unpack_rows(device const uchar *packed [[buffer(0)]], device ushort *out [[buffer(1)]], + constant Dsv41Pack &p [[buffer(2)]], uint2 i [[thread_position_in_grid]]) { + if (i.x >= p.width || i.y >= p.rows) return; + uint row = p.ring ? (p.first + i.y) % p.ring : p.first + i.y; + out[i.y * p.stride + i.x] = dsv41_round(dsv41_unpack(packed + row * p.packed_stride, + i.x, p.width, p.block, p.kind)); +} diff --git a/zig/kernels/metal/deepseek_v41/mhc.metal b/zig/kernels/metal/deepseek_v41/mhc.metal new file mode 100644 index 0000000000..8b6203a4cf --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/mhc.metal @@ -0,0 +1,120 @@ +// Reference: DeepSeek-V4.1-Flash inference/model.py and kernel.py, revision 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +#include +using namespace metal; +inline float bf(ushort x) { return as_type(uint(x)<<16); } +inline ushort rb(float x) { + uint u=as_type(x); + if ((u&0x7fffffff)>0x7f800000) return ushort((u>>16)|0x40); + return ushort((u+0x7fff+((u>>16)&1))>>16); +} +// Torch CUDA materializes FP32 squares before MeanOps accumulation. +#pragma clang fp contract(off) +struct Rms { uint rows,width,stream_stride; float eps; }; +kernel void dsv41_hc_rms(const device ushort* x [[buffer(0)]], device float* inv [[buffer(1)]], + constant Rms& d [[buffer(2)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + // Fixed single-row CUDA tree: four vector accumulators, stride 512, independent of prefill chunks. + threadgroup float sums[512]; + float4 accum=0.0f; + for (uint i=4*t;i<4*d.width;i+=4*512) { + float4 a=float4(bf(x[ulong(row)*d.stream_stride+i]), + bf(x[ulong(row)*d.stream_stride+i+1]), + bf(x[ulong(row)*d.stream_stride+i+2]), + bf(x[ulong(row)*d.stream_stride+i+3])); + float4 square=a*a; + accum+=square; + } + float s=((accum.x+accum.y)+accum.z)+accum.w; + sums[t]=s; + for (uint offset=256;offset>=32;offset>>=1) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t0;offset>>=1) s+=simd_shuffle_down(s,offset); + // MeanOps projects with a rounded FP32 factor, not a division of the sum. + float factor=1.0f/float(4*d.width); + if (lane==0) inv[row]=rsqrt(s*factor+d.eps); + } +} +#pragma clang fp contract(on) +struct Split { uint rows,mix_stride,coefficient_stride,comb_stride; float eps; }; +kernel void dsv41_hc_split(const device float* mixes [[buffer(0)]], const device float* inv [[buffer(1)]], + const device float* scale [[buffer(2)]], const device float* base [[buffer(3)]], + device float* pre [[buffer(4)]], device float* post [[buffer(5)]], + device float* comb [[buffer(6)]], constant Split& d [[buffer(7)]], + uint row [[thread_position_in_grid]]) { + if (row>=d.rows) return; + float m[16]; + for (uint j=0;j<4;j++) { + float a=fma(mixes[ulong(row)*d.mix_stride+j]*inv[row],scale[0],base[j]); + float b=fma(mixes[ulong(row)*d.mix_stride+4+j]*inv[row],scale[1],base[4+j]); + pre[ulong(row)*d.coefficient_stride+j]=1.0f/(1.0f+exp(-a))+d.eps; + post[ulong(row)*d.coefficient_stride+j]=2.0f/(1.0f+exp(-b)); + for (uint k=0;k<4;k++) { + uint q=4*j+k; + m[q]=fma(mixes[ulong(row)*d.mix_stride+8+q]*inv[row],scale[2],base[8+q]); + } + } + // Initial row softmax + eps, then a column normalization. This counts as iteration one. + for (uint j=0;j<4;j++) { + float mx=max(max(m[4*j],m[4*j+1]),max(m[4*j+2],m[4*j+3])); + for (uint k=0;k<4;k++) m[4*j+k]=exp(m[4*j+k]-mx); + // Pinned TileLang HC row/column sums use FP32 butterfly pairs. + float s=(m[4*j]+m[4*j+2])+(m[4*j+1]+m[4*j+3]); + for (uint k=0;k<4;k++) m[4*j+k]=m[4*j+k]/s+d.eps; + } + for (uint it=0;it<20;it++) { + if (it>0) { + for (uint j=0;j<4;j++) { + float s=((m[4*j]+m[4*j+2])+(m[4*j+1]+m[4*j+3]))+d.eps; + for (uint k=0;k<4;k++) m[4*j+k]/=s; + } + } + for (uint k=0;k<4;k++) { + float s=((m[k]+m[8+k])+(m[4+k]+m[12+k]))+d.eps; + for (uint j=0;j<4;j++) m[4*j+k]/=s; + } + } + for (uint q=0;q<16;q++) comb[ulong(row)*d.comb_stride+q]=m[q]; +} +// PyTorch materializes these FP32 products before reduction and branch addition. +#pragma clang fp contract(off) +struct Collapse { uint rows,width,stream_stride,y_stride,pre_stride; }; +kernel void dsv41_hc_collapse(const device ushort* x [[buffer(0)]], const device float* previous_pre [[buffer(1)]], + device ushort* y [[buffer(2)]], constant Collapse& d [[buffer(3)]], + uint2 i [[thread_position_in_grid]]) { + if (i.x>=d.width || i.y>=d.rows) return; + float s=0; + for (uint c=0;c<4;c++) s+=previous_pre[ulong(i.y)*d.pre_stride+c]*bf(x[ulong(i.y)*d.stream_stride+ulong(c)*d.width+i.x]); + y[ulong(i.y)*d.y_stride+i.x]=rb(s); +} +struct Post { uint rows,width,branch_stride,residual_stride,y_stride,coefficient_stride,comb_stride; }; +kernel void dsv41_hc_post(const device ushort* branch [[buffer(0)]], const device ushort* residual [[buffer(1)]], + const device float* post [[buffer(2)]], const device float* comb [[buffer(3)]], + device ushort* y [[buffer(4)]], constant Post& d [[buffer(5)]], + uint2 i [[thread_position_in_grid]]) { + if (i.x>=d.width || i.y>=d.rows) return; + float r[4]; + for (uint j=0;j<4;j++) r[j]=bf(residual[ulong(i.y)*d.residual_stride+ulong(j)*d.width+i.x]); + float b=bf(branch[ulong(i.y)*d.branch_stride+i.x]); + for (uint k=0;k<4;k++) { + float s=0; + for (uint j=0;j<4;j++) s+=comb[ulong(i.y)*d.comb_stride+4*j+k]*r[j]; + y[ulong(i.y)*d.y_stride+ulong(k)*d.width+i.x]=rb(post[ulong(i.y)*d.coefficient_stride+k]*b+s); + } +} +#pragma clang fp contract(on) +struct Expand { uint rows,width,x_stride,stream_stride,pre_stride; }; +kernel void dsv41_hc_expand(const device ushort* x [[buffer(0)]], device ushort* stream [[buffer(1)]], + device float* pre [[buffer(2)]], constant Expand& d [[buffer(3)]], + uint2 i [[thread_position_in_grid]]) { + if (i.x>=d.width || i.y>=d.rows) return; + for (uint c=0;c<4;c++) stream[ulong(i.y)*d.stream_stride+ulong(c)*d.width+i.x]=x[ulong(i.y)*d.x_stride+i.x]; + if (i.x==0) for (uint c=0;c<4;c++) pre[ulong(i.y)*d.pre_stride+c]=c==0?1.0f:0.0f; +} diff --git a/zig/kernels/metal/deepseek_v41/ops.metal b/zig/kernels/metal/deepseek_v41/ops.metal new file mode 100644 index 0000000000..d831901dd5 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/ops.metal @@ -0,0 +1,124 @@ +// Math: deepseek-ai/DeepSeek-V4.1-Flash inference/model.py, revision 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +#include +using namespace metal; +inline float bf(ushort x) { return as_type(uint(x) << 16); } +inline ushort rb(float x) { + uint u = as_type(x); + if ((u & 0x7fffffff) > 0x7f800000) return ushort((u >> 16) | 0x40); + return ushort((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} +inline float rd(const device uchar* p, ulong i, uint dtype) { + if (dtype==1) return ((const device float*)p)[i]; + if (dtype==2) return float(((const device half*)p)[i]); + return bf(((const device ushort*)p)[i]); +} +struct Rms { uint rows, width, x_stride, y_stride; float eps; uint weight_f32; }; +#pragma clang fp contract(off) +kernel void dsv41_rms(const device ushort* x [[buffer(0)]], const device uchar* w [[buffer(1)]], + device ushort* y [[buffer(2)]], constant Rms& d [[buffer(3)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float sums[8]; + float s = 0; + for (uint i = t; i < d.width; i += 256) { float a = bf(x[ulong(row)*d.x_stride+i]); s += a*a; } + s = simd_sum(s); + if (lane == 0) sums[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float total = simd_sum(lane < 8 ? sums[lane] : 0.0f); + const float factor = 1.0f / float(d.width); + if (lane == 0) sums[0] = rsqrt(total * factor + d.eps); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = t; i < d.width; i += 256) + y[ulong(row)*d.y_stride+i] = rb((bf(x[ulong(row)*d.x_stride+i])*sums[0])*rd(w,i,d.weight_f32)); +} +#pragma clang fp contract(on) +// Each SIMD group emits four columns; compensated dots do not depend on token count. + +// Precision: 0=BF16, 1=FP32, 2=FP16; FP16 weights never round through BF16. +inline void dense_add(thread float& high, thread float& low, float value) { + float next = high + value; + float borrowed = next - high; + low += (high - (next - borrowed)) + (value - borrowed); + high = next; +} +struct Dense { uint rows, k, n, groups, x_stride, w_stride, y_stride, x_dtype, w_dtype, y_dtype; }; +kernel void dsv41_dense(const device uchar* x [[buffer(0)]], const device uchar* w [[buffer(1)]], + device uchar* y [[buffer(2)]], constant Dense& d [[buffer(3)]], + uint3 g [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + uint col = g.x*16 + sg*4; + float a[4] = {0,0,0,0}, residual[4] = {0,0,0,0}; + for (uint k = lane; k < d.k; k += 32) { + float v = rd(x,ulong(g.y)*d.x_stride+ulong(g.z)*d.k+k,d.x_dtype); + for (uint j=0;j<4;j++) if (col+j>= 1) { + float other = simd_shuffle_xor(a[j], offset); + float other_residual = simd_shuffle_xor(residual[j], offset); + dense_add(a[j], residual[j], other); + residual[j] += other_residual; + } + float v = a[j] + residual[j]; + if (lane==0 && col+j=d.width || i.y>=d.rows) return; + ushort v=table[ulong(ids[i.y])*d.table_stride+i.x]; + for (uint c=0;c=d.width || i.y>=d.rows) return; + ulong k=ulong(i.y)*d.input_stride+i.x; + float a=min(rd(gate,k,d.input_f32),10.0f), b=clamp(rd(up,k,d.input_f32),-10.0f,10.0f); + float v=(a/(1.0f+exp(-a)))*b; + if (d.weighted) v=weights[i.y]*v; + y[ulong(i.y)*d.output_stride+i.x]=rb(v); +} +struct Combine { uint rows,width,routed_row_stride,slot_stride,shared_stride,y_stride,input_f32,id_stride; }; +kernel void dsv41_combine(const device uchar* routed [[buffer(0)]], const device uchar* shared [[buffer(1)]], + const device uint* ids [[buffer(2)]], device ushort* y [[buffer(3)]], + constant Combine& d [[buffer(4)]], + uint2 i [[thread_position_in_grid]], uint t [[thread_index_in_threadgroup]]) { + threadgroup uint slots[6]; + if (t==0) { + for (uint a=0;a<6;a++) slots[a]=a; + for (uint a=1;a<6;a++) { + uint slot=slots[a], j=a; + while (j>0 && ids[ulong(i.y)*d.id_stride+slots[j-1]]>ids[ulong(i.y)*d.id_stride+slot]) { + slots[j]=slots[j-1]; j--; + } + slots[j]=slot; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i.x>=d.width || i.y>=d.rows) return; + for (uint slot=0;slot<6;slot++) if (ids[ulong(i.y)*d.id_stride+slot]>=384) { + y[ulong(i.y)*d.y_stride+i.x]=rb(NAN); + return; + } + float s=0; + for (uint rank=0;rank<6;rank++) + s+=rd(routed,ulong(i.y)*d.routed_row_stride+ulong(slots[rank])*d.slot_stride+i.x,d.input_f32); + s+=rd(shared,ulong(i.y)*d.shared_stride+i.x,d.input_f32); + y[ulong(i.y)*d.y_stride+i.x]=rb(s); +} diff --git a/zig/kernels/metal/deepseek_v41/rope.metal b/zig/kernels/metal/deepseek_v41/rope.metal new file mode 100644 index 0000000000..09c4d1cb2e --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/rope.metal @@ -0,0 +1,17 @@ +// Adjacent-pair RoPE and YaRN match DeepSeek V4.1's in-place BF16 rotation. +#include +using namespace metal; +struct Dsv41Rope { uint rows, heads, width, rotary, stride, position, step, inverse, spectrum; }; +kernel void dsv41_rope(device ushort *x [[buffer(0)]], constant Dsv41Rope &p [[buffer(1)]], + device const float2 *factors [[buffer(2)]], + uint3 i [[thread_position_in_grid]]) { + if (i.x >= p.rotary / 2 || i.y >= p.heads || i.z >= p.rows) return; + // The encoder binds the selected spectrum; rows are absolute source positions. + float2 rotation = factors[(ulong(p.position) + ulong(i.z) * p.step) * (p.rotary / 2) + i.x]; + float c = rotation.x, s = p.inverse ? -rotation.y : rotation.y; + uint at = i.z * p.stride + i.y * p.width + p.width - p.rotary + 2 * i.x; + float a = as_type(uint(x[at]) << 16), b = as_type(uint(x[at + 1]) << 16); + uint u = as_type(a * c - b * s), v = as_type(a * s + b * c); + x[at] = ushort((u + 0x7fff + ((u >> 16) & 1)) >> 16); + x[at + 1] = ushort((v + 0x7fff + ((v >> 16) & 1)) >> 16); +} diff --git a/zig/kernels/metal/deepseek_v41/route.metal b/zig/kernels/metal/deepseek_v41/route.metal new file mode 100644 index 0000000000..51f5454dd3 --- /dev/null +++ b/zig/kernels/metal/deepseek_v41/route.metal @@ -0,0 +1,53 @@ +// Reference: DeepSeek-V4.1-Flash inference/model.py Gate, revision 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +#include +using namespace metal; +inline float lp(float x) { float u=1.0f+x; return u==1.0f?x:x*(log(u)/(u-1.0f)); } +inline float score(float x) { return sqrt(x>20.0f?x:lp(exp(x))); } +struct Shape { uint rows,logit_stride,output_stride; float temperature; uint bias_dtype; }; +kernel void dsv41_route(const device float* logits [[buffer(0)]], const device uchar* bias [[buffer(1)]], + device uint* ids [[buffer(2)]], device float* weights [[buffer(3)]], + constant Shape& d [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float scores[384], select[384], best[4], picked[6]; + threadgroup uint best_id[4], winner; + for (uint e=t;e<384;e+=128) { + float b; + if (d.bias_dtype==1) b=((const device float*)bias)[e]; + else if (d.bias_dtype==2) b=float(((const device half*)bias)[e]); + else b=as_type(uint(((const device ushort*)bias)[e])<<16); + scores[e]=score(logits[ulong(row)*d.logit_stride+e]/d.temperature); + select[e]=scores[e]+b; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // Every expert participates in every rank selection; ties use the lower expert id deterministically. + for (uint slot=0;slot<6;slot++) { + float v=-INFINITY; + uint id=0xffffffffu; + for (uint e=t;e<384;e+=128) if (select[e]>v || (select[e]==v && e std.math.maxInt(u32) or experts > std.math.maxInt(u32)) + return error.InvalidExpertGrouping; + const active = @min(rows, experts); + const tiles = try std.math.add(usize, active, (rows - active) / 8); + const padded = try std.math.mul(usize, tiles, 8); + if (padded > std.math.maxInt(u32)) return error.ExpertGroupingTooLarge; + const sizes = [_]usize{ experts * 4, experts * 4, rows * 4, padded * 4, padded * 4, padded * 4 }; + var plan = Plan{ .rows = @intCast(rows), .experts = @intCast(experts), .padded = @intCast(padded), .offsets = undefined, .bytes = 0 }; + for (sizes, 0..) |size, i| { + plan.offsets[i] = std.mem.alignForward(usize, plan.bytes, 16); + plan.bytes = try std.math.add(usize, plan.offsets[i], size); + } + plan.bytes = std.mem.alignForward(usize, plan.bytes, 16); + return plan; + } +}; + +pub const Scratch = struct { + buffer: mtl.Buffer, + offset: usize = 0, + plan: Plan, + + pub fn view(self: Scratch, field: Field) exl3.View { + return .{ .buffer = self.buffer, .offset = self.offset + self.plan.offsets[@backingInt(field)] }; + } +}; + +pub const Pipelines = struct { + stages: [4]mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const lib = try mtl.Library.fromSource(device, sources.core_expert_groups, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var stages: [4]mtl.Pipeline = undefined; + var loaded: usize = 0; + errdefer for (stages[0..loaded]) |pipeline| pipeline.deinit(); + inline for (.{ "tf_expert_group_clear", "tf_expert_group_count", "tf_expert_group_prefix", "tf_expert_group_scatter" }, 0..) |name, i| { + stages[i] = try mtl.Pipeline.init(device, lib, name, false); + loaded += 1; + } + return .{ .stages = stages }; + } + pub fn deinit(self: Pipelines) void { + for (self.stages) |pipeline| pipeline.deinit(); + } +}; +const Args = extern struct { rows: u32, experts: u32, repeat: u32, padded: u32 }; + +pub fn encode(enc: mtl.ComputeEncoder, pipelines: *const Pipelines, ids: exl3.View, scratch: Scratch, repeat: u32) !void { + // Plan must come from Plan.init; route IDs remain GPU-resident and are bounds-checked by the kernels. + if (repeat == 0 or scratch.plan.rows == 0 or scratch.plan.experts == 0 or + scratch.plan.rows % repeat != 0) return error.InvalidExpertGrouping; + if (ids.offset % 4 != 0 or scratch.offset % 4 != 0 or + ids.offset > ids.buffer.length() or @as(usize, scratch.plan.rows) * 4 > ids.buffer.length() - ids.offset or + scratch.offset > scratch.buffer.length() or scratch.plan.bytes > scratch.buffer.length() - scratch.offset) + return error.InvalidExpertGroupingStorage; + const args = Args{ .rows = scratch.plan.rows, .experts = scratch.plan.experts, .repeat = repeat, .padded = scratch.plan.padded }; + for (pipelines.stages, 0..) |pipeline, phase| { + enc.setPipeline(pipeline); + enc.setBuffer(ids.buffer, ids.offset, 0); + inline for (std.enums.values(Field), 1..) |field, i| { + const view = scratch.view(field); + enc.setBuffer(view.buffer, view.offset, i); + } + enc.setValue(args, 7); + enc.dispatchThreads(.{ .width = switch (phase) { + 0 => @max(args.experts, args.padded), + 2 => 1, + else => args.rows, + } }, .{ .width = 256 }); + enc.barrier(); + } +} + +test "expert grouping bounds padding without context-sized allocation" { + const small = try Plan.init(6, 384); + try std.testing.expectEqual(@as(u32, 48), small.padded); + const full = try Plan.init(1536, 384); + try std.testing.expectEqual(@as(u32, 4224), full.padded); + try std.testing.expect(full.bytes < 64 * 1024); + try std.testing.expectError(error.InvalidExpertGrouping, Plan.init(0, 384)); + try std.testing.expectError(error.ExpertGroupingTooLarge, Plan.init(std.math.maxInt(u32), 384)); +} diff --git a/zig/src/core/gpu_counters.zig b/zig/src/core/gpu_counters.zig index 2e0b216165..72edee1343 100644 --- a/zig/src/core/gpu_counters.zig +++ b/zig/src/core/gpu_counters.zig @@ -1,4 +1,4 @@ -//! Same-command-buffer GPU timestamps, bounded resources and no per-span CPU waits. +//! Bounded same-command GPU timestamps; disabled init allocates nothing, no per-span CPU waits or CPU timing fallback. const std = @import("std"); const mtl = @import("metal"); const counters = mtl.counters; @@ -9,6 +9,8 @@ pub const max_span_capacity: usize = 65536; pub const Options = struct { enabled: bool = false, max_spans: usize = 512 }; pub const Span = struct { kind: Class = .dense, + /// Caller-owned diagnostic label; does not affect sampling or generic class aggregation. + metadata: u32 = 0, begin_timestamp: u64 = 0, end_timestamp: u64 = 0, gpu_seconds: f64 = 0, @@ -16,7 +18,7 @@ pub const Span = struct { pub const Aggregate = struct { calls: u64 = 0, gpu_seconds: f64 = 0 }; const State = enum { idle, recording, resolved, read, broken }; -/// Two reference samples convert GPU ticks to nanoseconds; never use CPU span time. +/// Apple's paired reference samples convert GPU ticks to CPU nanoseconds, not elapsed operator time. pub const Calibration = struct { before: counters.ClockSample, after: counters.ClockSample, @@ -29,7 +31,8 @@ pub const Calibration = struct { } pub fn duration(self: Calibration, begin: u64, end: u64) !f64 { - if (begin == 0 or end == 0 or begin == counters.error_value or end == counters.error_value) return error.InvalidCounterSample; + if (begin == 0 or end == 0 or begin == counters.error_value or end == counters.error_value) + return error.InvalidCounterSample; if (end < begin) return error.ReversedCounterSample; if (begin < self.before.gpu or end > self.after.gpu) return error.CounterOutsideCalibration; const ticks: f64 = @floatFromInt(end - begin); @@ -79,7 +82,7 @@ pub const Recorder = struct { calibration: ?Calibration = null, totals: [class_count]Aggregate = @splat(.{}), - /// Init-only bounded storage; disabled init allocates nothing. + /// Init-only storage, descriptors and fence; 1..65536 spans, at most 16 CPU-resolve bytes per span. pub fn init(allocator: std.mem.Allocator, device: mtl.Device, options: Options) !Recorder { var self = Recorder{ .allocator = allocator }; if (!options.enabled) return self; @@ -120,8 +123,8 @@ pub const Recorder = struct { self.state = .recording; } - /// Non-nested spans; stage mode rotates fenced encoders, requiring resource and pipeline rebinding. - pub fn beginSpan(self: *Recorder, encoder: *mtl.ComputeEncoder, kind: Class) !void { + /// No nesting; stage sampling rotates fenced serial encoders: callers must rebind and redeclare resources. + pub fn beginSpan(self: *Recorder, encoder: *mtl.ComputeEncoder, kind: Class, metadata: u32) !void { if (!self.enabled()) return; if (self.state != .recording or !self.encoder_open) return error.CounterNotRecording; if (self.active_span != null) return error.NestedCounterSpan; @@ -141,7 +144,7 @@ pub const Recorder = struct { self.encoder_open = true; resources.fence.waitCompute(encoder.*); } - self.spans[index] = .{ .kind = kind }; + self.spans[index] = .{ .kind = kind, .metadata = metadata }; self.span_count += 1; self.active_span = index; } @@ -172,7 +175,7 @@ pub const Recorder = struct { self.active_span = null; } - /// Ends the final encoder before the caller's existing commit/wait; resolution requires completion. + /// End the final encoder before existing commit/wait; resolve counters only after command completion. pub fn finishEncoding(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { if (!self.enabled()) { encoder.end(); @@ -186,7 +189,7 @@ pub const Recorder = struct { self.state = .resolved; } - /// Cancels unsubmitted work, closing any surviving encoder, including after encoder creation failure. + /// Close and discard failed, unsubmitted encoding without committing; safe after encoder creation failure. pub fn cancelEncoding(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { if (self.command_buffer) |cb| { if (cb.status() != .not_enqueued) return error.CounterCommandAlreadySubmitted; @@ -198,7 +201,7 @@ pub const Recorder = struct { try self.discardAfterCompletion(); } - /// No wait or fallback; results are borrowed until beginCommand, discardAfterCompletion or deinit. + /// No wait or fallback; results are borrowed until reuse/discard/deinit, errors and invalid calibration propagate. pub fn readAfterCompletion(self: *Recorder) ![]const Span { if (!self.enabled()) return &.{}; if (self.state == .read) return self.spans[0..self.span_count]; @@ -244,7 +247,7 @@ pub const Recorder = struct { return self.totals; } - /// Discards closed, unsubmitted or completed recordings; in-flight resources are never reused. + /// Discard completed/unsubmitted recordings after closing encoders; in-flight reuse is rejected, never waited. pub fn discardAfterCompletion(self: *Recorder) !void { if (self.encoder_open) return error.CounterEncoderStillOpen; if (self.command_buffer) |cb| { @@ -288,14 +291,14 @@ test "disabled recorder allocates nothing and never touches a device or encoder var recorder = try Recorder.init(std.testing.failing_allocator, undefined, .{ .enabled = false }); defer recorder.deinit() catch unreachable; var encoder: mtl.ComputeEncoder = undefined; - try recorder.beginSpan(&encoder, .attention); + try recorder.beginSpan(&encoder, .attention, 0); try recorder.endSpan(&encoder); try std.testing.expect(!recorder.enabled()); try std.testing.expectEqual(@as(usize, 0), (try recorder.readAfterCompletion()).len); for (recorder.aggregates()) |aggregate| try std.testing.expectEqual(@as(u64, 0), aggregate.calls); } -test "calibration rejects missing samples, both reversed clocks and both range boundaries" { +test "calibration rejects missing samples and both clock and range boundaries" { const c = Calibration{ .before = .{ .cpu = 1000, .gpu = 100 }, .after = .{ .cpu = 5000, .gpu = 300 } }; try std.testing.expectError(error.InvalidCounterSample, c.duration(0, 200)); try std.testing.expectError(error.InvalidCounterSample, c.duration(100, counters.error_value)); @@ -308,12 +311,45 @@ test "calibration rejects missing samples, both reversed clocks and both range b } test "enabled capacity validation occurs before any device access" { - try std.testing.expectError(error.InvalidSpanCapacity, Recorder.init(std.testing.failing_allocator, undefined, .{ - .enabled = true, - .max_spans = 0, - })); - try std.testing.expectError(error.InvalidSpanCapacity, Recorder.init(std.testing.failing_allocator, undefined, .{ - .enabled = true, - .max_spans = max_span_capacity + 1, - })); + try std.testing.expectError(error.InvalidSpanCapacity, Recorder.init(std.testing.failing_allocator, undefined, .{ .enabled = true, .max_spans = 0 })); + try std.testing.expectError(error.InvalidSpanCapacity, Recorder.init(std.testing.failing_allocator, undefined, .{ .enabled = true, .max_spans = max_span_capacity + 1 })); +} + +test "counter ownership rejects open and in-flight reuse without a CPU span wait" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + if (counters.query(device).mode() == .unsupported) return error.SkipZigTest; + const queue = try device.queue(); + defer queue.deinit(); + const gate = try device.sharedEvent(); + defer gate.deinit(); + var recorder = try Recorder.init(std.testing.allocator, device, .{ .enabled = true, .max_spans = 1 }); + defer recorder.deinit() catch unreachable; + const cb = queue.commandBuffer(); + cb.waitFor(gate, 1); + var encoder = cb.compute(.serial); + try recorder.beginCommand(cb); + var submitted = false; + defer { + if (submitted) { + gate.set(1); + cb.wait(); + recorder.discardAfterCompletion() catch unreachable; + } else recorder.cancelEncoding(&encoder) catch unreachable; + } + try std.testing.expectError(error.CounterEncoderStillOpen, recorder.deinit()); + try recorder.beginSpan(&encoder, .dense, 42); + try std.testing.expectError(error.NestedCounterSpan, recorder.beginSpan(&encoder, .copy, 43)); + try std.testing.expectError(error.UnclosedCounterSpan, recorder.finishEncoding(&encoder)); + try recorder.endSpan(&encoder); + try std.testing.expectError(error.CounterSpanCapacityExceeded, recorder.beginSpan(&encoder, .copy, 43)); + try recorder.finishEncoding(&encoder); + cb.commit(); + submitted = true; + try std.testing.expectError(error.CounterCommandNotCompleted, recorder.readAfterCompletion()); + try std.testing.expectError(error.CounterRecordingInFlight, recorder.beginCommand(cb)); + try std.testing.expectError(error.CounterCommandNotCompleted, recorder.discardAfterCompletion()); + try std.testing.expectError(error.CounterCommandNotCompleted, recorder.deinit()); } diff --git a/zig/src/families/deepseek_v41/attention.zig b/zig/src/families/deepseek_v41/attention.zig new file mode 100644 index 0000000000..6d7cfe5da7 --- /dev/null +++ b/zig/src/families/deepseek_v41/attention.zig @@ -0,0 +1,258 @@ +//! Pinned DeepSeek V4.1 attention, resident EXL3 projections and causal packed-cache chunk execution. +const mtl = @import("metal"); +const exl3 = @import("../../core/exl3_projection.zig"); +const checkpoint = @import("../../core/checkpoint_metal.zig"); +const weights = @import("weights.zig"); +const config = @import("config.zig"); +const ops = @import("ops.zig"); +const csa2 = @import("csa2.zig"); +const kv = @import("kv.zig"); +const index = @import("csa2_index.zig"); +const rope = @import("rope.zig"); +pub const View = @import("view.zig").View; +pub const State = kv.cache.State; + +/// Opt-in bounded intermediate capture; ordinary execution does not install an observer. +pub const DiagnosticStage = enum { + qr_projection, + qr_norm, + q_projection, + q, + kv_projection, + kv_norm, + kv_rope, + kv_quantized, + sparse_o, + o, + grouped, + wo_b, +}; +pub const DiagnosticObserver = struct { + ctx: *anyopaque, + capture: *const fn (*anyopaque, mtl.ComputeEncoder, DiagnosticStage, View) anyerror!void, + + fn observe(comptime enabled: bool, self: ?DiagnosticObserver, e: mtl.ComputeEncoder, stage: DiagnosticStage, source: View) !void { + if (enabled) if (self) |observer| try observer.capture(observer.ctx, e, stage, source); + } +}; + +pub const Pipelines = struct { + plain: ops.Pipelines, + cache: kv.Pipelines, + select: index.Pipelines, + rotary: rope.Pipelines, + sparse: csa2.Pipelines, + + pub fn load(device: mtl.Device, capacity: u32, c: config.Text) !Pipelines { + var plain = try ops.Pipelines.load(device); + errdefer plain.deinit(); + var cache = try kv.Pipelines.load(device); + errdefer cache.deinit(); + var select = try index.Pipelines.load(device); + errdefer select.deinit(); + var rotary = try rope.Pipelines.load(device, capacity, c.qk_rope_head_dim, .{ .theta = c.rope_theta, .original = 0 }, .{ .theta = c.compress_rope_theta, .factor = c.rope_scaling.factor, .original = c.rope_scaling.original_max_position_embeddings, .beta_fast = c.rope_scaling.beta_fast, .beta_slow = c.rope_scaling.beta_slow }); + errdefer rotary.deinit(); + const sparse = try csa2.Pipelines.load(device); + return .{ .plain = plain, .cache = cache, .select = select, .rotary = rotary, .sparse = sparse }; + } + pub fn residentBytes(capacity: u32, c: config.Text) usize { + return rope.Pipelines.residentBytes(capacity, c.qk_rope_head_dim); + } + pub fn registerResidency(p: *const Pipelines, set: mtl.ResidencySet) void { + p.rotary.registerResidency(set); + } + pub fn deinit(p: *Pipelines) void { + p.plain.deinit(); + p.cache.deinit(); + p.select.deinit(); + p.rotary.deinit(); + p.sparse.deinit(); + } +}; + +// Row-sized regions never contain dequantized weights or a query-by-context tensor. +const regions = [_]usize{ 1280 * 2, 32768 * 2, 512 * 2, 512 * 4, 512 * 4, 512 * 2, 128 * 2, 4096 * 2, 32 * 2, 8192 * 2 }; +const Region = enum(usize) { qr, query, raw_kv, values, gates, latent, index_k, index_q, index_weights, grouped }; + +pub const Workspace = struct { + buffer: mtl.Buffer, + selection: index.Scratch, + max_rows: u32, + selection_layer: ?u32 = null, + selection_position: u32 = 0, + selection_rows: u32 = 0, + candidate_position: u32 = 0, + candidate_rows: u32 = 0, + + pub fn bytes(max_rows: u32) usize { + var per_row: usize = 0; + for (regions) |size| per_row += size; + return @as(usize, max_rows) * per_row + index.Scratch.bytes(max_rows); + } + pub fn init(device: mtl.Device, max_rows: u32) !Workspace { + if (max_rows == 0 or max_rows > 4096) return error.InvalidAttentionWorkspace; + const buffer = try device.buffer(bytes(max_rows) - index.Scratch.bytes(max_rows), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer buffer.deinit(); + const selection = try index.Scratch.init(device, max_rows); + return .{ .buffer = buffer, .selection = selection, .max_rows = max_rows }; + } + pub fn deinit(w: *Workspace) void { + w.buffer.deinit(); + w.selection.deinit(); + } + pub fn residentBytes(w: *const Workspace) usize { + return bytes(w.max_rows); + } + pub fn registerResidency(w: *const Workspace, set: mtl.ResidencySet) void { + set.add(w.buffer); + set.add(w.selection.buffer); + } + fn view(w: *const Workspace, region: Region) View { + var offset: usize = 0; + for (regions[0..@backingInt(region)]) |size| offset += size * w.max_rows; + return .{ .buffer = w.buffer, .offset = offset }; + } +}; + +fn tensor(t: checkpoint.Tensor) View { + return .{ .buffer = t.buffer, .offset = t.offset }; +} +fn core(v: View) exl3.View { + return .{ .buffer = v.buffer, .offset = v.offset }; +} +fn projection(p: *const exl3.Pipelines, e: mtl.ComputeEncoder, m: exl3.Matrix, x: View, y: View, scratch: exl3.View, rows: u32, xs: u32, ys: u32, output: exl3.Output) !void { + try exl3.encode(p, e, m, core(x), core(y), scratch, .{ .count = rows, .x_stride = xs, .y_stride = ys, .output = output }); +} +fn norm(p: *const Pipelines, e: mtl.ComputeEncoder, x: View, weight: checkpoint.Tensor, rows: u32, width: u32, eps: f32) void { + ops.encodeRms(e, &p.plain, x, tensor(weight), x, .{ .rows = rows, .width = width, .x_stride = width, .y_stride = width, .eps = eps, .weight_f32 = weight.dtype == .f32 }); + e.barrier(); +} +fn rotate(p: *const Pipelines, e: mtl.ComputeEncoder, x: View, rows: u32, heads: u32, width: u32, position: u32, ratio: u32, step: u32, inverse: bool, c: config.Text) !void { + try p.rotary.encode(e, x, .{ .rows = rows, .heads = heads, .width = width, .rotary = c.qk_rope_head_dim, .stride = heads * width, .position = position, .step = step, .inverse = @intFromBool(inverse), .spectrum = if (ratio == 0) .normal else .compressed }); +} +fn precedingIndex(layer: u32) u32 { + var source: u32 = 2; + for (kv.cache.index_sources) |id| { + if (id > layer) break; + source = id; + } + return source; +} + +/// BF16 x/y [rows,5120]; layers 0..39 share the chunk-start cursor, appended by the caller after all layers. + +// No extra RMS after wq_b or index-query Hadamard outside EXL3. +pub fn encode(p: *const Pipelines, e: mtl.ComputeEncoder, w: *const weights.Attention, state: *State, work: *Workspace, x: View, y: View, ep: *const exl3.Pipelines, scratch: exl3.View, layer: u32, rows: u32, c: config.Text) !void { + return encodeDiagnostic(p, e, w, state, work, x, y, ep, scratch, layer, rows, c, false, null); +} + +/// Same encoder and arithmetic as encode, with optional copies at operator boundaries. +pub fn encodeDiagnostic(p: *const Pipelines, e: mtl.ComputeEncoder, w: *const weights.Attention, state: *State, work: *Workspace, x: View, y: View, ep: *const exl3.Pipelines, scratch: exl3.View, layer: u32, rows: u32, c: config.Text, comptime diagnostic: bool, observer: ?DiagnosticObserver) !void { + if (layer >= 40 or rows == 0 or rows > work.max_rows or + state.position > state.plan.capacity or rows > state.plan.capacity - state.position) + return error.InvalidAttentionChunk; + const position = state.position; + const ratio = c.compress_ratios[layer]; + const end = position + rows; + const q_width = c.num_attention_heads * c.head_dim; + const qr = work.view(.qr); + const q = work.view(.query); + const raw = work.view(.raw_kv); + const ring = try state.ring(layer); + const chunk = work.selection.windowChunk(); + e.barrier(); + try projection(ep, e, w.wq_a, x, qr, scratch, rows, c.hidden_size, c.q_lora_rank, .bf16); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .qr_projection, qr); + norm(p, e, qr, w.q_norm, rows, c.q_lora_rank, c.rms_norm_eps); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .qr_norm, qr); + try projection(ep, e, w.wq_b, qr, q, scratch, rows, c.q_lora_rank, q_width, .bf16); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .q_projection, q); + try rotate(p, e, q, rows, c.num_attention_heads, c.head_dim, position, ratio, 1, false, c); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .q, q); + try projection(ep, e, w.wkv, x, raw, scratch, rows, c.hidden_size, c.head_dim, .bf16); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_projection, raw); + norm(p, e, raw, w.kv_norm, rows, c.head_dim, c.rms_norm_eps); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_norm, raw); + try rotate(p, e, raw, rows, 1, c.head_dim, position, ratio, 1, false, c); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_rope, raw); + try p.cache.encodePack(e, raw, chunk, .{ .rows = rows, .width = c.head_dim, .stride = c.head_dim, .packed_stride = kv.cache.ring_stride, .block = 32, .kind = .window_fp8 }); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_quantized, chunk); + + var main = chunk; + var compressed_rows: u32 = 0; + if (ratio != 0) { + const source = try state.source(layer, ratio); + main = source.main; + compressed_rows = end / ratio; + const complete = compressed_rows - position / ratio; + const latent = work.view(.latent); + if (source.owner == layer) { + const compressor = if (w.compressor) |*value| value else return error.MissingAttentionCompressor; + const output: exl3.Output = if (ratio == 1) .bf16 else .f32; + const values = work.view(.values); + const gates = work.view(.gates); + try projection(ep, e, compressor.wkv, x, values, scratch, rows, c.hidden_size, c.head_dim, output); + if (ratio > 1) { + const gate = compressor.wgate orelse return error.MissingAttentionCompressorGate; + try projection(ep, e, gate, x, gates, scratch, rows, c.hidden_size, c.head_dim, .f32); + } + try p.cache.encodePool(e, values, gates, source.partial, tensor(compressor.norm), latent, .{ .rows = rows, .ratio = ratio, .position = position, .stride = c.head_dim, .input_f32 = @intFromBool(ratio > 1), .eps = c.rms_norm_eps }); + if (complete != 0) { + const selector = if (w.indexer) |*value| value else return error.MissingAttentionIndexer; + const wk = selector.wk orelse return error.MissingAttentionIndexKeyProjection; + const k_norm = selector.k_norm orelse return error.MissingAttentionIndexKeyNorm; + const key = work.view(.index_k); + try projection(ep, e, wk, latent, key, scratch, complete, c.head_dim, c.index_head_dim, .bf16); + norm(p, e, key, k_norm, complete, c.index_head_dim, c.rms_norm_eps); + try rotate(p, e, key, complete, 1, c.index_head_dim, position / ratio * ratio, ratio, ratio, false, c); + try p.cache.encodePack(e, key, source.index, .{ .rows = complete, .width = c.index_head_dim, .stride = c.index_head_dim, .packed_stride = kv.cache.index_stride, .block = 32, .kind = .index_fp4, .first = position / ratio }); + } + } + if (kv.cache.isIndexSource(layer)) { + const selector = if (w.indexer) |*value| value else return error.MissingAttentionIndexer; + const iq = work.view(.index_q); + const iw = work.view(.index_weights); + const iq_width = c.index_n_heads * c.index_head_dim; + if (compressed_rows != 0) { + try projection(ep, e, selector.wq_b, qr, iq, scratch, rows, c.q_lora_rank, iq_width, .bf16); + try rotate(p, e, iq, rows, c.index_n_heads, c.index_head_dim, position, ratio, 1, false, c); + try p.cache.encodePack(e, iq, work.selection.queryPacked(), .{ .rows = rows * c.index_n_heads, .width = c.index_head_dim, .stride = c.index_head_dim, .packed_stride = kv.cache.index_stride, .block = 32, .kind = .index_fp4, .inplace = 1 }); + ops.encodeDense(e, &p.plain, x, tensor(selector.weights_proj), iw, .{ .rows = rows, .k = c.hidden_size, .n = c.index_n_heads, .x_stride = c.hidden_size, .w_stride = c.hidden_size, .y_stride = c.index_n_heads, .w_dtype = .f16 }); + e.barrier(); + } + const make_candidates = layer == c.candidate_source_layer_id; + const use_candidates = layer > c.candidate_source_layer_id; + if (use_candidates and (work.candidate_position != position or work.candidate_rows != rows)) + return error.MissingAttentionCandidates; + try p.select.encode(e, iq, iw, source.index, work.selection, .{ .rows = rows, .q_stride = iq_width, .weight_stride = c.index_n_heads, .position = position, .ratio = ratio, .keys = compressed_rows, .make_candidates = @intFromBool(make_candidates), .use_candidates = @intFromBool(use_candidates) }); + work.selection_layer = layer; + work.selection_position = position; + work.selection_rows = rows; + if (make_candidates) { + work.candidate_position = position; + work.candidate_rows = rows; + } + } else if (work.selection_layer != precedingIndex(layer) or + work.selection_position != position or work.selection_rows != rows) return error.MissingAttentionSelection; + // Index K consumes the normalized latent before this in-place RoPE mutates it. + if (source.owner == layer and complete != 0) { + try rotate(p, e, latent, complete, 1, c.head_dim, position / ratio * ratio, ratio, ratio, false, c); + try p.cache.encodePack(e, latent, source.main, .{ .rows = complete, .width = c.head_dim, .stride = c.head_dim, .packed_stride = kv.cache.main_stride, .block = 16, .kind = .main_fp4, .first = position / ratio }); + } + } + try p.sparse.encode(e, q, ring, chunk, main, work.selection.positions(), tensor(w.sink), q, .{ .rows = rows, .q_stride = q_width, .position = position, .chunk_rows = rows, .ratio = ratio, .compressed_rows = compressed_rows, .topk = @min(c.index_topk, compressed_rows) }); + e.barrier(); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .sparse_o, q); + try p.cache.publishWindow(e, chunk, ring, rows, position); + try rotate(p, e, q, rows, c.num_attention_heads, c.head_dim, position, ratio, 1, true, c); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .o, q); + const grouped = work.view(.grouped); + const group_input = q_width / c.o_groups; + const group_output = c.o_groups * c.o_lora_rank; + for (w.wo_a, 0..) |matrix, group| { + try projection(ep, e, matrix, q.shifted(group * group_input * 2), grouped.shifted(group * c.o_lora_rank * 2), scratch, rows, q_width, group_output, .bf16); + } + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .grouped, grouped); + try projection(ep, e, w.wo_b, grouped, y, scratch, rows, group_output, c.hidden_size, .bf16); + if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .wo_b, y); +} diff --git a/zig/src/families/deepseek_v41/csa2.zig b/zig/src/families/deepseek_v41/csa2.zig new file mode 100644 index 0000000000..395de9107a --- /dev/null +++ b/zig/src/families/deepseek_v41/csa2.zig @@ -0,0 +1,77 @@ +//! Exact causal SWA plus shared compressed KV; fixed-buffer chunked prefill and decode use one path. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const View = @import("view.zig").View; +pub const index = @import("csa2_index.zig"); +pub const kv = @import("kv.zig"); +pub const rope = @import("rope.zig"); +pub const Params = extern struct { + rows: u32, + heads: u32 = 64, + width: u32 = 512, + q_stride: u32, + position: u32, + chunk_rows: u32, + window: u32 = 128, + ring_stride: u32 = 528, + chunk_stride: u32 = 528, + main_stride: u32 = 288, + main_block: u32 = 16, + ratio: u32, + compressed_rows: u32, + topk: u32 = 512, + index_stride: u32 = 512, + scale: f32 = 0.04419417382415922, + pub fn validate(p: Params) !void { + if (p.rows == 0 or p.heads == 0 or p.width == 0 or p.width > 512 or p.width % 32 != 0 or + @as(u64, p.q_stride) < @as(u64, p.heads) * p.width or + p.chunk_rows < p.rows or p.window != 128 or + p.ring_stride < p.width + p.width / 32 or p.chunk_stride < p.width + p.width / 32 or + p.main_stride < p.width / 2 + p.width / 16 or p.main_block != 16 or + (p.ratio != 0 and p.ratio != 1 and p.ratio != 2) or p.topk > 512 or p.index_stride < p.topk or + @as(u64, p.position) + p.rows > kv.cache.max_context or + p.compressed_rows > kv.cache.max_context or !std.math.isFinite(p.scale) or p.scale <= 0 or + @as(u64, p.rows -| 1) * p.q_stride + @as(u64, p.heads) * p.width > std.math.maxInt(u32) or + @as(u64, p.rows) * p.index_stride > std.math.maxInt(u32) or + @as(u64, p.rows) * p.chunk_stride > std.math.maxInt(u32) or + @as(u64, p.window) * p.ring_stride > std.math.maxInt(u32) or + @as(u64, p.compressed_rows) * p.main_stride > std.math.maxInt(u32)) + return error.InvalidSparseAttentionShape; + } +}; +pub const Pipelines = struct { + attention: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_kv ++ sources.deepseek_v41_csa2, options); + defer lib.deinit(); + return .{ .attention = try mtl.Pipeline.init(device, lib, "dsv41_sparse_attention", false) }; + } + pub fn deinit(p: *Pipelines) void { + p.attention.deinit(); + } + /// BF16 Q/output [rows,heads,width], FP32 sinks, I32 indices; publish this chunk to the ring only after dispatch. + pub fn encode(p: *const Pipelines, enc: mtl.ComputeEncoder, q: View, ring: View, chunk: View, compressed: View, positions: View, sink: View, out: View, args: Params) !void { + try args.validate(); + try q.requireRows(args.rows, args.q_stride, @as(usize, args.heads) * args.width, 2); + try out.requireRows(args.rows, args.q_stride, @as(usize, args.heads) * args.width, 2); + try sink.require(@as(usize, args.heads) * 4, 4); + if (args.position != 0) try ring.requireRows(@min(args.position, args.window), args.ring_stride, args.width + args.width / 32, 1); + try chunk.requireRows(args.rows, args.chunk_stride, args.width + args.width / 32, 1); + if (args.ratio != 0 and args.compressed_rows != 0 and args.topk != 0) { + try compressed.requireRows(args.compressed_rows, args.main_stride, args.width / 2 + args.width / 16, 1); + try positions.requireRows(args.rows, args.index_stride, args.topk, 4); + } + enc.setPipeline(p.attention); + for ([_]View{ q, ring, chunk, compressed, positions, sink, out }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 7); + enc.dispatchGroups(mtl.Size.of(args.heads, args.rows, 1), mtl.Size.of(32, 1, 1)); + } +}; + +test { + _ = @import("kv_test.zig"); +} diff --git a/zig/src/families/deepseek_v41/csa2_index.zig b/zig/src/families/deepseek_v41/csa2_index.zig new file mode 100644 index 0000000000..ed2ee3a833 --- /dev/null +++ b/zig/src/families/deepseek_v41/csa2_index.zig @@ -0,0 +1,139 @@ +//! Tiled GPU scores and bounded online heaps; the memory footprint does not grow with context length. +const std = @import("std"); +const max_context = @import("metal_kv.zig").max_context; +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const View = @import("view.zig").View; +pub const score_tile_width = 4096; +pub const Params = extern struct { + rows: u32, + heads: u32 = 32, + width: u32 = 128, + q_stride: u32, + weight_stride: u32, + position: u32, + ratio: u32, + keys: u32, + key_stride: u32 = 68, + tile_start: u32 = 0, + tile_width: u32 = score_tile_width, + topk: u32 = 512, + candidate_topk: u32 = 2048, + block_size: u32 = 8, + use_candidates: u32 = 0, + make_candidates: u32 = 0, + scale: f32 = 0.015625, + pub fn validate(p: Params) !void { + if (p.rows == 0 or p.heads == 0 or p.heads > 32 or p.width == 0 or p.width > 128 or p.width % 32 != 0 or + p.q_stride < p.heads * p.width or p.weight_stride < p.heads or + (p.ratio != 1 and p.ratio != 2) or p.key_stride < p.width / 2 + p.width / 32 or + p.topk != 512 or p.candidate_topk != 2048 or p.block_size != 8 or + p.tile_width != score_tile_width or p.use_candidates > 1 or p.make_candidates > 1 or + (p.use_candidates != 0 and p.make_candidates != 0) or + @as(u64, p.position) + p.rows > max_context or p.keys > max_context or + !std.math.isFinite(p.scale) or p.scale <= 0 or + @as(u64, p.rows) * p.q_stride > std.math.maxInt(u32) or + @as(u64, p.rows) * p.weight_stride > std.math.maxInt(u32) or + @as(u64, p.keys) * p.key_stride > std.math.maxInt(u32)) + return error.InvalidIndexerShape; + } +}; +pub const Scratch = struct { + buffer: mtl.Buffer, + rows: u32, + pub fn bytes(rows: u32) usize { + return @as(usize, rows) * (score_tile_width * 4 + 512 * 8 + 2048 * 8 + 512 * 4 + 528 + 32 * 68); + } + pub fn init(device: mtl.Device, rows: u32) !Scratch { + if (rows == 0 or rows > 4096) return error.InvalidIndexerScratch; + return .{ .buffer = try device.buffer(bytes(rows), mtl.ResourceOptions.shared), .rows = rows }; + } + pub fn deinit(s: *Scratch) void { + s.buffer.deinit(); + } + pub fn scores(s: Scratch) View { + return .{ .buffer = s.buffer }; + } + pub fn heap(s: Scratch) View { + return s.scores().shifted(@as(usize, s.rows) * score_tile_width * 4); + } + pub fn candidates(s: Scratch) View { + return s.heap().shifted(@as(usize, s.rows) * 512 * 8); + } + pub fn positions(s: Scratch) View { + return s.candidates().shifted(@as(usize, s.rows) * 2048 * 8); + } + pub fn windowChunk(s: Scratch) View { + return s.positions().shifted(@as(usize, s.rows) * 512 * 4); + } + pub fn queryPacked(s: Scratch) View { + return s.windowChunk().shifted(@as(usize, s.rows) * 528); + } +}; +pub const Pipelines = struct { + score: mtl.Pipeline, + clear: mtl.Pipeline, + merge: mtl.Pipeline, + finish: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_kv ++ sources.deepseek_v41_index, options); + defer lib.deinit(); + var p: Pipelines = undefined; + p.score = try mtl.Pipeline.init(device, lib, "dsv41_index_score", false); + errdefer p.score.deinit(); + p.clear = try mtl.Pipeline.init(device, lib, "dsv41_index_clear", false); + errdefer p.clear.deinit(); + p.merge = try mtl.Pipeline.init(device, lib, "dsv41_index_merge", false); + errdefer p.merge.deinit(); + p.finish = try mtl.Pipeline.init(device, lib, "dsv41_index_finish", false); + return p; + } + pub fn deinit(p: *Pipelines) void { + p.score.deinit(); + p.clear.deinit(); + p.merge.deinit(); + p.finish.deinit(); + } + /// Q has already undergone BF16 RoPE and in-place FP4/E8M0 quantization; weights are BF16 projection output. + pub fn encode(p: *const Pipelines, enc: mtl.ComputeEncoder, q: View, weights: View, keys: View, scratch: Scratch, args: Params) !void { + try args.validate(); + if (args.rows > scratch.rows) return error.IndexerScratchExceeded; + try (View{ .buffer = scratch.buffer }).require(Scratch.bytes(scratch.rows), 4); + if (args.keys != 0) { + try q.requireRows(args.rows, args.q_stride, @as(usize, args.heads) * args.width, 2); + try weights.requireRows(args.rows, args.weight_stride, args.heads, 2); + try keys.requireRows(args.keys, args.key_stride, args.width / 2 + args.width / 32, 1); + } + enc.setPipeline(p.clear); + for ([_]View{ scratch.heap(), scratch.candidates() }, 0..) |v, i| enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 2); + enc.dispatchThreads(mtl.Size.of(@max(args.topk, args.candidate_topk), args.rows, 1), mtl.Size.of(128, 1, 1)); + enc.barrier(); + const count = if (args.use_candidates != 0) args.candidate_topk * args.block_size else args.keys; + var at: u32 = 0; + while (at < count) : (at += score_tile_width) { + var tile = args; + tile.tile_start = at; + enc.setPipeline(p.score); + for ([_]View{ q, weights, keys, scratch.candidates(), scratch.scores() }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(tile, 5); + enc.dispatchGroups(mtl.Size.of(score_tile_width, args.rows, 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + enc.setPipeline(p.merge); + for ([_]View{ scratch.scores(), scratch.heap(), scratch.candidates() }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(tile, 3); + enc.dispatchThreads(mtl.Size.of(args.rows, 1, 1), mtl.Size.of(1, 1, 1)); + enc.barrier(); + } + enc.setPipeline(p.finish); + for ([_]View{ scratch.heap(), scratch.candidates(), scratch.positions() }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 3); + enc.dispatchThreads(mtl.Size.of(args.rows, 1, 1), mtl.Size.of(1, 1, 1)); + enc.barrier(); + } +}; diff --git a/zig/src/families/deepseek_v41/csa2_test.zig b/zig/src/families/deepseek_v41/csa2_test.zig new file mode 100644 index 0000000000..0ac81cacca --- /dev/null +++ b/zig/src/families/deepseek_v41/csa2_test.zig @@ -0,0 +1,309 @@ +const std = @import("std"); +const mtl = @import("metal"); +const csa2 = @import("csa2.zig"); +const index = @import("csa2_index.zig"); +const rope = @import("rope.zig"); +const View = @import("view.zig").View; +fn bf(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @intCast((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); +} +fn view(b: mtl.Buffer) View { + return .{ .buffer = b }; +} +fn submit(cb: mtl.CommandBuffer, enc: mtl.ComputeEncoder) !void { + enc.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); +} + +test "sparse attention preserves a causal window across ring wrap and masks future compressed keys" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try csa2.Pipelines.load(device); + defer p.deinit(); + const q = try device.buffer(3 * 512 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const ring = try device.buffer(128 * 528, mtl.ResourceOptions.shared); + defer ring.deinit(); + const chunk = try device.buffer(3 * 528, mtl.ResourceOptions.shared); + defer chunk.deinit(); + const main = try device.buffer(201 * 288, mtl.ResourceOptions.shared); + defer main.deinit(); + const positions = try device.buffer(6 * 4, mtl.ResourceOptions.shared); + defer positions.deinit(); + const sink = try device.buffer(4, mtl.ResourceOptions.shared); + defer sink.deinit(); + const out = try device.buffer(3 * 512 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memset(q.contents()[0 .. 3 * 512 * 2], 0); + sink.slice(f32, 1)[0] = 0; + for (0..128) |r| { + @memset(ring.contents()[r * 528 .. r * 528 + 512], 0x38); + @memset(ring.contents()[r * 528 + 512 .. (r + 1) * 528], 127); + } + for (0..3) |r| { + @memset(chunk.contents()[r * 528 .. r * 528 + 512], 0x40); + @memset(chunk.contents()[r * 528 + 512 .. (r + 1) * 528], 127); + positions.slice(i32, 6)[2 * r] = 0; + positions.slice(i32, 6)[2 * r + 1] = 200; + } + @memset(main.contents()[0..256], 0x66); + @memset(main.contents()[256..288], 0x38); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(q), view(ring), view(chunk), view(main), view(positions), view(sink), view(out), .{ .rows = 3, .heads = 1, .q_stride = 512, .position = 129, .chunk_rows = 3, .ratio = 2, .compressed_rows = 201, .topk = 2, .index_stride = 2 }); + try submit(cb, enc); + for (0..3) |r| { + const expected = bf((133 + @as(f32, @floatFromInt(r))) / 130); + for (out.slice(u16, 3 * 512)[r * 512 .. (r + 1) * 512]) |v| try std.testing.expectEqual(expected, v); + } +} + +test "sparse QK cancellation retains a small score before sink normalization" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try csa2.Pipelines.load(device); + defer p.deinit(); + const q = try device.buffer(512 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const chunk = try device.buffer(528, mtl.ResourceOptions.shared); + defer chunk.deinit(); + const sink = try device.buffer(4, mtl.ResourceOptions.shared); + defer sink.deinit(); + const out = try device.buffer(512 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memset(q.slice(u16, 512), 0); + q.slice(u16, 512)[0] = bf(0x1p24); + q.slice(u16, 512)[1] = bf(-0x1p24); + q.slice(u16, 512)[32] = bf(0x1p-6); + @memset(chunk.contents()[0..512], 0x38); + @memset(chunk.contents()[512..528], 127); + sink.slice(f32, 1)[0] = 0; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + // No prior window or compressed keys: their unused bindings share the immutable chunk. + try p.encode(enc, view(q), view(chunk), view(chunk), view(chunk), view(chunk), view(sink), view(out), .{ .rows = 1, .heads = 1, .q_stride = 512, .position = 0, .chunk_rows = 1, .ratio = 0, .compressed_rows = 0, .scale = 1 }); + try submit(cb, enc); + // Exact dot = 2^-6; 1 / (1 + exp(-2^-6)) rounds to this BF16 value. + for (out.slice(u16, 512)) |v| try std.testing.expectEqual(bf(0.50390625), v); +} + +test "the first query retains BF16 probability arithmetic regardless of future prefill rows" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try csa2.Pipelines.load(device); + defer p.deinit(); + const q = try device.buffer(32 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const ring = try device.buffer(128 * 33, mtl.ResourceOptions.shared); + defer ring.deinit(); + const chunk = try device.buffer(128 * 33, mtl.ResourceOptions.shared); + defer chunk.deinit(); + const main = try device.buffer(18, mtl.ResourceOptions.shared); + defer main.deinit(); + const positions = try device.buffer(4, mtl.ResourceOptions.shared); + defer positions.deinit(); + const sink = try device.buffer(4, mtl.ResourceOptions.shared); + defer sink.deinit(); + const out = try device.buffer(32 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memset(q.slice(u16, 32), 0); + q.slice(u16, 32)[0] = bf(1); + @memset(chunk.contents()[0..32], 0x38); + chunk.contents()[32] = 127; + @memset(main.contents()[0..16], 0x44); + @memset(main.contents()[16..18], 0x38); + positions.slice(i32, 1)[0] = 0; + sink.slice(f32, 1)[0] = 0.005; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(q), view(ring), view(chunk), view(main), view(positions), view(sink), view(out), .{ .rows = 1, .heads = 1, .width = 32, .q_stride = 32, .position = 0, .chunk_rows = 1, .ring_stride = 33, .chunk_stride = 33, .main_stride = 18, .ratio = 1, .compressed_rows = 1, .topk = 1, .index_stride = 1, .scale = 1 }); + try submit(cb, enc); + // BF16(exp(-1))*1 + 1*2; padded windows would move key 2 into a later online-softmax tile. + for (out.slice(u16, 32)) |v| try std.testing.expectEqual(bf(1.5703125), v); + const cb2 = queue.commandBuffer(); + const enc2 = cb2.compute(.serial); + try p.encode(enc2, view(q), view(ring), view(chunk), view(main), view(positions), view(sink), view(out), .{ .rows = 1, .heads = 1, .width = 32, .q_stride = 32, .position = 0, .chunk_rows = 128, .ring_stride = 33, .chunk_stride = 33, .main_stride = 18, .ratio = 1, .compressed_rows = 1, .topk = 1, .index_stride = 1, .scale = 1 }); + try submit(cb2, enc2); + for (out.slice(u16, 32)) |v| try std.testing.expectEqual(bf(1.5703125), v); +} + +test "hierarchical candidate selection pins the newest partial block and subsequent reindex uses those blocks" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try index.Pipelines.load(device); + defer p.deinit(); + var scratch = try index.Scratch.init(device, 1); + defer scratch.deinit(); + const q = try device.buffer(32 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const weights = try device.buffer(2, mtl.ResourceOptions.shared); + defer weights.deinit(); + const keys_count = 2048 * 8 + 3; + const keys = try device.buffer(keys_count * 17, mtl.ResourceOptions.shared); + defer keys.deinit(); + @memset(q.slice(u16, 32), bf(1)); + weights.slice(u16, 1)[0] = bf(1); + for (0..keys_count) |r| { + @memset(keys.contents()[r * 17 .. r * 17 + 16], 0x22); + keys.contents()[r * 17 + 16] = if (r >= 2048 * 8) 126 else 127; + } + var args = index.Params{ .rows = 1, .heads = 1, .width = 32, .q_stride = 32, .weight_stride = 1, .position = keys_count - 1, .ratio = 1, .keys = keys_count, .key_stride = 17, .make_candidates = 1, .scale = 1 }; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(q), view(weights), view(keys), scratch, args); + try submit(cb, enc); + const Pick = extern struct { score: f32, position: i32 }; + const candidates = @as([*]Pick, @ptrCast(@alignCast(scratch.candidates().buffer.contents() + scratch.candidates().offset)))[0..2048]; + try std.testing.expectEqual(@as(i32, 2048), candidates[2047].position); + for (candidates[0..2047], 0..) |pick, i| try std.testing.expectEqual(@as(i32, @intCast(i)), pick.position); + const cb2 = queue.commandBuffer(); + const enc2 = cb2.compute(.serial); + args.make_candidates = 0; + args.use_candidates = 1; + try p.encode(enc2, view(q), view(weights), view(keys), scratch, args); + try submit(cb2, enc2); + const selected = @as([*]i32, @ptrCast(@alignCast(scratch.positions().buffer.contents() + scratch.positions().offset)))[0..512]; + for (selected, 0..) |pos, i| try std.testing.expectEqual(@as(i32, @intCast(i)), pos); +} + +test "RoPE rotates adjacent tail pairs with BF16 assignment and compressed positions use group starts" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try rope.Pipelines.load(device, 3, 32, .{ .theta = 10000 }, .{ .theta = 160000 }); + defer p.deinit(); + const x = try device.buffer(2 * 64 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + @memset(x.slice(u16, 128), bf(1)); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try std.testing.expectError(error.InvalidRotaryShape, p.encode(enc, view(x), .{ .rows = 2, .heads = 1, .width = 64, .rotary = 32, .stride = 64, .position = 1, .step = 2 })); + try std.testing.expectError(error.InvalidRotaryShape, p.encode(enc, view(x), .{ .rows = 1, .heads = 1, .width = 64, .rotary = 16, .stride = 64, .position = 0 })); + try std.testing.expectError(error.InvalidRotaryShape, p.encode(enc, view(x), .{ .rows = 2, .heads = 1, .width = 64, .rotary = 32, .stride = 64, .position = std.math.maxInt(u32), .step = std.math.maxInt(u32) })); + try p.encode(enc, view(x), .{ .rows = 2, .heads = 1, .width = 64, .rotary = 32, .stride = 64, .position = 0, .step = 2 }); + try submit(cb, enc); + for (x.slice(u16, 128)[0..96]) |v| try std.testing.expectEqual(bf(1), v); + for (0..16) |pair| { + const power: f32 = @floatCast(std.math.pow(f64, 10000, @as(f32, @floatFromInt(pair * 2)) / 32)); + const frequency: f32 = 1 / power; + const angle: f32 = 2 * frequency; + const c: f32 = @floatCast(@cos(@as(f64, angle))); + const s: f32 = @floatCast(@sin(@as(f64, angle))); + const a = bf(c - s); + const b = bf(s + c); + try std.testing.expectEqual(a, x.slice(u16, 128)[96 + pair * 2]); + try std.testing.expectEqual(b, x.slice(u16, 128)[97 + pair * 2]); + } +} + +test "indexer visibility follows completed compressor groups for every query in a chunk" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try index.Pipelines.load(device); + defer p.deinit(); + var scratch = try index.Scratch.init(device, 3); + defer scratch.deinit(); + const q = try device.buffer(3 * 32 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const weights = try device.buffer(3 * 2, mtl.ResourceOptions.shared); + defer weights.deinit(); + const keys = try device.buffer(3 * 17, mtl.ResourceOptions.shared); + defer keys.deinit(); + @memset(q.slice(u16, 3 * 32), bf(1)); + @memset(weights.slice(u16, 3), bf(1)); + for (0..3) |r| { + @memset(keys.contents()[r * 17 .. r * 17 + 16], 0x22); + keys.contents()[r * 17 + 16] = 127; + } + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(q), view(weights), view(keys), scratch, .{ .rows = 3, .heads = 1, .width = 32, .q_stride = 32, .weight_stride = 1, .position = 0, .ratio = 2, .keys = 3, .key_stride = 17, .scale = 1 }); + try submit(cb, enc); + const selected = @as([*]i32, @ptrCast(@alignCast(scratch.positions().buffer.contents() + + scratch.positions().offset)))[0 .. 3 * 512]; + for (selected[0..512]) |pos| try std.testing.expectEqual(@as(i32, -1), pos); + for (1..3) |r| { + try std.testing.expectEqual(@as(i32, 0), selected[r * 512]); + for (selected[r * 512 + 1 .. (r + 1) * 512]) |pos| try std.testing.expectEqual(@as(i32, -1), pos); + } +} + +test "YaRN interpolates its correction band and inverse rotation conjugates the frequency" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try rope.Pipelines.load(device, 65537, 32, .{ .theta = 10000 }, .{ .theta = 160000, .original = 65536 }); + defer p.deinit(); + const x = try device.buffer(64 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + @memset(x.slice(u16, 64), bf(1)); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(x), .{ .rows = 1, .heads = 1, .width = 64, .rotary = 32, .stride = 64, .position = 65536, .inverse = 1, .spectrum = .compressed }); + try submit(cb, enc); + const low: f32 = @floatCast(@max(@floor(32 * @log(@as(f64, 65536) / (32 * 2 * std.math.pi)) / (2 * @log(@as(f64, 160000)))), 0)); + const high: f32 = @floatCast(@min(@ceil(32 * @log(@as(f64, 65536) / (2 * std.math.pi)) / (2 * @log(@as(f64, 160000)))), 31)); + for ([_]usize{ 0, 10, 15 }) |pair| { + const power: f32 = @floatCast(std.math.pow(f64, 160000, @as(f32, @floatFromInt(pair * 2)) / 32)); + const frequency: f32 = 1 / power; + const smooth: f32 = 1 - std.math.clamp((@as(f32, @floatFromInt(pair)) - low) / @max(high - low, 1e-3), 0, 1); + const slow: f32 = frequency / 16 * (1 - smooth); + const fast: f32 = frequency * smooth; + const angle: f32 = 65536 * (slow + fast); + const c: f32 = @floatCast(@cos(@as(f64, angle))); + const s: f32 = -@as(f32, @floatCast(@sin(@as(f64, angle)))); + try std.testing.expectEqual(bf(c - s), x.slice(u16, 64)[32 + pair * 2]); + try std.testing.expectEqual(bf(s + c), x.slice(u16, 64)[33 + pair * 2]); + } +} + +test "RoPE retained source pair rounds the position-five imaginary consumer to pinned CUDA BF16" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try rope.Pipelines.load(device, 6, 64, .{ .theta = 10000 }, .{ .theta = 160000, .original = 65536 }); + defer p.deinit(); + const x = try device.buffer(512 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + @memset(x.slice(u16, 512), bf(0)); + // Token-5/layer-0 rotary pair 21 crosses a BF16 tie under per-head cos; pinned source factors give .546875. + x.slice(u16, 512)[490] = bf(1.15625); + x.slice(u16, 512)[491] = bf(0.53125); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try p.encode(enc, view(x), .{ .rows = 1, .heads = 1, .width = 512, .rotary = 64, .stride = 512, .position = 5 }); + try submit(cb, enc); + try std.testing.expectEqual(bf(0.546875), x.slice(u16, 512)[491]); +} diff --git a/zig/src/families/deepseek_v41/engram.zig b/zig/src/families/deepseek_v41/engram.zig new file mode 100644 index 0000000000..0f5a0287c2 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram.zig @@ -0,0 +1,203 @@ +//! Resident Engram FP8 gathers and signed-sqrt gates; WKV is projected separately by the caller. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const View = @import("view.zig").View; +pub const hash = @import("engram_hash.zig"); + +pub const head_dim = 256; +pub const gather_dim = hash.columns * head_dim; +pub const history_bytes = 16; + +/// Borrowed full-resident E4M3FN rows and UE8M0 rowwise scale blocks, never a 32x32 scale matrix. +pub const Table = struct { + weight: View, + scale: View, + rows: u32, + weight_stride_bytes: usize = head_dim, + scale_stride_bytes: usize = head_dim / 32, + scale_block: u32 = 32, + + pub fn validate(self: Table) !void { + if (self.rows == 0 or self.weight_stride_bytes < head_dim or self.scale_stride_bytes < head_dim / 32 or self.scale_block != 32) return error.InvalidEngramTable; + try self.weight.requireRows(self.rows, self.weight_stride_bytes, head_dim, 1); + try self.scale.requireRows(self.rows, self.scale_stride_bytes, head_dim / 32, 1); + } +}; + +/// BF16 hidden [tokens,copies,dim], WKV [tokens,(copies+1)*dim]; norm weights contiguous FP32 [copies,dim]. +pub const Gate = struct { + q_weight: View, + k_weight: View, + dim: u32, + copies: u32, + hidden_stride: usize, + kv_stride: usize, + output_stride: usize, + eps: f32, + + pub fn validate(self: Gate) !void { + const width = try std.math.mul(usize, self.dim, self.copies); + const projected_width = try std.math.add(usize, width, self.dim); + if (self.dim == 0 or self.copies == 0 or self.hidden_stride < width or self.output_stride < width or self.kv_stride < projected_width or !std.math.isFinite(self.eps) or self.eps <= 0 or + self.dim > std.math.maxInt(u32) - 255) return error.InvalidEngramGate; + const weight_bytes = try std.math.mul(usize, width, 4); + try self.q_weight.require(weight_bytes, 4); + try self.k_weight.require(weight_bytes, 4); + } +}; + +pub const Pipelines = struct { + hash_ids: mtl.Pipeline, + commit: mtl.Pipeline, + reset: mtl.Pipeline, + gather: mtl.Pipeline, + gate: mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_engram, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const h = try mtl.Pipeline.init(device, lib, "tf_v41_engram_hash", false); + errdefer h.deinit(); + const c = try mtl.Pipeline.init(device, lib, "tf_v41_engram_commit", false); + errdefer c.deinit(); + const r = try mtl.Pipeline.init(device, lib, "tf_v41_engram_reset", false); + errdefer r.deinit(); + const g = try mtl.Pipeline.init(device, lib, "tf_v41_engram_gather", false); + errdefer g.deinit(); + const gate = try mtl.Pipeline.init(device, lib, "tf_v41_engram_gate", false); + return .{ .hash_ids = h, .commit = c, .reset = r, .gather = g, .gate = gate }; + } + + pub fn deinit(self: *Pipelines) void { + self.hash_ids.deinit(); + self.commit.deinit(); + self.reset.deinit(); + self.gather.deinit(); + self.gate.deinit(); + } +}; + +const HashArgs = extern struct { tokens: u32, pad: u32, masked: u32 }; +const GatherArgs = extern struct { + weight_stride: u64, + scale_stride: u64, + output_stride: u64, + tokens: u32, + columns: u32, + head_dim: u32, + scale_block: u32, +}; +const GateArgs = extern struct { + hidden_stride: u64, + kv_stride: u64, + output_stride: u64, + dim: u32, + copies: u32, + masked: u32, + eps: f32, +}; + +fn bind(e: mtl.ComputeEncoder, view: View, at: usize) void { + e.setBuffer(view.buffer, view.offset, at); +} + +/// Each history View is one sequence's 16-byte GPU record; initialize before its first hash. +pub fn encodeReset(e: mtl.ComputeEncoder, p: *const Pipelines, history: View) void { + e.setPipeline(p.reset); + bind(e, history, 0); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 1 }); + e.barrier(); +} + +/// Snapshot/rollback are GPU copies; caller ends compute encoding before beginning a blit encoder. +pub fn encodeSnapshot(e: mtl.BlitEncoder, history: View, saved: View) void { + e.copy(history.buffer, history.offset, saved.buffer, saved.offset, history_bytes); +} + +pub fn encodeRollback(e: mtl.BlitEncoder, saved: View, history: View) void { + encodeSnapshot(e, saved, history); +} + +/// Raw u32 IDs -> [tokens,24] hashes; map covers valid IDs, optional 0=DEAD mask, history is unchanged. +pub fn encodeHash(e: mtl.ComputeEncoder, p: *const Pipelines, meta: hash.Layer, ids: View, token_map: View, mask: ?View, history: View, out: View, tokens: u32) !void { + if (tokens == 0) return error.EmptyEngramChunk; + if (!std.mem.eql(u8, std.mem.asBytes(&meta), std.mem.asBytes(&hash.layers[0])) and + !std.mem.eql(u8, std.mem.asBytes(&meta), std.mem.asBytes(&hash.layers[1]))) return error.InvalidEngramMetadata; + try ids.requireRows(tokens, 1, 1, 4); + try token_map.require(4, 4); + try history.require(history_bytes, 4); + try out.requireRows(tokens, hash.columns, hash.columns, 4); + if (mask) |m| try m.require(tokens, 1); + e.setPipeline(p.hash_ids); + bind(e, ids, 0); + bind(e, token_map, 1); + bind(e, mask orelse ids, 2); + bind(e, history, 3); + bind(e, out, 4); + e.setValue(meta, 5); + e.setValue(HashArgs{ .tokens = tokens, .pad = hash.pad_id, .masked = @intFromBool(mask != null) }, 6); + e.dispatchThreads(.{ .width = tokens }, .{ .width = 64 }); + e.barrier(); +} + +/// Commit exactly once per chunk after both layers have hashed; remains bounded at three past tokens. +pub fn encodeCommit(e: mtl.ComputeEncoder, p: *const Pipelines, ids: View, token_map: View, mask: ?View, history: View, tokens: u32) !void { + if (tokens == 0) return error.EmptyEngramChunk; + try ids.requireRows(tokens, 1, 1, 4); + try token_map.require(4, 4); + try history.require(history_bytes, 4); + if (mask) |m| try m.require(tokens, 1); + e.barrier(); + e.setPipeline(p.commit); + bind(e, ids, 0); + bind(e, token_map, 1); + bind(e, mask orelse ids, 2); + bind(e, history, 3); + e.setValue(HashArgs{ .tokens = tokens, .pad = hash.pad_id, .masked = @intFromBool(mask != null) }, 4); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 1 }); + e.barrier(); +} + +/// Gather uint32 [tokens,24] row IDs into BF16 [tokens,6144]; IDs must be within the borrowed table. +pub fn encodeGather(e: mtl.ComputeEncoder, p: *const Pipelines, table: Table, ids: View, out: View, tokens: u32, output_stride: usize) !void { + try table.validate(); + if (tokens == 0 or output_stride < gather_dim) return error.InvalidEngramGather; + try ids.requireRows(tokens, hash.columns, hash.columns, 4); + try out.requireRows(tokens, output_stride, gather_dim, 2); + e.setPipeline(p.gather); + bind(e, ids, 0); + bind(e, table.weight, 1); + bind(e, table.scale, 2); + bind(e, out, 3); + e.setValue(GatherArgs{ .weight_stride = table.weight_stride_bytes, .scale_stride = table.scale_stride_bytes, .output_stride = output_stride, .tokens = tokens, .columns = hash.columns, .head_dim = head_dim, .scale_block = table.scale_block }, 4); + e.dispatchThreads(.{ .width = gather_dim, .height = tokens }, .{ .width = 256 }); + e.barrier(); +} + +/// Add per-copy gate times shared value to BF16 hidden; out may alias hidden, never the projected KV. +pub fn encodeGate(e: mtl.ComputeEncoder, p: *const Pipelines, gate: Gate, hidden: View, projected_kv: View, mask: ?View, out: View, tokens: u32) !void { + try gate.validate(); + if (tokens == 0) return error.EmptyEngramChunk; + const width = @as(usize, gate.dim) * gate.copies; + try hidden.requireRows(tokens, gate.hidden_stride, width, 2); + try projected_kv.requireRows(tokens, gate.kv_stride, width + gate.dim, 2); + try out.requireRows(tokens, gate.output_stride, width, 2); + if (mask) |m| try m.require(tokens, 1); + e.setPipeline(p.gate); + bind(e, hidden, 0); + bind(e, projected_kv, 1); + bind(e, gate.q_weight, 2); + bind(e, gate.k_weight, 3); + bind(e, mask orelse hidden, 4); + bind(e, out, 5); + e.setValue(GateArgs{ .hidden_stride = gate.hidden_stride, .kv_stride = gate.kv_stride, .output_stride = gate.output_stride, .dim = gate.dim, .copies = gate.copies, .masked = @intFromBool(mask != null), .eps = gate.eps }, 6); + e.dispatchGroups(.{ .width = tokens, .height = gate.copies }, .{ .width = 256 }); + e.barrier(); +} + +test { + _ = @import("engram_test.zig"); +} diff --git a/zig/src/families/deepseek_v41/engram_test.zig b/zig/src/families/deepseek_v41/engram_test.zig new file mode 100644 index 0000000000..dd085a8744 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_test.zig @@ -0,0 +1,247 @@ +//! Behavioral hash, rollback, selected-row dequantization and per-copy signed-sqrt gate regressions. +const std = @import("std"); +const hash = @import("engram_hash.zig"); +const engram = @import("engram.zig"); +const mtl = @import("metal"); +const View = engram.View; + +// Generated offline with the official NumPy multipliers and prime-bucket algorithm, not the native code. +const first = [_]u32{ + 8939594, 28321402, 40265357, 55817216, 71147020, 86924254, 110813108, 121427189, + 128454377, 156247504, 165979524, 188168291, 198564788, 211763979, 239563936, 241631034, + 263018334, 286373017, 297618161, 319241270, 333656732, 339279600, 364226324, 373889199, +}; +const third = [_]u32{ + 13176822, 30752650, 42532473, 56771564, 70133176, 85254602, 108815600, 115334181, + 138863141, 149049926, 167314083, 191090346, 204680718, 211476972, 237672992, 245604111, + 269210034, 282954889, 298312790, 313992267, 329778842, 345458689, 359433926, 374584261, +}; +const after_dead = [_]u32{ + 762556, 18016262, 38494153, 58317688, 64056565, 81970637, 106927994, 122902996, + 129060773, 153088767, 160206490, 178949865, 205667218, 211027556, 235070602, 251994813, + 263310652, 278975004, 292595948, 315407209, 325281398, 348093105, 353903080, 377405725, +}; + +test "official hash IDs preserve order, padding and globally unique layer buckets" { + var history: hash.History = .{}; + const a = try history.push(7, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &first, &a); + _ = try history.push(9, hash.layers[0]); + const c = try history.push(11, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &third, &c); + for (hash.layers) |meta| for (meta.primes, meta.offsets) |p, offset| { + try std.testing.expect(p >= 16000000); + try std.testing.expect(offset + p <= meta.rows); + }; + var high: hash.History = .{ .tail = .{ 1234, 5678, 42 }, .position = 3 }; + const result = try high.push(99091, hash.layers[1]); + const expected = [_]u32{ + 3315102, 30010869, 33829090, 51274514, 77267621, 84401643, 106522093, 127296527, + 137402994, 151593930, 162661867, 180852958, 195297507, 213361283, 228864003, 255457679, + 262256667, 282063670, 300635716, 304966933, 321424067, 350721784, 358839309, 370654406, + }; + try std.testing.expectEqualSlices(u32, &expected, &result); +} + +test "dead image positions block all older lookbacks and snapshot rollback restores the exact boundary" { + var history: hash.History = .{}; + _ = try history.push(7, hash.layers[0]); + _ = try history.push(9, hash.layers[0]); + const saved = history.snapshot(); + _ = try history.push(hash.dead, hash.layers[0]); + const got = try history.push(13, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &after_dead, &got); + history.rollback(saved); + const replay = try history.push(11, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &third, &replay); + history.reset(); + const reset = try history.push(7, hash.layers[0]); + try std.testing.expectEqualSlices(u32, &first, &reset); + try std.testing.expectError(error.InvalidEngramCompressedId, history.push(-2, hash.layers[0])); + try std.testing.expectEqual(@as(u64, 1), history.position); +} + +test "all prefill/decode partitions preserve hashes including a dead token at the chunk boundary" { + var bytes: [20 * 4]u8 = undefined; + for (0..20) |i| std.mem.writeInt(u32, bytes[i * 4 ..][0..4], @intCast(i), .little); + const map = hash.TokenMap{ .bytes = &bytes }; + const ids = [_]u32{ 7, 9, 11, 3, 13, 15, 17 }; + const mask = [_]u8{ 1, 1, 1, 0, 1, 1, 1 }; + for (hash.layers) |meta| { + var whole: hash.History = .{}; + var expected: [ids.len][hash.columns]u32 = undefined; + try hash.hashChunk(&whole, map, meta, &ids, &mask, &expected); + for (0..ids.len + 1) |split| { + var chunked: hash.History = .{}; + var actual: [ids.len][hash.columns]u32 = undefined; + try hash.hashChunk(&chunked, map, meta, ids[0..split], mask[0..split], actual[0..split]); + try hash.hashChunk(&chunked, map, meta, ids[split..], mask[split..], actual[split..]); + try std.testing.expectEqualDeep(expected, actual); + try std.testing.expectEqualDeep(whole, chunked); + } + const saved = whole.snapshot(); + var invalid: [2][hash.columns]u32 = undefined; + try std.testing.expectError(error.InvalidEngramToken, hash.hashChunk(&whole, map, meta, &.{ 1, 20 }, null, &invalid)); + try std.testing.expectEqualDeep(saved, whole); + } +} + +test "raw tokenizer IDs are compressed before hashing and normalization-equivalent tokens share hashes" { + var bytes: [4 * 4]u8 = undefined; + for ([_]u32{ 7, 9, 11, 9 }, 0..) |id, i| std.mem.writeInt(u32, bytes[i * 4 ..][0..4], id, .little); + const map = hash.TokenMap{ .bytes = &bytes }; + var a: hash.History = .{}; + var b: hash.History = .{}; + var out_a: [3][hash.columns]u32 = undefined; + var out_b: [3][hash.columns]u32 = undefined; + try hash.hashChunk(&a, map, hash.layers[0], &.{ 0, 1, 2 }, null, &out_a); + try hash.hashChunk(&b, map, hash.layers[0], &.{ 0, 3, 2 }, null, &out_b); + try std.testing.expectEqualDeep(out_a, out_b); + try std.testing.expectEqualSlices(u32, &first, &out_a[0]); + try std.testing.expectEqualSlices(u32, &third, &out_a[2]); +} + +fn view(buffer: mtl.Buffer) View { + return .{ .buffer = buffer }; +} +fn bf(value: f32) u16 { + const bits: u32 = @bitCast(value); + return @truncate((bits + 0x7fff + ((bits >> 16) & 1)) >> 16); +} +fn float(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} +fn completed(cb: mtl.CommandBuffer) !void { + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.EngramGPUFailure; +} + +test "Metal hashes chunks with image breaks and rolls back bounded GPU history" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + var p = try engram.Pipelines.load(device); + defer p.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const ids = try device.buffer(7 * 4, 0); + defer ids.deinit(); + @memcpy(ids.slice(u32, 7), &[_]u32{ 7, 9, 11, 3, 13, 15, 17 }); + const map = try device.buffer(20 * 4, 0); + defer map.deinit(); + for (map.slice(u32, 20), 0..) |*id, i| id.* = @intCast(i); + const mask = try device.buffer(7, 0); + defer mask.deinit(); + @memcpy(mask.slice(u8, 7), &[_]u8{ 1, 1, 1, 0, 1, 1, 1 }); + const history = try device.buffer(engram.history_bytes, 0); + defer history.deinit(); + const saved = try device.buffer(engram.history_bytes, 0); + defer saved.deinit(); + const out = try device.buffer(7 * 24 * 4, 0); + defer out.deinit(); + const replay = try device.buffer(24 * 4, 0); + defer replay.deinit(); + const cb = queue.commandBuffer(); + var e = cb.compute(.serial); + engram.encodeReset(e, &p, view(history)); + try engram.encodeHash(e, &p, hash.layers[0], view(ids), view(map), view(mask), view(history), view(out), 3); + try engram.encodeCommit(e, &p, view(ids), view(map), view(mask), view(history), 3); + e.end(); + var blit = cb.blit(); + engram.encodeSnapshot(blit, view(history), view(saved)); + blit.end(); + e = cb.compute(.serial); + try engram.encodeHash(e, &p, hash.layers[0], view(ids).shifted(12), view(map), view(mask).shifted(3), view(history), view(out).shifted(3 * 24 * 4), 4); + try engram.encodeCommit(e, &p, view(ids).shifted(12), view(map), view(mask).shifted(3), view(history), 4); + e.end(); + blit = cb.blit(); + engram.encodeRollback(blit, view(saved), view(history)); + blit.end(); + e = cb.compute(.serial); + try engram.encodeHash(e, &p, hash.layers[1], view(ids).shifted(16), view(map), null, view(history), view(replay), 1); + e.end(); + try completed(cb); + var cpu: hash.History = .{}; + for (ids.slice(u32, 7), mask.slice(u8, 7), 0..) |id, alive, i| { + const want = try cpu.push(if (alive == 0) hash.dead else @intCast(id), hash.layers[0]); + try std.testing.expectEqualSlices(u32, &want, out.slice(u32, 7 * 24)[i * 24 ..][0..24]); + } + cpu = .{ .tail = .{ 11, 9, 7 }, .position = 3 }; + const want = try cpu.push(13, hash.layers[1]); + try std.testing.expectEqualSlices(u32, &want, replay.slice(u32, 24)); + try std.testing.expectEqualSlices(i32, &.{ 11, 9, 7 }, history.slice(i32, 4)[0..3]); +} + +test "Metal gathers only selected FP8 rows with rowwise scales and gates each mHC copy independently" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + var p = try engram.Pipelines.load(device); + defer p.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + const table = try device.buffer(3 * 272, 0); + defer table.deinit(); + @memset(table.slice(u8, 3 * 272), 0x7f); + const scales = try device.buffer(3 * 12, 0); + defer scales.deinit(); + @memset(scales.slice(u8, 3 * 12), 255); + for (0..3) |r| { + @memset(table.slice(u8, 3 * 272)[r * 272 ..][0..256], if (r == 0) 0x38 else if (r == 1) 0xb8 else 0x7e); + for (0..8) |g| scales.slice(u8, 3 * 12)[r * 12 + g] = @intCast(127 + g); + } + table.slice(u8, 3 * 272)[2 * 272] = 0x01; + table.slice(u8, 3 * 272)[2 * 272 + 1] = 0x80; + scales.slice(u8, 3 * 12)[12 + 7] = 255; + const ids = try device.buffer(24 * 4, 0); + defer ids.deinit(); + for (ids.slice(u32, 24), 0..) |*id, i| id.* = @intCast((i * 2 + 1) % 3); + const gathered = try device.buffer((engram.gather_dim + 16) * 2, 0); + defer gathered.deinit(); + const hidden = try device.buffer(2 * 4 * 2, 0); + defer hidden.deinit(); + for (hidden.slice(u16, 8), 0..) |*v, i| v.* = bf(if (i < 4) 1 else 2); + const kv = try device.buffer(12 * 2, 0); + defer kv.deinit(); + for (kv.slice(u16, 12), 0..) |*v, i| v.* = bf(if (i < 4) 1 else if (i < 8) -2 else 3); + const q = try device.buffer(8 * 4, 0); + defer q.deinit(); + for (q.slice(f32, 8), 0..) |*v, i| v.* = if (i < 4) 1 else 0.5; + const k = try device.buffer(8 * 4, 0); + defer k.deinit(); + for (k.slice(f32, 8), 0..) |*v, i| v.* = if (i < 4) 1 else 2; + const output = try device.buffer(8 * 2, 0); + defer output.deinit(); + const dead_mask = try device.buffer(1, 0); + defer dead_mask.deinit(); + dead_mask.slice(u8, 1)[0] = 0; + const masked_output = try device.buffer(8 * 2, 0); + defer masked_output.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + try engram.encodeGather(e, &p, .{ .weight = view(table), .scale = view(scales), .rows = 3, .weight_stride_bytes = 272, .scale_stride_bytes = 12 }, view(ids), view(gathered), 1, engram.gather_dim + 16); + try engram.encodeGate(e, &p, .{ .q_weight = view(q), .k_weight = view(k), .dim = 4, .copies = 2, .hidden_stride = 8, .kv_stride = 12, .output_stride = 8, .eps = 1e-20 }, view(hidden), view(kv), null, view(output), 1); + try engram.encodeGate(e, &p, .{ .q_weight = view(q), .k_weight = view(k), .dim = 4, .copies = 2, .hidden_stride = 8, .kv_stride = 12, .output_stride = 8, .eps = 1e-20 }, view(hidden), view(kv), view(dead_mask), view(masked_output), 1); + e.end(); + try completed(cb); + for (ids.slice(u32, 24), 0..) |r, col| for (0..256) |d| { + const got = gathered.slice(u16, engram.gather_dim)[col * 256 + d]; + if (r == 1 and d >= 224) { + try std.testing.expect(std.math.isNan(float(got))); + continue; + } + const base: f32 = if (r == 2 and d == 0) 0x1p-9 else if (r == 2 and d == 1) -0.0 else if (r == 0) 1 else if (r == 1) -1 else 448; + const scale: f32 = @floatFromInt(@as(u32, 1) << @as(u5, @intCast(d / 32))); + try std.testing.expectEqual(bf(base * scale), got); + }; + const gate_positive = 1.0 / (1.0 + @exp(-@sqrt(@as(f32, 2)))); + const gate_negative = 1.0 / (1.0 + @exp(@sqrt(@as(f32, 2)))); + for (output.slice(u16, 8), 0..) |v, i| { + const want = bf(if (i < 4) 1 + 3 * gate_positive else 2 + 3 * gate_negative); + try std.testing.expectApproxEqAbs(float(want), float(v), 0.015625); + } + try std.testing.expectEqualSlices(u16, hidden.slice(u16, 8), masked_output.slice(u16, 8)); +} diff --git a/zig/src/families/deepseek_v41/fixtures/hc_attention_rms.json b/zig/src/families/deepseek_v41/fixtures/hc_attention_rms.json new file mode 100644 index 0000000000..d1f1c2b61d --- /dev/null +++ b/zig/src/families/deepseek_v41/fixtures/hc_attention_rms.json @@ -0,0 +1,25 @@ +{ + "embedding_bf16": "WT1bPAO857tZvTS9Cr2HvDY9MT2nvUI867vvPIG9Xr3wPMs8GLwLvCG91b1CPf87mTzuvC69rTzOvM28x7w+PSC6VL0YPOG8Fb1TvRE84zwNuw48pTwrvF89OTxEveU8SDybPBA9Yz0XPKU8Nj33PJK8jjtTvNW8C731PDk8Rz0SvAK9l7xnPB69ojvNuxM9y7yMvDC97Dx3PA89Ib2TO0k8j7orvDc7xbwwvCM9lLydvDS9OD1JPQC9JL2yvBu8ebyCvUE9Qz3+vDm8Bb0XPGG8gryUuAi8ZT1vvN+8wbyQPI47orz/vCE93jzVPHU90DwgvVe8rLzPvDe9Xr1tPZK8ojtCvcG84byJOzq9QD2APGE9Xzu3OwY90LzkPNY9iLsZPao8kz2Ivbs85LxKvZu9AjuiPQO9BL6RvAk9hL37u+08gjoavBw97rxDPXq5kz3vO5e7arp3vX+9hj2dO6m8B7sBPHS9jDywPGI907zYvAy8P7y9PDw8u7zrO1O8rjzLvD28Yz3yu5u8lLz0vCa9kjxuPE+9ED1XvJW7RbwfPDe9Mz1yvRy9Br0CPNc8OL2zPLI8tjwdPQw8Fb03vIM8kTzrvPU8DTwgvdM8YbwJPKY877vVPLo8Y7x6PAw9IbyOPEu9+TxuPfo8E7zVPO48ybzNPOe8EbtlvYi8Kb3VvCI8Wz2rvb88Wz1HvYe9oLzFPMY8pDyCPXO8k7yQPQq9IrxbvS89JjwKvQo7Jb0QvVa9oDxWPLM8rrgMvI46mjwWPdi8nzwevUG9Dz3eumy8K70bPZS9Jz2/vIO9xrxLvZi9W73GOsY9mL2OvCe8BL1DvOE8D70tvbU9mzxPvT082zwnvCy9Ar0+O1s9qTwbvIY9Sb2qvTY9KzvvvCA93TxAPeG83TtbPN07s7vXu2M8RjzyvAe9wb1fvUk9YrwwvE+8nbywPVG8WzsDPZI9XD2RvG09+DybPAo81zx8PdI8kr0JvJI7Xj1GPTS9irycPFS8gTzcvES9wLwvPY+9Bju2vRs9IDziPF09GD3dvHA7ML0MvEa9nroCPbo8yrtcvZg7QT3WvC47aLxmPKa7zbwVvGq7Dz30Oyc9Bz2FPBo9Ob14vVW8ELsDPUc8qL2kvCA8DbzlOz49EbwsvS8907wQO5w6sDzcuo085jt3Pdu6Hr0UPBm9OLzVvOq8LT0wO8y7Hz2/vFa9rLwePSC91LyXPa87rr0xveK8jDuMu8C54rx6uzU7HT3AvLE8Jj37O047/zz1PB+7Ab0QPDa9Oz2JvZW8Lb3guwQ9njtivey8vTyhOuE8Brx5vXC80jxdOw08i70OvcS8MD2qvbu8XbyZPLm9mz1aPV66iTx3vMQ7oTyyPIw9PbuOuzM9HjzgO4c7ubzivS69PDzTvLi8H712u767ZL2dPFk9gTzvvK48BTwOvYA9TbtlPIy8Vj0wPSu8D7x2vZQ88ruvPLy8tjwyPSk99rycOxk91DzgvL87pjyGPUi8Uz1ivBS9hb28OjM9izxevFq7Rb39vK07fL2QvcU8uD1vvPq6TLz2PCk9hLzhvIC7gzyBO2Q65DySux6967zbvA09Iz2cPVe8rrwmvbC8djwtPVe8Pz1qvUi9e71ZvYU9R7okPS49jbpIvaY8QbwiPee8rLwKvbE6WT2MPFa8YbxJvLo7Hr0iPVM94DzMu4s8Ojw5PRs9N730O568Fz27vFK8tL2uvEa8irrtO0K89TtAPVk8vbgmvTk9pzqDvLs7KTxtO6W8GjoWPCc9DjxBPZi9oryoPJS8O72jvBu9Fz3lPKA8bj1lPAS9UD0iPQ48KbxFPJy9J703Pbi8/rwqPVs7Er3Iu3q9qDsvvR28VrxzPVE8Y72UPB+9WbsfvS+8KL2BOu08B71ZPVK6Nbo0PZg9iT1pPYE8RD0IPB89Mr3NPEG8oz2wvLU8QjzavLK7qLwgve88ybxCPV68obxbvbu8f71xPGa87LwRPa08Pj1vvG89mz1WPdk8E711O5C89zwXPOy8hDzxPAc9ELwqvGI9I70SO7Q7Ab3AOzg8abtfvTe8zbzAvLS8ibyVPTu9tTyYPKy857wnPZi9JDx5uz07E70MPVq8ZLysOy67BT00Pda8bTx0PN082T0nvdS8Lztyuzm95DyoO1I9vLv5PBO9FD2kPWQ9kLuguQY8A73LPD27pryUvPq8Rb2Pvf68kjxzvGS9+LyDvOC8TDyRPAG9HDy0vLW89DyGPQg78LqMOza9ULt+POM44zzRvIU9xTyLPQc8zbxtvaY9bLuoPKw79zvMO7U8h72pPAa9ML0BvUK907xDvZC9Fbx3PfA8VDxsvW88Bz0TvDK91DwrPIk9ab2uu0a8iTtBvck8sjz7PCQ9D70JvQe94DrHOhG7hLzvvJ894DwxO2q7mb1jPDY8qTzPvEM9N7zZPHE9IboPPUM84zxbO0y9q7ylOh68JD3/u5c7ujwKPIu7yTzLvOO8Ej0YPdE9r726PR28C732PBW8wztrvZY8eLyVPLU8T7xhPea9i7qvPES7qbsvPI08jj2uPOM8T71bvOC8ijwrPY28G707PeO8JD0BvC85sbwPPA491bzNO5Y6J7wbPX88n71hPA+9wbyzPBi7kj2JPM+7j7sROgG8Zj2nPJ88uzwGvJE9j70PPay7yLz+vAC9Sj19vYe8hz3OvBW9DrswvYM8Eb2dvGy9/TuKPeI7JTwIvUk9mLuLvI68rrx2vUg9izxfPf08ez33vB+8L70lvdk8Sb0lvYY8Pry0POk7yDy/PCA8cDwNPQY8dT2DvYu6FT2dPWy8C7wcPUw9nzxhvKA9t7sWPci8Tr25vL47Vb02vBQ9j71vva+8tzyNPDe9lrm3PAK82jq1O2M7WbuXvCw9krsmvag7tDoIPXM86zzJPAU9aj2bvRG8uD3SOqa7XrxNuy09gDvHPF+93DwxPAu9Ej1FPU69hTuyPOu8L70avQ89zTqzvUk9tjvvPBw7u7zEO7k8Ez0hO0U7NL01PQM83zzJO3y8oDz8PO0797yOPIG9nLx9PLs8sbx2PZa7kT2NPCe8k7vUPJU6Gb0fPL48cDzDOzQ8RLxKPU89Wr36PJY8rDy+uya7E7wMPau6TryIO6O9kDyyvc87Er3hvAg8BLzWPEi8yzxqO5I9vLwWvRs9sryUPCa8trw6PA09EDuNvaa7QDxXvSq8U71fvHQ9M71QvU88TLuYvV89xrxFPXS8/jxAu2E9cLyZPL29+DvHPOq8QL0XvCi9TT0JPX+8prxPPSe9wr7FvG09OD2mPJM8g72NvLS80jytPLo8F73cOz+9CD2LO1Q9Mrx1vA+9WT2UO4K9MT1uPIk96L1Qu9i8cL1oPU69v7rePNg7Kz2ePLA7Dj0oPEU9Mj3/vDY9hL2SO488J7sGvIE8Kz1Xuxs9KDwCPpa8O725vGA8SrxAvSI8A73HvF47n7x7PEA9ejyWPUw8NT0TPXo9DT2ePGu9yjzXu4I9tTsAPee97rsEvSO9GDw/PEq96DvOPQS8Yz1dvSw9cz2vOxq71boZvVM8AD1vvYQ5UT08POm8zjuEvZy8JT0Puw466LyfPfg87zqFvMw6Ujzqulm6fz2muwM9t72pPc87+LzIPci8xLzzO2S95LtoPXu9bDx0vBG8/TzPPHC97zzWPDe9k7wXPJQ78rwavQW9gj2NOFw8/rxgvX28gjwRvei8F7ywvDW9z7s0PUE90rsePV2987yZO8o8Rj2TO8e7mT3oPM47zDs3vBS90jykPBa9zTyNvOM8kTwTvT87Wjz2PLi8rDs8ugS9CT11u3e8Wr0rvNs7ITz3vA09T73bvIu8qjyIvYC9oz21vIC9Z7uAPEs9Gr0MvYe84ztXPeq86TykvAW9FD0uvII8ELxAvK67Sj3GPIe9HrxcPAW9yjxdPSc9RL3kvC48gbysPBg9BL0SvUE83jzCPMI7Cryyu448+rr0PEM8Wj3mvAY8mLyGuyw8Mjp4vcS9AjwlPXw8Rz1BPq89HT2yu9i8Zbz4Ois9S701Oym9FrwIvaI9lz1nPMO8nb0OutA7rT2IPFa96ryrvTG9e73xvKS8ETzsvGs95zuzPC28dz2dPIc9/DxPvGG8tDw+PGo8h70HPCq93TwqPJG6erznO8C8qbxfPBK8iD0EvWw68bwpO2a9lbvSPHE73DoEPGo9YD3vOzY9Q7ouPCo9I7wjPeu65TysvBA7gD2VO5S9BD07vbo8hL2QPG28AD0vvWG9JL3EPFw9q72guw68qzwRPRy9rDybPPG8BL3svGK9nj17PWi7ETwEO+s7Czz8u/E80DyPvFA9fLyHPEI9vT1ovJe7Gb0aPRQ8br0IPVe8Wr04u8c7BTw8vT+9Kb3ZPJW9+DxQPFC8UrzBPDc9Jb2Dvbq8Sz2zPIa9aT17POO8lbzEPL88Ar66O4U8Zj0Xu5W8Vz0NPSQ9tr2dvIM8PL0VvMs8n7wBvXq9sTuCPbu9ezwBvQO9Ubw4PSA9ITyruwm70bwSPUA8hDwPvZ89F73tvOa7Ij3OuZU8AT1svb28P7xIOzu8jz1oPQo9HDwePcG8ozy+vSm9vjyJPEW9ezyOPO28tTzLu8E8Er1mvdK8TD25vCO9fz1VPJu9BbxxvGc8rbwJvYO8bTzKPIo9Rb0yPJo82LzAu9I8Bz1XvGI9RrwMvSE9IjvAvA08FzoUvbQ8PLxbvfm8Zz0zvIK7zTv/PF88Mz1PvJa9IT3JPKW8Sb0QvQW8cjyjPLs737z5O4+807w0vQs9ZD0LvUk9Gj20PPe8I7suvOm8l7xVvca8KrtzvQ89QzyHPWU8ejzUPA+9Aj2aPBY9ebwmPXy7bj1mPS09Z7sLPZ+8KL28PCS9z7xQvW47QTw3vdu7/DvhvPq8X7ySvTq6O72nPUY80LwJvR496TxXPIU7LrywPfq87zz5uwW9/DsEvUI8Qr2wPPE7pz1SvbS8Xjx5vTe8Qj0LPcM9JjwzPZa8HT0zvBw9KjyIPZQ9WLwTPIg8cTutvHm8+Lweveo8k72ivOM8Sz17vYO9VD0Zvek8wjzAvdg8cr0IPFo9mLygPZe9GTz+OYc9mD1NvBs97TyHvTq96bwovXE8zTtIPK08qz2XPCw8MT2/O0m9ezwEPE07lT1uPU29s7zHPAO9CD3Ru2g8K71yPUA8nD3Ku0I8YT0ePU48TbzYu1s8sDwtPY29BD0UPPI74joFPUs9rr0uPD88Cb3CvOa8rjx1vVe8xbzlO6S8yTyBPJq7YLpOvLe8IbyzPDW9C7vkO8s8izwtvQQ9gbwbvIW8HT5PPXo8ID0oPAu8dbstvRq7pzxkvMM87jymPcE6BDyOvSI9w7tiPXW8vbwJvAi8pryJvHa6oz3APPS8QD3wvIE94Ty8vNc8Qz2RPGs9LD3YO4+9r7yhvEO8nzutvDe88bxVPHc8tTwuvV68Ab0NO5w9ijz5O6W8S71UPQG87TyIvAA9KLyWuys9/buVPHG8dL3oPMg8oTy7vNO86DlLvQs94jwUPHK8rTsTvQ49SzwyvXo8XTzbut28aTwQvOS5KbzkvJs7Pbkcvdm8P7x9vek8V70KPb489zx8PVo8gz3lPAE9RD04vSs6Trx6u9M8uDwOPLq7OT0rvSm8Jz2svI08t7yiO5A9X73Au1i9zjxvPES9Qj3zvGw8Tb0+vSw9Vb1pOxu9zL00vYu9QL0fPcs8djy1vS88vD1PPcI8NzyFPIK97zxXvOE8hz3iOg49Gz3XvBA9sr0bvJ884jzkvNy8kbyzO++8U7s/vXo8hj2cvDQ9R72DPYI9P7p/u2m9z7w2vDK8GD02vbc97bwqPP28Yb0xvWY8RjsQvX47eDySvAU9nLwVPVG9Hz3SPIa8f723PEo8Gb28PIW7Nr1cO5M9+rxsPDA9ODz5vIE8Az3QPFu9jTw7PQs9Yj2IvBk9kL2GvAS9zLx6PI+8oLzMuwu9Jj3pvGu8Vr1SvOw7V7yJPLc8G72ovGA8dT1JvKu8KT2CPHU9ET2ZPKM84rn5ug49mD0sPfo8Uz1yu8u8T7kPvXs9qDyrPe27Lb2ePKo8/rswvPS7wb2pvLc7n72VvFe8BzyrPKC8CD2ovHA8oD0ivF07+ryLvTA8mTxtPBS9xTvXO/K8Qb2fPDu8ID2sPLO7Rj0Kve68ibzrus68jrxMOye9ybyFPF88TD2KvfU84Ty1vKM86Lw4vFm9iTzcPCq997wIPBo9TD3MPUg8W7xbPVA93jiwvGA9J70uPbu8Rz0hPM883LxNPR48Jb0HPCu9sDyNuwu8ED3RPP28gr0dPUo8Ej0pvM29tT3dvLy85rw6vZc9k7ynPM+7Q7zeu0u9Rr3tPLK8Fz2FvGs7kbxWvOS7LD2FvDC90jwjvIE8Kb2mvH48PDzlvNC5G7z7PBW9Kb1vPG48oTwkO6e9HT26OzW8QL0APEg9t7uRvA68l73zOaG87LwwvVy7ZLy+vNg7HT0sOyM8vDyevcg7xrx0vNE9wryNvYw7Nz0wvIy9W701vCi8qD2DPba83zx7PSU9+ju5PD69Lj2BvCU9hjlJPQO9sjwdPQS7Mb2ROzm9OjwYvWQ9Ur13O0I9gT3dPKU8xjuVPIk5gjyku1+97rwcvIq9dTsKvVQ88jvLPB894Tv0vAu8yL1XuqC9ibwivSY8RDw1vYw8oTwEvZI8C71xvEM9HztYPXY9FLwdu5G8dDxavDu8ir3ZvPo8tbyWvVu8Kr0pvZA8ZT3Euxe9e7youQ28ZjwHvVk9Dr1OPEQ8/ryfPD+9kb0RPFS91jurvL08pDsFPRS957wLPSa9QLw9PVu8J73+uFK89DxXve68mD2XO4Q9Bb1FvIg8Uz1pPWg9qbtsvI494TwPuzI9hLysu4k9QDvKvPM8bbsGPea7Jb30vBs8yzxavRa9STyWvOc8Sj2KvHy9ZTwqO4y8mLwhvWW8GTxZOzo9wzx+vVK8Srx6vJi8EL2nvCg9AjxwvSM8ZD0uuK88x7zuvDa9UL1SvIe8CD0hvZ088Lymu4I81D2TPIi6OD3aPGk857yOPaA8yTzdvJG8uTsmvcs8ijz8PNE8J70WvG68OL3lu9q8l732vY49Ab1sPG28bbw6vRm9zbxEvUA9+Lw7vH07Yr0QvAo8orxmvHi6FLyavCu9gDyNvBy8zjxAvQY9aT38ux49bTynvAY9ojycvQ080zy9uhi9Zz3bOxC9jT3tu0m8iLuRvYY8s7zQu5g9jT3nvM67LD0nPcG8fTwkuxc6tr11utk8l7uvPDq9Dr2kPVW9Lj2hPCO9Qz1TvY68cL21PD69mLw/O6q84T2JPYO9cz2mPZA8v7xeuXy8hz14vJ08ljxyPZY9Uz16Paq65TmKvGc7OLrou928Ib1mPcC8ub0tPSg7LTz3O6W8/LyIPUM7Trxkvbw9Dr3cvOO8C71aPeS8lrw4vDe8db07PXq5CLzDPMI96DywvAA9KrzrPJY8Dz2jPHG8k7xBvGi9OTyxPMi8QbyIPJu8k7t/POc6rbrWPCi9NTxOPZ49VT0UPSy9Sj1YPCO9JztUPQ68gjwwvLi8jDuRPEw9U70xvcW73btcPbE9/jw/PfW8ND0Kvay8W72wPFU7Iz2bPeG7OTyCPao8nzwUPTI8fLvdO/G6trt1vfy897tBvIa8JzwWPS49uzytPKM8k71VPN+8KT1bvdM8zrxaPD07VT2kPBg9TLyPPJa9xLsjPSe7GT1hvO88QjxSvds8bz2mPGw9xLx8PW89pb0Kuww9jzzAPW+7E7wxvLC8WDxnPAq8sLzNu0C9KD3jPNy8/jvqPKq9I70WPDs8UD27PNE8vjyAvWM9Xrx0vIM8Zjw3vAi9Bj5XvT09j7xvvDm9Rr15PJS8/DuXPKU9vrw0PYm8CzuYPOe8LT0zvQG9aL36uxI8RjxKu667Uj0fO6w8Lb29uqM8hrzmPMW4Xr2OPBm9/zyNPZ+7Lbw4PVi8j72/u4q7tjtJvDO9ED0QPHk9wzybPKI8i727PIe8RjwNPVo9ADwgu1U8ybxJPey8VztTOwc9KD3lO0+9sbwpvTM8gLwQvSQ8SL2ZPNC8ST2pPNC8aT0CPXe9dLuCPGw9YD1fPAY81TxIOye8wrwcPRe9SjwePdw8BT12vQe8Vb0JvYQ96zuTPS09PrzdPJE8rDzmvKe7/Dz9u/K8Zj2IPZ28Ez1SvCc8T7wEvb67R73mO/O8sTvXO3W8FL1Avfs7xTwavVY757x/PcW7GL30vBc91rz0uiY9i71kvIG8WzzQvMc9QrzHPDu9ULvwPGQ9XT26vPy6q7uGvQQ9DT07PCk9mbyHO0S7eD0BPSE9Aj3bOsQ8Fr1dvbO8IT0QPsG81L2UPN28JT0GPYG8lzwsPWy7kLyyPeE8Wru0PBq86btYu/68ibxsPGk8Jj13vMe8jLxqPGg90btGPIO8gT27PJK9P737vF28tro6vZO8Jr20O3A96zw3vZe8eT1yPYw8vDsqPRk7obxnPd87j73lPNs8Hzo2PGW82jxOvYo8vbxPPWU8HD1SPZm8Br3GvKG8ZTylPG89kD3xvH69Ir1+Ozm95DoEPIO8Pj1BPQo9hb05PHW96jxevaC8TT2BuoO807w8vNq8E71JPQy8gz0DvKS8lTyQvTC9vjyUvcG7Qr0WvK+8CL27PcK8Hj25PEo9KD2KvBo9ybv5u628abouvAW96bwpPVG8tzoOO+a7WrxKvcE8xLtQvWs95bqDPa889bsePaY8STkyPbI7FD2HO6W8u7yJvT89FjwOPcG84Tyhvey8arzGOtc81ryyvNs8fT0avYq7Kz26PNK9Szk5vfI8PT0ZvQc8vD1IvaK8Bj2mPAe94bvhPEA9TzyEPZE9zDsKvSI9s71APXS73TyBPKQ8Hr0APQs8bz0zu1e9qDx5vZq9fzybPXa9qr2OvKg9STz1vIs9cLzvPBQ96Tsmvam9/rsMvUo957o6PbA8tbslO3W7az0MPaw8Cr1gPMa8ab1ku828Gb23OwI9b7yYvCQ9aT0WvbG847wevFA7BjwbPRc8hLx5u4i8gDy4PCA9dD2MvFC8iD30PFc9SDyNPHi7xrwyvXW9kT2LPFy9+Tt+PKw8qLwUPe+55DtzvRo9HD0NupG9rjsIPCc7pjpDOwk7Nr2UvMO86zyHvTK8H70lvas9MTwvPZG8hL5UPXg7FzxrvCu9d71LvS69vbwKvVG9tzy2vZe7ED1kvFy8nDvgPIo9BL2DPQ89LL3avEA9Tj1TPYY99r0ZPR27w7tru7Q8O7sBPd+827wqPYM8WDxkPaO8Pjx1Pc26mrzXOyW9mj0Tvc29sjzWuzY9Arx4vUM81Twduhs9hjsCvS89qDudPDM6lj3NuxW7Gb1EPcK8xjzwPAi9Xz0SOw094rkXPdq8ab0/PfI7bbv7uh28GL1YPfw8K7xJPQS9mbxbvMa8xjt0PBK9a70qPSS9HDwHPTO9yTxNPMs8DbxkPCW9lL38O029nz1QvL08hb2wvJS8Zjx6vMC8JTs8PQ28aDwUu1A8Bb0GO8c8Dj1NPFi8G72EPZ48DzxrPbI7ljwYu9e5Nj0TvaC65zyLPU886j0OvZo907u1vCo9Jb0au1E7szwdPSc98Ls0Pdc8Bjz1PLK8FT0Jvmm8l7yzO1E8Uj1BPfu7MD2lvHK8i7wWvTu8TzzlvIm98zvZPBi5gb3ZO4w9xzz7O/C84Tv1vCE8PL1GvRS8Y7z+PMs8nbx3OxG9XrwsvRA927uQvJ+9U71zvDo9oTyVPTS8wTwRvbG8FD3CPYy9vb2QPUC9XD0iPC06PL2VvLa8C72VOz89ZL33vHi8Sr0KPQI9nrxUPQq9xDpOvaO7sjxWPUc8+jxbvco8jDwAPXs8TLyWvPe7vLz7u7y6Wz1Ju7Y8BD0uu2C8j7zbutM7IDz/Oki9YbsoPCC9F7ytvBO85Ls2vb28UDuWO5w8Br0TPSu9ZbwUvfe8dj2gvUA9pT2nPFC9bzy9vLa8+TqCPPY8ibx/PeY8kLw6vfK7vzzzOw88u7wnvcs6mzwwvbu8PLx2PVA6Zj2SvQu9G72rPIy93bxiPQA9Iz3KvJQ9t7zPPX49oTy3vDS7XDzZORY83rwRPeM8/zvVvC+9CLy2vBM9SD3gvB07qzzgvIS8qz2KvOs6iLtIPS+8hT1dvP88lTzUPJq6WL1KvY69WD3HvD+93rzbvAY8HLx3vXK8ujtMPNk5yrxqvKS8nLtevYO9or3lu5C97rwOOma8lT09Pci8iT2KvJI8Hr0Kvcw8Jz0OvPY8kT33O6m9mzubvOq8JD2XOje9Mj1FPZc8pbvYvCK93Tsxu9U8ELomPZo8wbyHPL87+LyHPJ+6c7x2va68i7wqvTa8mj1zvaS8lj23PDC9sryhPSi9sTy9u6C8t7ktvCI9VT6+O7o8vDvfPGG8ozyJPH68tDuQvE69AzwrvPA8YD3MveI74bwDPBO9YD0SO1e9Fj0qO2W8IbvEPAQ9WjwCvb+8s72uOQo6nL2YPKq8Zb2zPJk8Cr0WPGi9Jj2tuxW9izzGPTY81rwEvM28pbnjPbS8gDyIPQQ96bzwPFS9SrqWPJe90zugO+q7G714OkI8vDyeOzW9aLzCPcS8cj1NPci8/rycPLc8Qb2IPPg6irw1PQi9irzkPIk9brzuvAo9cD2pvA87vjyeunK8kL1RvUY8ebrEO/u83DwLvMM8jz2OPVI8FT3evB09PzwgvUs9JT3ovG+8yTvDuGq8urxfvJg9uT0uPY69YzxHvIo8AL0cPSw8trz9PMG9xL0QPdo8CL1bPNm9UjyzO4K907zKu5w8ID0yPBo8s7wwvYu8jrynu328TztyvAQ9YzygvAo9hDxMN9Y8DT2yPc0837yVuhU8kbw3PV07h7wPOw69XD2rPDk7azxXPRQ9Hz0LvVm8zzxkvX+9Gb1GvJE7oL0/uwS8WDxAvYO9JD3xu887Fj3LvEo94zwlPSi74rwnu/Y8ZD2/PAy9/zymOhQ9uDuJvEm8STwLvZS8jj1xPSK9oDygPTk8VT3bPFq7rDsWPJQ8iTwFPQa9ezztvCE9bLs1u/s8bz0fvQe4Q7sfvea8Db0dPca8MD2bPYM8Ir0OPaY6hjwkPSQ6fDsBu1C9zzxNPd65LLuQvVk9ibtMvDS81jxnPWK9SjzTvKi8Sjx3vSg9KT0NPV68Jr0Rvau8DL0IvIk83DxePYK8iz2RO9u8tb3/u6y8LT23PGY9x7wru5E9ojzlPJ28hLyDOx28JLzwvCs9zTzzPNO6gT1xO4m8nr0uu3I8FT1vPFu8Dry+vBG9tbxYvWu8Az1tPKK8HT3CPI49pzxUPQO8sTxmvXK9STxAPLQ9bTy2PMA8tj2GvKu8eLwWvSy9Kjz5uzQ8F733PLK8q7zuO0O9rDx5PYC9BL0kvRu827z8u5S9izz4O5A86Dv3PBE8Gz0dvKc7ozxWvHG9EL1DvZC9KTwNPci83jqevYE8mDzxPNq7Dr3bunA9Tr01vZc7UD2AOhU9dTp0PIE66LxUPFU8GLv4vA89YjxjvaQ8w7y+PAa8eDw1O049Ajw9vMS9R73ovB481zuTvPQ5az1lPcA8urzVu3c95zrSvA47ejzIvIO9WTuSu7Y8JD1sOza9FL2EPMO8nz3yuys8Yz2EOQG9f7zHPEO9LD3QPLA8Xr1MPU88+bxZPKK7dTwNPfS7kTyYu1Y8Sb3zPEA8NL0UPdc8nzqBvJ28gzwTPaY8y7z8vBM8+ryeOiI9vDwzPfc7zbsVvNu8QjxlPfw6ZLwwvJ28FbsCvWW9LT2rveE8HT2FPUg8v7xKPO29VL1dve08Qry3u1E8IbwOPaw86jyuO5g8nz3mvOO8jzy7vAO9BT1CPQq8H72hvNO7ijwvuz+8pL2QvVg9d7sCvtK8U70XPa07IL1rvea8PLugPRS9hDxqvP+8G7ywvHY7OztIvdc627sJvKg80DxFPEQ9dj0Zuwo9uzx2PG87kD1aPac7oj1xvHE7dDyfPMW8Ejucu5e7Ub3UOwc8qDwMPXu6uz3VvF+8j7zEPIk6tLyUPWY7LbypvPW8NbxGuuy6qD1TvWq707wWvfO8Q71yPTQ8Ej1sPU+7wzqdu9c4X7xKvQ88SrzTuxK9qjtZvQW9vjuJvMQ8Ir1gvRo9ojxfPSg9yjyFPIA8A71EO7a8ebzUvIi8FDwUPRe9yr1xPeK6JrsOveC8hzzmO3+8IDySPRa9Jz1MPKE8Dz2SvfQ8xjyWvMC85LyavB28lDwVvQC7TjuGvAK96zyxPIG8VDzkuzq6ZzqsPHY7Az2EvTU9B7zDvNC8mzzsPJE8Kj2JvCK8lDz/PLc8l7xJvFw9Bz1uPIA8sDzXvHA9Zz1Cvam90zyjvAM9qDmGvHo7nDsQPd28bLzZvKg9Cr0bvTE9qLdAPc86Hr2oPKu8drwNvVm8Hr32vKu8cb0ivZU6870HPJk9hbxsPaE8CD2AvF09sLyOPDG9tzz5vK+9nbpBPRc9jT1CPNW7i7xaPai8I72cvB07j7zOO0m9p7xvPYC7pLu3PEy93rwzPQ09EDzEPPM8XbwLO4w9JD1Hu2O9qDm1OxK9yTx8u7y8qryZvNA8FT0Ku5m9TT0LvOO8sbyMPTa9VTy2vKi7Cz2RPJK8JrsrvQ69iL3MPL68F70mPai9NT39OyI8xbyBPVo8yLuBvdm7gTxSO9E8JrwhPBE9RL2fvNe7Hzw+PJy9wTqzPGq7ET0mvYY8NL2TvHU9FD2WPO+88Dw/vWU9VrwavRo9h70RPbS7bz1HvQe727w7PQe7J70ZO0G9Kb4hPK681TswPDO8ozyBPB49Qj33vP47Pr2WvcW8I70UPfw8Vz0qPRW7vrwFOzu9YLmWPEG98bwUPS69Or1cvPc8lTz4OrG7Qb2IPIk9mjwVPLk8zryGvII8U7zdPIW8jLwxvHM8BbyOvPA6Dr3OOS+81rwNvF06tzzRvM08Vj3TPHU7jjsPPXC75rwlPQG88jzaPJS7Ij3ivKw9qjxTPeY8DT0OPTk9Lj1BvEq9RT1TPIy8lroSPdo7ez1AvYy8xzsbPYc9ND25vJQ9qz3mvBA99LtivMG8WLyePLc8xDyrvPC8AjzYPI49V7yuPJo9nT0mvUq8P70pvSi9hbwEvNS7Mz0pPQu7mbxWPIm8u7xrvXK7LzxkPDG9IL0aPS69g72jPD09ELwIPWw9ETytvEA8NTyYPM08xjwPPYe8hjyjO8G8bryGOhw8MTyIvPA7TD06vZy8Br3zPM48Dz1SvbM87b0wuwC9hrxvvTY9VroZPUq8WLyBuxw9sLyWPYi8Eb1SvRc9uzxWu628BL1GvHy9pDuMvAy9A7xPvIU8dbxVuwM9X7wevTC8jL2yvK29ObyYPbM82ryrPLQ6yrwVu068yD2evAu9Bz0qvYm7UjwNPTQ97TuhvFG9pLsHvl+6Yj00PQG817wsvf48jjwiPai7H72JvDM73DyovDI8o7yUvG69qLy0Pbg8VrzOPAs9qzw5vuM78jwRPTA9JLtZPAm8sL0PPRg9HjzoO2g9C7wJvea8VL3IvMa8Oz0Uvby89L28POa8Qz1KvZq8zbuPPDk97jk2vM88/zwnPRu9lLxCPTe9szyuPLQ6mz1vvGG9mb26PA==", + "mixes_f32": "5ertwO6atMABjNfAmaqHQE2ySsG204rBMOp7wJFAr8AAOxE/S0CVuhU2ssD8mqBA2IAxv1nkMT83yZu6uaNIuUHJwbqG2oG/SCKCPzAy3LloCZa7EL8+uv22bLroEMs7", + "scale_f32": "UyCBPWQ7RDxMEXw9", + "base_f32": "ChY2QHj3VcDr8U7AWOUNwHV/hcDK9HzApJWcwXxptb9nNFU+IZzLwZlLo8AJrcU9GSQBP0av4b4WAsrBt7OswV3ty8HInUA+BOflvQtYsMFBCZzB416twTi2s8G+BFdB", + "branch_bf16_bits": [ + 16053, + 16067 + ], + "residual_bf16_bits": [ + 48359, + 48377, + 48359, + 48377, + 48359, + 48377, + 48359, + 48377 + ], + "expected_source_tree_target_bf16_bits": [ + 14092, + 46834 + ], + "provenance": "Token 671, layer-0 attention HC, official model revision 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. Four copies of the actual BF16 embedding; mixes are the independently executed native FP32 projection boundary. Target columns 3093 and 4687 of copy 3 distinguish native old RMS from pinned Torch CUDA single-row mean. Base64 payloads are little-endian BF16 or FP32." +} diff --git a/zig/src/families/deepseek_v41/index_precision_test.zig b/zig/src/families/deepseek_v41/index_precision_test.zig new file mode 100644 index 0000000000..1822f7ad0f --- /dev/null +++ b/zig/src/families/deepseek_v41/index_precision_test.zig @@ -0,0 +1,57 @@ +//! BF16 tensor boundaries are observable through which compressed positions enter sparse attention. +const std = @import("std"); +const mtl = @import("metal"); +const index = @import("csa2_index.zig"); +const kv = @import("kv.zig"); + +// The last key's extra 2^-9 dot term disappears at the source BF16 boundary; all 513 scores must tie. +test "index BF16 dot product and reduction boundaries preserve top-512 tie selection" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var ip = try index.Pipelines.load(device); + defer ip.deinit(); + var kp = try kv.Pipelines.load(device); + defer kp.deinit(); + var scratch = try index.Scratch.init(device, 1); + defer scratch.deinit(); + const q = try device.buffer(4096 * 2, mtl.ResourceOptions.shared); + defer q.deinit(); + const w = try device.buffer(32 * 2, mtl.ResourceOptions.shared); + defer w.deinit(); + const raw = try device.buffer(513 * 128 * 2, mtl.ResourceOptions.shared); + defer raw.deinit(); + const storage = try device.buffer(513 * 68, mtl.ResourceOptions.shared); + defer storage.deinit(); + @memset(q.slice(u16, 4096), 0); + @memset(w.slice(u16, 32), 0); + @memset(raw.slice(u16, 513 * 128), 0); + q.slice(u16, 4096)[0] = 0x3f80; // 1 + q.slice(u16, 4096)[32] = 0x3b00; // 2^-9, in a separate FP4 scale block + q.slice(u16, 4096)[128] = 0x3f80; + w.slice(u16, 32)[0] = 0x3f40; // 3/4 + w.slice(u16, 32)[1] = 0x3b00; + for (0..513) |row| raw.slice(u16, 513 * 128)[row * 128] = 0x3f80; + raw.slice(u16, 513 * 128)[512 * 128 + 32] = 0x3f80; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try kp.encodePack(enc, .{ .buffer = raw }, .{ .buffer = storage }, .{ .rows = 513, .width = 128, .stride = 128, .packed_stride = 68, .block = 32, .kind = .index_fp4 }); + try ip.encode(enc, .{ .buffer = q }, .{ .buffer = w }, .{ .buffer = storage }, scratch, .{ .rows = 1, .q_stride = 4096, .weight_stride = 32, .position = 512, .ratio = 1, .keys = 513 }); + enc.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |message| { + std.log.err("index precision qualification: {s}", .{message}); + return error.GpuExecutionFailed; + } + const positions = scratch.positions(); + const ids: [*]const i32 = @ptrCast(@alignCast(positions.buffer.contents() + positions.offset)); + for (0..512) |i| try std.testing.expectEqual(@as(i32, @intCast(i)), ids[i]); + const scores = scratch.scores(); + const values: [*]const f32 = @ptrCast(@alignCast(scores.buffer.contents() + scores.offset)); + try std.testing.expectEqual(@as(f32, 0.01171875), values[0]); + try std.testing.expectEqual(values[0], values[512]); +} diff --git a/zig/src/families/deepseek_v41/kv.zig b/zig/src/families/deepseek_v41/kv.zig new file mode 100644 index 0000000000..c8882ff590 --- /dev/null +++ b/zig/src/families/deepseek_v41/kv.zig @@ -0,0 +1,162 @@ +//! Cache operators accept projected inputs; no source projection or hidden allocations live here. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const View = @import("view.zig").View; +pub const cache = @import("metal_kv.zig"); +pub const Kind = enum(u32) { window_fp8 = 0, main_fp4 = 1, index_fp4 = 2 }; +pub const Pack = extern struct { + rows: u32, + width: u32, + stride: u32, + packed_stride: u32, + block: u32, + kind: Kind, + first: u32 = 0, + ring: u32 = 0, + inplace: u32 = 0, + pub fn validate(p: Pack) !void { + const block: u32 = if (p.kind == .main_fp4) 16 else 32; + const bytes = if (p.kind == .window_fp8) p.width else p.width / 2; + if (p.rows == 0 or p.width == 0 or p.width % block != 0 or p.block != block or + p.stride < p.width or @as(u64, p.packed_stride) < @as(u64, bytes) + p.width / block or + p.inplace > 1 or @as(u64, p.first) + p.rows > std.math.maxInt(u32) or + @as(u64, p.rows -| 1) * p.stride + p.width > std.math.maxInt(u32) or + @as(u64, if (p.ring != 0) p.ring else p.first +| p.rows) * p.packed_stride > + std.math.maxInt(u32) or (p.ring != 0 and p.rows > p.ring)) + return error.InvalidCachePackShape; + } +}; +pub const Pool = extern struct { + rows: u32, + width: u32 = 512, + ratio: u32, + position: u32, + stride: u32, + input_f32: u32, + eps: f32 = 1e-20, + pub fn completed(p: Pool) u32 { + return (p.position + p.rows) / p.ratio - p.position / p.ratio; + } + pub fn validate(p: Pool) !void { + if (p.rows == 0 or p.width == 0 or p.width > 512 or (p.ratio != 1 and p.ratio != 2) or + p.stride < p.width or p.input_f32 != @intFromBool(p.ratio > 1) or + !std.math.isFinite(p.eps) or p.eps <= 0 or + @as(u64, p.position) + p.rows > cache.max_context or + @as(u64, p.rows -| 1) * p.stride + p.width > std.math.maxInt(u32)) + return error.InvalidCompressorShape; + } +}; +pub const Pipelines = struct { + pack: mtl.Pipeline, + pool: mtl.Pipeline, + tail: mtl.Pipeline, + copy: mtl.Pipeline, + publish: mtl.Pipeline, + unpack: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_kv, options); + defer lib.deinit(); + var p: Pipelines = undefined; + p.pack = try mtl.Pipeline.init(device, lib, "dsv41_pack", false); + errdefer p.pack.deinit(); + p.pool = try mtl.Pipeline.init(device, lib, "dsv41_pool", false); + errdefer p.pool.deinit(); + p.tail = try mtl.Pipeline.init(device, lib, "dsv41_pool_tail", false); + errdefer p.tail.deinit(); + p.copy = try mtl.Pipeline.init(device, lib, "dsv41_copy", false); + errdefer p.copy.deinit(); + p.publish = try mtl.Pipeline.init(device, lib, "dsv41_publish_window", false); + errdefer p.publish.deinit(); + p.unpack = try mtl.Pipeline.init(device, lib, "dsv41_unpack_rows", false); + return p; + } + pub fn deinit(p: *Pipelines) void { + p.pack.deinit(); + p.pool.deinit(); + p.tail.deinit(); + p.copy.deinit(); + p.publish.deinit(); + p.unpack.deinit(); + } + pub fn encodePack(p: *const Pipelines, enc: mtl.ComputeEncoder, x: View, storage: View, args: Pack) !void { + try args.validate(); + try x.requireRows(args.rows, args.stride, args.width, 2); + try storage.requireRows(if (args.ring != 0) args.ring else args.first + args.rows, args.packed_stride, args.width / (if (args.kind == .window_fp8) @as(u32, 1) else 2) + + args.width / args.block, 1); + enc.setPipeline(p.pack); + enc.setBuffer(x.buffer, x.offset, 0); + enc.setBuffer(storage.buffer, storage.offset, 1); + enc.setValue(args, 2); + enc.dispatchGroups(mtl.Size.of(args.width / args.block, args.rows, 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + } + pub fn encodeUnpack(p: *const Pipelines, enc: mtl.ComputeEncoder, storage: View, out: View, args: Pack) !void { + try args.validate(); + try out.requireRows(args.rows, args.stride, args.width, 2); + try storage.requireRows(if (args.ring != 0) args.ring else args.first + args.rows, args.packed_stride, args.width / (if (args.kind == .window_fp8) @as(u32, 1) else 2) + + args.width / args.block, 1); + enc.setPipeline(p.unpack); + enc.setBuffer(storage.buffer, storage.offset, 0); + enc.setBuffer(out.buffer, out.offset, 1); + enc.setValue(args, 2); + enc.dispatchThreads(mtl.Size.of(args.width, args.rows, 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + } + /// ratio2 values/gates are FP32; ratio1 values and RMS weight are BF16. Latent output is BF16. + pub fn encodePool(p: *const Pipelines, enc: mtl.ComputeEncoder, values: View, gates: View, partial: View, norm_weight: View, latent: View, args: Pool) !void { + try args.validate(); + try values.requireRows(args.rows, args.stride, args.width, if (args.input_f32 != 0) 4 else 2); + if (args.ratio > 1) { + try gates.requireRows(args.rows, args.stride, args.width, 4); + try partial.require(@as(usize, 2) * args.ratio * args.width * 4, 4); + } + if (args.completed() != 0) { + try norm_weight.require(@as(usize, args.width) * 2, 2); + try latent.requireRows(args.completed(), args.width, args.width, 2); + enc.setPipeline(p.pool); + for ([_]View{ values, gates, partial, norm_weight, latent }, 0..) |v, i| + enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 5); + enc.dispatchGroups(mtl.Size.of(args.completed(), 1, 1), mtl.Size.of(128, 1, 1)); + enc.barrier(); + } + if (args.ratio > 1) { + enc.setPipeline(p.tail); + for ([_]View{ values, gates, partial }, 0..) |v, i| enc.setBuffer(v.buffer, v.offset, i); + enc.setValue(args, 3); + enc.dispatchThreads(mtl.Size.of(args.width, 1, 1), mtl.Size.of(128, 1, 1)); + enc.barrier(); + } + } + pub fn encodeCopy(p: *const Pipelines, enc: mtl.ComputeEncoder, src: View, dst: View, bytes: u32) void { + enc.setPipeline(p.copy); + enc.setBuffer(src.buffer, src.offset, 0); + enc.setBuffer(dst.buffer, dst.offset, 1); + enc.setValue(bytes, 2); + enc.dispatchThreads(mtl.Size.of(bytes, 1, 1), mtl.Size.of(256, 1, 1)); + enc.barrier(); + } + /// Copy only the final window; each destination slot has one writer even for large prefill chunks. + pub fn publishWindow(p: *const Pipelines, enc: mtl.ComputeEncoder, chunk: View, ring: View, rows: u32, position: u32) !void { + if (rows == 0 or @as(u64, position) + rows > cache.max_context) return error.NoCacheRows; + try chunk.requireRows(rows, 528, 528, 1); + try ring.require(128 * 528, 1); + enc.setPipeline(p.publish); + enc.setBuffer(chunk.buffer, chunk.offset, 0); + enc.setBuffer(ring.buffer, ring.offset, 1); + enc.setValue([4]u32{ rows, position, 528, 528 }, 2); + enc.dispatchThreads(mtl.Size.of(528, @min(rows, 128), 1), mtl.Size.of(32, 1, 1)); + enc.barrier(); + } + pub fn snapshot(p: *const Pipelines, enc: mtl.ComputeEncoder, state: *cache.State) !void { + try state.begin(); + p.encodeCopy(enc, .{ .buffer = state.live }, .{ .buffer = state.backup }, @intCast(state.plan.mutableBytes())); + } + pub fn rollback(p: *const Pipelines, enc: mtl.ComputeEncoder, state: *cache.State) !void { + try state.rollback(); + p.encodeCopy(enc, .{ .buffer = state.backup }, .{ .buffer = state.live }, @intCast(state.plan.mutableBytes())); + } +}; diff --git a/zig/src/families/deepseek_v41/kv_test.zig b/zig/src/families/deepseek_v41/kv_test.zig new file mode 100644 index 0000000000..088f36a3fe --- /dev/null +++ b/zig/src/families/deepseek_v41/kv_test.zig @@ -0,0 +1,234 @@ +const std = @import("std"); +const mtl = @import("metal"); +const kv = @import("kv.zig"); +const cache = @import("metal_kv.zig"); +const View = @import("view.zig").View; +fn bf(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @intCast((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); +} +fn view(b: mtl.Buffer) View { + return .{ .buffer = b }; +} +fn submit(cb: mtl.CommandBuffer, enc: mtl.ComputeEncoder) !void { + enc.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); +} + +test "cache owners and one million token storage match the released layout" { + try std.testing.expectEqual(@as(usize, 0), try cache.ownerFor(7, 2)); + try std.testing.expectEqual(@as(usize, 1), try cache.ownerFor(13, 2)); + try std.testing.expectEqual(@as(usize, 2), try cache.ownerFor(19, 2)); + try std.testing.expectEqual(@as(usize, 3), try cache.ownerFor(39, 1)); + try std.testing.expectError(error.NoCompressedCache, cache.ownerFor(0, 0)); + const plan = cache.Plan{ .capacity = 1048576, .layers = 40 }; + try std.testing.expectEqual(@as(usize, 1048576 * 890), plan.packedBytes()); + try std.testing.expectEqual(@as(usize, 2732032), plan.mutableBytes()); + try plan.validate(); + try std.testing.expectError(error.InvalidCacheCapacity, (cache.Plan{ .capacity = 1, .layers = 41 }).validate()); + try std.testing.expectError(error.InvalidCacheLayer, cache.ownerFor(40, 1)); + try std.testing.expectError(error.InvalidCacheOwner, cache.ownerFor(39, 2)); +} + +test "FP4 E2M1 ties, signs, zero scales and BF16 in-place values" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + const x = try device.buffer(64, mtl.ResourceOptions.shared); + defer x.deinit(); + const storage = try device.buffer(36, mtl.ResourceOptions.shared); + defer storage.deinit(); + const values = [_]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.5, -1, -1.5, -2, 0.25, 0.75, 1.25, 1.75 }; + for (values, 0..) |v, i| x.slice(u16, 32)[i] = bf(v); + @memset(x.slice(u16, 32)[16..32], 0); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try pipes.encodePack(enc, view(x), view(storage), .{ .rows = 2, .width = 16, .stride = 16, .packed_stride = 9, .block = 16, .kind = .main_fp4, .inplace = 1 }); + try submit(cb, enc); + const expected = [_]u8{ 0x10, 0x32, 0x54, 0x76, 0xa9, 0xcb, 0x20, 0x42, 0x38 }; + try std.testing.expectEqualSlices(u8, &expected, storage.contents()[0..9]); + try std.testing.expectEqual(@as(u8, 1), storage.contents()[17]); + const rounded = [_]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.5, -1, -1.5, -2, 0, 1, 1, 2 }; + for (rounded, 0..) |v, i| try std.testing.expectEqual(bf(v), x.slice(u16, 32)[i]); +} + +test "compressor carries a partial group across chunks without changing the pooled latent" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + const values = try device.buffer(3 * 512 * 4, mtl.ResourceOptions.shared); + defer values.deinit(); + const gates = try device.buffer(3 * 512 * 4, mtl.ResourceOptions.shared); + defer gates.deinit(); + const partial = try device.buffer(4 * 512 * 4, mtl.ResourceOptions.shared); + defer partial.deinit(); + const one_shot = try device.buffer(4 * 512 * 4, mtl.ResourceOptions.shared); + defer one_shot.deinit(); + const weight = try device.buffer(512 * 2, mtl.ResourceOptions.shared); + defer weight.deinit(); + const out = try device.buffer(2 * 512 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memset(partial.contents()[0 .. 4 * 512 * 4], 0); + @memset(one_shot.contents()[0 .. 4 * 512 * 4], 0); + @memset(weight.slice(u16, 512), bf(1)); + for (0..3 * 512) |i| { + values.slice(f32, 3 * 512)[i] = @as(f32, @floatFromInt((i % 512) % 11)) - 5 + @as(f32, @floatFromInt(i / 512)); + gates.slice(f32, 3 * 512)[i] = @as(f32, @floatFromInt(i % 7)) * 0.25; + } + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try pipes.encodePool(enc, view(values), view(gates), view(partial), view(weight), view(out), .{ .rows = 1, .ratio = 2, .position = 0, .stride = 512, .input_f32 = 1 }); + try pipes.encodePool(enc, view(values).shifted(512 * 4), view(gates).shifted(512 * 4), view(partial), view(weight), view(out), .{ .rows = 2, .ratio = 2, .position = 1, .stride = 512, .input_f32 = 1 }); + try pipes.encodePool(enc, view(values), view(gates), view(one_shot), view(weight), view(out).shifted(512 * 2), .{ .rows = 3, .ratio = 2, .position = 0, .stride = 512, .input_f32 = 1 }); + try submit(cb, enc); + try std.testing.expectEqualSlices(u16, out.slice(u16, 1024)[0..512], out.slice(u16, 1024)[512..1024]); + try std.testing.expectEqualSlices(f32, partial.slice(f32, 2048)[0..512], one_shot.slice(f32, 2048)[0..512]); + try std.testing.expectEqualSlices(f32, partial.slice(f32, 2048)[1024..1536], one_shot.slice(f32, 2048)[1024..1536]); +} + +test "compressor rounds the pooled FP32 latent to BF16 before RMS normalization" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + const values = try device.buffer(4 * 4, mtl.ResourceOptions.shared); + defer values.deinit(); + const gates = try device.buffer(4 * 4, mtl.ResourceOptions.shared); + defer gates.deinit(); + const partial = try device.buffer(4 * 2 * 4, mtl.ResourceOptions.shared); + defer partial.deinit(); + const weight = try device.buffer(2 * 2, mtl.ResourceOptions.shared); + defer weight.deinit(); + const out = try device.buffer(2 * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + @memcpy(values.slice(f32, 4), &[_]f32{ 1.00390625, 0.5, 1.00390625, 0.5 }); + @memset(gates.slice(f32, 4), 0); + @memset(partial.contents()[0..32], 0); + @memset(weight.slice(u16, 2), bf(1)); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try pipes.encodePool(enc, view(values), view(gates), view(partial), view(weight), view(out), .{ .rows = 2, .width = 2, .ratio = 2, .position = 0, .stride = 2, .input_f32 = 1 }); + try submit(cb, enc); + // 1.00390625 ties to 1 in BF16; normalizing the unrounded value changes the second result. + try std.testing.expectEqual(bf(1.265625), out.slice(u16, 2)[0]); + try std.testing.expectEqual(bf(0.6328125), out.slice(u16, 2)[1]); +} + +test "ring publish handles chunks beyond the window and transaction rollback restores every mutable byte" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + var state = try cache.State.init(device, .{ .capacity = 1024, .layers = 40 }); + defer state.deinit(); + const chunk = try device.buffer(131 * 528, mtl.ResourceOptions.shared); + defer chunk.deinit(); + for (0..131) |r| @memset(chunk.contents()[r * 528 .. (r + 1) * 528], @intCast(r)); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + try pipes.publishWindow(enc, view(chunk), try state.ring(0), 131, 0); + try state.append(131); + try pipes.snapshot(enc, &state); + try pipes.publishWindow(enc, view(chunk), try state.ring(0), 131, 131); + try state.append(131); + try pipes.rollback(enc, &state); + try submit(cb, enc); + try std.testing.expectEqual(@as(u32, 131), state.position); + try std.testing.expectEqualSlices(u8, state.backup.contents()[0..state.plan.mutableBytes()], state.live.contents()[0..state.plan.mutableBytes()]); + for (3..131) |r| { + const slot = r % 128; + for (state.live.contents()[slot * 528 .. (slot + 1) * 528]) |byte| + try std.testing.expectEqual(@as(u8, @intCast(r)), byte); + } + const cb2 = queue.commandBuffer(); + const enc2 = cb2.compute(.serial); + try pipes.snapshot(enc2, &state); + try pipes.publishWindow(enc2, view(chunk), try state.ring(0), 131, 131); + try state.append(131); + try state.commit(); + try submit(cb2, enc2); + try std.testing.expectEqual(@as(u32, 262), state.position); + try std.testing.expect(state.saved_position == null); +} + +test "window FP8 quantization includes the RoPE tail and index FP4 uses E8M0 scale bytes" { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var pipes = try kv.Pipelines.load(device); + defer pipes.deinit(); + const x = try device.buffer(64 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + const storage = try device.buffer(66, mtl.ResourceOptions.shared); + defer storage.deinit(); + const decoded = try device.buffer(64 * 2, mtl.ResourceOptions.shared); + defer decoded.deinit(); + @memset(x.slice(u16, 64), 0); + @memset(x.slice(u16, 64)[32..64], bf(0.3)); + x.slice(u16, 64)[63] = bf(448); + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + const args = kv.Pack{ .rows = 1, .width = 64, .stride = 64, .packed_stride = 66, .block = 32, .kind = .window_fp8, .inplace = 1 }; + try pipes.encodePack(enc, view(x), view(storage), args); + try pipes.encodeUnpack(enc, view(storage), view(decoded), args); + try submit(cb, enc); + try std.testing.expectEqual(@as(u8, 127), storage.contents()[65]); + try std.testing.expectEqual(@as(u8, 126), storage.contents()[63]); + try std.testing.expectEqual(bf(0.3125), decoded.slice(u16, 64)[62]); + try std.testing.expectEqualSlices(u16, x.slice(u16, 64), decoded.slice(u16, 64)); + @memset(x.slice(u16, 64)[0..32], bf(6)); + const cb2 = queue.commandBuffer(); + const enc2 = cb2.compute(.serial); + try pipes.encodePack(enc2, view(x), view(storage), .{ .rows = 1, .width = 32, .stride = 32, .packed_stride = 17, .block = 32, .kind = .index_fp4, .inplace = 1 }); + try submit(cb2, enc2); + for (storage.contents()[0..16]) |byte| try std.testing.expectEqual(@as(u8, 0x77), byte); + try std.testing.expectEqual(@as(u8, 127), storage.contents()[16]); +} + +test { + _ = @import("csa2_test.zig"); +} + +test "cache pack rejects undersized borrowed input and output before encoding" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const small = try device.buffer(16, mtl.ResourceOptions.shared); + defer small.deinit(); + const input = try device.buffer(64, mtl.ResourceOptions.shared); + defer input.deinit(); + const args = kv.Pack{ .rows = 1, .width = 32, .stride = 32, .packed_stride = 33, .block = 32, .kind = .window_fp8 }; + const pipelines: *const kv.Pipelines = undefined; + try std.testing.expectError(error.InvalidTensorView, pipelines.encodePack(undefined, view(small), view(input), args)); + try std.testing.expectError(error.InvalidTensorView, pipelines.encodePack(undefined, view(input), view(small), args)); + try std.testing.expectError(error.InvalidTensorView, pipelines.encodePack(undefined, view(input).shifted(1), view(input), args)); +} + +test "cache shapes reject overflow and nonfinite compressor epsilon" { + try std.testing.expectError(error.InvalidCachePackShape, (kv.Pack{ .rows = 2, .width = 32, .stride = 32, .packed_stride = 33, .block = 32, .kind = .window_fp8, .first = std.math.maxInt(u32) }).validate()); + try std.testing.expectError(error.InvalidCompressorShape, (kv.Pool{ .rows = 1, .ratio = 1, .position = 0, .stride = 512, .input_f32 = 0, .eps = std.math.nan(f32) }).validate()); +} diff --git a/zig/src/families/deepseek_v41/metal_kv.zig b/zig/src/families/deepseek_v41/metal_kv.zig new file mode 100644 index 0000000000..f2b70175a7 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal_kv.zig @@ -0,0 +1,111 @@ +//! Shared cache ownership and exact transaction snapshots; appended rows are cursor-delimited. +const std = @import("std"); +const mtl = @import("metal"); +const View = @import("view.zig").View; +pub const owners = [_]u32{ 2, 8, 14, 20 }; +pub const ratios = [_]u32{ 2, 2, 2, 1 }; +pub const index_sources = [_]u32{ 2, 8, 14, 20, 24, 28, 32, 36 }; +pub const main_stride = 288; +pub const index_stride = 68; +pub const ring_stride = 528; +pub const window = 128; +pub const max_context = 1048576; + +pub fn ownerFor(layer: u32, ratio: u32) !usize { + if (ratio == 0) return error.NoCompressedCache; + if (layer >= 40) return error.InvalidCacheLayer; + if (layer < 2 or ratio != (if (layer < 20) @as(u32, 2) else 1)) return error.InvalidCacheOwner; + var found: ?usize = null; + for (owners, 0..) |owner, i| { + if (owner <= layer and ratios[i] == ratio) found = i; + } + return found orelse error.InvalidCacheOwner; +} +pub fn isIndexSource(layer: u32) bool { + return std.mem.indexOfScalar(u32, &index_sources, layer) != null; +} +pub const Source = struct { owner: u32, ratio: u32, rows: u32, main: View, index: View, partial: View }; +pub const Plan = struct { + capacity: u32, + layers: u32, + pub fn validate(p: Plan) !void { + if (p.capacity == 0 or p.capacity > max_context or p.layers != 40) + return error.InvalidCacheCapacity; + } + pub fn packedBytes(p: Plan) usize { + var bytes: usize = 0; + for (ratios) |r| bytes += @as(usize, p.capacity / r) * (main_stride + index_stride); + return bytes; + } + pub fn mutableBytes(p: Plan) usize { + return @as(usize, p.layers) * window * ring_stride + 7 * 2 * 512 * 4; + } + pub fn residentBytes(p: Plan) usize { + return p.packedBytes() + 2 * p.mutableBytes(); + } +}; +pub const State = struct { + plan: Plan, + storage: mtl.Buffer, + live: mtl.Buffer, + backup: mtl.Buffer, + position: u32 = 0, + saved_position: ?u32 = null, + + pub fn init(device: mtl.Device, plan: Plan) !State { + try plan.validate(); + const storage = try device.buffer(plan.packedBytes(), mtl.ResourceOptions.shared); + errdefer storage.deinit(); + const live = try device.buffer(plan.mutableBytes(), mtl.ResourceOptions.shared); + errdefer live.deinit(); + const backup = try device.buffer(plan.mutableBytes(), mtl.ResourceOptions.shared); + @memset(live.contents()[0..plan.mutableBytes()], 0); + return .{ .plan = plan, .storage = storage, .live = live, .backup = backup }; + } + pub fn deinit(s: *State) void { + s.storage.deinit(); + s.live.deinit(); + s.backup.deinit(); + } + pub fn residentBytes(s: *const State) usize { + return s.plan.residentBytes(); + } + pub fn registerResidency(s: *const State, set: mtl.ResidencySet) void { + set.add(s.storage); + set.add(s.live); + set.add(s.backup); + } + pub fn ring(s: *const State, layer: u32) !View { + if (layer >= s.plan.layers) return error.InvalidCacheLayer; + return .{ .buffer = s.live, .offset = @as(usize, layer) * window * ring_stride }; + } + pub fn source(s: *const State, layer: u32, ratio: u32) !Source { + if (layer >= s.plan.layers) return error.InvalidCacheLayer; + if (s.position > s.plan.capacity) return error.CacheCapacityExceeded; + const id = try ownerFor(layer, ratio); + var packed_off: usize = 0; + var partial_off: usize = @as(usize, s.plan.layers) * window * ring_stride; + for (ratios[0..id]) |r| { + packed_off += @as(usize, s.plan.capacity / r) * (main_stride + index_stride); + partial_off += 2 * r * 512 * 4; + } + return .{ .owner = owners[id], .ratio = ratio, .rows = s.position / ratio, .main = .{ .buffer = s.storage, .offset = packed_off }, .index = .{ .buffer = s.storage, .offset = packed_off + @as(usize, s.plan.capacity / ratio) * main_stride }, .partial = .{ .buffer = s.live, .offset = partial_off } }; + } + pub fn append(s: *State, rows: u32) !void { + if (s.position > s.plan.capacity or rows == 0 or rows > s.plan.capacity - s.position) + return error.CacheCapacityExceeded; + s.position += rows; + } + pub fn begin(s: *State) !void { + if (s.saved_position != null) return error.CacheTransactionActive; + s.saved_position = s.position; + } + pub fn commit(s: *State) !void { + if (s.saved_position == null) return error.NoCacheTransaction; + s.saved_position = null; + } + pub fn rollback(s: *State) !void { + s.position = s.saved_position orelse return error.NoCacheTransaction; + s.saved_position = null; + } +}; diff --git a/zig/src/families/deepseek_v41/metal_model.zig b/zig/src/families/deepseek_v41/metal_model.zig new file mode 100644 index 0000000000..0178b0660f --- /dev/null +++ b/zig/src/families/deepseek_v41/metal_model.zig @@ -0,0 +1,363 @@ +//! Plain V4.1 execution over resident EXL3 weights and source-owned packed caches. +const std = @import("std"); +const mtl = @import("metal"); +const checkpoint = @import("../../core/checkpoint_metal.zig"); +const exl3 = @import("../../core/exl3_projection.zig"); +const config = @import("config.zig"); +const load = @import("resident_load.zig"); +const weights = @import("weights.zig"); +const work = @import("workspace.zig"); +pub const attention = @import("attention.zig"); +const kv = @import("kv.zig"); +pub const ops = @import("ops.zig"); +pub const mhc = @import("mhc.zig"); +const route = @import("route.zig"); +const engram = @import("engram.zig"); +const View = @import("view.zig").View; +const gpu_counters = @import("profiler.zig"); +const expert_groups = @import("../../core/expert_groups.zig"); + +pub const prefill_step = 256; +pub const Logits = enum { none, last, all }; +pub const Options = struct { + dir: []const u8, + table_dir: ?[]const u8 = null, + token_map: []const u8, + context: u32 = 0, + chunk_rows: u32 = prefill_step, +}; + +const Pipelines = struct { + projection: exl3.Pipelines, + ops: ops.Pipelines, + mhc: mhc.Pipelines, + route: route.Pipelines, + engram: engram.Pipelines, + attention: attention.Pipelines, + groups: expert_groups.Pipelines, + + fn init(device: mtl.Device, capacity: u32, c: config.Text) !Pipelines { + var ep = try exl3.Pipelines.load(device); + errdefer ep.deinit(); + var op = try ops.Pipelines.load(device); + errdefer op.deinit(); + var hp = try mhc.Pipelines.load(device); + errdefer hp.deinit(); + var rp = try route.Pipelines.load(device); + errdefer rp.deinit(); + var gp = try engram.Pipelines.load(device); + errdefer gp.deinit(); + var groups = try expert_groups.Pipelines.load(device); + errdefer groups.deinit(); + const ap = try attention.Pipelines.load(device, capacity, c); + return .{ .projection = ep, .ops = op, .mhc = hp, .route = rp, .engram = gp, .attention = ap, .groups = groups }; + } + + fn deinit(p: *Pipelines) void { + p.attention.deinit(); + p.engram.deinit(); + p.route.deinit(); + p.groups.deinit(); + p.mhc.deinit(); + p.ops.deinit(); + p.projection.deinit(); + } +}; + +/// Serialized unique owner; submissions finish before return/release, cancellation is between bounded chunks. +pub const Model = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + residency: mtl.ResidencySet, + checkpoint: checkpoint.Checkpoint, + weights: weights.Model, + cfg: config.Config, + cfg_arena: std.heap.ArenaAllocator, + cache: kv.cache.State, + workspace: work.Workspace, + attention_workspace: attention.Workspace, + pipelines: Pipelines, + admitted_bytes: usize, + last_gpu_seconds: f64 = 0, + logits_rows: usize = 0, + valid: bool = true, + profile: ?*gpu_counters.Recorder = null, + grouped_prefill: bool = true, + observer: ?Observer = null, + pub const Observer = struct { + ctx: *anyopaque, + encode: *const fn (*anyopaque, mtl.ComputeEncoder, u32, gpu_counters.Stage, *const Model) void, + }; + + pub fn open(gpa: std.mem.Allocator, io: std.Io, o: Options) !Model { + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + const c = try config.Config.load(arena.allocator(), io, o.dir); + const capacity = try c.context(o.context); + const table_dir = o.table_dir orelse c.engram_table_dir orelse return error.MissingEngramTables; + var plan = try load.Plan.open(gpa, io, o.dir, table_dir); + defer plan.deinit(); + try weights.validate(&plan); + const map_blob = try std.Io.Dir.cwd().readFileAlloc(io, o.token_map, arena.allocator(), .limited(2 << 20)); + const map = try engram.hash.TokenMap.parse(map_blob); + const wp = try work.Plan.init(o.chunk_rows); + const kp = kv.cache.Plan{ .capacity = capacity, .layers = 40 }; + const workspace_bytes = try std.math.add(usize, wp.bytes, attention.Workspace.bytes(o.chunk_rows)); + const cache_bytes = try std.math.add(usize, kp.residentBytes(), weights.sizeofDescriptors()); + const execution_bytes = try std.math.add(usize, workspace_bytes, cache_bytes); + const rotary_bytes = attention.Pipelines.residentBytes(capacity, c.text_config); + const needed = try std.math.add(usize, plan.bytes, try std.math.add(usize, execution_bytes, rotary_bytes)); + const device = try mtl.Device.init(); + errdefer device.deinit(); + const margin: u64 = 8 << 30; + const available = device.maxWorkingSet() -| device.allocated() -| margin; + if (needed > available) { + std.log.err("full residency needs {d} bytes plus {d} margin; Metal admission allows {d}", .{ needed, margin, available }); + return error.InsufficientResidentMemory; + } + const max_buffer = mtl.objc.msg(u64, device.id, "maxBufferLength", .{}); + for (plan.shards.items) |shard| if (shard.end - shard.begin > max_buffer) return error.ShardExceedsMetalBufferLimit; + std.log.info("full resident admission: weights={d}, total={d}, context={d}, headroom={d}", .{ plan.bytes, needed, capacity, available - needed }); + var pipelines = try Pipelines.init(device, capacity, c.text_config); + errdefer pipelines.deinit(); + const queue = try device.queue(); + errdefer queue.deinit(); + const residency = try device.residencySet(plan.shards.items.len + 128); + errdefer residency.deinit(); + var cp = try plan.materialize(gpa, device, residency, queue); + errdefer cp.deinit(); + var bound = try weights.bind(device, &cp); + errdefer bound.deinit(); + var cache = try kv.cache.State.init(device, kp); + errdefer cache.deinit(); + var workspace = try work.Workspace.init(device, wp, map.bytes); + errdefer workspace.deinit(); + var aw = try attention.Workspace.init(device, o.chunk_rows); + errdefer aw.deinit(); + bound.registerResidency(residency); + cache.registerResidency(residency); + residency.add(workspace.buffer); + aw.registerResidency(residency); + pipelines.attention.registerResidency(residency); + residency.commit(); + residency.requestResidency(); + queue.addResidencySet(residency); + return .{ .gpa = gpa, .device = device, .queue = queue, .residency = residency, .checkpoint = cp, .weights = bound, .cfg = c, .cfg_arena = arena, .cache = cache, .workspace = workspace, .attention_workspace = aw, .pipelines = pipelines, .admitted_bytes = needed }; + } + + pub fn deinit(m: *Model) void { + m.queue.removeResidencySet(m.residency); + m.residency.deinit(); + m.attention_workspace.deinit(); + m.workspace.deinit(); + m.cache.deinit(); + m.weights.deinit(); + m.checkpoint.deinit(); + m.pipelines.deinit(); + m.queue.deinit(); + m.device.deinit(); + m.cfg_arena.deinit(); + } + + pub fn reset(m: *Model) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + m.valid = false; + m.logits_rows = 0; + m.cache.position = 0; + m.attention_workspace.selection_layer = null; + m.attention_workspace.selection_rows = 0; + m.attention_workspace.candidate_rows = 0; + m.last_gpu_seconds = 0; + m.cache.saved_position = null; + @memset(m.cache.live.contents()[0..m.cache.plan.mutableBytes()], 0); + const cb = m.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + engram.encodeReset(enc, &m.pipelines.engram, m.workspace.view(.history)); + enc.end(); + cb.commit(); + cb.wait(); + try completed(cb); + m.valid = true; + m.logits_rows = 0; + } + + /// One command buffer per chunk; no operator waits or reads routes back to the CPU. + pub fn forward(m: *Model, ids: []const u32, logits: Logits) ![]const f32 { + if (!m.valid) return error.InvalidModelState; + if (ids.len == 0 or ids.len > m.workspace.plan.rows) return error.InvalidPrefillChunk; + if (m.cache.position > m.cache.plan.capacity or + ids.len > m.cache.plan.capacity - m.cache.position) return error.ContextFull; + for (ids) |id| if (id >= 129280) return error.UnsupportedToken; + m.logits_rows = 0; + @memcpy(m.workspace.tokenIds(ids.len), ids); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const rows: u32 = @intCast(ids.len); + const e = m.queue.commandBufferUnretained(); + var enc = e.compute(.serial); + for (m.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + var encoding = true; + var record_started = false; + errdefer { + if (m.profile) |recorder| { + if (record_started) { + if (encoding) recorder.cancelEncoding(&enc) catch {} else recorder.discardAfterCompletion() catch {}; + } else if (encoding) enc.end(); + } else if (encoding) enc.end(); + } + if (m.profile) |recorder| { + try recorder.beginCommand(e); + record_started = true; + } + try m.profileSpan(&enc, .other); + const p = &m.pipelines; + const ws = &m.workspace; + for (engram.hash.layers, 0..) |meta, index| { + try engram.encodeHash(enc, &p.engram, meta, ws.view(.ids), ws.view(.token_map), null, ws.view(.history), ws.view(if (index == 0) .engram_ids_0 else .engram_ids_1), rows); + } + try engram.encodeCommit(enc, &p.engram, ws.view(.ids), ws.view(.token_map), null, ws.view(.history), rows); + ops.encodeEmbedding(enc, &p.ops, ws.view(.ids), tensor(m.weights.embed), ws.view(.collapsed), .{ .rows = rows, .width = 5120, .table_stride = 5120, .y_stride = 5120, .copies = 1 }); + enc.barrier(); + mhc.encodeExpand(enc, &p.mhc, ws.view(.collapsed), ws.view(.stream), ws.view(.previous_pre), .{ .rows = rows, .width = 5120, .x_stride = 5120, .stream_stride = 20480 }); + enc.barrier(); + for (&m.weights.layers, 0..) |*layer, index| { + if (layer.engram) |g| { + try m.profileSpan(&enc, .engram); + try engram.encodeGather(enc, &p.engram, .{ .weight = tensor(g.table), .scale = tensor(g.scales), .rows = @intCast(g.table.shape[0]) }, ws.view(if (index == 1) .engram_ids_0 else .engram_ids_1), ws.view(.engram_gather), rows, 6144); + try m.project(enc, g.wkv, ws.view(.engram_gather), ws.view(.engram_kv), rows, 6144, 25600, .bf16); + try engram.encodeGate(enc, &p.engram, .{ .q_weight = tensor(g.q_weight), .k_weight = tensor(g.k_weight), .dim = 5120, .copies = 4, .hidden_stride = 20480, .kv_stride = 25600, .output_stride = 20480, .eps = m.cfg.text_config.rms_norm_eps }, ws.view(.stream), ws.view(.engram_kv), null, ws.view(.stream), rows); + } + try m.profileSpan(&enc, .mhc); + m.mixes(enc, layer.hc_attn, ws.view(.attention_pre), rows); + mhc.encodeCollapse(enc, &p.mhc, ws.view(.stream), ws.view(.previous_pre), ws.view(.collapsed), .{ .rows = rows, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + enc.barrier(); + m.normalize(enc, ws.view(.collapsed), layer.attn_norm, ws.view(.normalized), rows); + try m.profileSpan(&enc, .attention); + try attention.encode(&p.attention, enc, &layer.attn, &m.cache, &m.attention_workspace, ws.view(.normalized), ws.view(.branch), &p.projection, ws.projection(), @intCast(index), rows, m.cfg.text_config); + enc.barrier(); + try m.profileSpan(&enc, .mhc); + m.post(enc, rows); + if (m.observer) |observer| observer.encode(observer.ctx, enc, @intCast(index), .attention, m); + m.mixes(enc, layer.hc_ffn, ws.view(.previous_pre), rows); + mhc.encodeCollapse(enc, &p.mhc, ws.view(.stream), ws.view(.attention_pre), ws.view(.collapsed), .{ .rows = rows, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + enc.barrier(); + m.normalize(enc, ws.view(.collapsed), layer.ffn_norm, ws.view(.normalized), rows); + try m.profileSpan(&enc, .moe); + try m.moe(enc, &layer.moe, rows); + try m.profileSpan(&enc, .mhc); + m.post(enc, rows); + if (m.observer) |observer| observer.encode(observer.ctx, enc, @intCast(index), .moe, m); + } + if (logits != .none) { + try m.profileSpan(&enc, .head); + const output_rows: u32 = if (logits == .all) rows else 1; + const first: usize = if (logits == .all) 0 else rows - 1; + mhc.encodeCollapse(enc, &p.mhc, ws.view(.stream).shifted(first * 20480 * 2), ws.view(.previous_pre).shifted(first * 4 * 4), ws.view(.collapsed), .{ .rows = output_rows, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + enc.barrier(); + m.normalize(enc, ws.view(.collapsed), m.weights.norm, ws.view(.normalized), output_rows); + try m.project(enc, m.weights.head, ws.view(.normalized), ws.view(.logits), output_rows, 5120, 129280, .f32); + } + if (m.profile) |recorder| { + if (recorder.activeSpan()) try recorder.endSpan(&enc); + try recorder.finishEncoding(&enc); + } else enc.end(); + encoding = false; + m.valid = false; + e.commit(); + e.wait(); + try completed(e); + if (m.profile) |recorder| _ = try recorder.readAfterCompletion(); + try m.cache.append(rows); + m.last_gpu_seconds = e.gpuSeconds(); + m.valid = true; + m.logits_rows = if (logits == .all) rows else if (logits == .last) 1 else 0; + return ws.logits(m.logits_rows); + } + + fn profileSpan(m: *Model, encoder: *mtl.ComputeEncoder, stage: gpu_counters.Stage) !void { + const recorder = m.profile orelse return; + if (!recorder.enabled()) return; + if (recorder.activeSpan()) try recorder.endSpan(encoder); + try recorder.beginSpan(encoder, stage); + for (m.checkpoint.shards.items) |shard| encoder.useResource(shard.buffer, mtl.ResourceUsage.read); + } + + fn normalize(m: *Model, enc: mtl.ComputeEncoder, x: View, w: checkpoint.Tensor, y: View, rows: u32) void { + ops.encodeRms(enc, &m.pipelines.ops, x, tensor(w), y, .{ .rows = rows, .width = 5120, .x_stride = 5120, .y_stride = 5120, .eps = m.cfg.text_config.rms_norm_eps }); + enc.barrier(); + } + + fn mixes(m: *Model, enc: mtl.ComputeEncoder, w: weights.Hc, pre: View, rows: u32) void { + const ws = &m.workspace; + mhc.encodeRms(enc, &m.pipelines.mhc, ws.view(.stream), ws.view(.inv), .{ .rows = rows, .width = 5120, .stream_stride = 20480, .eps = m.cfg.text_config.rms_norm_eps }); + ops.encodeDense(enc, &m.pipelines.ops, ws.view(.stream), tensor(w.projection), ws.view(.mixes), .{ .rows = rows, .k = 20480, .n = 24, .x_stride = 20480, .w_stride = 20480, .y_stride = 24, .w_dtype = .f32, .y_dtype = .f32 }); + enc.barrier(); + mhc.encodeSplit(enc, &m.pipelines.mhc, ws.view(.mixes), ws.view(.inv), tensor(w.scale), tensor(w.base), pre, ws.view(.post), ws.view(.comb), .{ .rows = rows, .eps = m.cfg.text_config.hc_eps }); + enc.barrier(); + } + + fn post(m: *Model, enc: mtl.ComputeEncoder, rows: u32) void { + const ws = &m.workspace; + mhc.encodePost(enc, &m.pipelines.mhc, ws.view(.branch), ws.view(.stream), ws.view(.post), ws.view(.comb), ws.view(.stream), .{ .rows = rows, .width = 5120, .branch_stride = 5120, .residual_stride = 20480, .y_stride = 20480 }); + enc.barrier(); + } + + fn project(m: *Model, enc: mtl.ComputeEncoder, matrix: exl3.Matrix, x: View, y: View, rows: u32, xs: usize, ys: usize, output: exl3.Output) !void { + try exl3.encode(&m.pipelines.projection, enc, matrix, convert(x), convert(y), m.workspace.projection(), .{ .count = rows, .x_stride = xs, .y_stride = ys, .output = output }); + enc.barrier(); + } + + fn moe(m: *Model, enc: mtl.ComputeEncoder, w: *const weights.Moe, rows: u32) !void { + const ws = &m.workspace; + const p = &m.pipelines; + ops.encodeDense(enc, &p.ops, ws.view(.normalized), tensor(w.gate), ws.view(.router_logits), .{ .rows = rows, .k = 5120, .n = 384, .x_stride = 5120, .w_stride = 5120, .y_stride = 384, .w_dtype = .f16, .y_dtype = .f32 }); + enc.barrier(); + route.encode(enc, &p.route, ws.view(.router_logits), tensor(w.bias), ws.view(.route_ids), ws.view(.route_weights), .{ .rows = rows, .bias_dtype = .f16 }); + enc.barrier(); + const grouped: ?expert_groups.Scratch = if (m.grouped_prefill and rows >= 8) try ws.grouping(rows) else null; + var routed_rows = exl3.Rows{ .count = @as(usize, rows) * 6, .x_stride = 5120, .y_stride = 2304, .input_rows = rows, .output_rows = @as(usize, rows) * 6, .routes = .{ .ids = convert(ws.view(.route_ids)), .repeat = 6 } }; + if (grouped) |g| { + try expert_groups.encode(enc, &p.groups, convert(ws.view(.route_ids)), g, 6); + routed_rows.count = g.plan.padded; + routed_rows.input_ids = g.view(.token_ids); + routed_rows.output_ids = g.view(.slot_ids); + routed_rows.routes = .{ .ids = g.view(.expert_ids) }; + routed_rows.grouped = true; + } + for (0..2) |i| { + try exl3.encodeIndexed(&p.projection, enc, .{ .descriptors = .{ .buffer = w.descriptors[i] }, .count = 384, .k = 5120, .n = 2304 }, convert(ws.view(.normalized)), convert(ws.view(if (i == 0) .gate else .up)), ws.projection(), routed_rows); + enc.barrier(); + } + ops.encodeSwiGLU(enc, &p.ops, ws.view(.gate), ws.view(.up), ws.view(.route_weights), ws.view(.activation), .{ .rows = rows * 6, .width = 2304, .input_stride = 2304, .output_stride = 2304, .weighted = true }); + enc.barrier(); + routed_rows.x_stride = 2304; + routed_rows.y_stride = 5120; + routed_rows.routes.?.repeat = 1; + routed_rows.input_rows = @as(usize, rows) * 6; + if (grouped) |g| routed_rows.input_ids = g.view(.slot_ids); + try exl3.encodeIndexed(&p.projection, enc, .{ .descriptors = .{ .buffer = w.descriptors[2] }, .count = 384, .k = 2304, .n = 5120 }, convert(ws.view(.activation)), convert(ws.view(.routed)), ws.projection(), routed_rows); + enc.barrier(); + try m.project(enc, w.shared[0], ws.view(.normalized), ws.view(.shared_gate), rows, 5120, 2304, .bf16); + try m.project(enc, w.shared[1], ws.view(.normalized), ws.view(.shared_up), rows, 5120, 2304, .bf16); + ops.encodeSwiGLU(enc, &p.ops, ws.view(.shared_gate), ws.view(.shared_up), ws.view(.route_weights), ws.view(.shared_activation), .{ .rows = rows, .width = 2304, .input_stride = 2304, .output_stride = 2304 }); + enc.barrier(); + try m.project(enc, w.shared[2], ws.view(.shared_activation), ws.view(.shared_down), rows, 2304, 5120, .bf16); + ops.encodeCombine(enc, &p.ops, ws.view(.routed), ws.view(.shared_down), ws.view(.route_ids), ws.view(.branch), .{ .rows = rows, .width = 5120, .routed_row_stride = 6 * 5120, .slot_stride = 5120, .shared_stride = 5120, .y_stride = 5120 }); + enc.barrier(); + } +}; + +fn tensor(t: checkpoint.Tensor) View { + return .{ .buffer = t.buffer, .offset = t.offset }; +} +fn convert(v: View) exl3.View { + return .{ .buffer = v.buffer, .offset = v.offset }; +} +fn completed(cb: mtl.CommandBuffer) !void { + if (cb.failure()) |message| { + std.log.err("DeepSeek V4.1 GPU execution failed: {s}", .{message}); + return error.GpuExecutionFailed; + } +} diff --git a/zig/src/families/deepseek_v41/mhc.zig b/zig/src/families/deepseek_v41/mhc.zig new file mode 100644 index 0000000000..3ac499423f --- /dev/null +++ b/zig/src/families/deepseek_v41/mhc.zig @@ -0,0 +1,95 @@ +//! Four-copy mHC; callers validate shapes/storage and own previous-pre coefficients. +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const View = @import("view.zig").View; +pub const Pipelines = struct { + rms: mtl.Pipeline, + split: mtl.Pipeline, + collapse: mtl.Pipeline, + post: mtl.Pipeline, + expand: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_mhc, options); + defer lib.deinit(); + const rms = try mtl.Pipeline.init(device, lib, "dsv41_hc_rms", false); + errdefer rms.deinit(); + const split = try mtl.Pipeline.init(device, lib, "dsv41_hc_split", false); + errdefer split.deinit(); + const collapse = try mtl.Pipeline.init(device, lib, "dsv41_hc_collapse", false); + errdefer collapse.deinit(); + const post = try mtl.Pipeline.init(device, lib, "dsv41_hc_post", false); + errdefer post.deinit(); + const expand = try mtl.Pipeline.init(device, lib, "dsv41_hc_expand", false); + return .{ .rms = rms, .split = split, .collapse = collapse, .post = post, .expand = expand }; + } + pub fn deinit(p: *Pipelines) void { + p.rms.deinit(); + p.split.deinit(); + p.collapse.deinit(); + p.post.deinit(); + p.expand.deinit(); + } +}; +fn bind(e: mtl.ComputeEncoder, refs: anytype) void { + inline for (refs, 0..) |r, i| e.setBuffer(r.buffer, r.offset, i); +} +pub const Rms = extern struct { rows: u32, width: u32, stream_stride: u32, eps: f32 }; +/// BF16 flattened [rows,4*width] -> FP32 inv[rows]; fixed CUDA single-row tree, not every batched CUDA grouping. +pub fn encodeRms(e: mtl.ComputeEncoder, p: *const Pipelines, stream: View, inv: View, d: Rms) void { + e.setPipeline(p.rms); + bind(e, .{ stream, inv }); + e.setValue(d, 2); + e.dispatchGroups(mtl.Size.of(d.rows, 1, 1), mtl.Size.of(512, 1, 1)); +} +pub const Split = extern struct { + rows: u32, + mix_stride: u32 = 24, + coefficient_stride: u32 = 4, + comb_stride: u32 = 16, + eps: f32 = 1e-6, +}; +/// FP32 mixes[rows,24], inv, scale[3], base[24] -> pre/post[rows,4], comb[rows,4,4]; 20 iterations. +pub fn encodeSplit(e: mtl.ComputeEncoder, p: *const Pipelines, mixes: View, inv: View, scale: View, base: View, pre: View, post: View, comb: View, d: Split) void { + e.setPipeline(p.split); + bind(e, .{ mixes, inv, scale, base, pre, post, comb }); + e.setValue(d, 7); + e.dispatchThreads(mtl.Size.of(d.rows, 1, 1), mtl.Size.of(32, 1, 1)); +} +pub const Collapse = extern struct { rows: u32, width: u32, stream_stride: u32, y_stride: u32, pre_stride: u32 = 4 }; +/// Collapse BF16 [rows,4,width] with previous sublayer pre[rows,4]; head uses final FFN pre, initially one-hot. +pub fn encodeCollapse(e: mtl.ComputeEncoder, p: *const Pipelines, stream: View, previous_pre: View, y: View, d: Collapse) void { + e.setPipeline(p.collapse); + bind(e, .{ stream, previous_pre, y }); + e.setValue(d, 3); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} +pub const Post = extern struct { + rows: u32, + width: u32, + branch_stride: u32, + residual_stride: u32, + y_stride: u32, + coefficient_stride: u32 = 4, + comb_stride: u32 = 16, +}; +/// y[k,d]=post[k]*branch[d]+sum_j comb[j,k]*residual[j,d]; BF16 inputs/output, FP32 coefficients/math. + +// Element pitches; exact residual alias is safe, but output must not alias branch/post/comb. +pub fn encodePost(e: mtl.ComputeEncoder, p: *const Pipelines, branch: View, residual: View, post: View, comb: View, y: View, d: Post) void { + e.setPipeline(p.post); + bind(e, .{ branch, residual, post, comb, y }); + e.setValue(d, 5); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} +pub const Expand = extern struct { rows: u32, width: u32, x_stride: u32, stream_stride: u32, pre_stride: u32 = 4 }; +/// BF16 x[rows,width] copied into four streams and FP32 initial pre[rows,4] set to [1,0,0,0]. +pub fn encodeExpand(e: mtl.ComputeEncoder, p: *const Pipelines, x: View, stream: View, pre: View, d: Expand) void { + e.setPipeline(p.expand); + bind(e, .{ x, stream, pre }); + e.setValue(d, 3); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} diff --git a/zig/src/families/deepseek_v41/mhc_test.zig b/zig/src/families/deepseek_v41/mhc_test.zig new file mode 100644 index 0000000000..ee0d888e28 --- /dev/null +++ b/zig/src/families/deepseek_v41/mhc_test.zig @@ -0,0 +1,209 @@ +//! Consumer-visible mHC coefficient, reduction and BF16 boundary regressions. +const std = @import("std"); +const mtl = @import("metal"); +const mhc = @import("mhc.zig"); +const support = @import("test_support.zig"); +const bf = support.bf; +const fp = support.fp; +const view = support.view; +const Buffers = support.Buffers; +const done = support.done; +const close = support.close; +fn splitReference(mixes: []const f32, inv: f64, scales: []const f32, base: []const f32) [24]f64 { + var out: [24]f64 = undefined; + for (0..4) |j| { + out[j] = 1 / (1 + @exp(-(@as(f64, mixes[j]) * inv * scales[0] + base[j]))) + 1e-6; + out[4 + j] = 2 / (1 + @exp(-(@as(f64, mixes[4 + j]) * inv * scales[1] + base[4 + j]))); + } + for (0..4) |j| { + var max: f64 = -std.math.inf(f64); + for (0..4) |k| { + const q = 8 + 4 * j + k; + out[q] = @as(f64, mixes[q]) * inv * scales[2] + base[q]; + max = @max(max, out[q]); + } + var sum: f64 = 0; + for (0..4) |k| { + const q = 8 + 4 * j + k; + out[q] = @exp(out[q] - max); + sum += out[q]; + } + for (0..4) |k| out[8 + 4 * j + k] = out[8 + 4 * j + k] / sum + 1e-6; + } + for (0..20) |it| { + if (it > 0) for (0..4) |j| { + var s: f64 = 0; + for (0..4) |k| s += out[8 + 4 * j + k]; + for (0..4) |k| out[8 + 4 * j + k] /= s + 1e-6; + }; + for (0..4) |k| { + var s: f64 = 0; + for (0..4) |j| s += out[8 + 4 * j + k]; + for (0..4) |j| out[8 + 4 * j + k] /= s + 1e-6; + } + } + return out; +} + +test "V4.1 flattened mHC RMS, 20 Sinkhorn iterations, previous collapse, source-destination post and identity expand" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try mhc.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const stream = try b.get(u16, 16); + for (stream.slice(u16, 16), 0..) |*v, i| v.* = bf(@as(f32, @floatFromInt(i + 1)) * 0.25); + const inv = try b.get(f32, 1); + const mixes = try b.get(f32, 24); + for (mixes.slice(f32, 24), 0..) |*v, i| v.* = (@as(f32, @floatFromInt(i % 9)) - 4) * 3; + const scale = try b.get(f32, 3); + @memcpy(scale.slice(f32, 3), &[_]f32{ 0.3, -0.7, 1.2 }); + const base = try b.get(f32, 24); + for (base.slice(f32, 24), 0..) |*v, i| v.* = (@as(f32, @floatFromInt(i % 5)) - 2) * 0.3; + const pre = try b.get(f32, 4); + const post = try b.get(f32, 4); + const comb = try b.get(f32, 16); + const previous = try b.get(f32, 4); + @memcpy(previous.slice(f32, 4), &[_]f32{ 0.1, 0.2, 0.3, 0.4 }); + const collapse = try b.get(u16, 4); + const branch = try b.get(u16, 4); + @memset(branch.slice(u16, 4), bf(2)); + const y = try b.get(u16, 16); + const expanded = try b.get(u16, 16); + const initial = try b.get(f32, 4); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + mhc.encodeRms(e, &p, view(stream), view(inv), .{ .rows = 1, .width = 4, .stream_stride = 16, .eps = 1e-20 }); + mhc.encodeSplit(e, &p, view(mixes), view(inv), view(scale), view(base), view(pre), view(post), view(comb), .{ .rows = 1 }); + mhc.encodeCollapse(e, &p, view(stream), view(previous), view(collapse), .{ .rows = 1, .width = 4, .stream_stride = 16, .y_stride = 4 }); + mhc.encodePost(e, &p, view(branch), view(stream), view(post), view(comb), view(y), .{ .rows = 1, .width = 4, .branch_stride = 4, .residual_stride = 16, .y_stride = 16 }); + mhc.encodeExpand(e, &p, view(branch), view(expanded), view(initial), .{ .rows = 1, .width = 4, .x_stride = 4, .stream_stride = 16 }); + try done(cb, e); + var squares: f64 = 0; + for (stream.slice(u16, 16)) |v| squares += @as(f64, fp(v)) * fp(v); + const ir = 1 / @sqrt(squares / 16 + 1e-20); + try close(ir, inv.slice(f32, 1)[0], 1e-6); + const ref = splitReference(mixes.slice(f32, 24), ir, scale.slice(f32, 3), base.slice(f32, 24)); + for (0..4) |j| { + try close(ref[j], pre.slice(f32, 4)[j], 2e-6); + try close(ref[4 + j], post.slice(f32, 4)[j], 2e-6); + } + for (0..16) |q| try close(ref[8 + q], comb.slice(f32, 16)[q], 3e-6); + for (0..4) |i| { + var collapsed: f64 = 0; + for (0..4) |j| collapsed += @as(f64, previous.slice(f32, 4)[j]) * fp(stream.slice(u16, 16)[j * 4 + i]); + try close(collapsed, fp(collapse.slice(u16, 4)[i]), 0.02); + for (0..4) |k| { + var expected: f64 = 2 * ref[4 + k]; + for (0..4) |j| expected += ref[8 + 4 * j + k] * fp(stream.slice(u16, 16)[j * 4 + i]); + try close(expected, fp(y.slice(u16, 16)[k * 4 + i]), 0.04); + try std.testing.expectEqual(bf(2), expanded.slice(u16, 16)[k * 4 + i]); + } + try std.testing.expectEqual(@as(f32, if (i == 0) 1 else 0), initial.slice(f32, 4)[i]); + } +} + +test "V4.1 source single-row HC RMS reaches the retained attention BF16 consumer" { + const Fixture = struct { + embedding_bf16: []const u8, + mixes_f32: []const u8, + scale_f32: []const u8, + base_f32: []const u8, + branch_bf16_bits: [2]u16, + residual_bf16_bits: [8]u16, + expected_source_tree_target_bf16_bits: [2]u16, + }; + const fixture = try std.json.parseFromSlice(Fixture, std.testing.allocator, @embedFile("fixtures/hc_attention_rms.json"), .{ .ignore_unknown_fields = true }); + defer fixture.deinit(); + const f = fixture.value; + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try mhc.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const width = 5120; + const stream = try b.get(u16, 4 * width); + const xs = stream.slice(u16, 4 * width); + try std.base64.standard.Decoder.decode(std.mem.sliceAsBytes(xs[0..width]), f.embedding_bf16); + for (1..4) |copy| @memcpy(xs[copy * width ..][0..width], xs[0..width]); + const inv = try b.get(f32, 1); + const mixes = try b.get(f32, 24); + const scale = try b.get(f32, 3); + const base = try b.get(f32, 24); + try std.base64.standard.Decoder.decode(std.mem.sliceAsBytes(mixes.slice(f32, 24)), f.mixes_f32); + try std.base64.standard.Decoder.decode(std.mem.sliceAsBytes(scale.slice(f32, 3)), f.scale_f32); + try std.base64.standard.Decoder.decode(std.mem.sliceAsBytes(base.slice(f32, 24)), f.base_f32); + const pre = try b.get(f32, 4); + const post = try b.get(f32, 4); + const comb = try b.get(f32, 16); + const branch = try b.get(u16, 2); + const residual = try b.get(u16, 8); + const output = try b.get(u16, 8); + @memcpy(branch.slice(u16, 2), &f.branch_bf16_bits); + @memcpy(residual.slice(u16, 8), &f.residual_bf16_bits); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + // Recompute statistic, coefficients and post products from the retained projection; never inject an oracle inverse. + mhc.encodeRms(e, &p, view(stream), view(inv), .{ .rows = 1, .width = width, .stream_stride = 4 * width, .eps = 1e-20 }); + mhc.encodeSplit(e, &p, view(mixes), view(inv), view(scale), view(base), view(pre), view(post), view(comb), .{ .rows = 1 }); + mhc.encodePost(e, &p, view(branch), view(residual), view(post), view(comb), view(output), .{ .rows = 1, .width = 2, .branch_stride = 2, .residual_stride = 8, .y_stride = 8 }); + try done(cb, e); + // The former reduction produces [14091,46835], despite unchanged branch BF16 bits. + for (0..2) |i| try std.testing.expectEqual(f.expected_source_tree_target_bf16_bits[i], output.slice(u16, 8)[6 + i]); +} + +test "V4.1 hyperconnections preserve FP32 products and residual boundaries under cancellation" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try mhc.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const stream = try b.get(u16, 8); + @memcpy(stream.slice(u16, 8), &[_]u16{ bf(3), bf(3), bf(0), bf(0), bf(1), bf(1), bf(1), bf(1) }); + const previous = try b.get(f32, 8); + @memcpy(previous.slice(f32, 8), &[_]f32{ -1, 1 + 0x1p-23, 0, 0, 1, 0x1p-24, 0x1p-24, -1 }); + const collapsed = try b.get(u16, 2); + const branch = try b.get(u16, 2); + @memcpy(branch.slice(u16, 2), &[_]u16{ bf(0), bf(1) }); + const post = try b.get(f32, 8); + post.slice(f32, 8)[4] = -1; + const comb = try b.get(f32, 32); + comb.slice(f32, 32)[0] = -1; + comb.slice(f32, 32)[4] = 1 + 0x1p-23; + comb.slice(f32, 32)[16] = 1; + comb.slice(f32, 32)[20] = 0x1p-24; + comb.slice(f32, 32)[24] = 0x1p-24; + const output = try b.get(u16, 8); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + mhc.encodeCollapse(e, &p, view(stream), view(previous), view(collapsed), .{ .rows = 2, .width = 1, .stream_stride = 4, .y_stride = 1 }); + mhc.encodePost(e, &p, view(branch), view(stream), view(post), view(comb), view(output), .{ .rows = 2, .width = 1, .branch_stride = 1, .residual_stride = 4, .y_stride = 4 }); + try done(cb, e); + try std.testing.expectEqual(bf(0x1p-21), collapsed.slice(u16, 2)[0]); + try std.testing.expectEqual(bf(0), collapsed.slice(u16, 2)[1]); + try std.testing.expectEqual(bf(0x1p-21), output.slice(u16, 8)[0]); + try std.testing.expectEqual(bf(0), output.slice(u16, 8)[4]); + branch.slice(u16, 2)[0] = bf(3); + post.slice(f32, 8)[0] = 1 + 0x1p-23; + comb.slice(f32, 32)[4] = 0; + const branch_cb = queue.commandBuffer(); + const branch_enc = branch_cb.compute(.serial); + mhc.encodePost(branch_enc, &p, view(branch), view(stream), view(post), view(comb), view(output), .{ .rows = 1, .width = 1, .branch_stride = 1, .residual_stride = 4, .y_stride = 4 }); + try done(branch_cb, branch_enc); + try std.testing.expectEqual(bf(0x1p-21), output.slice(u16, 8)[0]); +} diff --git a/zig/src/families/deepseek_v41/ops.zig b/zig/src/families/deepseek_v41/ops.zig new file mode 100644 index 0000000000..68fd169857 --- /dev/null +++ b/zig/src/families/deepseek_v41/ops.zig @@ -0,0 +1,110 @@ +//! Trusted plaintext encoders: callers validate nonzero shapes/storage; strides count elements, EXL3 is separate. +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const View = @import("view.zig").View; + +pub const Pipelines = struct { + rms: mtl.Pipeline, + dense: mtl.Pipeline, + embedding: mtl.Pipeline, + swiglu: mtl.Pipeline, + combine: mtl.Pipeline, + + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_ops, options); + defer lib.deinit(); + const rms = try mtl.Pipeline.init(device, lib, "dsv41_rms", false); + errdefer rms.deinit(); + const dense = try mtl.Pipeline.init(device, lib, "dsv41_dense", false); + errdefer dense.deinit(); + const embedding = try mtl.Pipeline.init(device, lib, "dsv41_embedding", false); + errdefer embedding.deinit(); + const swiglu = try mtl.Pipeline.init(device, lib, "dsv41_swiglu", false); + errdefer swiglu.deinit(); + const combine = try mtl.Pipeline.init(device, lib, "dsv41_combine", false); + return .{ .rms = rms, .dense = dense, .embedding = embedding, .swiglu = swiglu, .combine = combine }; + } + pub fn deinit(p: *Pipelines) void { + p.rms.deinit(); + p.dense.deinit(); + p.embedding.deinit(); + p.swiglu.deinit(); + p.combine.deinit(); + } +}; +fn bind(e: mtl.ComputeEncoder, refs: anytype) void { + inline for (refs, 0..) |r, i| e.setBuffer(r.buffer, r.offset, i); +} +pub const Rms = struct { rows: u32, width: u32, x_stride: u32, y_stride: u32, eps: f32, weight_f32: bool = false }; +/// BF16 x/y [rows,width], BF16 or FP32 weight [width]; sums, mean and affine product are FP32. +pub fn encodeRms(e: mtl.ComputeEncoder, p: *const Pipelines, x: View, weight: View, y: View, d: Rms) void { + e.setPipeline(p.rms); + bind(e, .{ x, weight, y }); + const D = extern struct { rows: u32, width: u32, x_stride: u32, y_stride: u32, eps: f32, weight_f32: u32 }; + e.setValue(D{ .rows = d.rows, .width = d.width, .x_stride = d.x_stride, .y_stride = d.y_stride, .eps = d.eps, .weight_f32 = @intFromBool(d.weight_f32) }, 3); + e.dispatchGroups(mtl.Size.of(d.rows, 1, 1), mtl.Size.of(256, 1, 1)); +} +pub const Dtype = enum(u32) { bf16 = 0, f32 = 1, f16 = 2 }; +pub const Dense = struct { + rows: u32, + k: u32, + n: u32, + groups: u32 = 1, + x_stride: u32, + w_stride: u32, + y_stride: u32, + x_dtype: Dtype = .bf16, + w_dtype: Dtype = .bf16, + y_dtype: Dtype = .bf16, +}; +/// x[rows,groups,k], w[groups,n,k], y[rows,groups,n]; packed groups, padded token pitches, FP32 sums, final BF16. +pub fn encodeDense(e: mtl.ComputeEncoder, p: *const Pipelines, x: View, w: View, y: View, d: Dense) void { + e.setPipeline(p.dense); + bind(e, .{ x, w, y }); + e.setValue([10]u32{ d.rows, d.k, d.n, d.groups, d.x_stride, d.w_stride, d.y_stride, @backingInt(d.x_dtype), @backingInt(d.w_dtype), @backingInt(d.y_dtype) }, 3); + e.dispatchGroups(mtl.Size.of((d.n + 15) / 16, d.rows, d.groups), mtl.Size.of(128, 1, 1)); +} +pub const Embedding = extern struct { rows: u32, width: u32, table_stride: u32, y_stride: u32, copies: u32 = 4 }; +/// ids[rows] u32 must be in vocabulary; table[vocab,width] BF16 -> y[rows,copies,width] BF16. +pub fn encodeEmbedding(e: mtl.ComputeEncoder, p: *const Pipelines, ids: View, table: View, y: View, d: Embedding) void { + e.setPipeline(p.embedding); + bind(e, .{ ids, table, y }); + e.setValue(d, 3); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} +pub const Activation = struct { + rows: u32, + width: u32, + input_stride: u32, + output_stride: u32, + weighted: bool = false, + input_f32: bool = false, +}; +/// Separate gate/up; apply FP32 weights before BF16 cast/down. Unweighted calls borrow but never read weights. +pub fn encodeSwiGLU(e: mtl.ComputeEncoder, p: *const Pipelines, gate: View, up: View, weights: View, y: View, d: Activation) void { + e.setPipeline(p.swiglu); + bind(e, .{ gate, up, weights, y }); + e.setValue([6]u32{ d.rows, d.width, d.input_stride, d.output_stride, @intFromBool(d.weighted), @intFromBool(d.input_f32) }, 4); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} +pub const Combine = struct { + rows: u32, + width: u32, + routed_row_stride: u32, + slot_stride: u32, + shared_stride: u32, + y_stride: u32, + input_f32: bool = false, + id_stride: u32 = 6, +}; +/// Sum six weighted down slots by ascending expert ID in FP32, then shared; final BF16, invalid IDs yield NaN. +pub fn encodeCombine(e: mtl.ComputeEncoder, p: *const Pipelines, routed: View, shared: View, ids: View, y: View, d: Combine) void { + e.setPipeline(p.combine); + bind(e, .{ routed, shared, ids, y }); + e.setValue([8]u32{ d.rows, d.width, d.routed_row_stride, d.slot_stride, d.shared_stride, d.y_stride, @intFromBool(d.input_f32), d.id_stride }, 4); + e.dispatchThreads(mtl.Size.of(d.width, d.rows, 1), mtl.Size.of(256, 1, 1)); +} diff --git a/zig/src/families/deepseek_v41/ops_check.zig b/zig/src/families/deepseek_v41/ops_check.zig new file mode 100644 index 0000000000..bfe08f4dd8 --- /dev/null +++ b/zig/src/families/deepseek_v41/ops_check.zig @@ -0,0 +1,191 @@ +//! Independent GPU consumer fixtures; pinned official inference equations and provenance in THIRD_PARTY_NOTICES.md. +const std = @import("std"); +const mtl = @import("metal"); +const ops = @import("ops.zig"); +const support = @import("test_support.zig"); +const bf = support.bf; +const fp = support.fp; +const view = support.view; +const Buffers = support.Buffers; +const done = support.done; +const close = support.close; + +test "V4.1 BF16 RMS, F16 dense row independence, grouped wo_a, embedding, clamped weighted SwiGLU and shared combine" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ops.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const x = try b.get(u16, 16); + const w = try b.get(u16, 4); + const norm = try b.get(u16, 16); + const xs = x.slice(u16, 16); + const ws = w.slice(u16, 4); + for (0..2) |r| for (0..4) |i| { + xs[r * 8 + i] = bf(@as(f32, @floatFromInt(1 + i + r))); + }; + for (ws, 0..) |*v, i| v.* = bf(@as(f32, @floatFromInt(i + 1)) * 0.5); + const f16w = try b.get(f16, 8); + const fws = f16w.slice(f16, 8); + for (fws, 0..) |*v, i| v.* = @floatCast(@as(f32, @floatFromInt(i + 1)) * 0.125 + 0.0009765625); + const dense = try b.get(f32, 8); + const dense_one = try b.get(f32, 2); + const grouped = try b.get(u16, 8); + const gw = try b.get(u16, 8); + for (gw.slice(u16, 8), 0..) |*v, i| v.* = bf(@as(f32, @floatFromInt(i + 1)) * 0.25); + const ids = try b.get(u32, 2); + ids.slice(u32, 2)[0] = 1; + ids.slice(u32, 2)[1] = 0; + const embed = try b.get(u16, 32); + const gate = try b.get(f32, 8); + const up = try b.get(f32, 8); + @memcpy(gate.slice(f32, 8), &[_]f32{ -100, -12, 1, 100, -2, 0, 8, 12 }); + @memcpy(up.slice(f32, 8), &[_]f32{ -100, -12, 1, 100, 12, -20, 0.5, 1 }); + const weights = try b.get(f32, 2); + @memcpy(weights.slice(f32, 2), &[_]f32{ 0.25, 1.25 }); + const act = try b.get(u16, 8); + const routed = try b.get(u16, 48); + const combine_ids = try b.get(u32, 12); + @memcpy(combine_ids.slice(u32, 12), &[_]u32{ 5, 3, 1, 4, 2, 0, 5, 3, 1, 4, 2, 0 }); + const shared = try b.get(u16, 8); + const combined = try b.get(u16, 8); + const values = [_]f32{ 1, 33554432, -33554432, 2, 4, 8 }; + for (routed.slice(u16, 48), 0..) |*v, i| v.* = bf(values[(i % 24) / 4]); + @memset(shared.slice(u16, 8), bf(2.5)); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + ops.encodeRms(e, &p, view(x), view(w), view(norm), .{ .rows = 2, .width = 4, .x_stride = 8, .y_stride = 8, .eps = 1e-20 }); + ops.encodeDense(e, &p, view(x), view(f16w), view(dense), .{ .rows = 2, .k = 4, .n = 2, .x_stride = 8, .w_stride = 4, .y_stride = 4, .w_dtype = .f16, .y_dtype = .f32 }); + ops.encodeDense(e, &p, view(x).shifted(16), view(f16w), view(dense_one), .{ .rows = 1, .k = 4, .n = 2, .x_stride = 8, .w_stride = 4, .y_stride = 2, .w_dtype = .f16, .y_dtype = .f32 }); + ops.encodeDense(e, &p, view(x), view(gw), view(grouped), .{ .rows = 2, .k = 2, .n = 2, .groups = 2, .x_stride = 8, .w_stride = 2, .y_stride = 4 }); + ops.encodeEmbedding(e, &p, view(ids), view(x), view(embed), .{ .rows = 2, .width = 4, .table_stride = 8, .y_stride = 16 }); + ops.encodeSwiGLU(e, &p, view(gate), view(up), view(weights), view(act), .{ .rows = 2, .width = 4, .input_stride = 4, .output_stride = 4, .weighted = true, .input_f32 = true }); + ops.encodeCombine(e, &p, view(routed), view(shared), view(combine_ids), view(combined), .{ .rows = 2, .width = 4, .routed_row_stride = 24, .slot_stride = 4, .shared_stride = 4, .y_stride = 4 }); + try done(cb, e); + for (0..2) |r| { + var sum: f64 = 0; + for (0..4) |i| sum += @as(f64, fp(xs[r * 8 + i])) * fp(xs[r * 8 + i]); + for (0..4) |i| try close(@as(f64, fp(xs[r * 8 + i])) / @sqrt(sum / 4 + 1e-20) * fp(ws[i]), fp(norm.slice(u16, 16)[r * 8 + i]), 0.04); + for (0..2) |j| { + var ref: f64 = 0; + for (0..4) |k| ref += @as(f64, fp(xs[r * 8 + k])) * @as(f64, fws[j * 4 + k]); + try close(ref, dense.slice(f32, 8)[r * 4 + j], 1e-6); + } + for (0..2) |g| for (0..2) |j| { + var ref: f32 = 0; + for (0..2) |k| ref += fp(xs[r * 8 + g * 2 + k]) * fp(gw.slice(u16, 8)[g * 4 + j * 2 + k]); + try std.testing.expectEqual(bf(ref), grouped.slice(u16, 8)[r * 4 + g * 2 + j]); + }; + for (0..4) |copy| for (0..4) |i| try std.testing.expectEqual(xs[(1 - r) * 8 + i], embed.slice(u16, 32)[r * 16 + copy * 4 + i]); + for (0..4) |i| { + const k = r * 4 + i; + const a: f64 = @min(gate.slice(f32, 8)[k], 10); + const u: f64 = @max(-10, @min(up.slice(f32, 8)[k], 10)); + const activated = a / (1 + @exp(-a)) * u * weights.slice(f32, 2)[r]; + try close(activated, fp(act.slice(u16, 8)[k]), @abs(activated) * 0.004 + 1e-35); + var ref: f32 = 0; + for ([_]usize{ 5, 2, 4, 1, 3, 0 }) |slot| ref += fp(routed.slice(u16, 48)[r * 24 + slot * 4 + i]); + try std.testing.expectEqual(bf(ref + 2.5), combined.slice(u16, 8)[k]); + } + } + for (0..2) |j| try std.testing.expectEqual(dense.slice(f32, 8)[4 + j], dense_one.slice(f32, 2)[j]); +} + +test "V4.1 CUDA RMS mean projection preserves the BF16 affine tie" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ops.Pipelines.load(device); + defer p.deinit(); + var buffers = Buffers{ .device = device }; + defer buffers.deinit(); + const width = 5120; + const x = try buffers.get(u16, width); + const weight = try buffers.get(f32, width); + const output = try buffers.get(u16, width); + x.slice(u16, width)[0] = bf(3); + @memset(weight.slice(f32, width), @bitCast(@as(u32, 0x3c65de29))); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + ops.encodeRms(e, &p, view(x), view(weight), view(output), .{ .rows = 1, .width = width, .x_stride = width, .y_stride = width, .eps = 1e-20, .weight_f32 = true }); + try done(cb, e); + // Pinned CUDA RMSNorm.forward returns 1; sum / width crosses to the next BF16 value. + try std.testing.expectEqual(bf(1), output.slice(u16, width)[0]); + for (output.slice(u16, width)[1..]) |v| try std.testing.expectEqual(@as(u16, 0), v); +} + +test "V4.1 FP32 dense retains product and reduction cancellation" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ops.Pipelines.load(device); + defer p.deinit(); + var buffers = Buffers{ .device = device }; + defer buffers.deinit(); + const k = 96; + const n = 3; + const x = try buffers.get(u16, k); + const w = try buffers.get(f32, k * n); + const y = try buffers.get(f32, n); + @memset(x.slice(u16, k), bf(3)); + const weights = w.slice(f32, k * n); + weights[0] = 0x1p25; + weights[32] = 1; + weights[64] = -0x1p25; + weights[k] = 1 + 0x1p-23; + weights[k + 32] = -1; + weights[2 * k] = 0x1p25; + weights[2 * k + 1] = 1; + weights[2 * k + 2] = -0x1p25; + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + ops.encodeDense(enc, &p, view(x), view(w), view(y), .{ .rows = 1, .k = k, .n = n, .x_stride = k, .w_stride = k, .y_stride = n, .w_dtype = .f32, .y_dtype = .f32 }); + try done(cb, enc); + for (0..n) |column| { + var expected: f64 = 0; + for (0..k) |index| expected += @as(f64, fp(x.slice(u16, k)[index])) * weights[column * k + index]; + try std.testing.expectEqual(@as(f32, @floatCast(expected)), y.slice(f32, n)[column]); + } +} + +test "V4.1 weighted SiLU preserves division rounding before BF16 activation" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ops.Pipelines.load(device); + defer p.deinit(); + var buffers = Buffers{ .device = device }; + defer buffers.deinit(); + const gate = try buffers.get(u16, 1); + const up = try buffers.get(u16, 1); + const weights = try buffers.get(f32, 1); + const output = try buffers.get(u16, 1); + gate.slice(u16, 1)[0] = bf(0.78515625); + up.slice(u16, 1)[0] = bf(-6.96875); + weights.slice(f32, 1)[0] = 0.25; + const cb = queue.commandBuffer(); + const encoder = cb.compute(.serial); + ops.encodeSwiGLU(encoder, &p, view(gate), view(up), view(weights), view(output), .{ .rows = 1, .width = 1, .input_stride = 1, .output_stride = 1, .weighted = true }); + try done(cb, encoder); + // Source SiLU divides before up/route multiplication; reciprocal multiplication crosses this BF16 tie. + try std.testing.expectEqual(bf(-0.94140625), output.slice(u16, 1)[0]); +} + +test { + _ = @import("mhc_test.zig"); + _ = @import("route_test.zig"); +} diff --git a/zig/src/families/deepseek_v41/profiler.zig b/zig/src/families/deepseek_v41/profiler.zig new file mode 100644 index 0000000000..fe3a85b5fc --- /dev/null +++ b/zig/src/families/deepseek_v41/profiler.zig @@ -0,0 +1,93 @@ +//! Family stage labels over the generic same-command-buffer counter recorder; no extra storage or GPU waits. +const std = @import("std"); +const mtl = @import("metal"); +const counters = @import("../../core/gpu_counters.zig"); + +pub const Stage = enum { engram, mhc, attention, moe, head, other }; +pub const stage_count = @typeInfo(Stage).@"enum".field_names.len; +pub const Options = counters.Options; +pub const Aggregate = counters.Aggregate; +pub const Span = counters.Span; + +pub fn class(stage: Stage) counters.Class { + return switch (stage) { + .engram => .embedding, + .mhc => .norm, + .attention => .attention, + .moe, .head => .dense, + .other => .copy, + }; +} +pub fn stageOf(span: Span) !Stage { + if (span.metadata >= stage_count) return error.InvalidProfileStage; + return @fromBackingInt(@intCast(span.metadata)); +} + +/// Unique owner of the recorder; stage metadata and class totals share the recorder's bounded span storage. +pub const Recorder = struct { + recorder: counters.Recorder, + totals: [stage_count]Aggregate = @splat(.{}), + + pub fn init(allocator: std.mem.Allocator, device: mtl.Device, options: Options) !Recorder { + return .{ .recorder = try counters.Recorder.init(allocator, device, options) }; + } + pub fn enabled(self: *const Recorder) bool { + return self.recorder.enabled(); + } + pub fn capability(self: *const Recorder) mtl.counters.Capability { + return self.recorder.capability(); + } + pub fn activeSpan(self: *const Recorder) bool { + return self.recorder.active_span != null; + } + pub fn beginCommand(self: *Recorder, cb: mtl.CommandBuffer) !void { + try self.recorder.beginCommand(cb); + self.totals = @splat(.{}); + } + pub fn beginSpan(self: *Recorder, encoder: *mtl.ComputeEncoder, stage: Stage) !void { + try self.recorder.beginSpan(encoder, class(stage), @intCast(@backingInt(stage))); + } + pub fn endSpan(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { + try self.recorder.endSpan(encoder); + } + pub fn finishEncoding(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { + try self.recorder.finishEncoding(encoder); + } + pub fn cancelEncoding(self: *Recorder, encoder: *mtl.ComputeEncoder) !void { + try self.recorder.cancelEncoding(encoder); + self.totals = @splat(.{}); + } + pub fn discardAfterCompletion(self: *Recorder) !void { + try self.recorder.discardAfterCompletion(); + self.totals = @splat(.{}); + } + pub fn readAfterCompletion(self: *Recorder) ![]const Span { + const spans = try self.recorder.readAfterCompletion(); + var totals: [stage_count]Aggregate = @splat(.{}); + for (spans) |span| { + const total = &totals[@backingInt(try stageOf(span))]; + total.calls += 1; + total.gpu_seconds += span.gpu_seconds; + } + self.totals = totals; + return spans; + } + pub fn aggregates(self: *const Recorder) [stage_count]Aggregate { + return self.totals; + } + pub fn deinit(self: *Recorder) !void { + try self.recorder.deinit(); + self.totals = @splat(.{}); + } +}; + +test "disabled family profiling allocates nothing and leaves no stage receipts" { + var recorder = try Recorder.init(std.testing.failing_allocator, undefined, .{}); + defer recorder.deinit() catch unreachable; + var encoder: mtl.ComputeEncoder = undefined; + try recorder.beginSpan(&encoder, .mhc); + try recorder.endSpan(&encoder); + try std.testing.expect(!recorder.activeSpan()); + try std.testing.expectEqual(@as(usize, 0), (try recorder.readAfterCompletion()).len); + for (recorder.aggregates()) |total| try std.testing.expectEqual(@as(u64, 0), total.calls); +} diff --git a/zig/src/families/deepseek_v41/rope.zig b/zig/src/families/deepseek_v41/rope.zig new file mode 100644 index 0000000000..cc3419f7c3 --- /dev/null +++ b/zig/src/families/deepseek_v41/rope.zig @@ -0,0 +1,134 @@ +//! Adjacent-pair BF16 RoPE using resident, source-style complex rotation factors. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const View = @import("view.zig").View; + +pub const Spec = struct { + theta: f32, + factor: f32 = 16, + original: u32 = 0, + beta_fast: f32 = 32, + beta_slow: f32 = 1, + + fn validate(s: Spec) !void { + if (!std.math.isFinite(s.theta) or s.theta <= 1 or + !std.math.isFinite(s.factor) or s.factor <= 0 or + !std.math.isFinite(s.beta_fast) or s.beta_fast <= 0 or + !std.math.isFinite(s.beta_slow) or s.beta_slow <= 0) + return error.InvalidRotarySpec; + } +}; +pub const Spectrum = enum(u32) { normal, compressed }; +pub const Params = extern struct { + rows: u32, + heads: u32, + width: u32, + rotary: u32, + stride: u32, + position: u32, + step: u32 = 1, + inverse: u32 = 0, + spectrum: Spectrum = .normal, + + pub fn validate(p: Params) !void { + if (p.rows == 0 or p.heads == 0 or p.rotary == 0 or p.rotary % 2 != 0 or + p.rotary > p.width or @as(u64, p.stride) < @as(u64, p.heads) * p.width or + p.step == 0 or p.inverse > 1 or + @as(u64, p.rows - 1) * p.stride + @as(u64, p.heads) * p.width > std.math.maxInt(u32)) + return error.InvalidRotaryShape; + } +}; + +// Keep FP32 tensor/angle boundaries, FP64 corrected_dim logs, and FP32-rounded polar unaries; never per-head recompute. +fn fillFactors(out: []f32, capacity: u32, rotary: u32, spec: Spec) void { + @setFloatMode(.strict); + const dim: f32 = @floatFromInt(rotary); + const pairs: usize = rotary / 2; + var low: f32 = 0; + var high: f32 = 0; + if (spec.original > 0) { + const original: f64 = @floatFromInt(spec.original); + const dim64: f64 = @floatFromInt(rotary); + const denominator = 2 * @log(@as(f64, spec.theta)); + low = @floatCast(@max(@floor(dim64 * @log(original / + (@as(f64, spec.beta_fast) * 2 * std.math.pi)) / denominator), 0)); + high = @floatCast(@min(@ceil(dim64 * @log(original / + (@as(f64, spec.beta_slow) * 2 * std.math.pi)) / denominator), dim64 - 1)); + } + for (0..pairs) |pair| { + const exponent: f32 = @as(f32, @floatFromInt(pair * 2)) / dim; + const power: f32 = @floatCast(std.math.pow(f64, spec.theta, exponent)); + var frequency: f32 = 1 / power; + if (spec.original > 0) { + const ramp: f32 = std.math.clamp((@as(f32, @floatFromInt(pair)) - low) / + @max(high - low, @as(f32, 1e-3)), 0, 1); + const smooth: f32 = 1 - ramp; + const scaled: f32 = frequency / spec.factor; + const slow: f32 = scaled * (1 - smooth); + const fast: f32 = frequency * smooth; + frequency = slow + fast; + } + for (0..capacity) |position| { + const angle: f32 = @as(f32, @floatFromInt(position)) * frequency; + const at = (position * pairs + pair) * 2; + out[at] = @floatCast(@cos(@as(f64, angle))); + out[at + 1] = @floatCast(@sin(@as(f64, angle))); + } + } +} + +pub const Pipelines = struct { + rotate: mtl.Pipeline, + factors: mtl.Buffer, + capacity: u32, + rotary: u32, + + pub fn residentBytes(capacity: u32, rotary: u32) usize { + return 2 * @as(usize, capacity) * (rotary / 2) * 2 * @sizeOf(f32); + } + + pub fn load(device: mtl.Device, capacity: u32, rotary: u32, normal: Spec, compressed: Spec) !Pipelines { + if (capacity == 0 or rotary == 0 or rotary % 2 != 0 or + @as(u64, capacity) * rotary > std.math.maxInt(usize) / (2 * @sizeOf(f32))) + return error.InvalidRotaryShape; + try normal.validate(); + try compressed.validate(); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_rope, options); + defer lib.deinit(); + const rotate = try mtl.Pipeline.init(device, lib, "dsv41_rope", false); + errdefer rotate.deinit(); + const factors = try device.buffer(residentBytes(capacity, rotary), mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer factors.deinit(); + const count = @as(usize, capacity) * rotary; + const values = factors.slice(f32, 2 * count); + fillFactors(values[0..count], capacity, rotary, normal); + fillFactors(values[count..], capacity, rotary, compressed); + return .{ .rotate = rotate, .factors = factors, .capacity = capacity, .rotary = rotary }; + } + pub fn deinit(p: *Pipelines) void { + p.factors.deinit(); + p.rotate.deinit(); + } + pub fn registerResidency(p: *const Pipelines, set: mtl.ResidencySet) void { + set.add(p.factors); + } + pub fn encode(p: *const Pipelines, enc: mtl.ComputeEncoder, x: View, args: Params) !void { + try args.validate(); + const last = @as(u64, args.position) + @as(u64, args.rows - 1) * args.step; + if (args.rotary != p.rotary or last >= p.capacity) return error.InvalidRotaryShape; + try x.requireRows(args.rows, args.stride, @as(usize, args.heads) * args.width, 2); + const offset = switch (args.spectrum) { + .normal => 0, + .compressed => @as(usize, p.capacity) * p.rotary * @sizeOf(f32), + }; + enc.setPipeline(p.rotate); + enc.setBuffer(x.buffer, x.offset, 0); + enc.setValue(args, 1); + enc.setBuffer(p.factors, offset, 2); + enc.dispatchThreads(mtl.Size.of(args.rotary / 2, args.heads, args.rows), mtl.Size.of(32, 1, 1)); + enc.barrier(); + } +}; diff --git a/zig/src/families/deepseek_v41/route.zig b/zig/src/families/deepseek_v41/route.zig new file mode 100644 index 0000000000..00c1a1ff10 --- /dev/null +++ b/zig/src/families/deepseek_v41/route.zig @@ -0,0 +1,39 @@ +//! Exhaustive 384-expert sqrt-softplus top-six routing. Bias affects selection, never returned weights. +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +pub const View = @import("view.zig").View; +pub const Dtype = @import("ops.zig").Dtype; +pub const Pipelines = struct { + select: mtl.Pipeline, + pub fn load(device: mtl.Device) !Pipelines { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, sources.deepseek_v41_route, options); + defer lib.deinit(); + return .{ .select = try mtl.Pipeline.init(device, lib, "dsv41_route", false) }; + } + pub fn deinit(p: *Pipelines) void { + p.select.deinit(); + } +}; +pub const Shape = struct { + rows: u32, + logit_stride: u32 = 384, + output_stride: u32 = 6, + temperature: f32 = 1, + bias_dtype: Dtype = .f32, +}; +/// FP32 logits[rows,384], bias F16/BF16/FP32 -> u32 IDs and FP32 weights[rows,6]; lower ID wins biased-score ties. + +// Nonzero rows, logit/output pitches >=384/6, finite positive temperature, complete storage are caller invariants. + +// All six execute; weights use unbiased scores / (sum+1e-20) * 1.5. Invalid comparisons emit invalid IDs/NaNs. +pub fn encode(e: mtl.ComputeEncoder, p: *const Pipelines, logits: View, bias: View, ids: View, weights: View, d: Shape) void { + e.setPipeline(p.select); + inline for (.{ logits, bias, ids, weights }, 0..) |r, i| e.setBuffer(r.buffer, r.offset, i); + const D = extern struct { rows: u32, logit_stride: u32, output_stride: u32, temperature: f32, bias_dtype: u32 }; + e.setValue(D{ .rows = d.rows, .logit_stride = d.logit_stride, .output_stride = d.output_stride, .temperature = d.temperature, .bias_dtype = @backingInt(d.bias_dtype) }, 4); + e.dispatchGroups(mtl.Size.of(d.rows, 1, 1), mtl.Size.of(128, 1, 1)); +} diff --git a/zig/src/families/deepseek_v41/route_test.zig b/zig/src/families/deepseek_v41/route_test.zig new file mode 100644 index 0000000000..f653fa9241 --- /dev/null +++ b/zig/src/families/deepseek_v41/route_test.zig @@ -0,0 +1,152 @@ +//! Consumer-visible expert IDs, weights and nonfinite-input routing boundaries. +const std = @import("std"); +const mtl = @import("metal"); +const route = @import("route.zig"); +const ops = @import("ops.zig"); +const support = @import("test_support.zig"); +const view = support.view; +const Buffers = support.Buffers; +const done = support.done; +const close = support.close; +fn scoreReference(x: f64) f64 { + return @sqrt(if (x > 20) x else std.math.log1p(@exp(x))); +} +test "V4.1 exhaustive router F16 correction bias, unbiased normalization, extreme logits and row independence" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try route.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + const logits = try b.get(f32, 2 * 389); + const bias = try b.get(f16, 384); + for (0..384) |i| { + logits.slice(f32, 778)[i] = (@as(f32, @floatFromInt(i)) - 190) * 0.5; + logits.slice(f32, 778)[389 + i] = 0; + bias.slice(f16, 384)[i] = if (i >= 17 and i < 23) 100 else 0; + } + const ids = try b.get(u32, 16); + const weights = try b.get(f32, 16); + const ids_one = try b.get(u32, 6); + const weights_one = try b.get(f32, 6); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + route.encode(e, &p, view(logits), view(bias), view(ids), view(weights), .{ .rows = 2, .logit_stride = 389, .output_stride = 8, .bias_dtype = .f16 }); + route.encode(e, &p, view(logits).shifted(389 * 4), view(bias), view(ids_one), view(weights_one), .{ .rows = 1, .bias_dtype = .f16 }); + try done(cb, e); + for (0..2) |r| { + var used: [384]bool = @splat(false); + var picks: [6]usize = undefined; + var raw: [6]f64 = undefined; + var sum: f64 = 0; + for (0..6) |slot| { + var best: f64 = -std.math.inf(f64); + var pick: usize = 0; + for (0..384) |i| { + const sc = scoreReference(logits.slice(f32, 778)[r * 389 + i]) + @as(f64, bias.slice(f16, 384)[i]); + if (!used[i] and sc > best) { + best = sc; + pick = i; + } + } + used[pick] = true; + picks[slot] = pick; + raw[slot] = scoreReference(logits.slice(f32, 778)[r * 389 + pick]); + sum += raw[slot]; + } + for (0..6) |slot| { + try std.testing.expectEqual(@as(u32, @intCast(picks[slot])), ids.slice(u32, 16)[r * 8 + slot]); + try close(raw[slot] / (sum + 1e-20) * 1.5, weights.slice(f32, 16)[r * 8 + slot], 3e-6); + } + } + for (0..6) |i| { + try std.testing.expectEqual(ids.slice(u32, 16)[8 + i], ids_one.slice(u32, 6)[i]); + try std.testing.expectEqual(weights.slice(f32, 16)[8 + i], weights_one.slice(f32, 6)[i]); + } +} + +test "V4.1 router normalization matches source CUDA six-score FP32 reduction boundaries" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try route.Pipelines.load(device); + defer p.deinit(); + var b = Buffers{ .device = device }; + defer b.deinit(); + // Power-of-two sqrt(softplus(score^2)) with logits >20 eliminates unary-rounding ambiguity. + const scores = [2][6]f32{ + .{ 268435456, 16, 16, 16, 16, 16 }, + .{ 2097152, 8, 16, 134217728, 8388608, 8192 }, + }; + // Independent pinned Torch CUDA Gate expected bits; the second row also distinguishes wide normalization. + const expected_bits = [2][6]u32{ + .{ 1069547517, 868220925, 868220925, 868220925, 868220925, 868220925 }, + .{ 1018303405, 867308461, 875697069, 1068635053, 1035080621, 951194541 }, + }; + const logits = try b.get(f32, scores.len * 384); + const bias = try b.get(f32, 384); + const ids = try b.get(u32, scores.len * 6); + const weights = try b.get(f32, scores.len * 6); + @memset(logits.slice(f32, scores.len * 384), -100); + for (scores, 0..) |row, r| { + for (row, 0..) |s, slot| logits.slice(f32, scores.len * 384)[r * 384 + slot] = s * s; + } + // Distinct biases force slot order without topk ties; weights still use unbiased scores. + for (0..6) |slot| bias.slice(f32, 384)[slot] = 1099511627776 - @as(f32, @floatFromInt(slot)) * 68719476736; + const cb = queue.commandBuffer(); + const encoder = cb.compute(.serial); + route.encode(encoder, &p, view(logits), view(bias), view(ids), view(weights), .{ .rows = scores.len, .bias_dtype = .f32 }); + try done(cb, encoder); + for (expected_bits, 0..) |row, r| { + for (row, 0..) |expected, slot| { + try std.testing.expectEqual(@as(u32, @intCast(slot)), ids.slice(u32, scores.len * 6)[r * 6 + slot]); + try std.testing.expectEqual(expected, @as(u32, @bitCast(weights.slice(f32, scores.len * 6)[r * 6 + slot]))); + } + } +} + +test "nonfinite router inputs propagate invalid routes without reading outside the expert arrays" { + const device = mtl.Device.init() catch return error.SkipZigTest; + defer device.deinit(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try route.Pipelines.load(device); + defer p.deinit(); + var buffers = Buffers{ .device = device }; + defer buffers.deinit(); + const logits = try buffers.get(f32, 384); + const bias = try buffers.get(f32, 384); + const ids = try buffers.get(u32, 6); + const weights = try buffers.get(f32, 6); + @memset(logits.slice(f32, 384), std.math.nan(f32)); + const cb = queue.commandBuffer(); + const encoder = cb.compute(.serial); + route.encode(encoder, &p, view(logits), view(bias), view(ids), view(weights), .{ .rows = 1 }); + try done(cb, encoder); + for (ids.slice(u32, 6), weights.slice(f32, 6)) |id, weight| { + try std.testing.expectEqual(std.math.maxInt(u32), id); + try std.testing.expect(std.math.isNan(weight)); + } + var combine = try ops.Pipelines.load(device); + defer combine.deinit(); + const routed = try buffers.get(u16, 6 * 4); + const shared = try buffers.get(u16, 4); + const output = try buffers.get(u16, 4); + // Invalid routes skip expert writes; finite slots from a prior chunk must not hide that failure. + @memset(routed.slice(u16, 24), support.bf(2)); + @memset(shared.slice(u16, 4), support.bf(3)); + const combine_cb = queue.commandBuffer(); + const combine_encoder = combine_cb.compute(.serial); + ops.encodeCombine(combine_encoder, &combine, view(routed), view(shared), view(ids), view(output), .{ .rows = 1, .width = 4, .routed_row_stride = 24, .slot_stride = 4, .shared_stride = 4, .y_stride = 4 }); + try done(combine_cb, combine_encoder); + for (output.slice(u16, 4)) |value| try std.testing.expect(std.math.isNan(support.fp(value))); +} diff --git a/zig/src/families/deepseek_v41/test_support.zig b/zig/src/families/deepseek_v41/test_support.zig new file mode 100644 index 0000000000..80bf87e8d2 --- /dev/null +++ b/zig/src/families/deepseek_v41/test_support.zig @@ -0,0 +1,40 @@ +const std = @import("std"); +const mtl = @import("metal"); +const View = @import("view.zig").View; + +pub fn bf(x: f32) u16 { + const u: u32 = @bitCast(x); + return @intCast((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} +pub fn fp(x: u16) f32 { + return @bitCast(@as(u32, x) << 16); +} +pub fn view(b: mtl.Buffer) View { + return .{ .buffer = b }; +} +pub const Buffers = struct { + device: mtl.Device, + items: [32]mtl.Buffer = undefined, + count: usize = 0, + + pub fn get(b: *Buffers, comptime T: type, n: usize) !mtl.Buffer { + if (b.count == b.items.len) return error.TooManyTestBuffers; + const x = try b.device.buffer(try std.math.mul(usize, @sizeOf(T), n), mtl.ResourceOptions.shared); + b.items[b.count] = x; + b.count += 1; + @memset(x.slice(T, n), 0); + return x; + } + pub fn deinit(b: *Buffers) void { + for (b.items[0..b.count]) |x| x.deinit(); + } +}; +pub fn done(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void { + e.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); +} +pub fn close(expected: f64, actual: f32, tolerance: f64) !void { + try std.testing.expect(@abs(expected - @as(f64, actual)) <= tolerance); +} diff --git a/zig/src/families/deepseek_v41/workspace.zig b/zig/src/families/deepseek_v41/workspace.zig new file mode 100644 index 0000000000..fcef57cd6e --- /dev/null +++ b/zig/src/families/deepseek_v41/workspace.zig @@ -0,0 +1,107 @@ +//! Fixed chunk-sized activations: no layer history or context-by-query score matrices. +const std = @import("std"); +const mtl = @import("metal"); +const exl3 = @import("../../core/exl3_projection.zig"); +const View = @import("view.zig").View; +const hash = @import("engram_hash.zig"); +const groups = @import("../../core/expert_groups.zig"); + +pub const Field = enum { + ids, + token_map, + history, + stream, + collapsed, + normalized, + branch, + previous_pre, + attention_pre, + post, + comb, + mixes, + inv, + engram_ids_0, + engram_ids_1, + engram_gather, + engram_kv, + router_logits, + route_ids, + route_weights, + gate, + up, + activation, + routed, + shared_gate, + shared_up, + shared_activation, + shared_down, + logits, + projection, + grouping, +}; +const count = std.enums.values(Field).len; + +pub const Plan = struct { + rows: u32, + offsets: [count]usize, + bytes: usize, + + pub fn init(rows: u32) !Plan { + if (rows == 0 or rows > 4096) return error.InvalidPrefillChunk; + const t: usize = rows; + const grouped = if (rows > 1) try groups.Plan.init(t * 6, 384) else null; + const projection_bytes = @max(try exl3.scratchBytes(5120, 129280, t), if (grouped) |g| @max(try exl3.scratchBytes(5120, 2304, g.padded), try exl3.scratchBytes(2304, 5120, g.padded)) else try exl3.scratchBytes(2304, 5120, t * 6)); + const sizes = [_]usize{ + t * 4, hash.raw_vocab * 4, 16, t * 20480 * 2, t * 5120 * 2, t * 5120 * 2, t * 5120 * 2, + t * 4 * 4, t * 4 * 4, t * 4 * 4, t * 16 * 4, t * 24 * 4, t * 4, t * 24 * 4, + t * 24 * 4, t * 6144 * 2, t * 25600 * 2, t * 384 * 4, t * 6 * 4, t * 6 * 4, t * 6 * 2304 * 2, + t * 6 * 2304 * 2, t * 6 * 2304 * 2, t * 6 * 5120 * 2, t * 2304 * 2, t * 2304 * 2, t * 2304 * 2, t * 5120 * 2, + t * 129280 * 4, projection_bytes, if (grouped) |g| g.bytes else 0, + }; + var out = Plan{ .rows = rows, .offsets = undefined, .bytes = 0 }; + for (sizes, 0..) |size, i| { + out.offsets[i] = std.mem.alignForward(usize, out.bytes, 16); + out.bytes = try std.math.add(usize, out.offsets[i], size); + } + out.bytes = std.mem.alignForward(usize, out.bytes, 16384); + return out; + } +}; + +pub const Workspace = struct { + buffer: mtl.Buffer, + plan: Plan, + + pub fn init(device: mtl.Device, plan: Plan, token_map: []const u8) !Workspace { + if (token_map.len != hash.raw_vocab * 4) return error.InvalidTokenMap; + const buffer = try device.buffer(plan.bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + const w = Workspace{ .buffer = buffer, .plan = plan }; + @memcpy(buffer.contents()[plan.offsets[@backingInt(Field.token_map)]..][0..token_map.len], token_map); + @memset(buffer.contents()[plan.offsets[@backingInt(Field.history)]..][0..16], 0); + return w; + } + + pub fn deinit(w: *Workspace) void { + w.buffer.deinit(); + } + pub fn view(w: *const Workspace, field: Field) View { + return .{ .buffer = w.buffer, .offset = w.plan.offsets[@backingInt(field)] }; + } + pub fn projection(w: *const Workspace) exl3.View { + const v = w.view(.projection); + return .{ .buffer = v.buffer, .offset = v.offset }; + } + pub fn grouping(w: *const Workspace, rows: u32) !groups.Scratch { + if (rows < 2 or rows > w.plan.rows) return error.InvalidPrefillChunk; + const v = w.view(.grouping); + return .{ .buffer = v.buffer, .offset = v.offset, .plan = try groups.Plan.init(@as(usize, rows) * 6, 384) }; + } + pub fn tokenIds(w: *const Workspace, rows: usize) []u32 { + const v = w.view(.ids); + return @as([*]u32, @ptrCast(@alignCast(v.buffer.contents() + v.offset)))[0..rows]; + } + pub fn logits(w: *const Workspace, rows: usize) []const f32 { + const v = w.view(.logits); + return @as([*]const f32, @ptrCast(@alignCast(v.buffer.contents() + v.offset)))[0 .. rows * 129280]; + } +}; diff --git a/zig/src/metal/counters.zig b/zig/src/metal/counters.zig index 610e22086b..0bbd65f186 100644 --- a/zig/src/metal/counters.zig +++ b/zig/src/metal/counters.zig @@ -1,4 +1,4 @@ -//! Metal counter sampling; owned resources must outlive every command buffer using them. +//! Metal counter sampling; owned descriptors, samples and fences must outlive their command buffers. const std = @import("std"); const objc = @import("objc.zig"); const types = @import("types.zig"); @@ -107,7 +107,7 @@ pub const SampleBuffer = struct { objc.msg(void, encoder.id, "sampleCountersInBuffer:atSampleIndex:withBarrier:", .{ self.id, index, true }); } - /// Caller must ensure GPU completion; borrows a bounded NSData result without a GPU blit or second copy. + /// Borrow bounded NSData bytes after completion; same-command blit resolution is unsupported, no result copy. pub fn resolveCompleted(self: SampleBuffer, count: usize) !Resolved { if (count == 0 or count > self.count) return error.InvalidSampleCount; const data = objc.msg(?objc.Id, self.id, "resolveCounterRange:", .{ diff --git a/zig/src/metal/metal.zig b/zig/src/metal/metal.zig index 6404c4227d..71bed247d2 100644 --- a/zig/src/metal/metal.zig +++ b/zig/src/metal/metal.zig @@ -29,6 +29,7 @@ pub const keepalive = @import("keepalive.zig"); pub const clock = @import("clock.zig"); pub const counters = @import("counters.zig"); pub const ane = @import("ane.zig"); +pub const counters = @import("counters.zig"); test { std.testing.refAllDecls(@This()); diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 0cf0c9f5a4..9461cc627f 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -28,6 +28,10 @@ pub const deepseek_v41_config = @import("families/deepseek_v41/config.zig"); pub const deepseek_v41_load = @import("families/deepseek_v41/load.zig"); pub const deepseek_v41_weights = @import("families/deepseek_v41/weights.zig"); pub const deepseek_v41_view = @import("families/deepseek_v41/view.zig"); +pub const deepseek_v41 = @import("families/deepseek_v41/model.zig"); +pub const deepseek_v41_profiler = @import("families/deepseek_v41/profiler.zig"); +pub const gpu_counters = @import("core/gpu_counters.zig"); +pub const expert_groups = @import("core/expert_groups.zig"); pub const deepseek_v41_engram_hash = @import("families/deepseek_v41/engram_hash.zig"); pub const exl3_format = @import("core/exl3_format.zig"); pub const exl3_projection = @import("core/exl3_projection.zig"); diff --git a/zig/tests/deepseek_v41_reference/attention_math.zig b/zig/tests/deepseek_v41_reference/attention_math.zig new file mode 100644 index 0000000000..d3df4d0ac7 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/attention_math.zig @@ -0,0 +1,298 @@ +//! CPU CSA2 from DeepSeek AI inference/{model,kernel}.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +const std = @import("std"); +const tf = @import("tensorfold"); +const m = @import("math.zig"); +const projection = @import("projection.zig"); +const Config = tf.deepseek_v41_config.Text; +const Attention = tf.deepseek_v41_weights.Attention; + +pub const max_capacity = 4096; +pub const State = struct { + gpa: std.mem.Allocator, + capacity: usize, + window: []u16, + sources: []u16, + scores: []f32, + candidates: []bool, + blocks: []f32, + active_owner: usize = 0, + selected: [512]u32 = undefined, + selected_len: usize = 0, + partial_kv: [4][2][512]f32 = undefined, + partial_gate: [4][2][512]f32 = undefined, + qr: [1280]u16 = undefined, + q: [32768]u16 = undefined, + index_q: [4096]u16 = undefined, + index_weights: [32]f32 = undefined, + latent: [512]u16 = undefined, + index_key: [128]u16 = undefined, + kv: [512]u16 = undefined, + o: [32768]u16 = undefined, + grouped: [8192]u16 = undefined, + observer: ?Observer = null, + sparse_snapshot: ?[]u16 = null, + pub const Observer = struct { + ctx: *anyopaque, + capture: *const fn (*anyopaque, tf.deepseek_v41.attention.DiagnosticStage, []const u16) void, + }; + fn observe(self: *State, stage: tf.deepseek_v41.attention.DiagnosticStage, values: []const u16) void { + if (self.observer) |observer| observer.capture(observer.ctx, stage, values); + } + + pub fn init(gpa: std.mem.Allocator, capacity: usize) !State { + if (capacity == 0 or capacity > max_capacity) return error.InvalidAttentionReferenceCapacity; + const window = try gpa.alloc(u16, 40 * 128 * 512); + errdefer gpa.free(window); + const sources = try gpa.alloc(u16, 4 * capacity * 640); + errdefer gpa.free(sources); + const scores = try gpa.alloc(f32, capacity); + errdefer gpa.free(scores); + const candidates = try gpa.alloc(bool, capacity); + errdefer gpa.free(candidates); + const blocks = try gpa.alloc(f32, (capacity + 7) / 8); + @memset(window, 0); + return .{ .gpa = gpa, .capacity = capacity, .window = window, .sources = sources, .scores = scores, .candidates = candidates, .blocks = blocks }; + } + pub fn deinit(self: *State) void { + self.gpa.free(self.window); + self.gpa.free(self.sources); + self.gpa.free(self.scores); + self.gpa.free(self.candidates); + self.gpa.free(self.blocks); + self.* = undefined; + } + fn mainRow(self: *State, owner: usize, row: usize) []u16 { + return self.sources[owner * self.capacity * 640 + row * 640 ..][0..512]; + } + fn indexRow(self: *State, owner: usize, row: usize) []u16 { + return self.sources[owner * self.capacity * 640 + row * 640 + 512 ..][0..128]; + } + fn publish(self: *State, a: Attention, layer: usize, x: []const u16, position: usize, c: Config, scratch: *projection.Scratch) !bool { + @setFloatMode(.strict); + const compressor = a.compressor orelse return false; + const owner = (layer - 2) / 6; + self.active_owner = owner; + const ratio: usize = c.compress_ratios[layer]; + if (ratio == 1) { + try projection.apply(compressor.wkv, x, &self.latent, scratch); + } else { + const slot = position % ratio; + try projection.applyF32(compressor.wkv, x, &self.partial_kv[owner][slot], scratch); + try projection.applyF32(compressor.wgate.?, x, &self.partial_gate[owner][slot], scratch); + if ((position + 1) % ratio != 0) return false; + for (0..512) |d| { + const high = @max(self.partial_gate[owner][0][d], self.partial_gate[owner][1][d]); + const p0 = @exp(self.partial_gate[owner][0][d] - high); + const p1 = @exp(self.partial_gate[owner][1][d] - high); + const denominator = p0 + p1; + const weighted0 = self.partial_kv[owner][0][d] * (p0 / denominator); + const weighted1 = self.partial_kv[owner][1][d] * (p1 / denominator); + const pooled = weighted0 + weighted1; + self.latent[d] = m.b(pooled); + } + } + m.rms(&self.latent, compressor.norm, c.rms_norm_eps, &self.latent); + const indexer = a.indexer.?; + try projection.apply(indexer.wk.?, &self.latent, &self.index_key, scratch); + m.rms(&self.index_key, indexer.k_norm.?, c.rms_norm_eps, &self.index_key); + const group_position = position + 1 - ratio; + m.rope(&self.index_key, group_position, true, false, c); + m.quantize(&self.index_key, .index); + const row = position / ratio; + @memcpy(self.indexRow(owner, row), &self.index_key); + m.rope(&self.latent, group_position, true, false, c); + m.quantize(&self.latent, .main); + @memcpy(self.mainRow(owner, row), &self.latent); + return true; + } + fn chooseCandidates(self: *State, count: usize, c: Config) void { + @memset(self.candidates[0..count], false); + const nblocks = (count + 7) / 8; + for (0..nblocks) |block| { + var score: f32 = -std.math.inf(f32); + for (block * 8..@min(block * 8 + 8, count)) |row| score = @max(score, self.scores[row]); + self.blocks[block] = score; + } + if (nblocks == 0) return; + self.blocks[nblocks - 1] = std.math.inf(f32); + for (0..@min(nblocks, c.candidate_topk_blocks)) |_| { + var best: usize = 0; + for (0..nblocks) |i| if (self.blocks[i] > self.blocks[best]) { + best = i; + }; + if (self.blocks[best] == -std.math.inf(f32)) break; + for (best * 8..@min(best * 8 + 8, count)) |row| self.candidates[row] = true; + self.blocks[best] = -std.math.inf(f32); + } + } + fn selectPositions(self: *State, a: Attention, layer: usize, x: []const u16, position: usize, count: usize, c: Config, scratch: *projection.Scratch) !void { + const indexer = a.indexer orelse return; + self.selected_len = @min(count, c.index_topk); + if (count == 0) return; + try projection.apply(indexer.wq_b, &self.qr, &self.index_q, scratch); + for (0..32) |head| { + const q = self.index_q[head * 128 ..][0..128]; + m.rope(q, position, true, false, c); + m.quantize(q, .index); + } + m.dense(indexer.weights_proj, x, &self.index_weights); + const scale: f32 = 1 / @sqrt(@as(f32, 128 * 32)); + // Python scalar multiplication preserves the source weights projection's BF16 dtype. + for (&self.index_weights) |*v| v.* = m.f(m.b(m.f(m.b(v.*)) * scale)); + for (0..count) |row| { + const key = self.indexRow(self.active_owner, row); + var score: f64 = 0; + for (0..32) |head| { + var dot: f64 = 0; + for (key, 0..) |v, d| dot += @as(f64, m.f(v)) * m.f(self.index_q[head * 128 + d]); + // einsum produces BF16, then BF16 multiply and head sum, in the pinned source. + const rectified = @max(m.f(m.b(@floatCast(dot))), 0); + const weighted = m.f(m.b(rectified * self.index_weights[head])); + score += weighted; + } + self.scores[row] = m.f(m.b(@floatCast(score))); + } + if (layer == c.candidate_source_layer_id) self.chooseCandidates(count, c); + if (layer > c.candidate_source_layer_id) for (0..count) |row| { + if (!self.candidates[row]) self.scores[row] = -std.math.inf(f32); + }; + for (0..self.selected_len) |pick| { + var best: usize = 0; + for (0..count) |i| if (self.scores[i] > self.scores[best]) { + best = i; + }; + self.selected[pick] = @intCast(best); + self.scores[best] = -std.math.inf(f32); + } + std.mem.sort(u32, self.selected[0..self.selected_len], {}, std.sort.asc(u32)); + } + pub fn apply(self: *State, a: Attention, layer: usize, x: []const u16, position: usize, c: Config, out: []u16, scratch: *projection.Scratch) !void { + const compressed = c.compress_ratios[layer] != 0; + try projection.apply(a.wq_a, x, &self.qr, scratch); + self.observe(.qr_projection, &self.qr); + m.rms(&self.qr, a.q_norm, c.rms_norm_eps, &self.qr); + self.observe(.qr_norm, &self.qr); + try projection.apply(a.wq_b, &self.qr, &self.q, scratch); + self.observe(.q_projection, &self.q); + for (0..64) |head| m.rope(self.q[head * 512 ..][0..512], position, compressed, false, c); + self.observe(.q, &self.q); + try projection.apply(a.wkv, x, &self.kv, scratch); + self.observe(.kv_projection, &self.kv); + m.rms(&self.kv, a.kv_norm, c.rms_norm_eps, &self.kv); + self.observe(.kv_norm, &self.kv); + m.rope(&self.kv, position, compressed, false, c); + self.observe(.kv_rope, &self.kv); + m.quantize(&self.kv, .window); + self.observe(.kv_quantized, &self.kv); + const window = self.window[layer * 128 * 512 ..][0 .. 128 * 512]; + @memcpy(window[(position % 128) * 512 ..][0..512], &self.kv); + var count: usize = 0; + if (compressed) { + _ = try self.publish(a, layer, x, position, c, scratch); + count = (position + 1) / c.compress_ratios[layer]; + try self.selectPositions(a, layer, x, position, count, c, scratch); + } + var indices: [640]i32 = undefined; + var slots: usize = 0; + if (position == 0) { + indices[0] = 0; + slots = 1; + } else { + const oldest = position % 128 + 1; + for (0..128) |i| { + const at = (oldest + i) % 128; + indices[i] = if (at > position) -1 else @intCast(at); + } + slots = 128; + } + if (compressed) for (self.selected[0..self.selected_len]) |row| { + indices[slots] = @intCast(128 + row); + slots += 1; + }; + for (0..64) |head| { + self.sparseHead(self.q[head * 512 ..][0..512], window, indices[0..slots], m.value(a.sink, head), self.o[head * 512 ..][0..512]); + self.observeHead(head); + m.rope(self.o[head * 512 ..][0..512], position, compressed, true, c); + } + if (self.sparse_snapshot) |snapshot| self.observe(.sparse_o, snapshot); + self.observe(.o, &self.o); + for (0..8) |group| try projection.apply(a.wo_a[group], self.o[group * 4096 ..][0..4096], self.grouped[group * 1024 ..][0..1024], scratch); + self.observe(.grouped, &self.grouped); + try projection.apply(a.wo_b, &self.grouped, out, scratch); + self.observe(.wo_b, out); + } + fn observeHead(self: *State, head: usize) void { + if (self.sparse_snapshot) |snapshot| + @memcpy(snapshot[head * 512 ..][0..512], self.o[head * 512 ..][0..512]); + } + fn keyRow(self: *State, window: []const u16, key_id: i32) []const u16 { + const at: usize = @intCast(key_id); + return if (at < 128) window[at * 512 ..][0..512] else self.mainRow(self.active_owner, at - 128); + } + pub fn sparseHead(self: *State, q: []const u16, window: []const u16, indices: []const i32, sink: f32, out: []u16) void { + @setFloatMode(.strict); + var acc: [512]f32 = @splat(0); + var high: f32 = -1e30; + var denominator: f32 = 0; + var base: usize = 0; + while (base < indices.len) : (base += 64) { + var probabilities: [64]f32 = undefined; + const width = @min(64, indices.len - base); + var next_high = high; + for (0..width) |j| { + const index = indices[base + j]; + var dot: f64 = 0; + if (index >= 0) { + const key = self.keyRow(window, index); + for (q, key) |a, v| dot += @as(f64, m.f(a)) * m.f(v); + } + probabilities[j] = if (index < 0) -std.math.inf(f32) else @as(f32, @floatCast(dot)) * (1 / @sqrt(@as(f32, 512))); + next_high = @max(next_high, probabilities[j]); + } + const correction = @exp(high - next_high); + var sum: f64 = 0; + for (0..width) |j| { + probabilities[j] = @exp(probabilities[j] - next_high); + sum += probabilities[j]; + } + denominator = denominator * correction + @as(f32, @floatCast(sum)); + for (0..512) |d| { + var accumulated: f64 = acc[d] * correction; + for (0..width) |j| if (indices[base + j] >= 0) { + // Online 64-column source kernel rounds unnormalized probabilities to BF16 before P*V. + accumulated += @as(f64, m.f(m.b(probabilities[j]))) * m.f(self.keyRow(window, indices[base + j])[d]); + }; + acc[d] = @floatCast(accumulated); + } + high = next_high; + } + denominator += @exp(sink - high); + for (out, acc) |*v, a| v.* = m.b(a / denominator); + } +}; + +test "online 64-column sparse attention rounds probabilities before value GEMM" { + var state = try State.init(std.testing.allocator, 2); + defer state.deinit(); + var q: [512]u16 = @splat(0); + q[0] = m.b(8); + state.window[1] = m.b(1); + state.window[512] = m.b(1); + var indices: [65]i32 = @splat(-1); + indices[0] = 0; + indices[1] = 1; + indices[64] = 0; + var out: [512]u16 = undefined; + state.sparseHead(&q, state.window[0..1024], &indices, 0, &out); + const probability = @exp(-@as(f32, 8) * (1 / @sqrt(@as(f32, 512)))); + var denominator = @as(f32, @floatCast(@as(f64, probability) + 1)); + denominator += probability; + denominator += probability; + const expected = m.b((2 * m.f(m.b(probability))) / denominator); + try std.testing.expectEqual(expected, out[1]); +} + +test "attention reference capacity rejects unbounded history before allocation" { + try std.testing.expectError(error.InvalidAttentionReferenceCapacity, State.init(std.testing.allocator, 0)); + try std.testing.expectError(error.InvalidAttentionReferenceCapacity, State.init(std.testing.allocator, max_capacity + 1)); +} diff --git a/zig/tests/deepseek_v41_reference/engram_math.zig b/zig/tests/deepseek_v41_reference/engram_math.zig new file mode 100644 index 0000000000..d57567957f --- /dev/null +++ b/zig/tests/deepseek_v41_reference/engram_math.zig @@ -0,0 +1,106 @@ +//! CPU Engram from DeepSeek AI inference/engram.py and model.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +const std = @import("std"); +const tf = @import("tensorfold"); +const m = @import("math.zig"); +const projection = @import("projection.zig"); +const Engram = tf.deepseek_v41_weights.Engram; +const Config = tf.deepseek_v41_config.Text; +const multipliers = [2][4]u64{ + .{ 76632096046245, 4839876093313, 35959672319349, 73987337458391 }, + .{ 67716810739261, 51510806800915, 30921347202721, 82619226485591 }, +}; +const primes = [2][24]u32{ + .{ 16000057, 16000079, 16000081, 16000097, 16000121, 16000129, 16000133, 16000183, 16000189, 16000207, 16000211, 16000253, 16000277, 16000289, 16000307, 16000321, 16000339, 16000381, 16000393, 16000399, 16000403, 16000409, 16000447, 16000463 }, + .{ 16000477, 16000487, 16000499, 16000507, 16000511, 16000573, 16000609, 16000627, 16000667, 16000669, 16000693, 16000697, 16000711, 16000729, 16000759, 16000769, 16000781, 16000799, 16000813, 16000819, 16000841, 16000877, 16000879, 16000889 }, +}; +fn word(raw: []const u8, at: usize) u32 { + return std.mem.readInt(u32, raw[at..][0..4], .little); +} +pub const State = struct { + map: []const u8, + history: [3]u32 = .{ 2, 2, 2 }, + gathered: [6144]u16 = undefined, + kv: [25600]u16 = undefined, + + pub fn init(blob: []const u8) !State { + const at = 96 + 2 * 232; + if (blob.len != at + 129280 * 4 or !std.mem.eql(u8, blob[0..8], "DSENGR01")) + return error.InvalidEngramMetadata; + for ([_]u32{ 129280, 99092, 2, 4, 8, 2 }, 0..) |v, i| + if (word(blob, 8 + i * 4) != v) return error.InvalidEngramMetadata; + for (0..2) |layer| { + const base = 96 + layer * 232; + if (word(blob, base) != (if (layer == 0) @as(u32, 1) else @as(u32, 14))) + return error.InvalidEngramMetadata; + var offset: u32 = 0; + for (multipliers[layer], 0..) |v, j| + if (std.mem.readInt(u64, blob[base + 8 + 8 * j ..][0..8], .little) != v) + return error.InvalidEngramMetadata; + for (primes[layer], 0..) |v, j| { + if (word(blob, base + 40 + 4 * j) != v or word(blob, base + 136 + 4 * j) != offset) + return error.InvalidEngramMetadata; + offset += v; + } + if (word(blob, base + 4) != offset) return error.InvalidEngramMetadata; + } + const map = blob[at..]; + const expected = [_]u8{ 0xc6, 0x0a, 0x86, 0x32, 0x2e, 0xc1, 0x7b, 0x41, 0x42, 0xbf, 0xef, 0x3c, 0x57, 0xa8, 0xd8, 0x1f, 0xb4, 0x28, 0x55, 0x0c, 0xdf, 0x48, 0x7f, 0x88, 0xa4, 0x32, 0x0c, 0xb5, 0x9f, 0xe4, 0x64, 0x81 }; + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(map, &digest, .{}); + if (!std.mem.eql(u8, &expected, &digest)) return error.InvalidEngramTokenMap; + return .{ .map = map }; + } + pub fn compressed(self: *const State, id: u32) !u32 { + if (id >= 129280) return error.InvalidToken; + return word(self.map, @as(usize, id) * 4); + } + pub fn advance(self: *State, id: u32) void { + self.history[2] = self.history[1]; + self.history[1] = self.history[0]; + self.history[0] = id; + } + pub fn apply(self: *State, e: Engram, layer: usize, id: u32, h: []u16, c: Config, scratch: *projection.Scratch) !void { + @setFloatMode(.strict); + const tokens = [4]u32{ id, self.history[0], self.history[1], self.history[2] }; + var rolling = @as(u64, tokens[0]) * multipliers[layer][0]; + var offset: usize = 0; + for (1..4) |lookback| { + rolling ^= @as(u64, tokens[lookback]) * multipliers[layer][lookback]; + for (0..8) |head| { + const column = (lookback - 1) * 8 + head; + const row = offset + @as(usize, @intCast(rolling % primes[layer][column])); + for (0..256) |d| self.gathered[column * 256 + d] = + m.b(m.value(e.table, row * 256 + d) * m.value(e.scales, row * 8 + d / 32)); + offset += primes[layer][column]; + } + } + try projection.apply(e.wkv, &self.gathered, &self.kv, scratch); + for (0..4) |copy| { + var hh: f64 = 0; + var kk: f64 = 0; + var dot: f64 = 0; + for (0..5120) |d| { + const at = copy * 5120 + d; + const x = m.f(h[at]); + const k = m.f(self.kv[at]); + const weight = m.value(e.q_weight, at) * m.value(e.k_weight, at); + hh += @as(f64, x) * x; + kk += @as(f64, k) * k; + const weighted: f32 = (x * weight) * k; + dot += weighted; + } + const h_mean = m.meanFromWide(hh, 5120); + const k_mean = m.meanFromWide(kk, 5120); + const h_inverse = m.idealRsqrt(h_mean + c.rms_norm_eps); + const k_inverse = m.idealRsqrt(k_mean + c.rms_norm_eps); + const rstd: f32 = h_inverse * k_inverse; + const dot_f32: f32 = @floatCast(dot); + const normalized: f32 = (dot_f32 * rstd) * (1 / @sqrt(@as(f32, 5120))); + const gate = m.sigmoid(std.math.copysign(@sqrt(@max(@abs(normalized), 1e-6)), normalized)); + for (0..5120) |d| { + const at = copy * 5120 + d; + h[at] = m.b(m.f(h[at]) + gate * m.f(self.kv[20480 + d])); + } + } + } +}; diff --git a/zig/tests/deepseek_v41_reference/math.zig b/zig/tests/deepseek_v41_reference/math.zig new file mode 100644 index 0000000000..a6dccb2b93 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/math.zig @@ -0,0 +1,289 @@ +//! Independent CPU equations: DeepSeek AI inference/{model,kernel}.py, rev 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. +const std = @import("std"); +const tf = @import("tensorfold"); +const Tensor = tf.checkpoint_metal.Tensor; +const Hc = tf.deepseek_v41_weights.Hc; +const Config = tf.deepseek_v41_config.Text; + +pub fn f(x: u16) f32 { + return @bitCast(@as(u32, x) << 16); +} +pub fn b(x: f32) u16 { + const bits: u32 = @bitCast(x); + if ((bits & 0x7fffffff) > 0x7f800000) return @truncate((bits >> 16) | 0x40); + return @truncate((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); +} +pub fn bytes(t: Tensor) []const u8 { + return t.buffer.contents()[t.offset..][0..t.bytes]; +} +pub fn value(t: Tensor, index: usize) f32 { + const raw = bytes(t); + return switch (t.dtype) { + .bf16 => f(std.mem.readInt(u16, raw[index * 2 ..][0..2], .little)), + .f16 => @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, raw[index * 2 ..][0..2], .little)))), + .f32 => @bitCast(std.mem.readInt(u32, raw[index * 4 ..][0..4], .little)), + .f8_e4m3 => e4(raw[index]), + .f8_e8m0 => e8(raw[index]), + else => unreachable, + }; +} +pub fn e4(code: u8) f32 { + const magnitude = code & 127; + if (magnitude == 127) return std.math.nan(f32); + const exponent = magnitude >> 3; + const mantissa = magnitude & 7; + const power: f32 = @bitCast((@as(u32, exponent) + 120) << 23); + const v: f32 = if (exponent == 0) @as(f32, @floatFromInt(mantissa)) / 512 else power * (1 + @as(f32, @floatFromInt(mantissa)) / 8); + return if ((code & 128) != 0) -v else v; +} +pub fn e8(code: u8) f32 { + if (code == 255) return std.math.nan(f32); + if (code == 0) return 0x1p-127; + return @bitCast(@as(u32, code) << 23); +} +pub fn roundE4(x: f32) f32 { + const a = @min(@abs(x), 448); + var low: u8 = 0; + var high: u8 = 126; + while (low < high) { + const mid = low + (high - low + 1) / 2; + if (e4(mid) <= a) low = mid else high = mid - 1; + } + var best = low; + if (low < 126) { + const below = a - e4(low); + const above = e4(low + 1) - a; + if (above < below or (above == below and (low + 1) % 2 == 0)) best += 1; + } + return std.math.copysign(e4(best), x); +} +pub fn roundE2(x: f32) f32 { + const levels = [_]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + var best: usize = 0; + var distance: f32 = std.math.inf(f32); + for (levels, 0..) |v, i| { + const d = @abs(@abs(x) - v); + if (d < distance or (d == distance and i % 2 == 0)) { + best = i; + distance = d; + } + } + return std.math.copysign(levels[best], x); +} +fn powerScale(x: f32) f32 { + const bits: u32 = @bitCast(x); + const exponent = (bits >> 23) & 255; + const rounded = exponent + @as(u32, @intFromBool((bits & 0x7fffff) != 0)); + return @bitCast(rounded << 23); +} +pub const Quant = enum { window, main, index }; +pub fn quantize(x: []u16, kind: Quant) void { + const group: usize = if (kind == .main) 16 else 32; + var start: usize = 0; + while (start < x.len) : (start += group) { + var a: f32 = if (kind == .window) 1e-4 else if (kind == .main) 6 * 0x1p-9 else 6 * 0x1p-126; + for (x[start..][0..group]) |v| a = @max(a, @abs(f(v))); + const scale = if (kind == .main) roundE4(a / 6) else powerScale(a * (if (kind == .window) @as(f32, 1.0 / 448.0) else @as(f32, 1.0 / 6.0))); + for (x[start..][0..group]) |*v| { + const z = f(v.*) / scale; + v.* = b((if (kind == .window) roundE4(z) else roundE2(@max(-6, @min(6, z)))) * scale); + } + } +} +pub fn dense(t: Tensor, x: []const u16, out: []f32) void { + for (out, 0..) |*y, row| { + var sum: f64 = 0; + for (x, 0..) |v, k| sum += @as(f64, f(v)) * value(t, row * x.len + k); + y.* = @floatCast(sum); + } +} +/// CUDA MeanOps scalar projection; the internal reduction remains independently wide. +pub fn meanFromWide(sum: f64, count: usize) f32 { + @setFloatMode(.strict); + const accumulator: f32 = @floatCast(sum); + const factor: f32 = 1 / @as(f32, @floatFromInt(count)); + return accumulator * factor; +} +/// Independent ideal FP32 unary result, not an emulation of CUDA's rsqrt approximation. +pub fn idealRsqrt(x: f32) f32 { + @setFloatMode(.strict); + return @floatCast(1 / @sqrt(@as(f64, x))); +} +pub fn rms(x: []const u16, weight: Tensor, eps: f32, out: []u16) void { + @setFloatMode(.strict); + var sum: f64 = 0; + for (x) |v| sum += @as(f64, f(v)) * f(v); + const mean = meanFromWide(sum, x.len); + const r = idealRsqrt(mean + eps); + for (x, out, 0..) |v, *y, i| y.* = b(f(v) * r * value(weight, i)); +} +pub fn sigmoid(x: f32) f32 { + return 1 / (1 + @exp(-x)); +} +pub const Mix = struct { pre: [4]f32, post: [4]f32, comb: [16]f32 }; +pub fn mixes(x: []const u16, hc: Hc, c: Config) Mix { + @setFloatMode(.strict); + var projected: [24]f32 = undefined; + dense(hc.projection, x, &projected); + var squares: f64 = 0; + for (x) |v| squares += @as(f64, f(v)) * f(v); + const mean = meanFromWide(squares, x.len); + const r = idealRsqrt(mean + c.rms_norm_eps); + var m: Mix = undefined; + for (0..4) |i| { + m.pre[i] = sigmoid(@mulAdd(f32, projected[i] * r, value(hc.scale, 0), value(hc.base, i))) + c.hc_eps; + m.post[i] = 2 * sigmoid(@mulAdd(f32, projected[i + 4] * r, value(hc.scale, 1), value(hc.base, i + 4))); + } + for (0..4) |i| { + var high: f32 = -std.math.inf(f32); + for (0..4) |j| { + const at = i * 4 + j; + m.comb[at] = @mulAdd(f32, projected[at + 8] * r, value(hc.scale, 2), value(hc.base, at + 8)); + high = @max(high, m.comb[at]); + } + for (0..4) |j| { + m.comb[i * 4 + j] = @exp(m.comb[i * 4 + j] - high); + } + // TileLang 0.1.8 emits FP32 XOR reductions with offsets 2, then 1. + const denominator = (m.comb[i * 4] + m.comb[i * 4 + 2]) + + (m.comb[i * 4 + 1] + m.comb[i * 4 + 3]); + for (0..4) |j| m.comb[i * 4 + j] = m.comb[i * 4 + j] / denominator + c.hc_eps; + } + for (0..c.hc_sinkhorn_iters) |iteration| { + if (iteration != 0) for (0..4) |i| { + const denominator = ((m.comb[i * 4] + m.comb[i * 4 + 2]) + + (m.comb[i * 4 + 1] + m.comb[i * 4 + 3])) + c.hc_eps; + for (0..4) |j| m.comb[i * 4 + j] /= denominator; + }; + for (0..4) |j| { + const denominator = ((m.comb[j] + m.comb[8 + j]) + + (m.comb[4 + j] + m.comb[12 + j])) + c.hc_eps; + for (0..4) |i| m.comb[i * 4 + j] /= denominator; + } + } + return m; +} +pub fn collapse(h: []const u16, pre: [4]f32, out: []u16) void { + @setFloatMode(.strict); + for (out, 0..) |*v, d| { + var sum: f32 = 0; + for (0..4) |i| { + const product: f32 = pre[i] * f(h[i * out.len + d]); + sum += product; + } + v.* = b(sum); + } +} +pub fn expand(h: []u16, x: []const u16, m: Mix) void { + @setFloatMode(.strict); + for (0..x.len) |d| { + var old: [4]f32 = undefined; + for (0..4) |i| old[i] = f(h[i * x.len + d]); + for (0..4) |i| { + var residual: f32 = 0; + for (0..4) |j| { + const product: f32 = m.comb[j * 4 + i] * old[j]; + residual += product; + } + h[i * x.len + d] = b(m.post[i] * f(x[d]) + residual); + } + } +} +pub fn ropeFrequency(index: usize, theta: f32, original: u32, factor: f32, beta_fast: f32, beta_slow: f32) f32 { + @setFloatMode(.strict); + const exponent: f32 = @as(f32, @floatFromInt(index * 2)) / 64; + const power: f32 = @floatCast(std.math.pow(f64, theta, exponent)); + var frequency: f32 = 1 / power; + if (original > 0) { + const original_f64: f64 = @floatFromInt(original); + const theta_f64: f64 = theta; + const low: f32 = @floatCast(@max(@floor(64 * @log(original_f64 / + (@as(f64, beta_fast) * 2 * std.math.pi)) / (2 * @log(theta_f64))), 0)); + const high: f32 = @floatCast(@min(@ceil(64 * @log(original_f64 / + (@as(f64, beta_slow) * 2 * std.math.pi)) / (2 * @log(theta_f64))), 63)); + const ramp: f32 = @max(0, @min(1, (@as(f32, @floatFromInt(index)) - low) / @max(high - low, 1e-3))); + const smooth: f32 = 1 - ramp; + frequency = frequency / factor * (1 - smooth) + frequency * smooth; + } + return frequency; +} +pub fn rope(x: []u16, position: usize, compressed: bool, inverse: bool, c: Config) void { + for (0..32) |i| { + const frequency = ropeFrequency(i, if (compressed) c.compress_rope_theta else c.rope_theta, if (compressed) c.rope_scaling.original_max_position_embeddings else 0, c.rope_scaling.factor, c.rope_scaling.beta_fast, c.rope_scaling.beta_slow); + const angle: f32 = @as(f32, @floatFromInt(position)) * frequency; + const cosine = @cos(angle); + const sine = @sin(angle) * (if (inverse) @as(f32, -1) else @as(f32, 1)); + const at = x.len - 64 + 2 * i; + const a = f(x[at]); + const z = f(x[at + 1]); + x[at] = b(a * cosine - z * sine); + x[at + 1] = b(a * sine + z * cosine); + } +} + +test "BF16 probability ties and FP4 ties retain source boundaries" { + try std.testing.expectEqual(@as(u16, 0x3f80), b(@bitCast(@as(u32, 0x3f808000)))); + try std.testing.expectEqual(@as(f32, 2), roundE2(2.5)); + try std.testing.expectEqual(@as(f32, 4), roundE2(5)); + try std.testing.expectEqual(@as(f32, 448), roundE4(500)); +} + +test "delayed hyperconnection residual uses comb columns, not rows" { + var h = [_]u16{ b(1), b(2), b(3), b(4) }; + const mix = Mix{ .pre = .{ 1, 0, 0, 0 }, .post = .{ 0, 0, 0, 0 }, .comb = .{ 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0 } }; + expand(&h, &.{b(0)}, mix); + try std.testing.expectEqualSlices(u16, &.{ b(4), b(1), b(2), b(3) }, &h); +} + +test "hyperconnection FP32 accumulation preserves collapse and branch cancellation" { + var h = [_]u16{ b(1), b(1), b(1), b(1) }; + var collapsed: [1]u16 = undefined; + collapse(&h, .{ 1, 0x1p-24, 0x1p-24, -1 }, &collapsed); + try std.testing.expectEqual(@as(u16, 0), collapsed[0]); + const mix = Mix{ .pre = .{ 1, 0, 0, 0 }, .post = .{ -1, 0, 0, 0 }, .comb = .{ 1, 0, 0, 0, 0x1p-24, 0, 0, 0, 0x1p-24, 0, 0, 0, 0, 0, 0, 0 } }; + expand(&h, &.{b(1)}, mix); + try std.testing.expectEqual(@as(u16, 0), h[0]); +} + +test "hyperconnection rounds elementwise products before cancelling reductions" { + var h = [_]u16{ b(3), b(3), b(0), b(0) }; + const coefficient: f32 = 1 + 0x1p-23; + var collapsed: [1]u16 = undefined; + collapse(&h, .{ -1, coefficient, 0, 0 }, &collapsed); + try std.testing.expectEqual(b(0x1p-21), collapsed[0]); + const mix = Mix{ .pre = .{ 1, 0, 0, 0 }, .post = .{ 0, 0, 0, 0 }, .comb = .{ -1, 0, 0, 0, coefficient, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 } }; + expand(&h, &.{b(0)}, mix); + try std.testing.expectEqual(b(0x1p-21), h[0]); + h = .{ b(3), b(0), b(0), b(0) }; + const branch_mix = Mix{ .pre = .{ 1, 0, 0, 0 }, .post = .{ coefficient, 0, 0, 0 }, .comb = .{ -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 } }; + expand(&h, &.{b(3)}, branch_mix); + try std.testing.expectEqual(b(0x1p-21), h[0]); +} + +test "RoPE and YaRN frequency vectors retain official FP32 elementwise boundaries" { + const parsed = try std.json.parseFromSlice(std.json.Value, std.testing.allocator, @embedFile("rope_frequencies.json"), .{}); + defer parsed.deinit(); + for (parsed.value.object.get("vectors").?.array.items) |vector| { + const theta: f32 = @floatFromInt(vector.object.get("theta").?.integer); + const original: u32 = @intCast(vector.object.get("original").?.integer); + for (vector.object.get("bits").?.array.items, 0..) |expected, index| { + const actual = ropeFrequency(index, theta, original, 16, 32, 1); + const bits: u32 = @intCast(expected.integer); + const source: f32 = @bitCast(bits); + try std.testing.expectApproxEqAbs(source, actual, source * std.math.floatEps(f32)); + } + } +} + +test "ideal FP32 inverse square root avoids an intermediate square-root rounding" { + // Actual CUDA agrees on this HC RMS argument; CPU/MPS's two-round result is one ULP higher. + try std.testing.expectEqual(@as(f32, 0.3651483356952667), idealRsqrt(7.500000953674316)); +} + +test "CUDA mean projection preserves the BF16 RMS affine tie" { + @setFloatMode(.strict); + // Constructed exact-sum row isolates canonical CUDA RMSNorm mean projection through the BF16 affine tie. + const gamma: f32 = @bitCast(@as(u32, 0x3c65de29)); + const inverse = idealRsqrt(meanFromWide(9, 5120) + @as(f32, 1e-20)); + try std.testing.expectEqual(b(1), b((@as(f32, 3) * inverse) * gamma)); +} diff --git a/zig/tests/deepseek_v41_reference/operator_attention.zig b/zig/tests/deepseek_v41_reference/operator_attention.zig new file mode 100644 index 0000000000..73e94641c3 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_attention.zig @@ -0,0 +1,155 @@ +//! Single-position retained-input attention isolation. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const diagnostic = @import("operator_diagnostic.zig"); +const diagnosticRow = diagnostic.diagnosticRow; +const attentionMetric = diagnostic.attentionMetric; +const attention_math = @import("attention_math.zig"); +const AttentionStage = diagnostic.AttentionStage; +const AttentionCapture = diagnostic.AttentionCapture; +const attention_elements = diagnostic.attention_elements; +const attentionProjection = diagnostic.attentionProjection; +const hyperconnections = @import("operator_hc.zig").hyperconnections; + +/// Local isolation receipt, not a replacement for the full eight-token oracle or its 1% gate. +pub fn runAttention(gpa: std.mem.Allocator, io: std.Io, model: *Model, snapshot: []const u8, dir: []const u8, layer: usize) !void { + const blob = try std.Io.Dir.cwd().readFileAlloc(io, snapshot, gpa, .limited(8 << 20)); + defer gpa.free(blob); + const cfg = model.cfg.text_config; + if (layer >= model.weights.layers.len) return error.InvalidAttentionLayer; + const snapshot_offset = layer * 2 * 61488; + if (blob.len < snapshot_offset + 40960 + 10240) return error.InvalidAttentionSnapshot; + try model.reset(); + const input = model.workspace.view(.normalized); + @memcpy(input.buffer.contents()[input.offset..][0..10240], blob[snapshot_offset + 40960 ..][0..10240]); + const input_row = @as([*]const u16, @ptrCast(@alignCast(input.buffer.contents() + input.offset)))[0..5120]; + const native = try model.device.buffer(attention_elements * 2, mtl.ResourceOptions.shared); + defer native.deinit(); + const reference_rows = try gpa.alloc(u16, attention_elements); + defer gpa.free(reference_rows); + var capture = AttentionCapture{ .model = model, .native = native, .reference = reference_rows }; + model.residency.add(native); + model.residency.commit(); + model.residency.requestResidency(); + defer { + model.residency.remove(native); + model.residency.commit(); + } + var scratch = try projection.Scratch.init(gpa, 8192, 32768); + defer scratch.deinit(); + var state = try attention_math.State.init(gpa, 1); + defer state.deinit(); + var output: [5120]u16 = undefined; + // Replay prerequisite normalized layers; complete each GPU command before reusing the CPU input buffer. + for (0..layer) |prior| { + const prior_offset = prior * 2 * 61488 + 40960; + @memcpy(input.buffer.contents()[input.offset..][0..10240], blob[prior_offset..][0..10240]); + const warm_cb = model.queue.commandBufferUnretained(); + const warm_enc = warm_cb.compute(.serial); + var warm_encoding = true; + errdefer if (warm_encoding) warm_enc.end(); + for (model.checkpoint.shards.items) |shard| warm_enc.useResource(shard.buffer, mtl.ResourceUsage.read); + try attention.encode(&model.pipelines.attention, warm_enc, &model.weights.layers[prior].attn, &model.cache, &model.attention_workspace, input, model.workspace.view(.branch), &model.pipelines.projection, model.workspace.projection(), @intCast(prior), 1, cfg); + warm_enc.end(); + warm_encoding = false; + warm_cb.commit(); + warm_cb.wait(); + if (warm_cb.failure() != null) return error.OperatorGpuFailure; + try state.apply(model.weights.layers[prior].attn, prior, input_row, 0, cfg, &output, &scratch); + } + @memcpy(input.buffer.contents()[input.offset..][0..10240], blob[snapshot_offset + 40960 ..][0..10240]); + const weights = &model.weights.layers[layer].attn; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + for (model.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + try attention.encodeDiagnostic(&model.pipelines.attention, enc, weights, &model.cache, &model.attention_workspace, input, model.workspace.view(.branch), &model.pipelines.projection, model.workspace.projection(), @intCast(layer), 1, cfg, true, .{ .ctx = &capture, .capture = AttentionCapture.gpu }); + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + const sparse_snapshot = try gpa.alloc(u16, 32768); + defer gpa.free(sparse_snapshot); + state.sparse_snapshot = sparse_snapshot; + state.observer = .{ .ctx = &capture, .capture = AttentionCapture.cpu }; + try state.apply(weights.*, layer, input_row, 0, cfg, &output, &scratch); + state.observer = null; + state.sparse_snapshot = null; + + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + var count: usize = 0; + var header: [512]u8 = undefined; + try report.appendSlice(gpa, try std.fmt.bufPrint(&header, "{{\"mode\":\"attention-sameinput\",\"token\":0,\"layer\":{d},\"input_offset_bytes\":{d},\"input_length_bytes\":10240,\"identical_normalized_input\":true,\"retained_normalized_prefix_layers\":{d},\"compression_ratio\":{d},\"reference_dot\":\"independent_EXL3_FP64\",\"qualification_gate_applied\":false,\"cases\":[\n", .{ layer, snapshot_offset + 40960, layer, cfg.compress_ratios[layer] })); + inline for (std.enums.values(AttentionStage)) |stage| + try attentionMetric(u16, gpa, &report, &count, @tagName(stage), "same_normalized_input_chain", capture.row(stage), capture.cpuRow(stage)); + + const expected = try gpa.alloc(f32, 32768); + defer gpa.free(expected); + try attentionProjection(gpa, model, weights.wq_a, input, "wq_a", &scratch, expected, &report, &count); + try attentionProjection(gpa, model, weights.wq_b, capture.view(.qr_norm), "wq_b", &scratch, expected, &report, &count); + try attentionProjection(gpa, model, weights.wkv, input, "wkv", &scratch, expected, &report, &count); + for (weights.wo_a, 0..) |matrix, group| { + var name: [32]u8 = undefined; + try attentionProjection(gpa, model, matrix, capture.view(.o).shifted(group * 4096 * 2), try std.fmt.bufPrint(&name, "wo_a_group{d}", .{group}), &scratch, expected, &report, &count); + } + try attentionProjection(gpa, model, weights.wo_b, capture.view(.grouped), "wo_b", &scratch, expected, &report, &count); + + var same_input: [32768]u16 = undefined; + math.rms(capture.row(.qr_projection), weights.q_norm, cfg.rms_norm_eps, same_input[0..1280]); + try attentionMetric(u16, gpa, &report, &count, "qr_norm", "same_native_input_reference_rms", capture.row(.qr_norm), same_input[0..1280]); + math.rms(capture.row(.kv_projection), weights.kv_norm, cfg.rms_norm_eps, same_input[0..512]); + try attentionMetric(u16, gpa, &report, &count, "kv_norm", "same_native_input_reference_rms", capture.row(.kv_norm), same_input[0..512]); + @memcpy(same_input[0..512], capture.row(.kv_rope)); + math.quantize(same_input[0..512], .window); + try attentionMetric(u16, gpa, &report, &count, "kv_quantized", "same_native_input_reference_quantize", capture.row(.kv_quantized), same_input[0..512]); + if (cfg.compress_ratios[layer] == 1) { + const source = try model.cache.source(@intCast(layer), 1); + const compressed = model.workspace.view(.collapsed); + const unpack_cb = model.queue.commandBufferUnretained(); + const unpack_enc = unpack_cb.compute(.serial); + var unpack_encoding = true; + errdefer if (unpack_encoding) unpack_enc.end(); + try model.pipelines.attention.cache.encodeUnpack(unpack_enc, source.main, compressed, .{ .rows = 1, .width = 512, .stride = 512, .packed_stride = 288, .block = 16, .kind = .main_fp4 }); + unpack_enc.end(); + unpack_encoding = false; + unpack_cb.commit(); + unpack_cb.wait(); + if (unpack_cb.failure() != null) return error.OperatorGpuFailure; + const actual_compressed = diagnosticRow(u16, compressed, 512); + const reference_compressed = state.sources[state.active_owner * state.capacity * 640 ..][0..512]; + try attentionMetric(u16, gpa, &report, &count, "compressed_kv", "same_retained_normalized_prefix", actual_compressed, reference_compressed); + var sinks: [64]f32 = undefined; + for (&sinks, 0..) |*v, head| v.* = math.value(weights.sink, head); + inline for (.{ "attention-compressed-native.bf16", "attention-compressed-reference.bf16", "attention-sinks.f32" }, .{ std.mem.sliceAsBytes(actual_compressed), std.mem.sliceAsBytes(reference_compressed), std.mem.sliceAsBytes(&sinks) }) |name, data| { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); + } + @memcpy(reference_compressed, actual_compressed); + const indices = [_]i32{ 0, 128 }; + for (0..64) |head| + state.sparseHead(capture.row(.q)[head * 512 ..][0..512], capture.row(.kv_quantized), &indices, sinks[head], same_input[head * 512 ..][0..512]); + try attentionMetric(u16, gpa, &report, &count, "sparse_o", "same_native_q_window_and_compressed_kv_reference_sparse", capture.row(.sparse_o), &same_input); + } + if (cfg.compress_ratios[layer] == 0) { + const indices = [_]i32{0}; + for (0..64) |head| + state.sparseHead(capture.row(.q)[head * 512 ..][0..512], capture.row(.kv_quantized), &indices, math.value(weights.sink, head), same_input[head * 512 ..][0..512]); + try attentionMetric(u16, gpa, &report, &count, "sparse_o", "same_native_q_and_quantized_kv_reference_sparse", capture.row(.sparse_o), &same_input); + } + try hyperconnections(gpa, model, blob[snapshot_offset..], capture.view(.wo_b), &report, &count, layer); + try report.appendSlice(gpa, "\n]}\n"); + inline for (.{ "attention-sameinput.json", "attention-native.bf16", "attention-reference.bf16", "attention-input.bf16" }, .{ report.items, std.mem.sliceAsBytes(native.slice(u16, attention_elements)), std.mem.sliceAsBytes(reference_rows), std.mem.sliceAsBytes(input_row) }) |name, data| { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); + } +} diff --git a/zig/tests/deepseek_v41_reference/operator_diagnostic.zig b/zig/tests/deepseek_v41_reference/operator_diagnostic.zig new file mode 100644 index 0000000000..c37682bf8e --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_diagnostic.zig @@ -0,0 +1,124 @@ +//! Shared retained-input diagnostic capture and metrics. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const ex = tf.exl3_projection; + +pub const AttentionStage = attention.DiagnosticStage; +pub const attention_widths = [_]usize{ 1280, 1280, 32768, 32768, 512, 512, 512, 512, 32768, 32768, 8192, 5120 }; +pub const attention_elements = blk: { + var total: usize = 0; + for (attention_widths) |width| total += width; + break :blk total; +}; +pub fn attentionOffset(stage: AttentionStage) usize { + var offset: usize = 0; + for (attention_widths[0..@backingInt(stage)]) |width| offset += width; + return offset; +} +pub const AttentionCapture = struct { + model: *Model, + native: mtl.Buffer, + reference: []u16, + + pub fn view(self: *const AttentionCapture, stage: AttentionStage) attention.View { + return .{ .buffer = self.native, .offset = attentionOffset(stage) * 2 }; + } + pub fn row(self: *const AttentionCapture, stage: AttentionStage) []u16 { + return self.native.slice(u16, attention_elements)[attentionOffset(stage)..][0..attention_widths[@backingInt(stage)]]; + } + pub fn cpuRow(self: *const AttentionCapture, stage: AttentionStage) []u16 { + return self.reference[attentionOffset(stage)..][0..attention_widths[@backingInt(stage)]]; + } + pub fn gpu(ctx: *anyopaque, enc: mtl.ComputeEncoder, stage: AttentionStage, source: attention.View) !void { + const self: *AttentionCapture = @ptrCast(@alignCast(ctx)); + const cache = &self.model.pipelines.attention.cache; + if (stage == .kv_quantized) { + try cache.encodeUnpack(enc, source, self.view(stage), .{ .rows = 1, .width = 512, .stride = 512, .packed_stride = 528, .block = 32, .kind = .window_fp8 }); + } else { + cache.encodeCopy(enc, source, self.view(stage), @intCast(attention_widths[@backingInt(stage)] * 2)); + } + } + pub fn cpu(ctx: *anyopaque, stage: AttentionStage, values: []const u16) void { + const self: *AttentionCapture = @ptrCast(@alignCast(ctx)); + @memcpy(self.cpuRow(stage), values); + } +}; + +pub fn attentionMetric(comptime T: type, gpa: std.mem.Allocator, report: *std.ArrayList(u8), count: *usize, name: []const u8, comparison: []const u8, actual: []const T, expected: []const T) !void { + if (actual.len == 0 or actual.len != expected.len) return error.InvalidAttentionMetricShape; + var squared: f64 = 0; + var magnitude: f64 = 0; + var max_abs: f64 = 0; + var worst: usize = 0; + var unequal: usize = 0; + var first: ?usize = null; + for (actual, expected, 0..) |a, r, i| { + const x = if (T == u16) math.f(a) else a; + const y = if (T == u16) math.f(r) else r; + if (!std.math.isFinite(x) or !std.math.isFinite(y)) return error.NonfiniteAttentionDiagnostic; + const delta = @as(f64, x) - y; + squared += delta * delta; + magnitude += @as(f64, y) * y; + if (@abs(delta) > max_abs) { + max_abs = @abs(delta); + worst = i; + } + const ab = if (T == u16) a else math.b(a); + const rb = if (T == u16) r else math.b(r); + if (ab != rb) { + unequal += 1; + if (first == null) first = i; + } + } + const line = try std.json.Stringify.valueAlloc(gpa, .{ + .name = name, + .comparison = comparison, + .elements = actual.len, + .output_dtype = if (T == u16) "bf16" else "f32", + .normalized_rms_error = if (magnitude == 0) @sqrt(squared) else @sqrt(squared / magnitude), + .max_absolute_error = max_abs, + .worst_index = worst, + .bf16_unequal_count = unequal, + .bf16_first_unequal_index = first, + .worst_native = if (T == u16) math.f(actual[worst]) else actual[worst], + .worst_reference = if (T == u16) math.f(expected[worst]) else expected[worst], + }, .{}); + defer gpa.free(line); + if (count.* != 0) try report.appendSlice(gpa, ",\n"); + try report.appendSlice(gpa, line); + count.* += 1; + std.debug.print("{s}\n", .{line}); +} + +pub fn attentionProjection(gpa: std.mem.Allocator, model: *Model, matrix: ex.Matrix, input: attention.View, name: []const u8, scratch: *projection.Scratch, expected: []f32, report: *std.ArrayList(u8), count: *usize) !void { + const input_row = @as([*]const u16, @ptrCast(@alignCast(input.buffer.contents() + input.offset)))[0..matrix.k]; + try projection.applyF32(matrix, input_row, expected[0..matrix.n], scratch); + const output = model.workspace.view(.logits); + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + ex.useMatrix(enc, matrix); + try ex.encode(&model.pipelines.projection, enc, matrix, .{ .buffer = input.buffer, .offset = input.offset }, .{ .buffer = output.buffer, .offset = output.offset }, model.workspace.projection(), .{ .count = 1, .x_stride = matrix.k, .y_stride = matrix.n, .output = .f32 }); + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + const actual = @as([*]const f32, @ptrCast(@alignCast(output.buffer.contents() + output.offset)))[0..matrix.n]; + try attentionMetric(f32, gpa, report, count, name, "same_native_input_fp64_dot", actual, expected[0..matrix.n]); +} + +pub fn diagnosticRow(comptime T: type, view: attention.View, len: usize) []T { + return @as([*]T, @ptrCast(@alignCast(view.buffer.contents() + view.offset)))[0..len]; +} +pub fn attentionWrite(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8, data: []const u8) !void { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); +} diff --git a/zig/tests/deepseek_v41_reference/operator_hc.zig b/zig/tests/deepseek_v41_reference/operator_hc.zig new file mode 100644 index 0000000000..4aec23362a --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_hc.zig @@ -0,0 +1,98 @@ +//! Retained-input hyperconnection isolation, separate from qualification. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const diagnostic = @import("operator_diagnostic.zig"); +const diagnosticRow = diagnostic.diagnosticRow; +const attentionMetric = diagnostic.attentionMetric; + +pub fn hyperconnections(gpa: std.mem.Allocator, model: *Model, stream_bytes: []const u8, branch: attention.View, report: *std.ArrayList(u8), count: *usize, layer: usize) !void { + const ws = &model.workspace; + const hc = model.weights.layers[layer].hc_ffn; + const cfg = model.cfg.text_config; + const stream = ws.view(.stream); + @memcpy(stream.buffer.contents()[stream.offset..][0..40960], stream_bytes[0..40960]); + var input: [20480]u16 = undefined; + @memcpy(&input, diagnosticRow(u16, stream, input.len)); + var projected: [24]f32 = undefined; + math.dense(hc.projection, &input, &projected); + const expected_mix = math.mixes(&input, hc, cfg); + var squares: f64 = 0; + for (input) |v| squares += @as(f64, math.f(v)) * math.f(v); + const mean = math.meanFromWide(squares, input.len); + const expected_inv = [_]f32{math.idealRsqrt(mean + cfg.rms_norm_eps)}; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + for (model.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + tf.deepseek_v41.mhc.encodeRms(enc, &model.pipelines.mhc, stream, ws.view(.inv), .{ .rows = 1, .width = 5120, .stream_stride = 20480, .eps = cfg.rms_norm_eps }); + tf.deepseek_v41.ops.encodeDense(enc, &model.pipelines.ops, stream, .{ .buffer = hc.projection.buffer, .offset = hc.projection.offset }, ws.view(.mixes), .{ .rows = 1, .k = 20480, .n = 24, .x_stride = 20480, .w_stride = 20480, .y_stride = 24, .w_dtype = .f32, .y_dtype = .f32 }); + enc.barrier(); + tf.deepseek_v41.mhc.encodeSplit(enc, &model.pipelines.mhc, ws.view(.mixes), ws.view(.inv), .{ .buffer = hc.scale.buffer, .offset = hc.scale.offset }, .{ .buffer = hc.base.buffer, .offset = hc.base.offset }, ws.view(.attention_pre), ws.view(.post), ws.view(.comb), .{ .rows = 1, .eps = cfg.hc_eps }); + enc.barrier(); + tf.deepseek_v41.mhc.encodeCollapse(enc, &model.pipelines.mhc, stream, ws.view(.attention_pre), ws.view(.collapsed), .{ .rows = 1, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + enc.barrier(); + tf.deepseek_v41.mhc.encodePost(enc, &model.pipelines.mhc, branch, stream, ws.view(.post), ws.view(.comb), stream, .{ .rows = 1, .width = 5120, .branch_stride = 5120, .residual_stride = 20480, .y_stride = 20480 }); + enc.end(); + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + try attentionMetric(f32, gpa, report, count, "hc_projection", "same_retained_stream_fp64_dot", diagnosticRow(f32, ws.view(.mixes), 24), &projected); + try attentionMetric(f32, gpa, report, count, "hc_inverse", "same_retained_stream_reference_statistic", diagnosticRow(f32, ws.view(.inv), 1), &expected_inv); + try attentionMetric(f32, gpa, report, count, "hc_pre", "same_retained_stream_chain", diagnosticRow(f32, ws.view(.attention_pre), 4), &expected_mix.pre); + try attentionMetric(f32, gpa, report, count, "hc_post", "same_retained_stream_chain", diagnosticRow(f32, ws.view(.post), 4), &expected_mix.post); + try attentionMetric(f32, gpa, report, count, "hc_comb", "same_retained_stream_chain", diagnosticRow(f32, ws.view(.comb), 16), &expected_mix.comb); + var native_mix: math.Mix = undefined; + @memcpy(&native_mix.pre, diagnosticRow(f32, ws.view(.attention_pre), 4)); + @memcpy(&native_mix.post, diagnosticRow(f32, ws.view(.post), 4)); + @memcpy(&native_mix.comb, diagnosticRow(f32, ws.view(.comb), 16)); + var collapsed: [5120]u16 = undefined; + math.collapse(&input, native_mix.pre, &collapsed); + try attentionMetric(u16, gpa, report, count, "hc_collapse", "same_retained_stream_and_native_coefficients", diagnosticRow(u16, ws.view(.collapsed), 5120), &collapsed); + math.expand(&input, diagnosticRow(u16, branch, 5120), native_mix); + try attentionMetric(u16, gpa, report, count, "hc_expand", "same_retained_stream_branch_and_native_coefficients", diagnosticRow(u16, stream, 20480), &input); + if (layer + 1 < model.weights.layers.len and stream_bytes.len >= 61488 + 40960) { + @memcpy(stream.buffer.contents()[stream.offset..][0..40960], stream_bytes[61488..][0..40960]); + @memcpy(&input, diagnosticRow(u16, stream, input.len)); + const norm = model.weights.layers[layer + 1].attn_norm; + const next_cb = model.queue.commandBufferUnretained(); + const next_enc = next_cb.compute(.serial); + next_enc.useResource(norm.buffer, mtl.ResourceUsage.read); + tf.deepseek_v41.mhc.encodeCollapse(next_enc, &model.pipelines.mhc, stream, ws.view(.attention_pre), ws.view(.collapsed), .{ .rows = 1, .width = 5120, .stream_stride = 20480, .y_stride = 5120 }); + next_enc.barrier(); + tf.deepseek_v41.ops.encodeRms(next_enc, &model.pipelines.ops, ws.view(.collapsed), .{ .buffer = norm.buffer, .offset = norm.offset }, ws.view(.branch), .{ .rows = 1, .width = 5120, .x_stride = 5120, .y_stride = 5120, .eps = cfg.rms_norm_eps }); + next_enc.end(); + next_cb.commit(); + next_cb.wait(); + if (next_cb.failure() != null) return error.OperatorGpuFailure; + math.collapse(&input, native_mix.pre, &collapsed); + try attentionMetric(u16, gpa, report, count, "next_collapse", "same_retained_ffn_stream_and_native_pre", diagnosticRow(u16, ws.view(.collapsed), 5120), &collapsed); + math.collapse(&input, expected_mix.pre, &collapsed); + try attentionMetric(u16, gpa, report, count, "next_collapse", "same_retained_ffn_stream_reference_pre", diagnosticRow(u16, ws.view(.collapsed), 5120), &collapsed); + math.rms(&collapsed, norm, cfg.rms_norm_eps, &collapsed); + try attentionMetric(u16, gpa, report, count, "next_normalized", "same_retained_ffn_stream_reference_pre_and_norm", diagnosticRow(u16, ws.view(.branch), 5120), &collapsed); + } +} + +pub fn runHyperconnections(gpa: std.mem.Allocator, io: std.Io, model: *Model, snapshot: []const u8, dir: []const u8, layer: usize) !void { + if (layer >= model.weights.layers.len) return error.InvalidHyperconnectionLayer; + const blob = try std.Io.Dir.cwd().readFileAlloc(io, snapshot, gpa, .limited(8 << 20)); + defer gpa.free(blob); + const offset = layer * 2 * 61488; + if (blob.len < offset + 2 * 61488) return error.InvalidHyperconnectionSnapshot; + try model.reset(); + const branch = model.workspace.view(.branch); + @memcpy(branch.buffer.contents()[branch.offset..][0..10240], blob[offset + 61488 + 51200 ..][0..10240]); + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + var count: usize = 0; + var header: [256]u8 = undefined; + try report.appendSlice(gpa, try std.fmt.bufPrint(&header, "{{\"mode\":\"hc-sameinput\",\"token\":null,\"token_identity_source\":\"unknown_headerless_snapshot\",\"layer\":{d},\"qualification_gate_applied\":false,\"cases\":[\n", .{layer})); + try hyperconnections(gpa, model, blob[offset..], branch, &report, &count, layer); + try report.appendSlice(gpa, "\n]}\n"); + const path = try std.fs.path.join(gpa, &.{ dir, "hc-sameinput.json" }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = report.items }); +} diff --git a/zig/tests/deepseek_v41_reference/operator_history.zig b/zig/tests/deepseek_v41_reference/operator_history.zig new file mode 100644 index 0000000000..499a1cd463 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_history.zig @@ -0,0 +1,298 @@ +//! Retained-input attention history and source-layout receipts. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const diagnostic = @import("operator_diagnostic.zig"); +const diagnosticRow = diagnostic.diagnosticRow; +const attentionMetric = diagnostic.attentionMetric; +const attention_math = @import("attention_math.zig"); +const AttentionStage = diagnostic.AttentionStage; +const AttentionCapture = diagnostic.AttentionCapture; +const attention_widths = diagnostic.attention_widths; +const attention_elements = diagnostic.attention_elements; +const attentionOffset = diagnostic.attentionOffset; +const attentionProjection = diagnostic.attentionProjection; +const attentionWrite = diagnostic.attentionWrite; + +// Source indices address physical ring slots then compressed rows; reference reserves 128 window rows even at zero. +fn attentionHistoryIndices(position: usize, selected: []const i32, compressed_offset: usize, out: *[640]i32) []const i32 { + const window_columns: usize = if (position == 0) 1 else 128; + if (position == 0) { + out[0] = 0; + } else { + for (0..128) |slot| { + const at = (position % 128 + 1 + slot) % 128; + out[slot] = if (at > position) -1 else @intCast(at); + } + } + for (selected, 0..) |row, slot| + out[window_columns + slot] = if (row < 0) -1 else @intCast(compressed_offset + @as(usize, @intCast(row))); + return out[0 .. window_columns + selected.len]; +} + +/// Replay prerequisite layers at each position; native/reference cursors advance without borrowed intermediates. +pub fn runAttentionHistory(gpa: std.mem.Allocator, io: std.Io, model: *Model, trace_dir: []const u8, dir: []const u8, token: usize, layer: usize) !void { + if (layer >= model.weights.layers.len) return error.InvalidAttentionLayer; + if (token >= model.cache.plan.capacity) return error.InvalidAttentionHistory; + if (token >= attention_math.max_capacity) return error.DiagnosticHistoryTooLarge; + const prefix_elements = try std.math.mul(usize, token + 1, (layer + 1) * 5120); + if (prefix_elements > (256 << 20) / @sizeOf(u16)) return error.DiagnosticHistoryTooLarge; + const cfg = model.cfg.text_config; + const history_count = token + 1; + const ratio = cfg.compress_ratios[layer]; + const compressed_rows = if (ratio == 0) 0 else history_count / ratio; + const topk = @min(compressed_rows, cfg.index_topk); + const window_rows: usize = if (token == 0) 1 else 128; + const snapshot_bytes = 61488; + const snapshot_offset = layer * 2 * snapshot_bytes; + const weights = &model.weights.layers[layer].attn; + const input = model.workspace.view(.normalized); + const native = try model.device.buffer(attention_elements * 2, mtl.ResourceOptions.shared); + defer native.deinit(); + const decoded = try model.device.buffer((128 * 512 + compressed_rows * 640) * 2, mtl.ResourceOptions.shared); + defer decoded.deinit(); + const reference_rows = try gpa.alloc(u16, attention_elements); + defer gpa.free(reference_rows); + const sparse_snapshot = try gpa.alloc(u16, 32768); + defer gpa.free(sparse_snapshot); + const normalized_prefix = try gpa.alloc(u16, prefix_elements); + defer gpa.free(normalized_prefix); + var retained_branch: [5120]u16 = undefined; + var capture = AttentionCapture{ .model = model, .native = native, .reference = reference_rows }; + for ([_]mtl.Buffer{ native, decoded }) |buffer| model.residency.add(buffer); + model.residency.commit(); + model.residency.requestResidency(); + defer { + for ([_]mtl.Buffer{ native, decoded }) |buffer| model.residency.remove(buffer); + model.residency.commit(); + } + var scratch = try projection.Scratch.init(gpa, 8192, 32768); + defer scratch.deinit(); + var state = try attention_math.State.init(gpa, history_count); + defer state.deinit(); + var output: [5120]u16 = undefined; + try model.reset(); + for (0..history_count) |position| { + var name: [64]u8 = undefined; + const trace_name = try std.fmt.bufPrint(&name, "trace-{d}-gpu.bin", .{position}); + const path = try std.fs.path.join(gpa, &.{ trace_dir, trace_name }); + defer gpa.free(path); + const blob = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(8 << 20)); + defer gpa.free(blob); + if (blob.len != model.weights.layers.len * 2 * snapshot_bytes) return error.InvalidAttentionSnapshot; + if (model.cache.position != position) return error.InvalidAttentionHistoryCursor; + for (0..layer + 1) |prior| { + const offset = prior * 2 * snapshot_bytes + 40960; + const row = normalized_prefix[(position * (layer + 1) + prior) * 5120 ..][0..5120]; + @memcpy(std.mem.sliceAsBytes(row), blob[offset..][0..10240]); + @memcpy(input.buffer.contents()[input.offset..][0..10240], std.mem.sliceAsBytes(row)); + const target = position == token and prior == layer; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + for (model.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + if (target) { + try attention.encodeDiagnostic(&model.pipelines.attention, enc, &model.weights.layers[prior].attn, &model.cache, &model.attention_workspace, input, model.workspace.view(.branch), &model.pipelines.projection, model.workspace.projection(), @intCast(prior), 1, cfg, true, .{ .ctx = &capture, .capture = AttentionCapture.gpu }); + @memcpy(std.mem.sliceAsBytes(&retained_branch), blob[snapshot_offset + 51200 ..][0..10240]); + } else { + try attention.encode(&model.pipelines.attention, enc, &model.weights.layers[prior].attn, &model.cache, &model.attention_workspace, input, model.workspace.view(.branch), &model.pipelines.projection, model.workspace.projection(), @intCast(prior), 1, cfg); + } + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + // Wait before the next retained row replaces the native workspace input. + if (target) { + state.observer = .{ .ctx = &capture, .capture = AttentionCapture.cpu }; + state.sparse_snapshot = sparse_snapshot; + } + try state.apply(model.weights.layers[prior].attn, prior, row, position, cfg, &output, &scratch); + state.observer = null; + state.sparse_snapshot = null; + } + try model.cache.append(1); + } + + // Export post-insertion physical ring slots, masking never-written slots in source indices. + const decoded_view = attention.View{ .buffer = decoded }; + const native_window = decoded.slice(u16, 128 * 512 + compressed_rows * 640)[0 .. 128 * 512]; + const native_compressed = decoded.slice(u16, 128 * 512 + compressed_rows * 640)[128 * 512 ..][0 .. compressed_rows * 512]; + const native_index_keys = decoded.slice(u16, 128 * 512 + compressed_rows * 640)[128 * 512 + compressed_rows * 512 ..][0 .. compressed_rows * 128]; + const unpack_cb = model.queue.commandBufferUnretained(); + const unpack_enc = unpack_cb.compute(.serial); + var unpack_encoding = true; + errdefer if (unpack_encoding) unpack_enc.end(); + const ring = try model.cache.ring(@intCast(layer)); + try model.pipelines.attention.cache.encodeUnpack(unpack_enc, ring, decoded_view, .{ .rows = 128, .width = 512, .stride = 512, .packed_stride = 528, .block = 32, .kind = .window_fp8 }); + if (compressed_rows != 0) { + const source = try model.cache.source(@intCast(layer), ratio); + try model.pipelines.attention.cache.encodeUnpack(unpack_enc, source.main, decoded_view.shifted(128 * 512 * 2), .{ .rows = @intCast(compressed_rows), .width = 512, .stride = 512, .packed_stride = 288, .block = 16, .kind = .main_fp4 }); + try model.pipelines.attention.cache.encodeUnpack(unpack_enc, source.index, decoded_view.shifted((128 * 512 + compressed_rows * 512) * 2), .{ .rows = @intCast(compressed_rows), .width = 128, .stride = 128, .packed_stride = 68, .block = 32, .kind = .index_fp4 }); + } + unpack_enc.end(); + unpack_encoding = false; + unpack_cb.commit(); + unpack_cb.wait(); + if (unpack_cb.failure() != null) return error.OperatorGpuFailure; + const reference_window = state.window[layer * 128 * 512 ..][0 .. 128 * 512]; + const reference_compressed = try gpa.alloc(u16, compressed_rows * 640); + defer gpa.free(reference_compressed); + const reference_main = reference_compressed[0 .. compressed_rows * 512]; + const reference_index_keys = reference_compressed[compressed_rows * 512 ..]; + for (0..compressed_rows) |row| { + const source_row = state.sources[state.active_owner * state.capacity * 640 + row * 640 ..][0..640]; + @memcpy(reference_main[row * 512 ..][0..512], source_row[0..512]); + @memcpy(reference_index_keys[row * 128 ..][0..128], source_row[512..640]); + } + const native_selected = diagnosticRow(i32, model.attention_workspace.selection.positions(), 512)[0..topk]; + var reference_selected: [512]i32 = undefined; + for (state.selected[0..topk], 0..) |row, i| reference_selected[i] = @intCast(row); + var native_indices_storage: [640]i32 = undefined; + var reference_indices_storage: [640]i32 = undefined; + const native_indices = attentionHistoryIndices(token, native_selected, window_rows, &native_indices_storage); + const reference_indices = attentionHistoryIndices(token, reference_selected[0..topk], window_rows, &reference_indices_storage); + var sinks: [64]f32 = undefined; + for (&sinks, 0..) |*value, head| value.* = math.value(weights.sink, head); + + var first_stage: ?[]const u8 = null; + inline for (std.enums.values(AttentionStage)) |stage| + if (first_stage == null and !std.mem.eql(u16, capture.row(stage), capture.cpuRow(stage))) { + first_stage = @tagName(stage); + }; + const header = try std.json.Stringify.valueAlloc(gpa, .{ + .mode = "attention-history-sameinput", + .token = token, + .layer = layer, + .history_count = history_count, + .token_identity_source = "caller_supplied_TOKEN_and_trace_POSITION_gpu_filenames_raw_snapshots_have_no_token_header", + .normalized_prefix_replay = .{ .file_pattern = "trace-{position}-gpu.bin", .origin = "retained_native_GPU_normalized_attention_inputs", .first_position = 0, .last_position = token, .first_layer = 0, .last_layer = layer, .order = "position_major_then_layer_ascending_inclusive", .snapshot_count_per_file = model.weights.layers.len * 2, .snapshot_length_bytes = snapshot_bytes, .attention_snapshot_index = "2*layer", .normalized_input_offset_within_snapshot_bytes = 40960, .normalized_input_length_bytes = 10240, .native_cache_reset = true, .native_cache_cursor_matches_each_position = true, .native_cache_append_rows_per_position = 1, .reference_position_passed_explicitly = true, .reference_capacity = state.capacity }, + .identical_retained_normalized_inputs = true, + .independent_reference_history = true, + .native_history_replay_branch_matches_retained_trace = std.mem.eql(u16, capture.row(.wo_b), &retained_branch), + .first_differing_same_history_input_stage = first_stage, + .window_history_equal = std.mem.eql(u16, native_window, reference_window), + .compressed_main_equal = std.mem.eql(u16, native_compressed, reference_main), + .compressed_index_keys_equal = std.mem.eql(u16, native_index_keys, reference_index_keys), + .selection_equal = std.mem.eql(i32, native_selected, reference_selected[0..topk]), + .compression_ratio = ratio, + .compressed_rows = compressed_rows, + .selection_count = topk, + .reference_dot = "independent_EXL3_FP64", + .reference_sparse_dot = "independent_FP64", + .qualification_gate_applied = false, + .qualification_oracle_modified_to_reuse_native_values = false, + .same_native_query_cache_diagnostic = "separate_reference_State_with_captured_native_window_main_selection_and_q_not_used_in_independent_chain", + .official_revision = "2cba9e42aa026125f3ed06c6d98c1db82f7ca027", + .receipt_dtype_endianness = "little_endian", + .receipt_strides = "contiguous_element_strides", + .files = .{ + .stage_chain = .{ .native = "attention-native.bf16", .reference = "attention-reference.bf16", .dtype = "bf16", .shape = .{attention_elements}, .strides = .{1}, .layout = "stages_table_below" }, + .normalized_prefix = .{ .file = "attention-normalized-prefix.bf16", .dtype = "bf16", .shape = .{ history_count, layer + 1, 5120 }, .strides = .{ (layer + 1) * 5120, 5120, 1 } }, + .target_input = .{ .file = "attention-input.bf16", .dtype = "bf16", .shape = .{5120}, .strides = .{1} }, + .retained_branch = .{ .file = "attention-retained-branch.bf16", .dtype = "bf16", .shape = .{5120}, .strides = .{1} }, + .query = .{ .native = "attention-query-native.bf16", .reference = "attention-query-reference.bf16", .dtype = "bf16", .shape = .{ 1, 1, 64, 512 }, .strides = .{ 32768, 32768, 512, 1 } }, + .window = .{ .native = "attention-window-native.bf16", .reference = "attention-window-reference.bf16", .dtype = "bf16", .shape = .{ 1, 128, 512 }, .strides = .{ 65536, 512, 1 }, .layout = "physical_ring_slot_position_modulo_128_includes_current_quantized_KV" }, + .compressed_main = .{ .native = "attention-compressed-native.bf16", .reference = "attention-compressed-reference.bf16", .dtype = "bf16", .shape = .{ 1, compressed_rows, 512 }, .strides = .{ compressed_rows * 512, 512, 1 } }, + .compressed_index_keys = .{ .native = "attention-index-keys-native.bf16", .reference = "attention-index-keys-reference.bf16", .dtype = "bf16", .shape = .{ 1, compressed_rows, 128 }, .strides = .{ compressed_rows * 128, 128, 1 } }, + .selection = .{ .native = "attention-selection-native.i32", .reference = "attention-selection-reference.i32", .dtype = "i32", .shape = .{ 1, topk }, .strides = .{ topk, 1 }, .layout = "actual_compressed_row_ids_compact_export_native_workspace_row_stride_is_512" }, + .source_kv = .{ .native = "attention-source-kv-native.bf16", .reference = "attention-source-kv-reference.bf16", .dtype = "bf16", .shape = .{ 1, window_rows + compressed_rows, 512 }, .strides = .{ (window_rows + compressed_rows) * 512, 512, 1 }, .layout = "source_sparse_attn_KV_window_then_compressed_rows_no_requantization" }, + .source_indices = .{ .native = "attention-source-indices-native.i32", .reference = "attention-source-indices-reference.i32", .dtype = "i32", .shape = .{ 1, 1, native_indices.len }, .strides = .{ native_indices.len, native_indices.len, 1 }, .layout = "source_get_window_topk_idxs_then_actual_selection_plus_window_rows_minus_one_remains_masked" }, + .sinks = .{ .file = "attention-sinks.f32", .dtype = "f32", .shape = .{64}, .strides = .{1} }, + .sparse_output = .{ .native = "attention-sparse-native.bf16", .reference = "attention-sparse-reference.bf16", .same_native_input_reference = "attention-sparse-same-native-input-reference.bf16", .dtype = "bf16", .shape = .{ 1, 1, 64, 512 }, .strides = .{ 32768, 32768, 512, 1 } }, + }, + .source_sparse_call = "sparse_attn(query,source_kv,sinks,source_indices,1/sqrt(512)); output_before_inverse_RoPE", + .source_sparse_replay_is_operator_evidence_only = true, + }, .{}); + defer gpa.free(header); + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + try report.appendSlice(gpa, header[0 .. header.len - 1]); + try report.appendSlice(gpa, ",\"cases\":[\n"); + var count: usize = 0; + inline for (std.enums.values(AttentionStage)) |stage| + try attentionMetric(u16, gpa, &report, &count, @tagName(stage), "same_retained_normalized_history_input_chain", capture.row(stage), capture.cpuRow(stage)); + try attentionMetric(u16, gpa, &report, &count, "retained_branch", "native_history_replay_vs_retained_native_trace", capture.row(.wo_b), &retained_branch); + try attentionMetric(u16, gpa, &report, &count, "window_history", "same_retained_normalized_history_input_chain", native_window, reference_window); + if (compressed_rows != 0) { + try attentionMetric(u16, gpa, &report, &count, "compressed_kv", "same_retained_normalized_history_input_chain", native_compressed, reference_main); + try attentionMetric(u16, gpa, &report, &count, "compressed_index_keys", "same_retained_normalized_history_input_chain", native_index_keys, reference_index_keys); + } + + // This separate diagnostic State never overwrites independent retained-history reference data. + var native_input_state = try attention_math.State.init(gpa, history_count); + defer native_input_state.deinit(); + for (0..compressed_rows) |row| + @memcpy(native_input_state.sources[row * 640 ..][0..512], native_compressed[row * 512 ..][0..512]); + var helper_indices_storage: [640]i32 = undefined; + const helper_indices = attentionHistoryIndices(token, native_selected, 128, &helper_indices_storage); + var same_input: [32768]u16 = undefined; + for (0..64) |head| + native_input_state.sparseHead(capture.row(.q)[head * 512 ..][0..512], native_window, helper_indices, sinks[head], same_input[head * 512 ..][0..512]); + try attentionMetric(u16, gpa, &report, &count, "sparse_o", "same_native_query_window_compressed_cache_selection_reference_sparse_diagnostic", capture.row(.sparse_o), &same_input); + try attentionWrite(gpa, io, dir, "attention-sparse-same-native-input-reference.bf16", std.mem.sliceAsBytes(&same_input)); + math.rms(capture.row(.qr_projection), weights.q_norm, cfg.rms_norm_eps, same_input[0..1280]); + try attentionMetric(u16, gpa, &report, &count, "qr_norm", "same_native_input_reference_rms_diagnostic", capture.row(.qr_norm), same_input[0..1280]); + math.rms(capture.row(.kv_projection), weights.kv_norm, cfg.rms_norm_eps, same_input[0..512]); + try attentionMetric(u16, gpa, &report, &count, "kv_norm", "same_native_input_reference_rms_diagnostic", capture.row(.kv_norm), same_input[0..512]); + @memcpy(&same_input, capture.row(.q_projection)); + for (0..64) |head| math.rope(same_input[head * 512 ..][0..512], token, ratio != 0, false, cfg); + try attentionMetric(u16, gpa, &report, &count, "q", "same_native_input_reference_rope_diagnostic", capture.row(.q), &same_input); + @memcpy(same_input[0..512], capture.row(.kv_norm)); + math.rope(same_input[0..512], token, ratio != 0, false, cfg); + try attentionMetric(u16, gpa, &report, &count, "kv_rope", "same_native_input_reference_rope_diagnostic", capture.row(.kv_rope), same_input[0..512]); + @memcpy(same_input[0..512], capture.row(.kv_rope)); + math.quantize(same_input[0..512], .window); + try attentionMetric(u16, gpa, &report, &count, "kv_quantized", "same_native_input_reference_quantize_diagnostic", capture.row(.kv_quantized), same_input[0..512]); + const expected = try gpa.alloc(f32, 32768); + defer gpa.free(expected); + try attentionProjection(gpa, model, weights.wq_a, input, "wq_a", &scratch, expected, &report, &count); + try attentionProjection(gpa, model, weights.wq_b, capture.view(.qr_norm), "wq_b", &scratch, expected, &report, &count); + try attentionProjection(gpa, model, weights.wkv, input, "wkv", &scratch, expected, &report, &count); + for (weights.wo_a, 0..) |matrix, group| { + var name: [32]u8 = undefined; + try attentionProjection(gpa, model, matrix, capture.view(.o).shifted(group * 4096 * 2), try std.fmt.bufPrint(&name, "wo_a_group{d}", .{group}), &scratch, expected, &report, &count); + } + try attentionProjection(gpa, model, weights.wo_b, capture.view(.grouped), "wo_b", &scratch, expected, &report, &count); + try report.appendSlice(gpa, "\n],\"stages\":["); + inline for (std.enums.values(AttentionStage), 0..) |stage, i| { + const descriptor = try std.json.Stringify.valueAlloc(gpa, .{ .name = @tagName(stage), .offset_bytes = attentionOffset(stage) * 2, .elements = attention_widths[i] }, .{}); + defer gpa.free(descriptor); + if (i != 0) try report.appendSlice(gpa, ","); + try report.appendSlice(gpa, descriptor); + } + try report.appendSlice(gpa, "]}\n"); + const source_kv = try gpa.alloc(u16, 2 * (window_rows + compressed_rows) * 512); + defer gpa.free(source_kv); + const source_native = source_kv[0 .. (window_rows + compressed_rows) * 512]; + const source_reference = source_kv[(window_rows + compressed_rows) * 512 ..]; + @memcpy(source_native[0 .. window_rows * 512], native_window[0 .. window_rows * 512]); + @memcpy(source_native[window_rows * 512 ..], native_compressed); + @memcpy(source_reference[0 .. window_rows * 512], reference_window[0 .. window_rows * 512]); + @memcpy(source_reference[window_rows * 512 ..], reference_main); + const input_row = normalized_prefix[(history_count * (layer + 1) - 1) * 5120 ..][0..5120]; + inline for (.{ + "attention-history-sameinput.json", "attention-native.bf16", "attention-reference.bf16", + "attention-normalized-prefix.bf16", "attention-input.bf16", "attention-retained-branch.bf16", + "attention-query-native.bf16", "attention-query-reference.bf16", "attention-window-native.bf16", + "attention-window-reference.bf16", "attention-compressed-native.bf16", "attention-compressed-reference.bf16", + "attention-index-keys-native.bf16", "attention-index-keys-reference.bf16", "attention-selection-native.i32", + "attention-selection-reference.i32", "attention-source-kv-native.bf16", "attention-source-kv-reference.bf16", + "attention-source-indices-native.i32", "attention-source-indices-reference.i32", "attention-sinks.f32", + "attention-sparse-native.bf16", "attention-sparse-reference.bf16", + }, .{ + report.items, std.mem.sliceAsBytes(native.slice(u16, attention_elements)), std.mem.sliceAsBytes(reference_rows), + std.mem.sliceAsBytes(normalized_prefix), std.mem.sliceAsBytes(input_row), std.mem.sliceAsBytes(&retained_branch), + std.mem.sliceAsBytes(capture.row(.q)), std.mem.sliceAsBytes(capture.cpuRow(.q)), std.mem.sliceAsBytes(native_window), + std.mem.sliceAsBytes(reference_window), std.mem.sliceAsBytes(native_compressed), std.mem.sliceAsBytes(reference_main), + std.mem.sliceAsBytes(native_index_keys), std.mem.sliceAsBytes(reference_index_keys), std.mem.sliceAsBytes(native_selected), + std.mem.sliceAsBytes(reference_selected[0..topk]), std.mem.sliceAsBytes(source_native), std.mem.sliceAsBytes(source_reference), + std.mem.sliceAsBytes(native_indices), std.mem.sliceAsBytes(reference_indices), std.mem.sliceAsBytes(&sinks), + std.mem.sliceAsBytes(capture.row(.sparse_o)), std.mem.sliceAsBytes(capture.cpuRow(.sparse_o)), + }) |name, data| try attentionWrite(gpa, io, dir, name, data); +} diff --git a/zig/tests/deepseek_v41_reference/operator_moe.zig b/zig/tests/deepseek_v41_reference/operator_moe.zig new file mode 100644 index 0000000000..eb9780db96 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operator_moe.zig @@ -0,0 +1,409 @@ +//! Retained-input MoE routing and expert isolation. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = tf.deepseek_v41.attention; +const Model = tf.deepseek_v41.Model; +const ex = tf.exl3_projection; +const diagnosticRow = @import("operator_diagnostic.zig").diagnosticRow; + +const moe_reference = @import("reference.zig"); +const MoeRows = struct { + gate: [7][2304]u16 = undefined, + up: [7][2304]u16 = undefined, + activation: [7][2304]u16 = undefined, + down: [7][5120]u16 = undefined, + branch: [5120]u16 = undefined, +}; +const MoeFirst = struct { + source_chain: ?[]const u8 = null, + source_expert_stage: ?[]const u8 = null, + native_input_diagnostic: ?[]const u8 = null, + native_weights_reference_chain_diagnostic: ?[]const u8 = null, + retained_weights_reference_chain_diagnostic: ?[]const u8 = null, + native_down_reference_combine_diagnostic_differs: bool = false, +}; + +fn moeMetric(comptime T: type, gpa: std.mem.Allocator, report: *std.ArrayList(u8), count: *usize, name: []const u8, comparison: []const u8, diagnostic_reuse: bool, actual: []const T, expected: []const T) !bool { + if (actual.len == 0 or actual.len != expected.len) return error.InvalidMoeMetricShape; + var squared: f64 = 0; + var magnitude: f64 = 0; + var max_abs: f64 = 0; + var worst: usize = 0; + var unequal: usize = 0; + var first: ?usize = null; + for (actual, expected, 0..) |a, r, i| { + const x = if (T == u16) math.f(a) else a; + const y = if (T == u16) math.f(r) else r; + if (!std.math.isFinite(x) or !std.math.isFinite(y)) return error.NonfiniteMoeDiagnostic; + const delta = @as(f64, x) - y; + squared += delta * delta; + magnitude += @as(f64, y) * y; + if (@abs(delta) > max_abs) { + max_abs = @abs(delta); + worst = i; + } + const ab = if (T == u16) a else @as(u32, @bitCast(a)); + const rb = if (T == u16) r else @as(u32, @bitCast(r)); + if (ab != rb) { + unequal += 1; + if (first == null) first = i; + } + } + const line = try std.json.Stringify.valueAlloc(gpa, .{ + .name = name, + .comparison = comparison, + .diagnostic_native_value_reuse = diagnostic_reuse, + .qualification_gate_applied = false, + .elements = actual.len, + .output_dtype = if (T == u16) "bf16" else "f32", + .normalized_rms_error = if (magnitude == 0) @sqrt(squared) else @sqrt(squared / magnitude), + .max_absolute_error = max_abs, + .worst_index = worst, + .bit_unequal_count = unequal, + .first_bit_unequal_index = first, + .worst_native = if (T == u16) math.f(actual[worst]) else actual[worst], + .worst_reference = if (T == u16) math.f(expected[worst]) else expected[worst], + }, .{}); + defer gpa.free(line); + if (count.* != 0) try report.appendSlice(gpa, ",\n"); + try report.appendSlice(gpa, line); + count.* += 1; + std.debug.print("{s}\n", .{line}); + return unequal != 0; +} + +fn moeWrite(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8, bytes: []const u8) !void { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = bytes }); +} + +fn moeWriteRows(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, prefix: []const u8, rows: *const MoeRows) !void { + inline for (.{ "gate", "up", "activation", "down", "branch" }) |stage| { + const name = try std.fmt.allocPrint(gpa, "{s}-{s}.bf16", .{ prefix, stage }); + defer gpa.free(name); + try moeWrite(gpa, io, dir, name, std.mem.asBytes(&@field(rows.*, stage))); + } +} + +fn moeMatrices(model: *Model, layer: usize, expert: ?u32) ![3]ex.Matrix { + if (expert == null) return model.weights.layers[layer].moe.shared; + var matrices: [3]ex.Matrix = undefined; + var name: [128]u8 = undefined; + inline for (.{ "w1", "w3", "w2" }, 0..) |part, i| { + const key = try std.fmt.bufPrint(&name, "layers.{d}.ffn.experts.{d}.{s}", .{ layer, expert.?, part }); + matrices[i] = try tf.deepseek_v41_weights.lookup(&model.checkpoint, key, if (i == 2) 2304 else 5120, if (i == 2) 5120 else 2304); + } + return matrices; +} + +fn moeCombineDown(ids: *const [6]u32, down: *const [7][5120]u16, branch: *[5120]u16) void { + @setFloatMode(.strict); + var sum: [5120]f32 = @splat(0); + // Keep slots in top-six order in receipts, but match the official ascending expert traversal. + for (0..384) |expert| { + for (ids, 0..) |id, slot| { + if (id == expert) + for (&sum, down[slot]) |*s, value| { + s.* += math.f(value); + }; + } + } + for (branch, sum, down[6]) |*v, routed, shared| v.* = math.b(routed + math.f(shared)); +} + +fn moeCombine(ids: *const [6]u32, rows: *MoeRows) void { + moeCombineDown(ids, &rows.down, &rows.branch); +} + +fn moeReferenceRows(model: *Model, layer: usize, input: []const u16, ids: *const [6]u32, routing: *const [6]f32, rows: *MoeRows, scratch: *projection.Scratch, previous: ?*const MoeRows, previous_ids: ?*const [6]u32) !void { + for (0..7) |slot| { + const matrices = try moeMatrices(model, layer, if (slot == 6) null else ids[slot]); + var reuse: ?usize = if (slot == 6 and previous != null) 6 else null; + if (slot != 6 and previous_ids != null) + for (previous_ids.?, 0..) |id, prior| { + if (id == ids[slot]) { + reuse = prior; + break; + } + }; + if (reuse) |prior| { + rows.gate[slot] = previous.?.gate[prior]; + rows.up[slot] = previous.?.up[prior]; + } else { + try projection.apply(matrices[0], input, &rows.gate[slot], scratch); + try projection.apply(matrices[1], input, &rows.up[slot], scratch); + } + moe_reference.moeActivation(&rows.gate[slot], &rows.up[slot], if (slot == 6) 1 else routing[slot], model.cfg.text_config, &rows.activation[slot]); + // A shared or identical-weight expert result can be reused only when the BF16 input is identical. + if (reuse != null and std.mem.eql(u16, &rows.activation[slot], &previous.?.activation[reuse.?])) { + rows.down[slot] = previous.?.down[reuse.?]; + } else { + try projection.apply(matrices[2], &rows.activation[slot], &rows.down[slot], scratch); + } + } + moeCombine(ids, rows); +} + +fn moeEncodeNative(model: *Model, layer: usize, select_routes: bool) !void { + const ws = &model.workspace; + const p = &model.pipelines; + const w = &model.weights.layers[layer].moe; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + for (model.checkpoint.shards.items) |shard| enc.useResource(shard.buffer, mtl.ResourceUsage.read); + for (w.descriptors) |buffer| enc.useResource(buffer, mtl.ResourceUsage.read); + if (select_routes) { + tf.deepseek_v41.ops.encodeDense(enc, &p.ops, ws.view(.normalized), .{ .buffer = w.gate.buffer, .offset = w.gate.offset }, ws.view(.router_logits), .{ .rows = 1, .k = 5120, .n = 384, .x_stride = 5120, .w_stride = 5120, .y_stride = 384, .w_dtype = .f16, .y_dtype = .f32 }); + enc.barrier(); + // Bind the existing production router; no new native interface or diagnostic math kernel. + enc.setPipeline(p.route.select); + inline for (.{ ws.view(.router_logits), attention.View{ .buffer = w.bias.buffer, .offset = w.bias.offset }, ws.view(.route_ids), ws.view(.route_weights) }, 0..) |view, i| + enc.setBuffer(view.buffer, view.offset, i); + const Shape = extern struct { rows: u32, logit_stride: u32, output_stride: u32, temperature: f32, bias_dtype: u32 }; + enc.setValue(Shape{ .rows = 1, .logit_stride = 384, .output_stride = 6, .temperature = 1, .bias_dtype = 2 }, 4); + enc.dispatchGroups(mtl.Size.of(1, 1, 1), mtl.Size.of(128, 1, 1)); + enc.barrier(); + } + var routed_rows = ex.Rows{ .count = 6, .x_stride = 5120, .y_stride = 2304, .routes = .{ .ids = .{ .buffer = ws.view(.route_ids).buffer, .offset = ws.view(.route_ids).offset }, .repeat = 6 } }; + for (0..2) |i| { + const target = ws.view(if (i == 0) .gate else .up); + const input = ws.view(.normalized); + try ex.encodeIndexed(&p.projection, enc, .{ .descriptors = .{ .buffer = w.descriptors[i] }, .count = 384, .k = 5120, .n = 2304 }, .{ .buffer = input.buffer, .offset = input.offset }, .{ .buffer = target.buffer, .offset = target.offset }, ws.projection(), routed_rows); + enc.barrier(); + } + tf.deepseek_v41.ops.encodeSwiGLU(enc, &p.ops, ws.view(.gate), ws.view(.up), ws.view(.route_weights), ws.view(.activation), .{ .rows = 6, .width = 2304, .input_stride = 2304, .output_stride = 2304, .weighted = true }); + enc.barrier(); + routed_rows.x_stride = 2304; + routed_rows.y_stride = 5120; + routed_rows.routes.?.repeat = 1; + const activation = ws.view(.activation); + const down = ws.view(.routed); + try ex.encodeIndexed(&p.projection, enc, .{ .descriptors = .{ .buffer = w.descriptors[2] }, .count = 384, .k = 2304, .n = 5120 }, .{ .buffer = activation.buffer, .offset = activation.offset }, .{ .buffer = down.buffer, .offset = down.offset }, ws.projection(), routed_rows); + enc.barrier(); + for (0..2) |i| { + const target = ws.view(if (i == 0) .shared_gate else .shared_up); + const input = ws.view(.normalized); + ex.useMatrix(enc, w.shared[i]); + try ex.encode(&p.projection, enc, w.shared[i], .{ .buffer = input.buffer, .offset = input.offset }, .{ .buffer = target.buffer, .offset = target.offset }, ws.projection(), .{ .count = 1, .x_stride = 5120, .y_stride = 2304 }); + enc.barrier(); + } + tf.deepseek_v41.ops.encodeSwiGLU(enc, &p.ops, ws.view(.shared_gate), ws.view(.shared_up), ws.view(.route_weights), ws.view(.shared_activation), .{ .rows = 1, .width = 2304, .input_stride = 2304, .output_stride = 2304 }); + enc.barrier(); + const shared_activation = ws.view(.shared_activation); + const shared_down = ws.view(.shared_down); + ex.useMatrix(enc, w.shared[2]); + try ex.encode(&p.projection, enc, w.shared[2], .{ .buffer = shared_activation.buffer, .offset = shared_activation.offset }, .{ .buffer = shared_down.buffer, .offset = shared_down.offset }, ws.projection(), .{ .count = 1, .x_stride = 2304, .y_stride = 5120 }); + enc.barrier(); + tf.deepseek_v41.ops.encodeCombine(enc, &p.ops, ws.view(.routed), ws.view(.shared_down), ws.view(.route_ids), ws.view(.branch), .{ .rows = 1, .width = 5120, .routed_row_stride = 6 * 5120, .slot_stride = 5120, .shared_stride = 5120, .y_stride = 5120 }); + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; +} + +fn moeCaptureNative(model: *Model, rows: *MoeRows) void { + const ws = &model.workspace; + inline for (.{ "gate", "up", "activation", "down" }, .{ .gate, .up, .activation, .routed }, .{ .shared_gate, .shared_up, .shared_activation, .shared_down }) |stage, routed, shared| { + const width = if (comptime std.mem.eql(u8, stage, "down")) 5120 else 2304; + const source = diagnosticRow(u16, ws.view(routed), 6 * width); + for (0..6) |slot| @memcpy(&@field(rows.*, stage)[slot], source[slot * width ..][0..width]); + @memcpy(&@field(rows.*, stage)[6], diagnosticRow(u16, ws.view(shared), width)); + } + @memcpy(&rows.branch, diagnosticRow(u16, ws.view(.branch), 5120)); +} + +fn moeCompareRows(gpa: std.mem.Allocator, report: *std.ArrayList(u8), count: *usize, comparison: []const u8, reuse: bool, actual: *const MoeRows, actual_ids: *const [6]u32, expected: *const MoeRows, expected_ids: *const [6]u32, first: *?[]const u8) !void { + // Compare like expert IDs rather than assuming independent top-six orders are identical. + inline for (.{ "gate", "up", "activation", "down" }) |stage| { + var differs = false; + for (0..7) |slot| { + var reference_slot: ?usize = if (slot == 6) 6 else null; + if (slot != 6) + for (expected_ids, 0..) |id, i| { + if (id == actual_ids[slot]) { + reference_slot = i; + break; + } + }; + if (reference_slot == null) continue; + var name: [64]u8 = undefined; + const label = if (slot == 6) try std.fmt.bufPrint(&name, "shared_{s}", .{stage}) else try std.fmt.bufPrint(&name, "expert_{d}_{s}", .{ actual_ids[slot], stage }); + differs = (try moeMetric(u16, gpa, report, count, label, comparison, reuse, &@field(actual.*, stage)[slot], &@field(expected.*, stage)[reference_slot.?])) or differs; + } + if (differs and first.* == null) first.* = stage; + } + if (try moeMetric(u16, gpa, report, count, "combined_branch", comparison, reuse, &actual.branch, &expected.branch)) + if (first.* == null) { + first.* = "combined_branch"; + }; +} + +/// Retained local snapshot isolation. MoE has no history; qualification computes every value independently. +pub fn runMoe(gpa: std.mem.Allocator, io: std.Io, model: *Model, snapshot: []const u8, dir: []const u8, layer: usize) !void { + if (layer >= model.weights.layers.len) return error.InvalidMoeLayer; + const blob = try std.Io.Dir.cwd().readFileAlloc(io, snapshot, gpa, .limited(8 << 20)); + defer gpa.free(blob); + const snapshot_offset = (layer * 2 + 1) * 61488; + if (blob.len != 80 * 61488 or blob.len < snapshot_offset + 61488) return error.InvalidMoeSnapshot; + try model.reset(); + const input = model.workspace.view(.normalized); + @memcpy(input.buffer.contents()[input.offset..][0..10240], blob[snapshot_offset + 40960 ..][0..10240]); + const input_row = diagnosticRow(u16, input, 5120); + var retained_ids: [6]u32 = undefined; + var retained_weights: [6]f32 = undefined; + var retained_branch: [5120]u16 = undefined; + for (&retained_ids, 0..) |*id, i| id.* = std.mem.readInt(u32, blob[snapshot_offset + 61440 + 4 * i ..][0..4], .little); + for (&retained_weights, 0..) |*weight, i| + weight.* = @bitCast(std.mem.readInt(u32, blob[snapshot_offset + 61464 + 4 * i ..][0..4], .little)); + @memcpy(std.mem.asBytes(&retained_branch), blob[snapshot_offset + 51200 ..][0..10240]); + for (retained_ids, retained_weights, 0..) |id, weight, i| { + if (id >= 384 or !std.math.isFinite(weight)) return error.InvalidMoeRetainedRouting; + for (retained_ids[0..i]) |prior| if (id == prior) return error.InvalidMoeRetainedRouting; + } + const native = try gpa.create(MoeRows); + defer gpa.destroy(native); + const source = try gpa.create(MoeRows); + defer gpa.destroy(source); + const native_weights_reference = try gpa.create(MoeRows); + defer gpa.destroy(native_weights_reference); + const retained_reference = try gpa.create(MoeRows); + defer gpa.destroy(retained_reference); + const native_inputs_reference = try gpa.create(MoeRows); + defer gpa.destroy(native_inputs_reference); + const retained_native = try gpa.create(MoeRows); + defer gpa.destroy(retained_native); + var scratch = try projection.Scratch.init(gpa, 5120, 5120); + defer scratch.deinit(); + + try moeEncodeNative(model, layer, true); + moeCaptureNative(model, native); + var native_ids: [6]u32 = undefined; + var native_weights: [6]f32 = undefined; + var native_logits_routing: moe_reference.MoeRouting = .{}; + @memcpy(&native_ids, diagnosticRow(u32, model.workspace.view(.route_ids), 6)); + @memcpy(&native_weights, diagnosticRow(f32, model.workspace.view(.route_weights), 6)); + @memcpy(&native_logits_routing.logits, diagnosticRow(f32, model.workspace.view(.router_logits), 384)); + moe_reference.moeSelect(model.weights.layers[layer].moe, model.cfg.text_config, &native_logits_routing.logits, &native_logits_routing.scores, &native_logits_routing.ids, &native_logits_routing.weights); + var source_routing: moe_reference.MoeRouting = .{}; + math.dense(model.weights.layers[layer].moe.gate, input_row, &source_routing.logits); + moe_reference.moeSelect(model.weights.layers[layer].moe, model.cfg.text_config, &source_routing.logits, &source_routing.scores, &source_routing.ids, &source_routing.weights); + try moeReferenceRows(model, layer, input_row, &source_routing.ids, &source_routing.weights, source, &scratch, null, null); + try moeReferenceRows(model, layer, input_row, &native_ids, &native_weights, native_weights_reference, &scratch, source, &source_routing.ids); + try moeReferenceRows(model, layer, input_row, &retained_ids, &retained_weights, retained_reference, &scratch, native_weights_reference, &native_ids); + for (0..7) |slot| { + native_inputs_reference.gate[slot] = native.gate[slot]; + native_inputs_reference.up[slot] = native.up[slot]; + moe_reference.moeActivation(&native.gate[slot], &native.up[slot], if (slot == 6) 1 else native_weights[slot], model.cfg.text_config, &native_inputs_reference.activation[slot]); + // Deliberately use the captured native activation, not the preceding reference activation. + if (std.mem.eql(u16, &native.activation[slot], &native_weights_reference.activation[slot])) { + native_inputs_reference.down[slot] = native_weights_reference.down[slot]; + } else { + const matrices = try moeMatrices(model, layer, if (slot == 6) null else native_ids[slot]); + try projection.apply(matrices[2], &native.activation[slot], &native_inputs_reference.down[slot], &scratch); + } + } + moeCombine(&native_ids, native_inputs_reference); + @memcpy(diagnosticRow(u32, model.workspace.view(.route_ids), 6), &retained_ids); + @memcpy(diagnosticRow(f32, model.workspace.view(.route_weights), 6), &retained_weights); + try moeEncodeNative(model, layer, false); + moeCaptureNative(model, retained_native); + + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + const header = try std.json.Stringify.valueAlloc(gpa, .{ + .mode = "moe-sameinput", + .token = @as(?usize, null), + .layer = layer, + .token_identity_source = "caller_supplied_snapshot_argument_raw_snapshot_has_no_token_header", + .retained_route_origin = "native_GPU_snapshot_command_contract", + .snapshot_count = 80, + .snapshot_length_bytes = 61488, + .snapshot_index = layer * 2 + 1, + .snapshot_stage = "moe", + .input_offset_bytes = snapshot_offset + 40960, + .input_length_bytes = 10240, + .retained_ids_offset_bytes = snapshot_offset + 61440, + .retained_weights_offset_bytes = snapshot_offset + 61464, + .identical_normalized_input = true, + .reference_dot = "independent_EXL3_FP64", + .qualification_gate_applied = false, + .qualification_oracle_modified_to_reuse_native_values = false, + .official_revision = "2cba9e42aa026125f3ed06c6d98c1db82f7ca027", + .source_ids = source_routing.ids, + .source_weights = source_routing.weights, + .native_ids = native_ids, + .native_weights = native_weights, + .retained_ids = retained_ids, + .retained_weights = retained_weights, + .source_ids_equal_native_order = std.mem.eql(u32, &source_routing.ids, &native_ids), + .native_ids_equal_retained_order = std.mem.eql(u32, &native_ids, &retained_ids), + .source_select_on_native_logits_ids = native_logits_routing.ids, + .source_select_on_native_logits_weights = native_logits_routing.weights, + .score_receipt_origin = "source_sqrt_softplus_applied_to_captured_native_logits_not_GPU_score_capture", + .expert_receipt_layout = "slot_major_top6_ids_order_then_shared_slot6", + .expert_sum_order = "ascending_expert_id_FP32_shared_last_final_BF16", + .receipt_prefixes = .{ "moe-native", "moe-source", "moe-native-weights-reference-diagnostic", "moe-native-inputs-reference-diagnostic", "moe-retained-weights-reference-diagnostic", "moe-retained-route-native-diagnostic" }, + .receipt_shapes = .{ .gate = .{ 7, 2304 }, .up = .{ 7, 2304 }, .activation = .{ 7, 2304 }, .down = .{ 7, 5120 }, .branch = .{5120} }, + .native_inputs_reference_contract = "activation:reference_on_native_gate_up_and_weight;down:FP64_on_native_activation;branch:reference_combine_of_these_down_outputs", + }, .{}); + defer gpa.free(header); + try report.appendSlice(gpa, header[0 .. header.len - 1]); + try report.appendSlice(gpa, ",\"cases\":[\n"); + var count: usize = 0; + var first: MoeFirst = .{}; + if (try moeMetric(f32, gpa, &report, &count, "route_logits", "source_chain", false, &native_logits_routing.logits, &source_routing.logits)) first.source_chain = "route_logits"; + if (try moeMetric(f32, gpa, &report, &count, "route_scores_source_transform", "source_score_transform_on_captured_native_logits_diagnostic", true, &native_logits_routing.scores, &source_routing.scores)) + if (first.source_chain == null) { + first.source_chain = "route_scores_source_transform"; + }; + if (!std.mem.eql(u32, &source_routing.ids, &native_ids) and first.source_chain == null) first.source_chain = "route_ids"; + if (std.mem.eql(u32, &source_routing.ids, &native_ids)) { + if (try moeMetric(f32, gpa, &report, &count, "route_weights", "source_chain", false, &native_weights, &source_routing.weights)) + if (first.source_chain == null) { + first.source_chain = "route_weights"; + }; + } + if (std.mem.eql(u32, &native_logits_routing.ids, &native_ids)) + _ = try moeMetric(f32, gpa, &report, &count, "route_weights", "source_select_on_same_native_logits_diagnostic", true, &native_weights, &native_logits_routing.weights); + if (std.mem.eql(u32, &retained_ids, &native_ids)) + _ = try moeMetric(f32, gpa, &report, &count, "route_weights", "retained_native_weights_vs_reexecuted_native", true, &native_weights, &retained_weights); + try moeCompareRows(gpa, &report, &count, "source_chain", false, native, &native_ids, source, &source_routing.ids, &first.source_expert_stage); + if (first.source_chain == null) first.source_chain = first.source_expert_stage; + try moeCompareRows(gpa, &report, &count, "same_native_ids_weights_reference_chain_diagnostic", true, native, &native_ids, native_weights_reference, &native_ids, &first.native_weights_reference_chain_diagnostic); + inline for (.{ "activation", "down" }) |stage| { + var differs = false; + for (0..7) |slot| { + var name: [64]u8 = undefined; + const label = if (slot == 6) try std.fmt.bufPrint(&name, "shared_{s}", .{stage}) else try std.fmt.bufPrint(&name, "expert_{d}_{s}", .{ native_ids[slot], stage }); + differs = (try moeMetric(u16, gpa, &report, &count, label, if (comptime std.mem.eql(u8, stage, "activation")) "same_native_gate_up_weight_reference_activation_diagnostic" else "same_native_activation_FP64_down_diagnostic", true, &@field(native.*, stage)[slot], &@field(native_inputs_reference.*, stage)[slot])) or differs; + } + if (differs and first.native_input_diagnostic == null) first.native_input_diagnostic = stage; + } + _ = try moeMetric(u16, gpa, &report, &count, "combined_branch", "same_native_activation_FP64_down_then_reference_combine_diagnostic", true, &native.branch, &native_inputs_reference.branch); + // Isolate combination arithmetic independently of every projection and route-weight difference. + var native_down_reference_branch: [5120]u16 = undefined; + moeCombineDown(&native_ids, &native.down, &native_down_reference_branch); + first.native_down_reference_combine_diagnostic_differs = + try moeMetric(u16, gpa, &report, &count, "combined_branch", "same_native_down_reference_combine_diagnostic", true, &native.branch, &native_down_reference_branch); + try moeWrite(gpa, io, dir, "moe-native-down-reference-combine-diagnostic.bf16", std.mem.asBytes(&native_down_reference_branch)); + try moeCompareRows(gpa, &report, &count, "same_retained_native_ids_weights_reference_chain_diagnostic", true, retained_native, &retained_ids, retained_reference, &retained_ids, &first.retained_weights_reference_chain_diagnostic); + _ = try moeMetric(u16, gpa, &report, &count, "combined_branch", "retained_native_branch_vs_reexecuted_native_routing", true, &native.branch, &retained_branch); + _ = try moeMetric(u16, gpa, &report, &count, "combined_branch", "retained_native_branch_vs_reexecuted_retained_routes_diagnostic", true, &retained_native.branch, &retained_branch); + const finding = try std.json.Stringify.valueAlloc(gpa, first, .{}); + defer gpa.free(finding); + try report.appendSlice(gpa, "\n],\"first_differing_intermediate\":"); + try report.appendSlice(gpa, finding); + try report.appendSlice(gpa, "}\n"); + try moeWrite(gpa, io, dir, "moe-sameinput.json", report.items); + inline for (.{ "moe-native", "moe-source", "moe-native-weights-reference-diagnostic", "moe-native-inputs-reference-diagnostic", "moe-retained-weights-reference-diagnostic", "moe-retained-route-native-diagnostic" }, .{ native, source, native_weights_reference, native_inputs_reference, retained_reference, retained_native }) |prefix, rows| + try moeWriteRows(gpa, io, dir, prefix, rows); + inline for (.{ "moe-input.bf16", "moe-retained-branch.bf16", "moe-native-logits.f32", "moe-source-logits.f32", "moe-native-logits-source-scores.f32", "moe-source-scores.f32", "moe-native-ids.u32", "moe-source-ids.u32", "moe-retained-ids.u32", "moe-native-weights.f32", "moe-source-weights.f32", "moe-retained-weights.f32", "moe-native-logits-source-select-ids.u32", "moe-native-logits-source-select-weights.f32" }, .{ std.mem.sliceAsBytes(input_row), std.mem.asBytes(&retained_branch), std.mem.asBytes(&native_logits_routing.logits), std.mem.asBytes(&source_routing.logits), std.mem.asBytes(&native_logits_routing.scores), std.mem.asBytes(&source_routing.scores), std.mem.asBytes(&native_ids), std.mem.asBytes(&source_routing.ids), std.mem.asBytes(&retained_ids), std.mem.asBytes(&native_weights), std.mem.asBytes(&source_routing.weights), std.mem.asBytes(&retained_weights), std.mem.asBytes(&native_logits_routing.ids), std.mem.asBytes(&native_logits_routing.weights) }) |name, bytes| + try moeWrite(gpa, io, dir, name, bytes); +} diff --git a/zig/tests/deepseek_v41_reference/operators.zig b/zig/tests/deepseek_v41_reference/operators.zig new file mode 100644 index 0000000000..8d050300e2 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/operators.zig @@ -0,0 +1,92 @@ +//! Actual checkpoint projections on identical deterministic BF16 rows, against independent FP64 dot products. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const Model = tf.deepseek_v41.Model; +const ex = tf.exl3_projection; +const Case = struct { name: []const u8, matrix: ex.Matrix }; + +pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, dir: []const u8) !void { + const first = &model.weights.layers[0]; + const cases = [_]Case{ + .{ .name = "wq_a", .matrix = first.attn.wq_a }, + .{ .name = "wq_b", .matrix = first.attn.wq_b }, + .{ .name = "wkv", .matrix = first.attn.wkv }, + .{ .name = "wo_a_group0", .matrix = first.attn.wo_a[0] }, + .{ .name = "wo_b", .matrix = first.attn.wo_b }, + .{ .name = "engram_wkv", .matrix = model.weights.layers[1].engram.?.wkv }, + .{ .name = "routed_gate", .matrix = try tf.deepseek_v41_weights.lookup(&model.checkpoint, "layers.0.ffn.experts.0.w1", 5120, 2304) }, + .{ .name = "routed_up", .matrix = try tf.deepseek_v41_weights.lookup(&model.checkpoint, "layers.0.ffn.experts.0.w3", 5120, 2304) }, + .{ .name = "routed_down", .matrix = try tf.deepseek_v41_weights.lookup(&model.checkpoint, "layers.0.ffn.experts.0.w2", 2304, 5120) }, + .{ .name = "head", .matrix = model.weights.head }, + }; + const x = try model.device.buffer(8192 * 2, mtl.ResourceOptions.shared); + defer x.deinit(); + const y = try model.device.buffer(2 * 129280 * 4, mtl.ResourceOptions.shared); + defer y.deinit(); + const scratch = try model.device.buffer(try ex.scratchBytes(8192, 129280, 1), mtl.ResourceOptions.shared); + defer scratch.deinit(); + for (x.slice(u16, 8192), 0..) |*value, i| + value.* = math.b(@as(f32, @floatFromInt(@as(i32, @intCast((i * 17) % 97)) - 48)) / 32); + for ([_]mtl.Buffer{ x, y, scratch }) |buffer| model.residency.add(buffer); + model.residency.commit(); + model.residency.requestResidency(); + defer { + for ([_]mtl.Buffer{ x, y, scratch }) |buffer| model.residency.remove(buffer); + model.residency.commit(); + } + var reference = try projection.Scratch.init(gpa, 8192, 129280); + defer reference.deinit(); + const expected = try gpa.alloc(f32, 129280); + defer gpa.free(expected); + var report: std.ArrayList(u8) = .empty; + defer report.deinit(gpa); + try report.appendSlice(gpa, "{\"input\":\"bf16((17*i%97-48)/32)\",\"reference_dot\":\"FP64\",\"cases\":["); + var all_pass = true; + for (cases, 0..) |case, index| { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const matrix = case.matrix; + const cb = model.queue.commandBufferUnretained(); + const enc = cb.compute(.serial); + var encoding = true; + errdefer if (encoding) enc.end(); + ex.useMatrix(enc, matrix); + try ex.encode(&model.pipelines.projection, enc, matrix, .{ .buffer = x }, .{ .buffer = y }, .{ .buffer = scratch }, .{ .count = 1, .x_stride = matrix.k, .y_stride = matrix.n, .output = .f32 }); + enc.end(); + encoding = false; + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.OperatorGpuFailure; + const actual = y.slice(f32, 129280); + try projection.applyF32(matrix, x.slice(u16, matrix.k), expected[0..matrix.n], &reference); + var max_abs: f64 = 0; + var worst: usize = 0; + var squared: f64 = 0; + var magnitude: f64 = 0; + for (actual[0..matrix.n], expected[0..matrix.n], 0..) |a, b, column| { + if (!std.math.isFinite(a) or !std.math.isFinite(b)) return error.NonfiniteOperatorOutput; + const delta: f64 = @as(f64, a) - b; + if (@abs(delta) > max_abs) { + max_abs = @abs(delta); + worst = column; + } + squared += delta * delta; + magnitude += @as(f64, b) * b; + } + const nrms = if (magnitude == 0) @sqrt(squared) else @sqrt(squared / magnitude); + all_pass = all_pass and nrms <= 1e-4; + const line = try std.fmt.allocPrint(gpa, "{{\"name\":\"{s}\",\"k\":{d},\"n\":{d},\"normalized_rms_error\":{d},\"max_absolute_error\":{d},\"worst_column\":{d},\"worst_gpu\":{d},\"worst_reference\":{d}}}", .{ case.name, matrix.k, matrix.n, nrms, max_abs, worst, actual[worst], expected[worst] }); + defer gpa.free(line); + if (index > 0) try report.appendSlice(gpa, ","); + try report.appendSlice(gpa, line); + std.debug.print("{s}\n", .{line}); + } + try report.appendSlice(gpa, "]}\n"); + const path = try std.fs.path.join(gpa, &.{ dir, "operators.json" }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = report.items }); + if (!all_pass) return error.OperatorReferenceMismatch; +} diff --git a/zig/tests/deepseek_v41_reference/profile.zig b/zig/tests/deepseek_v41_reference/profile.zig new file mode 100644 index 0000000000..0fb15a8b90 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/profile.zig @@ -0,0 +1,52 @@ +//! Real-checkpoint calibrated GPU stage timings, with profile-on/off logits checked before reporting. +const std = @import("std"); +const tf = @import("tensorfold"); +const counters = tf.deepseek_v41_profiler; +const Model = tf.deepseek_v41.Model; + +pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, ids: []const u32, dir: []const u8) !void { + if (ids.len == 0 or ids.len > model.workspace.plan.rows) return error.InvalidProfileWindow; + const previous_grouping = model.grouped_prefill; + defer model.grouped_prefill = previous_grouping; + model.grouped_prefill = false; + try model.reset(); + const baseline = try model.forward(ids, .all); + const baseline_gpu_seconds = model.last_gpu_seconds; + const expected = try gpa.dupe(f32, baseline); + defer gpa.free(expected); + model.grouped_prefill = true; + try model.reset(); + const candidate = try model.forward(ids, .all); + const candidate_gpu_seconds = model.last_gpu_seconds; + try check(expected, candidate); + var recorder = try counters.Recorder.init(gpa, model.device, .{ .enabled = true, .max_spans = 512 }); + defer recorder.deinit() catch unreachable; + const previous_profile = model.profile; + model.profile = &recorder; + defer model.profile = previous_profile; + try model.reset(); + const actual = try model.forward(ids, .last); + try check(expected[expected.len - 129280 ..], actual); + const totals = recorder.aggregates(); + var sampled_gpu_seconds: f64 = 0; + for (totals) |value| sampled_gpu_seconds += value.gpu_seconds; + var text: std.ArrayList(u8) = .empty; + defer text.deinit(gpa); + try text.print(gpa, "{{\"tokens\":{d},\"comparison_logits\":\"all\",\"candidate_grouping\":true,\"same_logits\":true,\"baseline_command_gpu_seconds\":{d},\"candidate_command_gpu_seconds\":{d},\"profiled_logits\":\"last\",\"sampled_span_gpu_seconds\":{d},\"stages\":[", .{ ids.len, baseline_gpu_seconds, candidate_gpu_seconds, sampled_gpu_seconds }); + for (std.enums.values(counters.Stage), totals, 0..) |stage, value, i| { + try text.print(gpa, "{s}{{\"stage\":\"{s}\",\"calls\":{d},\"gpu_seconds\":{d}}}", .{ if (i == 0) "" else ",", @tagName(stage), value.calls, value.gpu_seconds }); + } + try text.appendSlice(gpa, "]}\n"); + const path = try std.fs.path.join(gpa, &.{ dir, "gpu-profile.json" }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = text.items }); + std.debug.print("{s}", .{text.items}); +} + +fn check(expected: []const f32, actual: []const f32) !void { + if (expected.len != actual.len) return error.ProfileLogitShapeMismatch; + for (actual, expected) |a, b| { + if (!std.math.isFinite(a) or !std.math.isFinite(b)) return error.NonfiniteProfileLogits; + if (@as(u32, @bitCast(a)) != @as(u32, @bitCast(b))) return error.ProfileChangedLogits; + } +} diff --git a/zig/tests/deepseek_v41_reference/projection.zig b/zig/tests/deepseek_v41_reference/projection.zig new file mode 100644 index 0000000000..ca0f932c22 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -0,0 +1,261 @@ +//! Independent scalar EXL3 oracle over resident bytes; FP64 dots and FP32 rotation/output boundaries, no GPU work. + +// EXL3 attribution: docs/recipes/exl3.md; pinned checkpoint provenance in THIRD_PARTY_NOTICES.md. + +const std = @import("std"); +const exl3 = @import("tensorfold").exl3_projection; +const format = exl3.format; +const block_size = format.hadamard_dim; +const inverse: f32 = 0.08838834764831845; +const states_count = 65536; + +/// Reuse bounded vector/codebook storage for heterogeneous matrices; never allocate decoded W[K,N]. +pub const Scratch = struct { + gpa: std.mem.Allocator, + rotated: []f64, + sums: []f64, + codebooks: []f16, + + pub fn init(gpa: std.mem.Allocator, max_k: usize, max_n: usize) !Scratch { + if (max_k == 0 or max_n == 0) return error.InvalidReferenceProjectionCapacity; + const rotated = try gpa.alloc(f64, max_k); + errdefer gpa.free(rotated); + const sums = try gpa.alloc(f64, max_n); + errdefer gpa.free(sums); + const codebooks = try gpa.alloc(f16, 3 * states_count); + errdefer gpa.free(codebooks); + for ([_]format.Codebook{ .inst3, .mcg, .mul1 }, 0..) |book, bi| { + for (0..states_count) |state| { + codebooks[bi * states_count + state] = format.codebookValue(book, @intCast(state)); + } + } + return .{ .gpa = gpa, .rotated = rotated, .sums = sums, .codebooks = codebooks }; + } + + pub fn deinit(s: *Scratch) void { + s.gpa.free(s.rotated); + s.gpa.free(s.sums); + s.gpa.free(s.codebooks); + s.* = undefined; + } +}; + +/// One BF16 input row and BF16 output row; caller must finish GPU writes before borrowing shared CPU bytes. +pub fn apply(matrix: exl3.Matrix, x: []const u16, y: []u16, scratch: *Scratch) !void { + try project(matrix, x, y.len, scratch); + try finish(u16, matrix, y, scratch.sums[0..matrix.n]); +} + +/// The same arithmetic with an FP32 output boundary, for compressor projections and full-vocabulary logits. +pub fn applyF32(matrix: exl3.Matrix, x: []const u16, y: []f32, scratch: *Scratch) !void { + try project(matrix, x, y.len, scratch); + try finish(f32, matrix, y, scratch.sums[0..matrix.n]); +} + +fn borrowed(view: exl3.View, bytes: usize) ![]const u8 { + const len = view.buffer.length(); + if (view.offset % 2 != 0 or view.offset > len or bytes > len - view.offset) + return error.InvalidReferenceProjectionView; + const address = view.buffer.contents(); + if (@intFromPtr(address) == 0) return error.NonCpuAccessibleExl3Buffer; + return address[view.offset..][0..bytes]; +} + +fn scaleBytes(count: usize, kind: exl3.ScaleFormat) !usize { + return std.math.mul(usize, if (kind == .packed_signs) count / 16 else count, 2); +} + +fn half(bytes: []const u8, i: usize) f32 { + const value: f16 = @bitCast(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); + return @floatCast(value); +} + +fn scale(bytes: []const u8, i: usize, kind: exl3.ScaleFormat) f32 { + if (kind == .f16) return half(bytes, i); + const word = std.mem.readInt(u16, bytes[2 * (i / 16) ..][0..2], .little); + return if (((word >> @as(u4, @intCast(i % 16))) & 1) != 0) -1.0 else 1.0; +} + +fn bf16(value: f32) u16 { + const bits: u32 = @bitCast(value); + // Preserve a quiet NaN instead of letting rounding turn a tiny NaN payload into infinity. + if ((bits & 0x7fffffff) > 0x7f800000) return @as(u16, @truncate(bits >> 16)) | 0x0040; + return @truncate((bits +% (0x7fff + ((bits >> 16) & 1))) >> 16); +} + +/// Seven ordinary Sylvester butterflies, with explicit FP32 stage boundaries and normalization afterward. +fn hadamard(values: *[block_size]f32) void { + @setFloatMode(.strict); + var stride: usize = 1; + while (stride < block_size) : (stride *= 2) { + var base: usize = 0; + while (base < block_size) : (base += 2 * stride) { + for (0..stride) |j| { + const a = values[base + j]; + const b = values[base + stride + j]; + values[base + j] = a + b; + values[base + stride + j] = a - b; + } + } + } +} + +fn prepare(x: []const u16, su: []const u8, kind: exl3.ScaleFormat, rotated: []f64) void { + @setFloatMode(.strict); + var base: usize = 0; + while (base < x.len) : (base += block_size) { + var values: [block_size]f32 = undefined; + for (&values, 0..) |*value, i| { + const input: f32 = @bitCast(@as(u32, x[base + i]) << 16); + value.* = input * scale(su, base + i, kind); + } + hadamard(&values); + for (values, 0..) |value, i| { + const rounded: f16 = @floatCast(value * inverse); + rotated[base + i] = @floatCast(rounded); + } + } +} + +/// CPU stream reader consumes MSB-first bits from little-endian int16 pairs, not the GPU indexed decoder. +const Stream = struct { + bytes: []const u8, + cursor: usize = 0, + pending: u32 = 0, + available: u6 = 0, + + fn take(s: *Stream, count: u5) u16 { + if (s.available == 0) { + s.pending = std.mem.readInt(u32, s.bytes[s.cursor..][0..4], .little); + s.cursor += 4; + s.available = 32; + } + const first: u5 = @intCast(@min(s.available, count)); + const head: u32 = s.pending >> @as(u5, @intCast(32 - @as(u6, first))); + s.pending <<= first; + s.available -= first; + if (first == count) return @intCast(head); + const rest = count - first; + s.pending = std.mem.readInt(u32, s.bytes[s.cursor..][0..4], .little); + s.cursor += 4; + const tail = s.pending >> @as(u5, @intCast(32 - @as(u6, rest))); + s.pending <<= rest; + s.available = 32 - @as(u6, rest); + return @intCast((head << rest) | tail); + } +}; + +/// Circular initial state uses final 16 stream bits; decode only a sixteen-column accumulator and live stream. +fn accumulateTile(comptime halves: u16, tile: []const u8, table: []const f16, x: []const f64, acc: *[16]f64) void { + @setFloatMode(.strict); + var stream: Stream = .{ .bytes = tile }; + var state: u16 = @truncate(std.mem.readInt(u32, tile[tile.len - 4 ..][0..4], .little)); + var p: usize = 0; + while (p < 256) : (p += 1) { + // Odd half-widths alternate floor(bits),ceil(bits), exactly as format.streamEnds specifies. + const advance: u4 = @intCast(halves / 2 + (if (halves % 2 == 1) p % 2 else 0)); + state = (state << advance) | stream.take(advance); + const lane = p / 8; + const j = p % 8; + const row = 2 * (lane % 4) + (j & 1) + 8 * ((j >> 1) & 1); + const col = lane / 4 + 8 * (j >> 2); + const weight: f64 = @floatCast(table[state]); + acc[col] += x[row] * weight; + } +} + +fn product(comptime halves: u16, matrix: exl3.Matrix, trellis: []const u8, table: []const f16, rotated: []const f64, sums: []f64) !void { + const ks = try matrix.kStride(); + const ns = matrix.nStride(); + const tile_bytes = @as(usize, halves) * 16; + @memset(sums, 0); + for (0..matrix.k / 16) |tk| { + const input = rotated[tk * 16 ..][0..16]; + for (0..matrix.n / 16) |tn| { + const at = tk * ks + tn * ns; + accumulateTile(halves, trellis[at..][0..tile_bytes], table, input, sums[tn * 16 ..][0..16]); + } + } +} + +fn project(matrix: exl3.Matrix, x: []const u16, output_len: usize, scratch: *Scratch) !void { + try matrix.validate(); + if (x.len != matrix.k or output_len != matrix.n) return error.InvalidReferenceProjectionLength; + if (matrix.k > scratch.rotated.len or matrix.n > scratch.sums.len) + return error.ReferenceProjectionScratchTooSmall; + const trellis = try borrowed(matrix.trellis, try matrix.storageBytes()); + const su = try borrowed(matrix.suh, try scaleBytes(matrix.k, matrix.suh_format)); + // Check every borrowed view before doing arithmetic, including the output-only scale and optional bias. + _ = try borrowed(matrix.svh, try scaleBytes(matrix.n, matrix.svh_format)); + if (matrix.bias) |bias| _ = try borrowed(bias, try std.math.mul(usize, matrix.n, 2)); + const rotated = scratch.rotated[0..matrix.k]; + const sums = scratch.sums[0..matrix.n]; + prepare(x, su, matrix.suh_format, rotated); + const book_index: usize = @intCast(@backingInt(matrix.codebook)); + const table = scratch.codebooks[book_index * states_count ..][0..states_count]; + switch (matrix.width.halves) { + inline 2...16 => |halves| try product(halves, matrix, trellis, table, rotated, sums), + else => unreachable, + } +} + +fn finish(comptime T: type, matrix: exl3.Matrix, y: []T, sums: []const f64) !void { + @setFloatMode(.strict); + const sv = try borrowed(matrix.svh, try scaleBytes(matrix.n, matrix.svh_format)); + const bias = if (matrix.bias) |view| try borrowed(view, try std.math.mul(usize, matrix.n, 2)) else null; + var base: usize = 0; + while (base < matrix.n) : (base += block_size) { + var values: [block_size]f32 = undefined; + // One FP64->FP32 boundary replaces GPU FP32 dot-product accumulation, not the output rotation. + for (&values, 0..) |*value, i| value.* = @floatCast(sums[base + i]); + hadamard(&values); + for (values, 0..) |value, i| { + var out: f32 = (value * inverse) * scale(sv, base + i, matrix.svh_format); + if (bias) |bytes| out += half(bytes, base + i); + y[base + i] = if (T == u16) bf16(out) else out; + } + } +} + +test "CPU circular streaming matches scalar format for every supported width and codebook" { + var scratch = try Scratch.init(std.testing.allocator, 128, 128); + defer scratch.deinit(); + var words: [128]i16 = undefined; + for (&words, 0..) |*word, i| word.* = @bitCast(@as(u16, @truncate(i * 40503 + 17329))); + var x: [16]f64 = undefined; + for (&x, 0..) |*value, i| value.* = @as(f64, @floatFromInt(i + 1)) * 0.0625; + var placement: [256]format.Placement = undefined; + format.tilePlacement(&placement); + inline for (2..17) |halves| { + const width: format.Width = .{ .halves = halves }; + const tile = words[0..width.tileWords()]; + var states: [256]u32 = undefined; + format.tileStates(width, tile, &states); + for ([_]format.Codebook{ .inst3, .mcg, .mul1 }, 0..) |book, bi| { + if (width.isHalf() and (book != .mul1 or halves > 7)) continue; + const table = scratch.codebooks[bi * states_count ..][0..states_count]; + var expected: [16]f64 = @splat(0); + for (states, placement) |state, where| { + const weight: f64 = @floatCast(format.codebookValue(book, @intCast(state))); + expected[where.col] += x[where.row] * weight; + } + var actual: [16]f64 = @splat(0); + accumulateTile(halves, std.mem.sliceAsBytes(tile), table, &x, &actual); + try std.testing.expectEqualSlices(f64, &expected, &actual); + } + } +} + +test "FP32 Hadamard then FP16 and BF16 boundaries are deterministic" { + var x: [128]u16 = @splat(0); + x[0] = 0x3f80; + var signs: [16]u8 = @splat(0); + signs[0] = 1; + var rotated: [128]f64 = undefined; + prepare(&x, &signs, .packed_signs, &rotated); + const rounded: f16 = @floatCast(-inverse); + for (rotated) |value| try std.testing.expectEqual(@as(f64, @floatCast(rounded)), value); + try std.testing.expectEqual(@as(u16, 0x3f80), bf16(@bitCast(@as(u32, 0x3f808000)))); + try std.testing.expectEqual(@as(u16, 0x3f82), bf16(@bitCast(@as(u32, 0x3f818000)))); + try std.testing.expectEqual(@as(u16, 0x7fc0), bf16(@bitCast(@as(u32, 0x7f800001)))); +} diff --git a/zig/tests/deepseek_v41_reference/qualify.zig b/zig/tests/deepseek_v41_reference/qualify.zig new file mode 100644 index 0000000000..e5190ee546 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/qualify.zig @@ -0,0 +1,97 @@ +//! Actual checkpoint teacher comparison: independent CPU arithmetic versus the resident GPU decode path. +const std = @import("std"); +const mtl = @import("metal"); +const Model = @import("tensorfold").deepseek_v41.Model; +const Reference = @import("reference.zig").Reference; +const nrms_limit: f64 = 0.01; + +pub const Result = struct { + tokens: usize, + max_absolute_error: f64, + rms_error: f64, + normalized_rms_error: f64, + greedy_disagreements: usize, + gpu_seconds: f64, + cpu_seconds: f64, + normalized_rms_limit: f64 = nrms_limit, + numeric_gate_passed: bool = false, + greedy_gate_passed: bool = false, +}; + +fn maximum(row: []const f32) !u32 { + var id: u32 = 0; + var best = -std.math.inf(f32); + for (row, 0..) |v, i| { + if (!std.math.isFinite(v)) return error.NonfiniteTeacherLogits; + if (v > best) { + best = v; + id = @intCast(i); + } + } + return id; +} + +pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, ids: []const u32, token_map: []const u8, dir: []const u8) !Result { + if (ids.len == 0 or ids.len > 16) return error.InvalidOracleWindow; + var reference = try Reference.init(gpa, model.cfg.text_config, &model.weights, &model.checkpoint, token_map, @intCast(ids.len)); + defer reference.deinit(); + const cpu = try gpa.alloc(f32, ids.len * 129280); + defer gpa.free(cpu); + const gpu = try gpa.alloc(f32, cpu.len); + defer gpa.free(gpu); + try model.reset(); + var result = Result{ .tokens = ids.len, .max_absolute_error = 0, .rms_error = 0, .normalized_rms_error = 0, .greedy_disagreements = 0, .gpu_seconds = 0, .cpu_seconds = 0 }; + var squared_error: f64 = 0; + var squared_reference: f64 = 0; + for (ids, 0..) |id, i| { + const actual = try model.forward(&.{id}, .all); + result.gpu_seconds += model.last_gpu_seconds; + @memcpy(gpu[i * 129280 ..][0..129280], actual); + const start = mtl.clock.seconds(); + const expected = try reference.step(id); + result.cpu_seconds += mtl.clock.seconds() - start; + @memcpy(cpu[i * 129280 ..][0..129280], expected); + if (try maximum(actual) != try maximum(expected)) result.greedy_disagreements += 1; + for (actual, expected) |a, b| { + if (!std.math.isFinite(a) or !std.math.isFinite(b)) return error.NonfiniteTeacherLogits; + const error_: f64 = @as(f64, a) - b; + result.max_absolute_error = @max(result.max_absolute_error, @abs(error_)); + squared_error += error_ * error_; + squared_reference += @as(f64, b) * b; + } + std.debug.print("oracle teacher {d}/{d}: GPU token={d}, CPU token={d}\n", .{ i + 1, ids.len, try maximum(actual), try maximum(expected) }); + } + result.rms_error = @sqrt(squared_error / @as(f64, @floatFromInt(cpu.len))); + result.normalized_rms_error = if (squared_reference == 0) result.rms_error else @sqrt(squared_error / squared_reference); + result.numeric_gate_passed = result.normalized_rms_error <= result.normalized_rms_limit; + result.greedy_gate_passed = result.greedy_disagreements == 0; + const cp = try std.fs.path.join(gpa, &.{ dir, "oracle-logits.f32" }); + defer gpa.free(cp); + const gp = try std.fs.path.join(gpa, &.{ dir, "gpu-logits.f32" }); + defer gpa.free(gp); + const report = try std.fs.path.join(gpa, &.{ dir, "oracle.json" }); + defer gpa.free(report); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = cp, .data = std.mem.sliceAsBytes(cpu) }); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = gp, .data = std.mem.sliceAsBytes(gpu) }); + const json = try std.json.Stringify.valueAlloc(gpa, result, .{}); + defer gpa.free(json); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = report, .data = json }); + std.debug.print("{s}\n", .{json}); + try enforceGate(result.normalized_rms_error, result.greedy_disagreements); + return result; +} + +fn enforceGate(normalized_rms_error: f64, greedy_disagreements: usize) !void { + if (!std.math.isFinite(normalized_rms_error)) return error.NonfiniteTeacherLogits; + if (normalized_rms_error > nrms_limit) return error.TeacherReferenceMismatch; + if (greedy_disagreements != 0) return error.TeacherGreedyMismatch; +} + +test "teacher numeric gate remains independent of greedy agreement" { + try std.testing.expectError(error.TeacherReferenceMismatch, enforceGate(0.04, 0)); + try std.testing.expectError(error.TeacherReferenceMismatch, enforceGate(0.010000001, 1)); + try std.testing.expectError(error.TeacherGreedyMismatch, enforceGate(0.01, 1)); + try enforceGate(0.01, 0); + try std.testing.expectError(error.NonfiniteTeacherLogits, enforceGate(std.math.nan(f64), 0)); + try std.testing.expectError(error.NonfiniteTeacherLogits, enforceGate(std.math.inf(f64), 0)); +} diff --git a/zig/tests/deepseek_v41_reference/reference.zig b/zig/tests/deepseek_v41_reference/reference.zig new file mode 100644 index 0000000000..23d811a532 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/reference.zig @@ -0,0 +1,176 @@ +//! Independent wide-sum qualification with explicit FP32/BF16 boundaries; not GPU self-reference or CUDA emulation. + +// Equations: pinned DeepSeek AI inference/{model,engram,kernel}.py; provenance in THIRD_PARTY_NOTICES.md. +const std = @import("std"); +const tf = @import("tensorfold"); +const weights = tf.deepseek_v41_weights; +const config = tf.deepseek_v41_config; +const cp = tf.checkpoint_metal; +const exl3 = tf.exl3_projection; +const projection = @import("projection.zig"); +const math = @import("math.zig"); +const attention = @import("attention_math.zig"); +const engram = @import("engram_math.zig"); +const Stage = tf.deepseek_v41_profiler.Stage; + +/// Local routing isolation also uses these source equations; no native values enter qualification. +pub const MoeRouting = struct { + logits: [384]f32 = undefined, + scores: [384]f32 = undefined, + ids: [6]u32 = undefined, + weights: [6]f32 = undefined, +}; + +pub fn moeSelect(moe_weights: weights.Moe, c: config.Text, logits: *const [384]f32, scores: *[384]f32, ids: *[6]u32, routing_weights: *[6]f32) void { + @setFloatMode(.strict); + for (logits, scores) |z, *score| { + const softplus = if (z > 20) z else std.math.log1p(@exp(z)); + score.* = @sqrt(softplus); + } + var used: [384]bool = @splat(false); + for (0..6) |pick| { + var best: usize = 0; + var high: f32 = -std.math.inf(f32); + for (0..384) |expert| { + const biased = scores[expert] + math.value(moe_weights.bias, expert); + if (!used[expert] and biased > high) { + best = expert; + high = biased; + } + } + ids[pick] = @intCast(best); + routing_weights[pick] = scores[best]; + used[best] = true; + } + var normalization: f64 = 0; + for (routing_weights) |v| normalization += v; + const denominator = @as(f32, @floatCast(normalization)) + 1e-20; + for (routing_weights) |*v| v.* = (v.* / denominator) * c.routed_scaling_factor; +} + +pub fn moeActivation(gate: []const u16, up: []const u16, route: f32, c: config.Text, activation: []u16) void { + @setFloatMode(.strict); + for (activation, gate, up) |*v, g_value, u_value| { + const g = @min(math.f(g_value), c.swiglu_limit); + const u = @max(-c.swiglu_limit, @min(math.f(u_value), c.swiglu_limit)); + v.* = math.b(route * ((g / (1 + @exp(-g))) * u)); + } +} + +pub const Reference = struct { + gpa: std.mem.Allocator, + c: config.Text, + w: *const weights.Model, + checkpoint: *const cp.Checkpoint, + capacity: u32, + position: u32 = 0, + projection_scratch: projection.Scratch, + attention_state: attention.State, + engram_state: engram.State, + logits: []f32, + h: [20480]u16 = undefined, + input: [5120]u16 = undefined, + output: [5120]u16 = undefined, + gate: [2304]u16 = undefined, + up: [2304]u16 = undefined, + activation: [2304]u16 = undefined, + expert_output: [5120]u16 = undefined, + moe_sum: [5120]f32 = undefined, + route_logits: [384]f32 = undefined, + route_scores: [384]f32 = undefined, + last_routes: [6]u32 = undefined, + last_routing: [6]f32 = undefined, + observer: ?Observer = null, + pub const Observer = struct { + ctx: *anyopaque, + read: *const fn (*anyopaque, usize, Stage, *const Reference) void, + }; + + /// Borrowed model/checkpoint/config/map must outlive this instance; capacity is the teacher sequence length. + pub fn init(gpa: std.mem.Allocator, c: config.Text, w: *const weights.Model, checkpoint: *const cp.Checkpoint, token_map: []const u8, capacity: u32) !Reference { + if (capacity == 0 or capacity > c.max_position_embeddings or capacity > attention.max_capacity) + return error.InvalidReferenceCapacity; + if (c.hidden_size != 5120 or c.num_hidden_layers != 40 or c.vocab_size != 129280 or + c.num_experts_per_tok != 6 or c.n_routed_experts != 384 or c.hc_mult != 4) + return error.UnsupportedReferenceArchitecture; + const engram_state = try engram.State.init(token_map); + var projection_scratch = try projection.Scratch.init(gpa, 8192, 129280); + errdefer projection_scratch.deinit(); + var attention_state = try attention.State.init(gpa, capacity); + errdefer attention_state.deinit(); + const logits = try gpa.alloc(f32, 129280); + return .{ .gpa = gpa, .c = c, .w = w, .checkpoint = checkpoint, .capacity = capacity, .projection_scratch = projection_scratch, .attention_state = attention_state, .engram_state = engram_state, .logits = logits }; + } + pub fn deinit(self: *Reference) void { + self.projection_scratch.deinit(); + self.attention_state.deinit(); + self.gpa.free(self.logits); + self.* = undefined; + } + /// Consume exactly one teacher token and return all actual checkpoint logits, borrowed until the next step. + pub fn step(self: *Reference, id: u32) ![]const f32 { + if (self.position >= self.capacity) return error.ContextFull; + const compressed = try self.engram_state.compressed(id); + for (0..5120) |d| { + const v = math.b(math.value(self.w.embed, @as(usize, id) * 5120 + d)); + for (0..4) |copy| self.h[copy * 5120 + d] = v; + } + var pre = [4]f32{ 1, 0, 0, 0 }; + for (&self.w.layers, 0..) |*layer, layer_id| { + if (layer.engram) |e| try self.engram_state.apply(e, if (layer_id == 1) 0 else 1, compressed, &self.h, self.c, &self.projection_scratch); + const attn_mix = math.mixes(&self.h, layer.hc_attn, self.c); + math.collapse(&self.h, pre, &self.input); + math.rms(&self.input, layer.attn_norm, self.c.rms_norm_eps, &self.input); + try self.attention_state.apply(layer.attn, layer_id, &self.input, self.position, self.c, &self.output, &self.projection_scratch); + math.expand(&self.h, &self.output, attn_mix); + if (self.observer) |observer| observer.read(observer.ctx, layer_id, .attention, self); + const ffn_mix = math.mixes(&self.h, layer.hc_ffn, self.c); + math.collapse(&self.h, attn_mix.pre, &self.input); + math.rms(&self.input, layer.ffn_norm, self.c.rms_norm_eps, &self.input); + try self.moe(layer.moe, layer_id); + math.expand(&self.h, &self.output, ffn_mix); + if (self.observer) |observer| observer.read(observer.ctx, layer_id, .moe, self); + pre = ffn_mix.pre; + } + math.collapse(&self.h, pre, &self.input); + math.rms(&self.input, self.w.norm, self.c.rms_norm_eps, &self.input); + try projection.applyF32(self.w.head, &self.input, self.logits, &self.projection_scratch); + for (self.logits) |v| if (!std.math.isFinite(v)) return error.NonFiniteReferenceLogits; + self.engram_state.advance(compressed); + self.position += 1; + return self.logits; + } + fn moe(self: *Reference, moe_weights: weights.Moe, layer: usize) !void { + @setFloatMode(.strict); + math.dense(moe_weights.gate, &self.input, &self.route_logits); + moeSelect(moe_weights, self.c, &self.route_logits, &self.route_scores, &self.last_routes, &self.last_routing); + @memset(&self.moe_sum, 0); + // Official MoE visits expert ids in ascending order, accumulates FP32, then adds shared expert last. + for (0..384) |expert| { + var route: ?f32 = null; + for (self.last_routes, self.last_routing) |selected, weight| + if (selected == expert) { + route = weight; + }; + const expert_route = route orelse continue; + var matrices: [3]exl3.Matrix = undefined; + var name: [128]u8 = undefined; + inline for (.{ "w1", "w3", "w2" }, 0..) |part, i| { + const prefix = try std.fmt.bufPrint(&name, "layers.{d}.ffn.experts.{d}.{s}", .{ layer, expert, part }); + matrices[i] = try weights.lookup(self.checkpoint, prefix, if (i == 2) 2304 else 5120, if (i == 2) 5120 else 2304); + } + try self.applyExpert(matrices, expert_route); + for (&self.moe_sum, self.expert_output) |*sum, v| sum.* += math.f(v); + } + try self.applyExpert(moe_weights.shared, 1); + for (&self.output, self.moe_sum, self.expert_output) |*v, routed, shared| + v.* = math.b(routed + math.f(shared)); + } + fn applyExpert(self: *Reference, matrices: [3]exl3.Matrix, route: f32) !void { + @setFloatMode(.strict); + try projection.apply(matrices[0], &self.input, &self.gate, &self.projection_scratch); + try projection.apply(matrices[1], &self.input, &self.up, &self.projection_scratch); + moeActivation(&self.gate, &self.up, route, self.c, &self.activation); + try projection.apply(matrices[2], &self.activation, &self.expert_output, &self.projection_scratch); + } +}; diff --git a/zig/tests/deepseek_v41_reference/rope_frequencies.json b/zig/tests/deepseek_v41_reference/rope_frequencies.json new file mode 100644 index 0000000000..2ffe1d1556 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/rope_frequencies.json @@ -0,0 +1,14 @@ +{ + "reference": "deepseek-ai/DeepSeek-V4.1-Flash", + "revision": "2cba9e42aa026125f3ed06c6d98c1db82f7ca027", + "source_sha256": "4e9ae23620edc8028ccc5d5fef552ab7fdc7dcd6f79608754fe9f67644056f65", + "torch_version": "2.14.1", + "extraction": "precompute_freqs_cis FP32 frequency vector before outer/polar; CPU", + "factor": 16, + "beta_fast": 32, + "beta_slow": 1, + "vectors": [ + {"theta": 10000, "original": 0, "bits": [1065353216,1061157137,1058010522,1054337179,1050798235,1047712803,1043732615,1040747895,1036831949,1033475085,1030116802,1026341397,1023510243,1019364175,1016180025,1012562930,1008981770,1005959600,1001931932,998911607,995049016,991652109,988353832,984533390,981668463,977573004,974350870,970790693,967166816,964208661,960132947,957076594]}, + {"theta": 160000, "original": 65536, "bits": [1065353216,1060112954,1056054303,1051098369,1046804782,1042117749,1037601972,1033169251,1028443341,1024251131,1019326489,1015361742,1010249151,1006366088,1001209182,997081655,991197921,985076960,979287478,973418702,966505391,959469944,952012877,943627103,933484300,917783504,913828774,908753896,904558591,899759068,895336240,890797133]} + ] +} diff --git a/zig/tests/deepseek_v41_reference/trace.zig b/zig/tests/deepseek_v41_reference/trace.zig new file mode 100644 index 0000000000..79cdef0fa4 --- /dev/null +++ b/zig/tests/deepseek_v41_reference/trace.zig @@ -0,0 +1,123 @@ +//! Bounded investigation snapshots; file I/O and independent CPU work are not performance measurements. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const Model = tf.deepseek_v41.Model; +const Reference = @import("reference.zig").Reference; +const Stage = tf.deepseek_v41_profiler.Stage; +const widths = [_]usize{ 20480 * 2, 5120 * 2, 5120 * 2, 6 * 4, 6 * 4 }; +const stride = 61488; +const bytes = stride * 80; + +const Capture = struct { + buffer: mtl.Buffer, + pipeline: mtl.Pipeline, + cpu: []u8, + + fn gpu(ctx: *anyopaque, enc: mtl.ComputeEncoder, layer: u32, stage: Stage, model: *const Model) void { + const self: *Capture = @ptrCast(@alignCast(ctx)); + var destination: usize = (@as(usize, layer) * 2 + @intFromBool(stage == .moe)) * stride; + const ws = &model.workspace; + inline for (.{ ws.view(.stream), ws.view(.normalized), ws.view(.branch), ws.view(.route_ids), ws.view(.route_weights) }, widths, 0..) |source, count, field| { + if (stage == .moe or field < 3) { + enc.setPipeline(self.pipeline); + enc.setBuffer(source.buffer, source.offset, 0); + enc.setBuffer(self.buffer, destination, 1); + enc.setValue(@as(u32, @intCast(count / 4)), 2); + enc.dispatchThreads(.{ .width = count / 4 }, .{ .width = 256 }); + destination += count; + } + } + enc.barrier(); + } + fn reference(ctx: *anyopaque, layer: usize, stage: Stage, ref: *const Reference) void { + const self: *Capture = @ptrCast(@alignCast(ctx)); + const start = (layer * 2 + @intFromBool(stage == .moe)) * stride; + var destination: []u8 = self.cpu[start..][0..stride]; + inline for (.{ std.mem.sliceAsBytes(&ref.h), std.mem.sliceAsBytes(&ref.input), std.mem.sliceAsBytes(&ref.output), std.mem.sliceAsBytes(&ref.last_routes), std.mem.sliceAsBytes(&ref.last_routing) }, 0..) |source, field| { + if (stage == .moe or field < 3) { + @memcpy(destination[0..source.len], source); + destination = destination[source.len..]; + } + } + } +}; +fn normalizedError(a: []const u8, b: []const u8) !f64 { + var squared: f64 = 0; + var denominator: f64 = 0; + var i: usize = 0; + while (i < a.len) : (i += 2) { + const x: f32 = @bitCast(@as(u32, std.mem.readInt(u16, a[i..][0..2], .little)) << 16); + const y: f32 = @bitCast(@as(u32, std.mem.readInt(u16, b[i..][0..2], .little)) << 16); + if (!std.math.isFinite(x) or !std.math.isFinite(y)) return error.NonfiniteTraceValues; + squared += (@as(f64, x) - y) * (@as(f64, x) - y); + denominator += @as(f64, y) * y; + } + return if (denominator == 0) @sqrt(squared) else @sqrt(squared / denominator); +} + +pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, ids: []const u32, token_map: []const u8, dir: []const u8) !void { + if (ids.len == 0 or ids.len > 16) return error.InvalidOracleWindow; + const lib = try mtl.Library.fromSource(model.device, "#include \nusing namespace metal;\nkernel void trace_copy(const device uint* x [[buffer(0)]], device uint* y [[buffer(1)]], constant uint& n [[buffer(2)]], uint i [[thread_position_in_grid]]) { if(i 0.01 or !routes_match) std.debug.print("trace token={d} layer={d} stage={s} stream={d} input={d} branch={d} routes_match={}\n", .{ token, point / 2, if (point % 2 == 0) "attention" else "moe", stream_error, input_error, branch_error, routes_match }); + } + inline for (.{ "gpu", "cpu" }, .{ gpu, cpu }) |name, data| { + const path = try std.fmt.allocPrint(gpa, "{s}/trace-{d}-{s}.bin", .{ dir, token, name }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); + } + std.debug.print("trace captured token {d}/{d}, same_logits=true\n", .{ token + 1, ids.len }); + } + try report.appendSlice(gpa, "\n]\n"); + const path = try std.fs.path.join(gpa, &.{ dir, "layer-trace.json" }); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = report.items }); +} From ac719e365bd89035380cf91910a6e172071748e2 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:13:58 -0700 Subject: [PATCH 03/15] THIRD_PARTY_NOTICES: the native DS4.1 reference follows DeepSeek-V4.1-Flash's MIT inference --- THIRD_PARTY_NOTICES.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 018947342d..5b34839762 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -217,6 +217,13 @@ TensorFold ships no model weights. The `z-lab/Qwen3.8-27B-DFlash2` model card st The optional `incoai/GLM-5.3-Flash-DFlash2` model card states CC BY-NC-ND 4.0, for non-commercial use without derivatives; only the Python engine reads that drafter. Each checkpoint keeps its own license. +The native `zig/src/families/deepseek_v41/` implementation and `zig/tests/deepseek_v41_reference/` +follow [DeepSeek-V4.1-Flash's official inference](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/tree/2cba9e42aa026125f3ed06c6d98c1db82f7ca027/inference), +MIT License, Copyright (c) 2023 DeepSeek: the reference's equations, the Engram normalization, +PCG64 multipliers and prime bucket layout, and the pinned chat template come from that revision. +The EXL3 checkpoint reads through the ExLlamaV3 format covered above; no captured CUDA/Triton +kernel assets, full model shards or binary traces are bundled. + ## Structured output: libtfgrammar `lib/libtfgrammar` in the release archives is `zig/src/grammar/tf_grammar.cc` compiled with the C++ sources of From 2b1d49564a41a0145e1a553307239c5835d24728 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Tue, 13 Oct 2026 07:41:05 -0700 Subject: [PATCH 04/15] DeepSeek V4.1 load: a shard contributes only the tensors the index maps to it Rank-split packs (jayleaton's 2.8bpw) duplicate whole-load tensors in native.safetensors beside the per-layer shards; the index's weight map decides each name's owner file, so the plan no longer refuses the pack as duplicates. --- .../families/deepseek_v41/resident_load.zig | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/zig/src/families/deepseek_v41/resident_load.zig b/zig/src/families/deepseek_v41/resident_load.zig index 422af3555c..a84bf955b9 100644 --- a/zig/src/families/deepseek_v41/resident_load.zig +++ b/zig/src/families/deepseek_v41/resident_load.zig @@ -44,7 +44,10 @@ pub const Plan = struct { var n: usize = 0; while (file_it.next()) |file| : (n += 1) ordered[n] = file.*; std.mem.sort([]const u8, ordered, {}, less); - for (ordered) |file| try self.add(try std.fs.path.join(a, &.{ dir, file }), null); + var owned: std.StringHashMapUnmanaged([]const u8) = .empty; + entries.reset(); + while (entries.next()) |item| try owned.put(a, item.key_ptr.*, item.value_ptr.string); + for (ordered) |file| try self.add(try std.fs.path.join(a, &.{ dir, file }), null, &owned); if (self.tensors.count() != index.weight_map.object.count()) return error.CheckpointIndexMismatch; entries = index.weight_map.object.iterator(); while (entries.next()) |item| { @@ -57,7 +60,7 @@ pub const Plan = struct { .{ .layer = 14, .file = "model-00048-of-00048.safetensors", .rows = 384016682 }, }) |table| { const prefix = try std.fmt.allocPrint(a, "layers.{d}.engram.embed.", .{table.layer}); - try self.add(try std.fs.path.join(a, &.{ table_dir, table.file }), prefix); + try self.add(try std.fs.path.join(a, &.{ table_dir, table.file }), prefix, null); const weight = try std.fmt.allocPrint(a, "{s}weight", .{prefix}); const scale = try std.fmt.allocPrint(a, "{s}scale", .{prefix}); try self.require(weight, .f8_e4m3, &.{ table.rows, 256 }); @@ -88,7 +91,7 @@ pub const Plan = struct { } } - fn add(self: *Plan, path: []const u8, prefix: ?[]const u8) !void { + fn add(self: *Plan, path: []const u8, prefix: ?[]const u8, owned: ?*const std.StringHashMapUnmanaged([]const u8)) !void { const a = self.arena.allocator(); const own_path = try a.dupeSentinel(u8, path, 0); const fd = std.c.open(own_path, .{ .ACCMODE = .RDONLY }); @@ -109,10 +112,18 @@ pub const Plan = struct { var begin: usize = if (prefix == null) 0 else file_bytes - data_start; var end: usize = if (prefix == null) file_bytes - data_start else 0; for (names.keys(), names.values()) |name, entry| { + // A rank-split pack duplicates whole-load tensors beside the shards; the index decides the owner file. + if (owned) |map| { + const owner = map.get(name) orelse continue; + if (!std.mem.eql(u8, owner, std.fs.path.basename(path))) continue; + } begin = @min(begin, entry.begin); end = @max(end, entry.end); const slot = try self.tensors.getOrPut(a, name); - if (slot.found_existing) return error.DuplicateTensor; + if (slot.found_existing) { + std.log.err("duplicate tensor {s} in {s}", .{ name, path }); + return error.DuplicateTensor; + } slot.value_ptr.* = .{ .shard = self.shards.items.len, .entry = entry }; } if (prefix != null) { From 4d6dfa2a8e9a841f8bee42e3a4ef44c2096d337e Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Tue, 13 Oct 2026 07:47:38 -0700 Subject: [PATCH 05/15] DeepSeek V4.1 load: the router gate and bias ride the pack's dtypes The released checkpoints store the gate f16 and its bias f16; rank-split repacks store the weight bf16 and the bias f32. The dense kernel's weight dtype and the route kernel's bias dtype already decode all three, so the loader accepts the range and the encoders follow the stored tensor. --- zig/src/families/deepseek_v41/metal_model.zig | 12 ++++++-- zig/src/families/deepseek_v41/weights.zig | 28 +++++++++++++++++-- 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/zig/src/families/deepseek_v41/metal_model.zig b/zig/src/families/deepseek_v41/metal_model.zig index 0178b0660f..6889ed13f5 100644 --- a/zig/src/families/deepseek_v41/metal_model.zig +++ b/zig/src/families/deepseek_v41/metal_model.zig @@ -312,9 +312,17 @@ pub const Model = struct { fn moe(m: *Model, enc: mtl.ComputeEncoder, w: *const weights.Moe, rows: u32) !void { const ws = &m.workspace; const p = &m.pipelines; - ops.encodeDense(enc, &p.ops, ws.view(.normalized), tensor(w.gate), ws.view(.router_logits), .{ .rows = rows, .k = 5120, .n = 384, .x_stride = 5120, .w_stride = 5120, .y_stride = 384, .w_dtype = .f16, .y_dtype = .f32 }); + ops.encodeDense(enc, &p.ops, ws.view(.normalized), tensor(w.gate), ws.view(.router_logits), .{ .rows = rows, .k = 5120, .n = 384, .x_stride = 5120, .w_stride = 5120, .y_stride = 384, .w_dtype = switch (w.gate.dtype) { + .bf16 => .bf16, + else => .f16, + }, .y_dtype = .f32 }); enc.barrier(); - route.encode(enc, &p.route, ws.view(.router_logits), tensor(w.bias), ws.view(.route_ids), ws.view(.route_weights), .{ .rows = rows, .bias_dtype = .f16 }); + route.encode(enc, &p.route, ws.view(.router_logits), tensor(w.bias), ws.view(.route_ids), ws.view(.route_weights), .{ .rows = rows, .bias_dtype = switch (w.bias.dtype) { + .f32 => .f32, + .f16 => .f16, + .bf16 => .bf16, + else => .f16, + } }); enc.barrier(); const grouped: ?expert_groups.Scratch = if (m.grouped_prefill and rows >= 8) try ws.grouping(rows) else null; var routed_rows = exl3.Rows{ .count = @as(usize, rows) * 6, .x_stride = 5120, .y_stride = 2304, .input_rows = rows, .output_rows = @as(usize, rows) * 6, .routes = .{ .ids = convert(ws.view(.route_ids)), .repeat = 6 } }; diff --git a/zig/src/families/deepseek_v41/weights.zig b/zig/src/families/deepseek_v41/weights.zig index 1377466c9c..4c5603d45b 100644 --- a/zig/src/families/deepseek_v41/weights.zig +++ b/zig/src/families/deepseek_v41/weights.zig @@ -250,8 +250,32 @@ fn populate(comptime binding: bool, source: anytype, model: ?*Model) !void { } } const moe: ?*Moe = if (binding) &out.?.moe else null; - try plain(binding, source, try named(&name, "layers.{d}.ffn.gate.weight", .{layer}), .f16, &.{ 384, 5120 }, if (binding) &moe.?.gate else null); - try plain(binding, source, try named(&name, "layers.{d}.ffn.gate.bias", .{layer}), .f16, &.{384}, if (binding) &moe.?.bias else null); + // The released checkpoints store the router gate f16; rank-split repacks store it bf16. + if (binding) { + const tensor = try source.get(try named(&name, "layers.{d}.ffn.gate.weight", .{layer})); + if ((tensor.dtype != .f16 and tensor.dtype != .bf16) or tensor.rank != 2 or + tensor.shape[0] != 384 or tensor.shape[1] != 5120) + { + std.log.err("unsupported router gate {s}: {s} {any}", .{ name, @tagName(tensor.dtype), tensor.shape[0..tensor.rank] }); + return error.UnsupportedTensor; + } + moe.?.gate = tensor; + } else { + const tensor = try source.get(try named(&name, "layers.{d}.ffn.gate.weight", .{layer})); + if ((tensor.dtype != .f16 and tensor.dtype != .bf16) or tensor.rank != 2 or + tensor.shape[0] != 384 or tensor.shape[1] != 5120) return error.UnsupportedTensor; + } + // The bias rides the same three dtypes the route kernel decodes (f32, f16, bf16). + if (binding) { + const tensor = try source.get(try named(&name, "layers.{d}.ffn.gate.bias", .{layer})); + if ((tensor.dtype != .f16 and tensor.dtype != .bf16 and tensor.dtype != .f32) or tensor.rank != 1 or tensor.shape[0] != 384) + return error.UnsupportedTensor; + moe.?.bias = tensor; + } else { + const tensor = try source.get(try named(&name, "layers.{d}.ffn.gate.bias", .{layer})); + if ((tensor.dtype != .f16 and tensor.dtype != .bf16 and tensor.dtype != .f32) or tensor.rank != 1 or tensor.shape[0] != 384) + return error.UnsupportedTensor; + } inline for (.{ "w1", "w3", "w2" }, 0..) |part, projection| { const k: usize = if (projection == 2) 2304 else 5120; const n: usize = if (projection == 2) 5120 else 2304; From eebd6b5dfaa6950759d4317c4ff5a5b446327f51 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Tue, 13 Oct 2026 07:51:07 -0700 Subject: [PATCH 06/15] DeepSeek V4.1 load: the Engram gate weights ride f32 or bf16 The gate kernel decodes its q/k weights by dtype, the Gate descriptor carries it, and the loader accepts the rank-split packs' bf16 projections beside the released f32 ones. --- zig/src/families/deepseek_v41/engram.zig | 6 ++++- zig/src/families/deepseek_v41/metal_model.zig | 5 +++- zig/src/families/deepseek_v41/weights.zig | 23 +++++++++++++++++-- 3 files changed, 30 insertions(+), 4 deletions(-) diff --git a/zig/src/families/deepseek_v41/engram.zig b/zig/src/families/deepseek_v41/engram.zig index 0f5a0287c2..867cc78d39 100644 --- a/zig/src/families/deepseek_v41/engram.zig +++ b/zig/src/families/deepseek_v41/engram.zig @@ -3,6 +3,7 @@ const std = @import("std"); const mtl = @import("metal"); const sources = @import("kernel_sources"); pub const View = @import("view.zig").View; +const Dtype = @import("ops.zig").Dtype; pub const hash = @import("engram_hash.zig"); pub const head_dim = 256; @@ -29,6 +30,8 @@ pub const Table = struct { pub const Gate = struct { q_weight: View, k_weight: View, + /// The released checkpoints store these f32; rank-split repacks store bf16. + weight_dtype: Dtype = .f32, dim: u32, copies: u32, hidden_stride: usize, @@ -97,6 +100,7 @@ const GateArgs = extern struct { dim: u32, copies: u32, masked: u32, + weight_dtype: u32, eps: f32, }; @@ -193,7 +197,7 @@ pub fn encodeGate(e: mtl.ComputeEncoder, p: *const Pipelines, gate: Gate, hidden bind(e, gate.k_weight, 3); bind(e, mask orelse hidden, 4); bind(e, out, 5); - e.setValue(GateArgs{ .hidden_stride = gate.hidden_stride, .kv_stride = gate.kv_stride, .output_stride = gate.output_stride, .dim = gate.dim, .copies = gate.copies, .masked = @intFromBool(mask != null), .eps = gate.eps }, 6); + e.setValue(GateArgs{ .hidden_stride = gate.hidden_stride, .kv_stride = gate.kv_stride, .output_stride = gate.output_stride, .dim = gate.dim, .copies = gate.copies, .masked = @intFromBool(mask != null), .weight_dtype = @backingInt(gate.weight_dtype), .eps = gate.eps }, 6); e.dispatchGroups(.{ .width = tokens, .height = gate.copies }, .{ .width = 256 }); e.barrier(); } diff --git a/zig/src/families/deepseek_v41/metal_model.zig b/zig/src/families/deepseek_v41/metal_model.zig index 6889ed13f5..a43485d79f 100644 --- a/zig/src/families/deepseek_v41/metal_model.zig +++ b/zig/src/families/deepseek_v41/metal_model.zig @@ -227,7 +227,10 @@ pub const Model = struct { try m.profileSpan(&enc, .engram); try engram.encodeGather(enc, &p.engram, .{ .weight = tensor(g.table), .scale = tensor(g.scales), .rows = @intCast(g.table.shape[0]) }, ws.view(if (index == 1) .engram_ids_0 else .engram_ids_1), ws.view(.engram_gather), rows, 6144); try m.project(enc, g.wkv, ws.view(.engram_gather), ws.view(.engram_kv), rows, 6144, 25600, .bf16); - try engram.encodeGate(enc, &p.engram, .{ .q_weight = tensor(g.q_weight), .k_weight = tensor(g.k_weight), .dim = 5120, .copies = 4, .hidden_stride = 20480, .kv_stride = 25600, .output_stride = 20480, .eps = m.cfg.text_config.rms_norm_eps }, ws.view(.stream), ws.view(.engram_kv), null, ws.view(.stream), rows); + try engram.encodeGate(enc, &p.engram, .{ .q_weight = tensor(g.q_weight), .k_weight = tensor(g.k_weight), .dim = 5120, .copies = 4, .weight_dtype = switch (g.q_weight.dtype) { + .bf16 => .bf16, + else => .f32, + }, .hidden_stride = 20480, .kv_stride = 25600, .output_stride = 20480, .eps = m.cfg.text_config.rms_norm_eps }, ws.view(.stream), ws.view(.engram_kv), null, ws.view(.stream), rows); } try m.profileSpan(&enc, .mhc); m.mixes(enc, layer.hc_attn, ws.view(.attention_pre), rows); diff --git a/zig/src/families/deepseek_v41/weights.zig b/zig/src/families/deepseek_v41/weights.zig index 4c5603d45b..aef0314b34 100644 --- a/zig/src/families/deepseek_v41/weights.zig +++ b/zig/src/families/deepseek_v41/weights.zig @@ -106,6 +106,25 @@ fn named(buffer: []u8, comptime fmt: []const u8, args: anytype) ![]const u8 { return std.fmt.bufPrint(buffer, fmt, args); } +/// A plain tensor the encoders decode as f32, f16 or bf16; the shape must match exactly. +fn plainEither(comptime binding: bool, source: anytype, name: []const u8, shape: []const usize, destination: ?*Tensor) !void { + const tensor = try source.get(name); + if ((tensor.dtype != .f32 and tensor.dtype != .f16 and tensor.dtype != .bf16) or tensor.rank != shape.len or + !std.mem.eql(usize, tensor.shape[0..tensor.rank], shape)) + { + std.log.err("unsupported tensor {s}: {s} {any}", .{ name, @tagName(tensor.dtype), tensor.shape[0..tensor.rank] }); + return error.UnsupportedTensor; + } + if (binding) { + var count: usize = 1; + for (shape) |dim| count = try std.math.mul(usize, count, dim); + if (tensor.bytes != try std.math.mul(usize, count, tensor.dtype.size())) return error.UnsupportedTensor; + if (tensor.offset % tensor.dtype.size() != 0 or tensor.offset > tensor.buffer.length() or + tensor.bytes > tensor.buffer.length() - tensor.offset) return error.InvalidTensorView; + destination.?.* = tensor; + } else if (tensor.begin % 4 != 0) return error.InvalidTensorView; +} + fn plain(comptime binding: bool, source: anytype, name: []const u8, dtype: st.DType, shape: []const usize, destination: ?*Tensor) !void { const tensor = try source.get(name); if (tensor.dtype != dtype or tensor.rank != shape.len or @@ -289,8 +308,8 @@ fn populate(comptime binding: bool, source: anytype, model: ?*Model) !void { const rows: usize = if (layer == 1) 384006168 else 384016682; try plain(binding, source, try named(&name, "layers.{d}.engram.embed.weight", .{layer}), .f8_e4m3, &.{ rows, 256 }, if (binding) &e.?.table else null); try plain(binding, source, try named(&name, "layers.{d}.engram.embed.scale", .{layer}), .f8_e8m0, &.{ rows, 8 }, if (binding) &e.?.scales else null); - try plain(binding, source, try named(&name, "layers.{d}.engram.q_weight", .{layer}), .f32, &.{ 4, 5120 }, if (binding) &e.?.q_weight else null); - try plain(binding, source, try named(&name, "layers.{d}.engram.k_weight", .{layer}), .f32, &.{ 4, 5120 }, if (binding) &e.?.k_weight else null); + try plainEither(binding, source, try named(&name, "layers.{d}.engram.q_weight", .{layer}), &.{ 4, 5120 }, if (binding) &e.?.q_weight else null); + try plainEither(binding, source, try named(&name, "layers.{d}.engram.k_weight", .{layer}), &.{ 4, 5120 }, if (binding) &e.?.k_weight else null); try matrix(binding, source, try named(&prefix, "layers.{d}.engram.wkv", .{layer}), 6144, 25600, if (binding) &e.?.wkv else null); } } From d332595fb8f6f8f2cea4454e26e9c3ae4cf65460 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Tue, 13 Oct 2026 08:15:39 -0700 Subject: [PATCH 07/15] DeepSeek V4.1: the kept-dense EXL3 form, for rank-split repacks Jayleaton's 2.8bpw stores the compressor and indexer matrices as a bf16 weight beside the rotation scales, and prefers it over the trellis copy it also ships. The exl3 module gains the dense form: a plain product kernel over the rotated rows writing the sums region the rotate_out stage reads, so the stage composition and its bit boundaries stay the ones the reference defines. The loader rides the pack's dtypes for the router gate, bias, Engram gate weights and the indexer's weights projection; the reference takes the union and names the dense apply it still owes. The 2.8bpw plan passes: 406.6 GB, 44 shards. --- zig/kernels/metal/core/exl3_projection.metal | 15 +++++ zig/kernels/metal/deepseek_v41/engram.metal | 13 ++-- zig/src/core/exl3_projection.zig | 64 ++++++++++++++++++- zig/src/families/deepseek_v41/attention.zig | 29 ++++++--- zig/src/families/deepseek_v41/weights.zig | 55 ++++++++++++---- .../deepseek_v41_reference/attention_math.zig | 10 +-- .../deepseek_v41_reference/engram_math.zig | 2 +- .../operator_diagnostic.zig | 2 +- .../deepseek_v41_reference/operator_moe.zig | 8 +-- .../deepseek_v41_reference/operators.zig | 2 +- .../deepseek_v41_reference/projection.zig | 22 ++++++- .../deepseek_v41_reference/reference.zig | 8 +-- 12 files changed, 187 insertions(+), 43 deletions(-) diff --git a/zig/kernels/metal/core/exl3_projection.metal b/zig/kernels/metal/core/exl3_projection.metal index 0eca0d80c5..51777a82a8 100644 --- a/zig/kernels/metal/core/exl3_projection.metal +++ b/zig/kernels/metal/core/exl3_projection.metal @@ -239,6 +239,21 @@ kernel void tf_exl3_state_window_check(const device uint4* inputs [[buffer(0)]], const device ushort* tile = reinterpret_cast(inputs + tid); const uint off = inputs[tid].z; outputs[tid] = tf_exl3_mapped_state(tile, (off << 6) | (1u << 11)); + +// A kept-dense matrix's plain product: the rotated rows times the bf16 weight, accumulated +// in f32 in the sums region the rotate_out stage reads. Sequential k, one thread an output +// element — the fp64 reference mirrors this order. +kernel void tf_exl3_dense(const device half* rotated [[buffer(0)]], + const device bfloat* weight [[buffer(1)]], + device float* sums [[buffer(2)]], + constant TfExl3Args& a [[buffer(3)]], + const device uint* input_ids [[buffer(4)]], + uint2 g [[thread_position_in_grid]]) { + const uint row = a.input_ids ? input_ids[g.y] : g.y; + float acc = 0; + for (uint k = 0; k < a.k; ++k) + acc += float(rotated[ulong(g.y) * a.k + k]) * float(weight[ulong(g.x) * a.k + k]); + sums[ulong(g.y) * a.n + g.x] = acc; } // Four SIMDgroups share decoded weights while each retains its original eight-row compensated dots. diff --git a/zig/kernels/metal/deepseek_v41/engram.metal b/zig/kernels/metal/deepseek_v41/engram.metal index 532bbd01ee..b3d454cf59 100644 --- a/zig/kernels/metal/deepseek_v41/engram.metal +++ b/zig/kernels/metal/deepseek_v41/engram.metal @@ -1,6 +1,11 @@ // Math/layout attribution: DeepSeek AI DeepSeek-V4.1-Flash inference/{engram,model}.py (MIT). #include using namespace metal; +inline float rd(const device uchar* p, ulong i, uint dtype) { + if (dtype == 1u) return ((const device float*)p)[i]; + if (dtype == 2u) return float(((const device half*)p)[i]); + return as_type(uint(((const device ushort*)p)[i]) << 16); +} struct EngramLayer { uint layer_id, rows; @@ -96,15 +101,15 @@ kernel void tf_v41_engram_gather(const device uint* ids [[buffer(0)]], struct EngramGateArgs { ulong hidden_stride, kv_stride, output_stride; - uint dim, copies, masked; + uint dim, copies, masked, weight_dtype; float eps; }; // Each copy gets its own RMS dot gate; all copies share the value following the flattened keys. kernel void tf_v41_engram_gate(const device bfloat* hidden [[buffer(0)]], const device bfloat* kv [[buffer(1)]], - const device float* q_weight [[buffer(2)]], - const device float* k_weight [[buffer(3)]], + const device uchar* q_weight [[buffer(2)]], + const device uchar* k_weight [[buffer(3)]], const device uchar* mask [[buffer(4)]], device bfloat* out [[buffer(5)]], constant EngramGateArgs& args [[buffer(6)]], @@ -120,7 +125,7 @@ kernel void tf_v41_engram_gate(const device bfloat* hidden [[buffer(0)]], float3 acc = 0; for (uint d = t; d < args.dim; d += 256) { const float h = float(hidden[hbase + d]), k = float(kv[kbase + d]); - const float weight = q_weight[wbase + d] * k_weight[wbase + d]; + const float weight = rd(q_weight, wbase + d, args.weight_dtype) * rd(k_weight, wbase + d, args.weight_dtype); acc += float3(h * h, k * k, (h * weight) * k); } acc = float3(simd_sum(acc.x), simd_sum(acc.y), simd_sum(acc.z)); diff --git a/zig/src/core/exl3_projection.zig b/zig/src/core/exl3_projection.zig index 0c2074a793..f99e51053a 100644 --- a/zig/src/core/exl3_projection.zig +++ b/zig/src/core/exl3_projection.zig @@ -138,6 +138,7 @@ pub const Pipelines = struct { rotate_in: mtl.Pipeline, matmul: mtl.Pipeline, rotate_out: mtl.Pipeline, + dense: mtl.Pipeline, prefill: ?mtl.Pipeline = null, pub fn load(device: mtl.Device) !Pipelines { @@ -151,6 +152,8 @@ pub const Pipelines = struct { errdefer mm.deinit(); const ro = try mtl.Pipeline.init(device, lib, "tf_exl3_rotate_out", false); errdefer ro.deinit(); + const dn = try mtl.Pipeline.init(device, lib, "tf_exl3_dense", false); + errdefer dn.deinit(); for ([_]mtl.Pipeline{ ri, mm, ro }) |p| if (p.simdWidth() != 32 or p.maxThreads() < 32) return error.UnsupportedExl3Simd; var prefill = mtl.Pipeline.init(device, lib, "tf_exl3_matmul_prefill", false) catch null; @@ -160,12 +163,13 @@ pub const Pipelines = struct { prefill = null; } } - return .{ .rotate_in = ri, .matmul = mm, .rotate_out = ro, .prefill = prefill }; + return .{ .rotate_in = ri, .matmul = mm, .rotate_out = ro, .dense = dn, .prefill = prefill }; } pub fn deinit(p: *Pipelines) void { p.rotate_in.deinit(); p.matmul.deinit(); p.rotate_out.deinit(); + p.dense.deinit(); if (p.prefill) |pf| pf.deinit(); } }; @@ -288,6 +292,64 @@ pub fn encode(p: *const Pipelines, e: mtl.ComputeEncoder, m: Matrix, x: View, y: stages(p, e, a, m.trellis, m.suh, m.svh, m.bias orelse m.svh, scratch, x, y, scratch, r); } +/// A kept-dense EXL3 matrix: a bf16 weight beside the input/output rotation scales, no trellis. +pub const DenseMatrix = struct { + weight: View, + suh: View, + svh: View, + bias: ?View = null, + k: usize, + n: usize, +}; + +/// The kept-dense path: the same rotate stages as the quantized encode, with the trellis +/// product replaced by a plain dense product over the rotated rows, accumulated in f32 in +/// the scratch sums region the rotate_out stage reads. Sequential k; the fp64 reference +/// mirrors this order. +pub fn encodeDense(p: *const Pipelines, e: mtl.ComputeEncoder, m: DenseMatrix, x: View, y: View, scratch: View, r: Rows) !void { + try shape(m.k, m.n); + var a = try rowsArgs(m.k, m.n, r, x, y, scratch); + if (r.grouped) return error.InvalidExl3GroupedRows; + a.su_signs = 0; + a.sv_signs = 0; + a.has_bias = @intFromBool(m.bias != null); + const sums = View{ .buffer = scratch.buffer, .offset = scratch.offset + try rotatedBytes(m.k, r.count) }; + const sums_bytes = try std.math.mul(usize, try std.math.mul(usize, m.n, r.count), 4); + try readonly(m.weight, try std.math.mul(usize, try std.math.mul(usize, m.n, m.k), 2), y, 0, scratch, sums_bytes); + try readonly(m.suh, m.k * 2, y, 0, scratch, sums_bytes); + try readonly(m.svh, m.n * 2, y, 0, scratch, sums_bytes); + if (m.bias) |b| try readonly(b, m.n * 2, y, 0, scratch, sums_bytes); + e.setPipeline(p.rotate_in); + bind(e, x, 0); + bind(e, m.suh, 1); + bind(e, scratch, 2); + e.setValue(a, 3); + bind(e, scratch, 4); + bind(e, scratch, 5); + bind(e, scratch, 6); + e.dispatchGroups(.{ .width = a.k / 128, .height = a.rows }, .{ .width = 32 }); + e.barrier(); + e.setPipeline(p.dense); + bind(e, scratch, 0); + bind(e, m.weight, 1); + bind(e, sums, 2); + e.setValue(a, 3); + bind(e, scratch, 4); + e.dispatchThreads(.{ .width = a.n, .height = a.rows }, .{ .width = 64 }); + e.barrier(); + e.setPipeline(p.rotate_out); + bind(e, sums, 0); + bind(e, m.svh, 1); + bind(e, m.bias orelse m.svh, 2); + bind(e, y, 3); + e.setValue(a, 4); + bind(e, scratch, 5); + bind(e, scratch, 6); + bind(e, scratch, 7); + e.dispatchGroups(.{ .width = a.n / 128, .height = a.rows }, .{ .width = 32 }); + e.barrier(); +} + /// GPU routes choose caller-validated resident descriptors; grouped padding uses maxInt(u32). pub fn encodeIndexed(p: *const Pipelines, e: mtl.ComputeEncoder, experts: Experts, x: View, y: View, scratch: View, r: Rows) !void { if (r.routes == null or experts.count == 0 or experts.count > std.math.maxInt(u32)) return error.InvalidExl3Experts; diff --git a/zig/src/families/deepseek_v41/attention.zig b/zig/src/families/deepseek_v41/attention.zig index 6d7cfe5da7..a1f76f6886 100644 --- a/zig/src/families/deepseek_v41/attention.zig +++ b/zig/src/families/deepseek_v41/attention.zig @@ -120,8 +120,18 @@ fn tensor(t: checkpoint.Tensor) View { fn core(v: View) exl3.View { return .{ .buffer = v.buffer, .offset = v.offset }; } -fn projection(p: *const exl3.Pipelines, e: mtl.ComputeEncoder, m: exl3.Matrix, x: View, y: View, scratch: exl3.View, rows: u32, xs: u32, ys: u32, output: exl3.Output) !void { - try exl3.encode(p, e, m, core(x), core(y), scratch, .{ .count = rows, .x_stride = xs, .y_stride = ys, .output = output }); +fn projection(p: *const exl3.Pipelines, e: mtl.ComputeEncoder, m: weights.AnyMatrix, x: View, y: View, scratch: exl3.View, rows: u32, xs: u32, ys: u32, output: exl3.Output) !void { + switch (m) { + .trellis => |t| try exl3.encode(p, e, t, core(x), core(y), scratch, .{ .count = rows, .x_stride = xs, .y_stride = ys, .output = output }), + // A kept-dense repack: the same rotate stages, the middle a plain dense product. + .dense => |d| try exl3.encodeDense(p, e, .{ + .weight = core(tensor(d.weight)), + .suh = core(tensor(d.suh)), + .svh = core(tensor(d.svh)), + .k = d.k, + .n = d.n, + }, core(x), core(y), scratch, .{ .count = rows, .x_stride = xs, .y_stride = ys, .output = output }), + } } fn norm(p: *const Pipelines, e: mtl.ComputeEncoder, x: View, weight: checkpoint.Tensor, rows: u32, width: u32, eps: f32) void { ops.encodeRms(e, &p.plain, x, tensor(weight), x, .{ .rows = rows, .width = width, .x_stride = width, .y_stride = width, .eps = eps, .weight_f32 = weight.dtype == .f32 }); @@ -161,15 +171,15 @@ pub fn encodeDiagnostic(p: *const Pipelines, e: mtl.ComputeEncoder, w: *const we const ring = try state.ring(layer); const chunk = work.selection.windowChunk(); e.barrier(); - try projection(ep, e, w.wq_a, x, qr, scratch, rows, c.hidden_size, c.q_lora_rank, .bf16); + try projection(ep, e, .{ .trellis = w.wq_a }, x, qr, scratch, rows, c.hidden_size, c.q_lora_rank, .bf16); if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .qr_projection, qr); norm(p, e, qr, w.q_norm, rows, c.q_lora_rank, c.rms_norm_eps); if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .qr_norm, qr); - try projection(ep, e, w.wq_b, qr, q, scratch, rows, c.q_lora_rank, q_width, .bf16); + try projection(ep, e, .{ .trellis = w.wq_b }, qr, q, scratch, rows, c.q_lora_rank, q_width, .bf16); if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .q_projection, q); try rotate(p, e, q, rows, c.num_attention_heads, c.head_dim, position, ratio, 1, false, c); if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .q, q); - try projection(ep, e, w.wkv, x, raw, scratch, rows, c.hidden_size, c.head_dim, .bf16); + try projection(ep, e, .{ .trellis = w.wkv }, x, raw, scratch, rows, c.hidden_size, c.head_dim, .bf16); if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_projection, raw); norm(p, e, raw, w.kv_norm, rows, c.head_dim, c.rms_norm_eps); if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .kv_norm, raw); @@ -217,7 +227,10 @@ pub fn encodeDiagnostic(p: *const Pipelines, e: mtl.ComputeEncoder, w: *const we try projection(ep, e, selector.wq_b, qr, iq, scratch, rows, c.q_lora_rank, iq_width, .bf16); try rotate(p, e, iq, rows, c.index_n_heads, c.index_head_dim, position, ratio, 1, false, c); try p.cache.encodePack(e, iq, work.selection.queryPacked(), .{ .rows = rows * c.index_n_heads, .width = c.index_head_dim, .stride = c.index_head_dim, .packed_stride = kv.cache.index_stride, .block = 32, .kind = .index_fp4, .inplace = 1 }); - ops.encodeDense(e, &p.plain, x, tensor(selector.weights_proj), iw, .{ .rows = rows, .k = c.hidden_size, .n = c.index_n_heads, .x_stride = c.hidden_size, .w_stride = c.hidden_size, .y_stride = c.index_n_heads, .w_dtype = .f16 }); + ops.encodeDense(e, &p.plain, x, tensor(selector.weights_proj), iw, .{ .rows = rows, .k = c.hidden_size, .n = c.index_n_heads, .x_stride = c.hidden_size, .w_stride = c.hidden_size, .y_stride = c.index_n_heads, .w_dtype = switch (selector.weights_proj.dtype) { + .bf16 => .bf16, + else => .f16, + } }); e.barrier(); } const make_candidates = layer == c.candidate_source_layer_id; @@ -250,9 +263,9 @@ pub fn encodeDiagnostic(p: *const Pipelines, e: mtl.ComputeEncoder, w: *const we const group_input = q_width / c.o_groups; const group_output = c.o_groups * c.o_lora_rank; for (w.wo_a, 0..) |matrix, group| { - try projection(ep, e, matrix, q.shifted(group * group_input * 2), grouped.shifted(group * c.o_lora_rank * 2), scratch, rows, q_width, group_output, .bf16); + try projection(ep, e, .{ .trellis = matrix }, q.shifted(group * group_input * 2), grouped.shifted(group * c.o_lora_rank * 2), scratch, rows, q_width, group_output, .bf16); } if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .grouped, grouped); - try projection(ep, e, w.wo_b, grouped, y, scratch, rows, group_output, c.hidden_size, .bf16); + try projection(ep, e, .{ .trellis = w.wo_b }, grouped, y, scratch, rows, group_output, c.hidden_size, .bf16); if (diagnostic) try DiagnosticObserver.observe(true, observer, e, .wo_b, y); } diff --git a/zig/src/families/deepseek_v41/weights.zig b/zig/src/families/deepseek_v41/weights.zig index aef0314b34..650c37bd1e 100644 --- a/zig/src/families/deepseek_v41/weights.zig +++ b/zig/src/families/deepseek_v41/weights.zig @@ -8,9 +8,13 @@ const load = @import("resident_load.zig"); pub const Tensor = cp.Tensor; pub const Matrix = exl3.Matrix; + +/// A kept-dense repack stores small matrices as a bf16 weight beside the rotation scales. +pub const Dense = struct { weight: Tensor, suh: Tensor, svh: Tensor, k: usize, n: usize }; +pub const AnyMatrix = union(enum) { trellis: Matrix, dense: Dense }; pub const Hc = struct { projection: Tensor, scale: Tensor, base: Tensor }; -pub const Compressor = struct { wkv: Matrix, wgate: ?Matrix, norm: Tensor }; -pub const Indexer = struct { wq_b: Matrix, wk: ?Matrix, weights_proj: Tensor, k_norm: ?Tensor }; +pub const Compressor = struct { wkv: AnyMatrix, wgate: ?AnyMatrix, norm: Tensor }; +pub const Indexer = struct { wq_b: AnyMatrix, wk: ?AnyMatrix, weights_proj: Tensor, k_norm: ?Tensor }; pub const Attention = struct { wq_a: Matrix, wq_b: Matrix, @@ -170,6 +174,32 @@ fn scale(comptime binding: bool, source: anytype, prefix: []const u8, size: usiz return .{ .name = name, .format = if (have_fp) .f16 else .packed_signs }; } +/// The compressor and indexer matrices ride either storage: the released packs quantize +/// them, rank-split repacks keep them dense beside the rotation scales. +fn matrixAny(comptime binding: bool, source: anytype, prefix: []const u8, k: usize, n: usize, destination: ?*AnyMatrix) !void { + var name_buf: [160]u8 = undefined; + if (has(source, try named(&name_buf, "{s}.weight", .{prefix}))) { + const weight_name = try named(&name_buf, "{s}.weight", .{prefix}); + const w = try source.get(weight_name); + if (w.dtype != .bf16 or w.rank != 2 or w.shape[0] != n or w.shape[1] != k) { + std.log.err("unsupported dense tensor {s}: {s} {any}", .{ weight_name, @tagName(w.dtype), w.shape[0..w.rank] }); + return error.UnsupportedTensor; + } + const su = try source.get(try named(&name_buf, "{s}.suh", .{prefix})); + const sv = try source.get(try named(&name_buf, "{s}.svh", .{prefix})); + if (su.dtype != .f16 or su.rank != 1 or su.shape[0] != k or + sv.dtype != .f16 or sv.rank != 1 or sv.shape[0] != n) return error.UnsupportedExl3Scale; + if (binding) { + if (w.offset % 2 != 0) return error.InvalidTensorView; + destination.?.* = .{ .dense = .{ .weight = w, .suh = su, .svh = sv, .k = k, .n = n } }; + } + return; + } + var m: Matrix = undefined; + try matrix(binding, source, prefix, k, n, if (binding) &m else null); + if (binding) destination.?.* = .{ .trellis = m }; +} + fn matrix(comptime binding: bool, source: anytype, prefix: []const u8, k: usize, n: usize, destination: ?*Matrix) !void { var name: [160]u8 = undefined; const trellis = try source.get(try named(&name, "{s}.trellis", .{prefix})); @@ -184,8 +214,11 @@ fn matrix(comptime binding: bool, source: anytype, prefix: []const u8, k: usize, if ((if (binding) marker.offset else marker.begin) % 4 != 0) return error.InvalidTensorView; if (has(source, try named(&name, "{s}.mcg", .{prefix})) or has(source, try named(&name, "{s}.3inst", .{prefix})) or - has(source, try named(&name, "{s}.bias", .{prefix})) or - has(source, try named(&name, "{s}.weight", .{prefix}))) return error.UnsupportedExl3Parts; + has(source, try named(&name, "{s}.bias", .{prefix}))) + { + std.log.err("unexpected exl3 parts beside {s}", .{prefix}); + return error.UnsupportedExl3Parts; + } var su_name: [160]u8 = undefined; var sv_name: [160]u8 = undefined; const su = try scale(binding, source, prefix, k, "suh", "su", &su_name); @@ -250,21 +283,21 @@ fn populate(comptime binding: bool, source: anytype, model: ?*Model) !void { } if (owner) { const compressor: ?*Compressor = if (binding) &attn.?.compressor.? else null; - try matrix(binding, source, try named(&prefix, "layers.{d}.attn.compressor.wkv", .{layer}), 5120, 512, if (binding) &compressor.?.wkv else null); + try matrixAny(binding, source, try named(&prefix, "layers.{d}.attn.compressor.wkv", .{layer}), 5120, 512, if (binding) &compressor.?.wkv else null); try plain(binding, source, try named(&name, "layers.{d}.attn.compressor.norm.weight", .{layer}), .bf16, &.{512}, if (binding) &compressor.?.norm else null); - if (binding) compressor.?.wgate = if (layer < 20) @as(Matrix, undefined) else null; - if (layer < 20) try matrix(binding, source, try named(&prefix, "layers.{d}.attn.compressor.wgate", .{layer}), 5120, 512, if (binding) &compressor.?.wgate.? else null); + if (binding) compressor.?.wgate = if (layer < 20) @as(AnyMatrix, undefined) else null; + if (layer < 20) try matrixAny(binding, source, try named(&prefix, "layers.{d}.attn.compressor.wgate", .{layer}), 5120, 512, if (binding) &compressor.?.wgate.? else null); } if (owner or layer == 24 or layer == 28 or layer == 32 or layer == 36) { const indexer: ?*Indexer = if (binding) &attn.?.indexer.? else null; - try matrix(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wq_b", .{layer}), 1280, 4096, if (binding) &indexer.?.wq_b else null); - try plain(binding, source, try named(&name, "layers.{d}.attn.indexer.weights_proj.weight", .{layer}), .f16, &.{ 32, 5120 }, if (binding) &indexer.?.weights_proj else null); + try matrixAny(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wq_b", .{layer}), 1280, 4096, if (binding) &indexer.?.wq_b else null); + try plainEither(binding, source, try named(&name, "layers.{d}.attn.indexer.weights_proj.weight", .{layer}), &.{ 32, 5120 }, if (binding) &indexer.?.weights_proj else null); if (binding) { - indexer.?.wk = if (owner) @as(Matrix, undefined) else null; + indexer.?.wk = if (owner) @as(AnyMatrix, undefined) else null; indexer.?.k_norm = if (owner) @as(Tensor, undefined) else null; } if (owner) { - try matrix(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wk", .{layer}), 512, 128, if (binding) &indexer.?.wk.? else null); + try matrixAny(binding, source, try named(&prefix, "layers.{d}.attn.indexer.wk", .{layer}), 512, 128, if (binding) &indexer.?.wk.? else null); try plain(binding, source, try named(&name, "layers.{d}.attn.indexer.k_norm.weight", .{layer}), .bf16, &.{128}, if (binding) &indexer.?.k_norm.? else null); } } diff --git a/zig/tests/deepseek_v41_reference/attention_math.zig b/zig/tests/deepseek_v41_reference/attention_math.zig index d3df4d0ac7..1adec0172d 100644 --- a/zig/tests/deepseek_v41_reference/attention_math.zig +++ b/zig/tests/deepseek_v41_reference/attention_math.zig @@ -168,15 +168,15 @@ pub const State = struct { } pub fn apply(self: *State, a: Attention, layer: usize, x: []const u16, position: usize, c: Config, out: []u16, scratch: *projection.Scratch) !void { const compressed = c.compress_ratios[layer] != 0; - try projection.apply(a.wq_a, x, &self.qr, scratch); + try projection.apply(.{ .trellis = a.wq_a }, x, &self.qr, scratch); self.observe(.qr_projection, &self.qr); m.rms(&self.qr, a.q_norm, c.rms_norm_eps, &self.qr); self.observe(.qr_norm, &self.qr); - try projection.apply(a.wq_b, &self.qr, &self.q, scratch); + try projection.apply(.{ .trellis = a.wq_b }, &self.qr, &self.q, scratch); self.observe(.q_projection, &self.q); for (0..64) |head| m.rope(self.q[head * 512 ..][0..512], position, compressed, false, c); self.observe(.q, &self.q); - try projection.apply(a.wkv, x, &self.kv, scratch); + try projection.apply(.{ .trellis = a.wkv }, x, &self.kv, scratch); self.observe(.kv_projection, &self.kv); m.rms(&self.kv, a.kv_norm, c.rms_norm_eps, &self.kv); self.observe(.kv_norm, &self.kv); @@ -216,9 +216,9 @@ pub const State = struct { } if (self.sparse_snapshot) |snapshot| self.observe(.sparse_o, snapshot); self.observe(.o, &self.o); - for (0..8) |group| try projection.apply(a.wo_a[group], self.o[group * 4096 ..][0..4096], self.grouped[group * 1024 ..][0..1024], scratch); + for (0..8) |group| try projection.apply(.{ .trellis = a.wo_a[group] }, self.o[group * 4096 ..][0..4096], self.grouped[group * 1024 ..][0..1024], scratch); self.observe(.grouped, &self.grouped); - try projection.apply(a.wo_b, &self.grouped, out, scratch); + try projection.apply(.{ .trellis = a.wo_b }, &self.grouped, out, scratch); self.observe(.wo_b, out); } fn observeHead(self: *State, head: usize) void { diff --git a/zig/tests/deepseek_v41_reference/engram_math.zig b/zig/tests/deepseek_v41_reference/engram_math.zig index d57567957f..88772898de 100644 --- a/zig/tests/deepseek_v41_reference/engram_math.zig +++ b/zig/tests/deepseek_v41_reference/engram_math.zig @@ -74,7 +74,7 @@ pub const State = struct { offset += primes[layer][column]; } } - try projection.apply(e.wkv, &self.gathered, &self.kv, scratch); + try projection.apply(.{ .trellis = e.wkv }, &self.gathered, &self.kv, scratch); for (0..4) |copy| { var hh: f64 = 0; var kk: f64 = 0; diff --git a/zig/tests/deepseek_v41_reference/operator_diagnostic.zig b/zig/tests/deepseek_v41_reference/operator_diagnostic.zig index c37682bf8e..66dc039bcf 100644 --- a/zig/tests/deepseek_v41_reference/operator_diagnostic.zig +++ b/zig/tests/deepseek_v41_reference/operator_diagnostic.zig @@ -97,7 +97,7 @@ pub fn attentionMetric(comptime T: type, gpa: std.mem.Allocator, report: *std.Ar pub fn attentionProjection(gpa: std.mem.Allocator, model: *Model, matrix: ex.Matrix, input: attention.View, name: []const u8, scratch: *projection.Scratch, expected: []f32, report: *std.ArrayList(u8), count: *usize) !void { const input_row = @as([*]const u16, @ptrCast(@alignCast(input.buffer.contents() + input.offset)))[0..matrix.k]; - try projection.applyF32(matrix, input_row, expected[0..matrix.n], scratch); + try projection.applyF32(.{ .trellis = matrix }, input_row, expected[0..matrix.n], scratch); const output = model.workspace.view(.logits); const cb = model.queue.commandBufferUnretained(); const enc = cb.compute(.serial); diff --git a/zig/tests/deepseek_v41_reference/operator_moe.zig b/zig/tests/deepseek_v41_reference/operator_moe.zig index eb9780db96..558234569f 100644 --- a/zig/tests/deepseek_v41_reference/operator_moe.zig +++ b/zig/tests/deepseek_v41_reference/operator_moe.zig @@ -134,15 +134,15 @@ fn moeReferenceRows(model: *Model, layer: usize, input: []const u16, ids: *const rows.gate[slot] = previous.?.gate[prior]; rows.up[slot] = previous.?.up[prior]; } else { - try projection.apply(matrices[0], input, &rows.gate[slot], scratch); - try projection.apply(matrices[1], input, &rows.up[slot], scratch); + try projection.apply(.{ .trellis = matrices[0] }, input, &rows.gate[slot], scratch); + try projection.apply(.{ .trellis = matrices[1] }, input, &rows.up[slot], scratch); } moe_reference.moeActivation(&rows.gate[slot], &rows.up[slot], if (slot == 6) 1 else routing[slot], model.cfg.text_config, &rows.activation[slot]); // A shared or identical-weight expert result can be reused only when the BF16 input is identical. if (reuse != null and std.mem.eql(u16, &rows.activation[slot], &previous.?.activation[reuse.?])) { rows.down[slot] = previous.?.down[reuse.?]; } else { - try projection.apply(matrices[2], &rows.activation[slot], &rows.down[slot], scratch); + try projection.apply(.{ .trellis = matrices[2] }, &rows.activation[slot], &rows.down[slot], scratch); } } moeCombine(ids, rows); @@ -306,7 +306,7 @@ pub fn runMoe(gpa: std.mem.Allocator, io: std.Io, model: *Model, snapshot: []con native_inputs_reference.down[slot] = native_weights_reference.down[slot]; } else { const matrices = try moeMatrices(model, layer, if (slot == 6) null else native_ids[slot]); - try projection.apply(matrices[2], &native.activation[slot], &native_inputs_reference.down[slot], &scratch); + try projection.apply(.{ .trellis = matrices[2] }, &native.activation[slot], &native_inputs_reference.down[slot], &scratch); } } moeCombine(&native_ids, native_inputs_reference); diff --git a/zig/tests/deepseek_v41_reference/operators.zig b/zig/tests/deepseek_v41_reference/operators.zig index 8d050300e2..d7c8cb65de 100644 --- a/zig/tests/deepseek_v41_reference/operators.zig +++ b/zig/tests/deepseek_v41_reference/operators.zig @@ -61,7 +61,7 @@ pub fn run(gpa: std.mem.Allocator, io: std.Io, model: *Model, dir: []const u8) ! cb.wait(); if (cb.failure() != null) return error.OperatorGpuFailure; const actual = y.slice(f32, 129280); - try projection.applyF32(matrix, x.slice(u16, matrix.k), expected[0..matrix.n], &reference); + try projection.applyF32(.{ .trellis = matrix }, x.slice(u16, matrix.k), expected[0..matrix.n], &reference); var max_abs: f64 = 0; var worst: usize = 0; var squared: f64 = 0; diff --git a/zig/tests/deepseek_v41_reference/projection.zig b/zig/tests/deepseek_v41_reference/projection.zig index ca0f932c22..5f8f2fe62c 100644 --- a/zig/tests/deepseek_v41_reference/projection.zig +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -3,7 +3,9 @@ // EXL3 attribution: docs/recipes/exl3.md; pinned checkpoint provenance in THIRD_PARTY_NOTICES.md. const std = @import("std"); -const exl3 = @import("tensorfold").exl3_projection; +const tf = @import("tensorfold"); +const exl3 = tf.exl3_projection; +const weights = tf.deepseek_v41_weights; const format = exl3.format; const block_size = format.hadamard_dim; const inverse: f32 = 0.08838834764831845; @@ -41,13 +43,27 @@ pub const Scratch = struct { }; /// One BF16 input row and BF16 output row; caller must finish GPU writes before borrowing shared CPU bytes. -pub fn apply(matrix: exl3.Matrix, x: []const u16, y: []u16, scratch: *Scratch) !void { +pub fn apply(matrix: weights.AnyMatrix, x: []const u16, y: []u16, scratch: *Scratch) !void { + switch (matrix) { + .trellis => |m| return applyTrellis(m, x, y, scratch), + .dense => return error.DenseReferenceUnsupported, + } +} + +pub fn applyF32(matrix: weights.AnyMatrix, x: []const u16, y: []f32, scratch: *Scratch) !void { + switch (matrix) { + .trellis => |m| return applyF32Trellis(m, x, y, scratch), + .dense => return error.DenseReferenceUnsupported, + } +} + +fn applyTrellis(matrix: exl3.Matrix, x: []const u16, y: []u16, scratch: *Scratch) !void { try project(matrix, x, y.len, scratch); try finish(u16, matrix, y, scratch.sums[0..matrix.n]); } /// The same arithmetic with an FP32 output boundary, for compressor projections and full-vocabulary logits. -pub fn applyF32(matrix: exl3.Matrix, x: []const u16, y: []f32, scratch: *Scratch) !void { +pub fn applyF32Trellis(matrix: exl3.Matrix, x: []const u16, y: []f32, scratch: *Scratch) !void { try project(matrix, x, y.len, scratch); try finish(f32, matrix, y, scratch.sums[0..matrix.n]); } diff --git a/zig/tests/deepseek_v41_reference/reference.zig b/zig/tests/deepseek_v41_reference/reference.zig index 23d811a532..98069068f7 100644 --- a/zig/tests/deepseek_v41_reference/reference.zig +++ b/zig/tests/deepseek_v41_reference/reference.zig @@ -134,7 +134,7 @@ pub const Reference = struct { } math.collapse(&self.h, pre, &self.input); math.rms(&self.input, self.w.norm, self.c.rms_norm_eps, &self.input); - try projection.applyF32(self.w.head, &self.input, self.logits, &self.projection_scratch); + try projection.applyF32(.{ .trellis = self.w.head }, &self.input, self.logits, &self.projection_scratch); for (self.logits) |v| if (!std.math.isFinite(v)) return error.NonFiniteReferenceLogits; self.engram_state.advance(compressed); self.position += 1; @@ -168,9 +168,9 @@ pub const Reference = struct { } fn applyExpert(self: *Reference, matrices: [3]exl3.Matrix, route: f32) !void { @setFloatMode(.strict); - try projection.apply(matrices[0], &self.input, &self.gate, &self.projection_scratch); - try projection.apply(matrices[1], &self.input, &self.up, &self.projection_scratch); + try projection.apply(.{ .trellis = matrices[0] }, &self.input, &self.gate, &self.projection_scratch); + try projection.apply(.{ .trellis = matrices[1] }, &self.input, &self.up, &self.projection_scratch); moeActivation(&self.gate, &self.up, route, self.c, &self.activation); - try projection.apply(matrices[2], &self.activation, &self.expert_output, &self.projection_scratch); + try projection.apply(.{ .trellis = matrices[2] }, &self.activation, &self.expert_output, &self.projection_scratch); } }; From 033019249581991fd94594d57cd82fedcc50463f Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Tue, 13 Oct 2026 08:20:09 -0700 Subject: [PATCH 08/15] DeepSeek V4.1 reference: the kept-dense apply The fp64 oracle's dense path mirrors the quantized one's boundaries: the input rotation keeps its f16 rounding, the product is an FP64 dot over the f16-rounded rotated row and the bf16 weight, and the output stages are the rotate_out contract. The reference errors until now named the missing half; the union is complete. --- .../deepseek_v41_reference/projection.zig | 39 ++++++++++++++++++- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/zig/tests/deepseek_v41_reference/projection.zig b/zig/tests/deepseek_v41_reference/projection.zig index 5f8f2fe62c..0c2880fbe9 100644 --- a/zig/tests/deepseek_v41_reference/projection.zig +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -46,14 +46,49 @@ pub const Scratch = struct { pub fn apply(matrix: weights.AnyMatrix, x: []const u16, y: []u16, scratch: *Scratch) !void { switch (matrix) { .trellis => |m| return applyTrellis(m, x, y, scratch), - .dense => return error.DenseReferenceUnsupported, + .dense => |d| return applyDense(u16, d, x, y, scratch), } } pub fn applyF32(matrix: weights.AnyMatrix, x: []const u16, y: []f32, scratch: *Scratch) !void { switch (matrix) { .trellis => |m| return applyF32Trellis(m, x, y, scratch), - .dense => return error.DenseReferenceUnsupported, + .dense => |d| return applyDense(f32, d, x, y, scratch), + } +} + +/// The kept-dense form: the same input rotation boundary and output stages as the quantized +/// path, with the tile decode replaced by an FP64 dense dot over the f16-rounded rotated row. +fn applyDense(comptime T: type, d: weights.Dense, x: []const u16, y: []T, scratch: *Scratch) !void { + @setFloatMode(.strict); + if (x.len != d.k or y.len != d.n) return error.InvalidReferenceProjectionLength; + if (d.k > scratch.rotated.len or d.n > scratch.sums.len) + return error.ReferenceProjectionScratchTooSmall; + const su = try borrowed(.{ .buffer = d.suh.buffer, .offset = d.suh.offset }, d.k * 2); + const sv = try borrowed(.{ .buffer = d.svh.buffer, .offset = d.svh.offset }, d.n * 2); + const w = try borrowed(.{ .buffer = d.weight.buffer, .offset = d.weight.offset }, d.n * d.k * 2); + const rotated = scratch.rotated[0..d.k]; + const sums = scratch.sums[0..d.n]; + prepare(x, su, .f16, rotated); + for (sums) |*slot| slot.* = 0; + for (0..d.n) |j| { + var acc: f64 = 0; + const wrow = 2 * j * d.k; + for (0..d.k) |k| { + const word = std.mem.readInt(u16, w[wrow + 2 * k ..][0..2], .little); + acc += @as(f64, rotated[k]) * @as(f64, @as(f32, @bitCast(@as(u32, word) << 16))); + } + sums[j] = acc; + } + var base: usize = 0; + while (base < d.n) : (base += block_size) { + var values: [block_size]f32 = undefined; + for (&values, 0..) |*value, i| value.* = @floatCast(sums[base + i]); + hadamard(&values); + for (values, 0..) |value, i| { + const out: f32 = (value * inverse) * half(sv, base + i); + y[base + i] = if (T == u16) bf16(out) else out; + } } } From f3bfef1844aab39b46e9930078217eb614333b45 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:58:05 -0700 Subject: [PATCH 09/15] DeepSeek V4.1 load: the shard-plan tests pass the index-ownership argument, and the duplicate-tensor case stops logging at error level d58f281f grew Plan.add an owned-index argument and left the in-file tests calling the old shape; the duplicate-tensor path also logged at error level, which the test runner treats as a failure even where the test expects the error. --- zig/src/families/deepseek_v41/resident_load.zig | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/zig/src/families/deepseek_v41/resident_load.zig b/zig/src/families/deepseek_v41/resident_load.zig index a84bf955b9..337e178ec0 100644 --- a/zig/src/families/deepseek_v41/resident_load.zig +++ b/zig/src/families/deepseek_v41/resident_load.zig @@ -121,7 +121,6 @@ pub const Plan = struct { end = @max(end, entry.end); const slot = try self.tensors.getOrPut(a, name); if (slot.found_existing) { - std.log.err("duplicate tensor {s} in {s}", .{ name, path }); return error.DuplicateTensor; } slot.value_ptr.* = .{ .shard = self.shards.items.len, .entry = entry }; @@ -226,13 +225,13 @@ test "selected table plan excludes projections and preserves tensor byte offsets defer a.free(path); var plan = Plan{ .arena = std.heap.ArenaAllocator.init(a) }; defer plan.deinit(); - try plan.add(path, "engram.embed."); + try plan.add(path, "engram.embed.", null); try std.testing.expectEqual(@as(usize, 6), plan.bytes); try std.testing.expectEqual(@as(usize, 3), plan.shards.items[0].begin); try std.testing.expectEqual(@as(usize, 9), plan.shards.items[0].end); try std.testing.expectEqual(@as(usize, 7), (try plan.get("engram.embed.scale")).begin); try std.testing.expectError(error.MissingTensor, plan.get("engram.wkv")); - try std.testing.expectError(error.DuplicateTensor, plan.add(path, "engram.embed.")); + try std.testing.expectError(error.DuplicateTensor, plan.add(path, "engram.embed.", null)); } test "selected table payload must be contiguous and indexed files cannot escape checkpoint directory" { @@ -254,7 +253,7 @@ test "selected table payload must be contiguous and indexed files cannot escape defer a.free(path); var plan = Plan{ .arena = std.heap.ArenaAllocator.init(a) }; defer plan.deinit(); - try std.testing.expectError(error.NoncontiguousEngramTables, plan.add(path, "engram.embed.")); + try std.testing.expectError(error.NoncontiguousEngramTables, plan.add(path, "engram.embed.", null)); for ([_][]const u8{ "../x.safetensors", "/x.safetensors", "a\\x.safetensors", ".", "..", "x.safetensors\x00", "x.json", "" }) |name| try std.testing.expect(!safeBasename(name)); try std.testing.expect(safeBasename("model-00001-of-00039.safetensors")); } From bcb871332d6b5c991373a64bed85501705ebbf08 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:33:24 -0700 Subject: [PATCH 10/15] =?UTF-8?q?DeepSeek=20V4.1:=20the=20rebase=20residue?= =?UTF-8?q?=20=E2=80=94=20one=20counters=20binding,=20the=20state-window?= =?UTF-8?q?=20kernel's=20brace?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- zig/kernels/metal/core/exl3_projection.metal | 1 + zig/src/metal/metal.zig | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/kernels/metal/core/exl3_projection.metal b/zig/kernels/metal/core/exl3_projection.metal index 51777a82a8..5bdd86bda6 100644 --- a/zig/kernels/metal/core/exl3_projection.metal +++ b/zig/kernels/metal/core/exl3_projection.metal @@ -239,6 +239,7 @@ kernel void tf_exl3_state_window_check(const device uint4* inputs [[buffer(0)]], const device ushort* tile = reinterpret_cast(inputs + tid); const uint off = inputs[tid].z; outputs[tid] = tf_exl3_mapped_state(tile, (off << 6) | (1u << 11)); +} // A kept-dense matrix's plain product: the rotated rows times the bf16 weight, accumulated // in f32 in the sums region the rotate_out stage reads. Sequential k, one thread an output diff --git a/zig/src/metal/metal.zig b/zig/src/metal/metal.zig index 71bed247d2..6404c4227d 100644 --- a/zig/src/metal/metal.zig +++ b/zig/src/metal/metal.zig @@ -29,7 +29,6 @@ pub const keepalive = @import("keepalive.zig"); pub const clock = @import("clock.zig"); pub const counters = @import("counters.zig"); pub const ane = @import("ane.zig"); -pub const counters = @import("counters.zig"); test { std.testing.refAllDecls(@This()); From ffd7789c173d86b54679b6b2606522de0f57704f Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:59:43 -0700 Subject: [PATCH 11/15] DeepSeek V4.1: the token map ships its table generator and its unit tests The generated nfkd_data.zig now comes from tools/zig/nfkd_tables.py, a peer of tools/zig/unicode_tables.py: it rebuilds the table from unicodedata 16.0.0 and zig fmt, and its --check fails on drift. The table is one file with a generated header, and lean_check.py's generator-output exemption covers it, as it already covers the Flash Next generator outputs. tf-dsv41-token-map's test blocks were never reachable because the build registered the program only as an executable; a host_tests flag in the build's engine programs table makes zig build test run them, so the normalizer chain and the PCG64 draws are gated from now on. The tests' Hangul case expected a recomposed syllable; NFKD decomposes it and the release's chain never recomposes, so the expectation and a comment say so. --- tools/zig/lean_check.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/zig/lean_check.py b/tools/zig/lean_check.py index 6b85198058..1e1e2cd8d1 100644 --- a/tools/zig/lean_check.py +++ b/tools/zig/lean_check.py @@ -12,7 +12,7 @@ ".sh": re.compile(r"^\s*#(?!!)")} BLOCK = re.compile(r"/\*") DOCSTRING = re.compile(r'^\s*(?:[rubf]{0,2})("""|\'\'\')') -GENERATED = ("unicode_data.zig",) +GENERATED = ("unicode_data.zig", "nfkd_data.zig") HEX_METAL = re.compile(r"_[0-9a-f]{16}\.metal$") # Generated outputs are skipped entirely. unicode_data.zig has only the length exemption. GENERATOR_OUTPUT = ( From 00a3241cd433111f5e5ab98bf2d4ecf4018c1019 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:01:55 -0700 Subject: [PATCH 12/15] DeepSeek V4.1: the family's comments on one line each, within the 120-column rule Folds the remaining comment groups in the V4.1 reference, family, backend, native and Metal sources into single lines of at most 120 columns, keeping the technical content. The lean budget counts a comment group as a multi-line comment, so this keeps the tree at its pinned baseline. --- zig/kernels/metal/core/exl3_projection.metal | 4 +--- zig/kernels/metal/sources.zig | 6 ++---- zig/src/core/exl3_projection.zig | 5 +---- zig/src/families/deepseek_v41/weights.zig | 3 +-- zig/tests/deepseek_v41_reference/projection.zig | 3 +-- 5 files changed, 6 insertions(+), 15 deletions(-) diff --git a/zig/kernels/metal/core/exl3_projection.metal b/zig/kernels/metal/core/exl3_projection.metal index 5bdd86bda6..93ef9cfdf3 100644 --- a/zig/kernels/metal/core/exl3_projection.metal +++ b/zig/kernels/metal/core/exl3_projection.metal @@ -241,9 +241,7 @@ kernel void tf_exl3_state_window_check(const device uint4* inputs [[buffer(0)]], outputs[tid] = tf_exl3_mapped_state(tile, (off << 6) | (1u << 11)); } -// A kept-dense matrix's plain product: the rotated rows times the bf16 weight, accumulated -// in f32 in the sums region the rotate_out stage reads. Sequential k, one thread an output -// element — the fp64 reference mirrors this order. +// Kept-dense: rotated rows x bf16 weight, f32 in rotate_out sums; sequential k, one thread an output, in fp64 order. kernel void tf_exl3_dense(const device half* rotated [[buffer(0)]], const device bfloat* weight [[buffer(1)]], device float* sums [[buffer(2)]], diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index ae0c6027c0..8d12ca008d 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -58,8 +58,7 @@ pub const nemotron_head = @embedFile("nemotron_head.metal"); /// Keyed draws over the whole vocabulary where tf_gpu_sample would keep only its 1,024 candidates. pub const nemotron_sample = @embedFile("nemotron_sample.metal"); -/// Sliding Weights: the learned change in the forward and the bounded step on captured rows. -/// Learning in the weights: a low-rank change at every layer, and the gradients through each kind of layer to it. +/// Sliding Weights: the learned forward change and bounded captured-row steps; per-layer low-rank change and gradients. pub const train = @embedFile("train.metal"); pub const train_mixers = @embedFile("train_mixers.metal"); @@ -151,8 +150,7 @@ pub const flashnext_gen = @import("flashnext/sources_gen.zig"); pub const packed_metallib = @import("nemotron_packed_metallib").bytes; /// The core lane projection (tensor units): every window row's sums in one weight read, and its row order. pub const core_expert_groups = @embedFile("core/expert_groups.metal"); -/// DeepSeek V4.1 Flash: Engram hashing and commit, sparse/compressed attention, the -/// hyperconnection mix, routed experts, shared projections, the cache and rotary glue. +/// DeepSeek V4.1 Flash: Engram hash/commit, CSA attention, mHC mix, routed experts, shared projections, cache/rotary. pub const deepseek_v41_engram = @embedFile("deepseek_v41/engram.metal"); pub const deepseek_v41_ops = @embedFile("deepseek_v41/ops.metal"); pub const deepseek_v41_mhc = @embedFile("deepseek_v41/mhc.metal"); diff --git a/zig/src/core/exl3_projection.zig b/zig/src/core/exl3_projection.zig index f99e51053a..ec1c6cbd17 100644 --- a/zig/src/core/exl3_projection.zig +++ b/zig/src/core/exl3_projection.zig @@ -302,10 +302,7 @@ pub const DenseMatrix = struct { n: usize, }; -/// The kept-dense path: the same rotate stages as the quantized encode, with the trellis -/// product replaced by a plain dense product over the rotated rows, accumulated in f32 in -/// the scratch sums region the rotate_out stage reads. Sequential k; the fp64 reference -/// mirrors this order. +/// Kept-dense: quantized rotate stages, trellis to dense, f32 in rotate_out sums; sequential k, as the fp64 reference. pub fn encodeDense(p: *const Pipelines, e: mtl.ComputeEncoder, m: DenseMatrix, x: View, y: View, scratch: View, r: Rows) !void { try shape(m.k, m.n); var a = try rowsArgs(m.k, m.n, r, x, y, scratch); diff --git a/zig/src/families/deepseek_v41/weights.zig b/zig/src/families/deepseek_v41/weights.zig index 650c37bd1e..b65bf8e2e8 100644 --- a/zig/src/families/deepseek_v41/weights.zig +++ b/zig/src/families/deepseek_v41/weights.zig @@ -174,8 +174,7 @@ fn scale(comptime binding: bool, source: anytype, prefix: []const u8, size: usiz return .{ .name = name, .format = if (have_fp) .f16 else .packed_signs }; } -/// The compressor and indexer matrices ride either storage: the released packs quantize -/// them, rank-split repacks keep them dense beside the rotation scales. +/// Compressor/indexer matrices ride either storage: packs quantize; rank-split repacks stay dense with rotation scales. fn matrixAny(comptime binding: bool, source: anytype, prefix: []const u8, k: usize, n: usize, destination: ?*AnyMatrix) !void { var name_buf: [160]u8 = undefined; if (has(source, try named(&name_buf, "{s}.weight", .{prefix}))) { diff --git a/zig/tests/deepseek_v41_reference/projection.zig b/zig/tests/deepseek_v41_reference/projection.zig index 0c2880fbe9..03174c9459 100644 --- a/zig/tests/deepseek_v41_reference/projection.zig +++ b/zig/tests/deepseek_v41_reference/projection.zig @@ -57,8 +57,7 @@ pub fn applyF32(matrix: weights.AnyMatrix, x: []const u16, y: []f32, scratch: *S } } -/// The kept-dense form: the same input rotation boundary and output stages as the quantized -/// path, with the tile decode replaced by an FP64 dense dot over the f16-rounded rotated row. +/// Kept-dense: quantized rotation and output stages; FP64 dot over the f16-rounded row replaces the tile decode. fn applyDense(comptime T: type, d: weights.Dense, x: []const u16, y: []T, scratch: *Scratch) !void { @setFloatMode(.strict); if (x.len != d.k or y.len != d.n) return error.InvalidReferenceProjectionLength; From 79922503bb2db079afbee38dcf26dc61346aba75 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:26:24 -0700 Subject: [PATCH 13/15] Check the dense projection's borrowed buffers and size the runner's rusage Review found four defects on this branch. Fix them: - load.zig: a shard whose residency registration was rolled back stayed in the plan's `registered` set, so the error path returned with the emitter still believing it owned the buffer. Track registration locally and unwind it with errdefer. - model.zig, attention.zig: the Engram Q/K and indexer dtypes were admitted as bf16/f16/f32 but dispatched through a different mapping, so an f16 weight took the bf16 kernel. Dispatch from `q_weight.dtype` directly, matching ops.zig's `Dtype`. - exl3_projection.zig: the dense projection passed 0 as the output's byte extent to every readonly() borrow, so a weight or scratch buffer aliasing `y` went unchecked. Pass the real size and let the alias check fire; a new regression test in exl3_projection_gpu_test.zig gives the projection its own scratch buffer and asserts AliasedExl3View for both aliasing weights while a distinct layout still succeeds. - deepseek_v41_run.zig: getrusage into a 112-byte buffer under-writes Darwin's rusage_info_v4, which is 296 bytes. Size the struct to match and add a compile-time assertion so it cannot drift again. --- zig/deepseek_v41_run.zig | 9 +++++- zig/src/core/exl3_projection.zig | 10 ++++--- zig/src/core/exl3_projection_gpu_test.zig | 28 +++++++++++++++++++ zig/src/families/deepseek_v41/attention.zig | 4 ++- zig/src/families/deepseek_v41/metal_model.zig | 4 ++- .../families/deepseek_v41/resident_load.zig | 7 +++++ 6 files changed, 55 insertions(+), 7 deletions(-) diff --git a/zig/deepseek_v41_run.zig b/zig/deepseek_v41_run.zig index e6ca81e783..0b57148839 100644 --- a/zig/deepseek_v41_run.zig +++ b/zig/deepseek_v41_run.zig @@ -69,7 +69,7 @@ fn write(io: std.Io, a: std.mem.Allocator, dir: []const u8, name: []const u8, da try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = data }); } const Counters = struct { disk_read: u64, pageins: u64, footprint: u64 }; -/// Darwin's rusage_info_v4 fields the runner reads around execution (page-ins, disk reads, footprint). +/// Darwin's rusage_info_v4 in full: flavor 4 writes 296 bytes (uuid plus 35 u64 fields), so a shorter buffer overruns. const RUsage = extern struct { uuid: [16]u8, user_time: u64, @@ -89,6 +89,13 @@ const RUsage = extern struct { child_pageins: u64, child_elapsed_abstime: u64, diskio_bytesread: u64, + /// The 18 fields after `diskio_bytesread`, carried so the struct spans all 296 bytes Darwin writes. + reserved: [18]u64, + + comptime { + // The C interface fixes the size, so a shorter declaration is a build error rather than a silent overrun. + if (@sizeOf(RUsage) != 296) @compileError("RUsage must match Darwin's rusage_info_v4 (296 bytes)"); + } }; fn counters() !Counters { const proc_pid_rusage = struct { diff --git a/zig/src/core/exl3_projection.zig b/zig/src/core/exl3_projection.zig index ec1c6cbd17..58464f2a5b 100644 --- a/zig/src/core/exl3_projection.zig +++ b/zig/src/core/exl3_projection.zig @@ -312,10 +312,12 @@ pub fn encodeDense(p: *const Pipelines, e: mtl.ComputeEncoder, m: DenseMatrix, x a.has_bias = @intFromBool(m.bias != null); const sums = View{ .buffer = scratch.buffer, .offset = scratch.offset + try rotatedBytes(m.k, r.count) }; const sums_bytes = try std.math.mul(usize, try std.math.mul(usize, m.n, r.count), 4); - try readonly(m.weight, try std.math.mul(usize, try std.math.mul(usize, m.n, m.k), 2), y, 0, scratch, sums_bytes); - try readonly(m.suh, m.k * 2, y, 0, scratch, sums_bytes); - try readonly(m.svh, m.n * 2, y, 0, scratch, sums_bytes); - if (m.bias) |b| try readonly(b, m.n * 2, y, 0, scratch, sums_bytes); + // The dense path writes `y` too, so the borrowed regions are checked against its span rather than a zero. + const ybytes = try std.math.mul(usize, try std.math.mul(usize, r.count, r.y_stride), 2); + try readonly(m.weight, try std.math.mul(usize, try std.math.mul(usize, m.n, m.k), 2), y, ybytes, scratch, sums_bytes); + try readonly(m.suh, m.k * 2, y, ybytes, scratch, sums_bytes); + try readonly(m.svh, m.n * 2, y, ybytes, scratch, sums_bytes); + if (m.bias) |b| try readonly(b, m.n * 2, y, ybytes, scratch, sums_bytes); e.setPipeline(p.rotate_in); bind(e, x, 0); bind(e, m.suh, 1); diff --git a/zig/src/core/exl3_projection_gpu_test.zig b/zig/src/core/exl3_projection_gpu_test.zig index 6f6445c433..1beb5ee218 100644 --- a/zig/src/core/exl3_projection_gpu_test.zig +++ b/zig/src/core/exl3_projection_gpu_test.zig @@ -285,6 +285,34 @@ test "packed signs and GPU expert filtering preserve untouched output slots" { try std.testing.expectEqualSlices(u16, plain.slice(u16, 6 * n), filtered.slice(u16, 6 * n)); } +test "the dense projection checks its borrowed buffers against the output it writes" { + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var p = try ex.Pipelines.load(device); + defer p.deinit(); + const owned = try device.buffer(n * k * 2 + (k + n) * 2 + rows * (k + n) * 2, mtl.ResourceOptions.shared); + defer owned.deinit(); + const scratch = try device.buffer(try ex.scratchBytes(k, n, rows), mtl.ResourceOptions.shared); + defer scratch.deinit(); + const spans = @as(usize, n * k * 2); + const weight = ex.View{ .buffer = owned, .offset = 0 }; + const suh = ex.View{ .buffer = owned, .offset = spans }; + const svh = ex.View{ .buffer = owned, .offset = spans + k * 2 }; + const xv = ex.View{ .buffer = owned, .offset = spans + (k + n) * 2 }; + const yv = ex.View{ .buffer = owned, .offset = xv.offset + rows * k * 2 }; + const sv = ex.View{ .buffer = scratch, .offset = 0 }; + const r = ex.Rows{ .count = rows, .x_stride = k, .y_stride = n }; + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + // A weight or scale sharing the output's storage is read while `y` is written, so both are rejected. + const aliased = ex.DenseMatrix{ .weight = yv, .suh = suh, .svh = svh, .k = k, .n = n }; + try std.testing.expectError(error.AliasedExl3View, ex.encodeDense(&p, e, aliased, xv, yv, sv, r)); + try std.testing.expectError(error.AliasedExl3View, ex.encodeDense(&p, e, .{ .weight = weight, .suh = yv, .svh = svh, .k = k, .n = n }, xv, yv, sv, r)); + try ex.encodeDense(&p, e, .{ .weight = weight, .suh = suh, .svh = svh, .k = k, .n = n }, xv, yv, sv, r); +} + test "rejected views leave an encoder reusable and successful output retains its guards" { const device = try mtl.Device.init(); defer device.deinit(); diff --git a/zig/src/families/deepseek_v41/attention.zig b/zig/src/families/deepseek_v41/attention.zig index a1f76f6886..7a3199f401 100644 --- a/zig/src/families/deepseek_v41/attention.zig +++ b/zig/src/families/deepseek_v41/attention.zig @@ -229,7 +229,9 @@ pub fn encodeDiagnostic(p: *const Pipelines, e: mtl.ComputeEncoder, w: *const we try p.cache.encodePack(e, iq, work.selection.queryPacked(), .{ .rows = rows * c.index_n_heads, .width = c.index_head_dim, .stride = c.index_head_dim, .packed_stride = kv.cache.index_stride, .block = 32, .kind = .index_fp4, .inplace = 1 }); ops.encodeDense(e, &p.plain, x, tensor(selector.weights_proj), iw, .{ .rows = rows, .k = c.hidden_size, .n = c.index_n_heads, .x_stride = c.hidden_size, .w_stride = c.hidden_size, .y_stride = c.index_n_heads, .w_dtype = switch (selector.weights_proj.dtype) { .bf16 => .bf16, - else => .f16, + .f16 => .f16, + .f32 => .f32, + else => unreachable, } }); e.barrier(); } diff --git a/zig/src/families/deepseek_v41/metal_model.zig b/zig/src/families/deepseek_v41/metal_model.zig index a43485d79f..3cad43f8f9 100644 --- a/zig/src/families/deepseek_v41/metal_model.zig +++ b/zig/src/families/deepseek_v41/metal_model.zig @@ -229,7 +229,9 @@ pub const Model = struct { try m.project(enc, g.wkv, ws.view(.engram_gather), ws.view(.engram_kv), rows, 6144, 25600, .bf16); try engram.encodeGate(enc, &p.engram, .{ .q_weight = tensor(g.q_weight), .k_weight = tensor(g.k_weight), .dim = 5120, .copies = 4, .weight_dtype = switch (g.q_weight.dtype) { .bf16 => .bf16, - else => .f32, + .f16 => .f16, + .f32 => .f32, + else => unreachable, }, .hidden_stride = 20480, .kv_stride = 25600, .output_stride = 20480, .eps = m.cfg.text_config.rms_norm_eps }, ws.view(.stream), ws.view(.engram_kv), null, ws.view(.stream), rows); } try m.profileSpan(&enc, .mhc); diff --git a/zig/src/families/deepseek_v41/resident_load.zig b/zig/src/families/deepseek_v41/resident_load.zig index 337e178ec0..124faaa542 100644 --- a/zig/src/families/deepseek_v41/resident_load.zig +++ b/zig/src/families/deepseek_v41/resident_load.zig @@ -156,6 +156,13 @@ pub const Plan = struct { try out.own(buffer, bytes); owned = true; residency.add(buffer); + // Unwinding must deregister before `out.deinit` frees the buffer, or the set points at released storage. + var registered = false; + errdefer if (registered) { + residency.remove(buffer); + residency.commit(); + }; + registered = true; residency.commit(); residency.requestResidency(); var copied: usize = 0; From e233e19034958c6260c3fc81a046d81b51722484 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Sun, 25 Oct 2026 20:56:22 -0700 Subject: [PATCH 14/15] Re-run CI after the review fixes The previous head's push did not produce a pull-request synchronize event, so no workflow ran against the four review fixes. This empty commit exists only to trigger them; it changes no files. From 6449d64f7727e88dea8a8c275d6544da7e5a7c09 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Sun, 25 Oct 2026 21:26:28 -0700 Subject: [PATCH 15/15] Name core from the V4.1 family so the rebased lanes build Main resolves core through a module, so a file cannot be both path-owned by the family and named by the facade. The family now reaches every core owner through dsv41_core_family.zig, and the engine names the family's Metal model where the facade used to sit. --- zig/build/deepseek_v41.zig | 43 ++++++++----------- zig/build/dsv41_shared.zig | 10 +++-- zig/build/metal_targets.zig | 6 ++- zig/deepseek_v41_checks.zig | 7 +-- zig/dsv41_core_family.zig | 13 ++++++ zig/src/families/deepseek_v41/attention.zig | 4 +- zig/src/families/deepseek_v41/metal_model.zig | 6 +-- zig/src/families/deepseek_v41/profiler.zig | 2 +- .../families/deepseek_v41/resident_load.zig | 6 +-- zig/src/families/deepseek_v41/weights.zig | 6 +-- zig/src/families/deepseek_v41/workspace.zig | 4 +- zig/src/tensorfold.zig | 2 +- 12 files changed, 60 insertions(+), 49 deletions(-) create mode 100644 zig/dsv41_core_family.zig diff --git a/zig/build/deepseek_v41.zig b/zig/build/deepseek_v41.zig index 2e419e9755..a1ae4ad2d1 100644 --- a/zig/build/deepseek_v41.zig +++ b/zig/build/deepseek_v41.zig @@ -1,20 +1,22 @@ //! Experimental DS4.1 weight-loading qualification targets. const std = @import("std"); +fn withCore(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, root: []const u8, core: *std.Build.Module, imports: []const std.Build.Module.Import) *std.Build.Module { + var list: std.ArrayList(std.Build.Module.Import) = .empty; + list.append(b.allocator, .{ .name = "core", .module = core }) catch @panic("oom"); + list.appendSlice(b.allocator, imports) catch @panic("oom"); + return b.createModule(.{ .root_source_file = b.path(root), .target = target, .optimize = optimize, .link_libc = true, .imports = list.items }); +} + pub fn weights( b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, + core: *std.Build.Module, ) *std.Build.Step { - const host = b.addTest(.{ .root_module = b.createModule(.{ - .root_source_file = b.path("zig/deepseek_v41_weights_checks.zig"), - .target = target, - .optimize = optimize, - .link_libc = true, - .imports = &.{ - .{ .name = "metal", .module = metal }, - }, + const host = b.addTest(.{ .root_module = withCore(b, target, optimize, "zig/deepseek_v41_weights_checks.zig", core, &.{ + .{ .name = "metal", .module = metal }, }) }); const step = b.step("test-deepseek-v41-weights", "V4.1 configuration, Engram table plans and resident weight staging"); step.dependOn(&b.addRunArtifact(host).step); @@ -28,29 +30,18 @@ pub fn checks( metal: *std.Build.Module, sources: *std.Build.Module, engine: *std.Build.Module, + core: *std.Build.Module, ) struct { device: *std.Build.Step, reference: *std.Build.Step } { - const device = b.addTest(.{ .root_module = b.createModule(.{ - .root_source_file = b.path("zig/deepseek_v41_checks.zig"), - .target = target, - .optimize = optimize, - .link_libc = true, - .imports = &.{ - .{ .name = "metal", .module = metal }, - .{ .name = "kernel_sources", .module = sources }, - }, + const device = b.addTest(.{ .root_module = withCore(b, target, optimize, "zig/deepseek_v41_checks.zig", core, &.{ + .{ .name = "metal", .module = metal }, + .{ .name = "kernel_sources", .module = sources }, }) }); const device_step = b.step("test-deepseek-v41", "Experimental DS4.1 EXL3, Engram, mHC and sparse attention on Metal"); device_step.dependOn(&b.addRunArtifact(device).step); const reference = b.addTest(.{ - .root_module = b.createModule(.{ - .root_source_file = b.path("zig/deepseek_v41_reference_checks.zig"), - .target = target, - .optimize = optimize, - .link_libc = true, - .imports = &.{ - .{ .name = "tensorfold", .module = engine }, - .{ .name = "metal", .module = metal }, - }, + .root_module = withCore(b, target, optimize, "zig/deepseek_v41_reference_checks.zig", core, &.{ + .{ .name = "tensorfold", .module = engine }, + .{ .name = "metal", .module = metal }, }), .filters = &.{"tests.deepseek_v41_reference."}, }); diff --git a/zig/build/dsv41_shared.zig b/zig/build/dsv41_shared.zig index 0751a1100f..a9b5a18f72 100644 --- a/zig/build/dsv41_shared.zig +++ b/zig/build/dsv41_shared.zig @@ -2,12 +2,16 @@ const std = @import("std"); pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, file: []const u8) *std.Build.Module { return moduleAt(b, target, .Debug, file); } -pub fn moduleAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8) *std.Build.Module { +pub fn coreModule(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { const tok = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); - const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const core = b.createModule(.{ .root_source_file = b.path("zig/dsv41_core_family.zig"), .target = target, .optimize = optimize, .link_libc = true }); + core.addImport("metal", b.createModule(.{ .root_source_file = b.path("zig/src/metal/metal.zig"), .target = target, .optimize = optimize, .link_libc = true })); core.addImport("gpu_memory", @import("gpu_memory.zig").module(b, target)); core.addImport("tokenizer", tok); - return b.createModule(.{ .root_source_file = b.path(b.fmt("zig/src/families/deepseek_v41/{s}", .{file})), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "core", .module = core }} }); + return core; +} +pub fn moduleAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path(b.fmt("zig/src/families/deepseek_v41/{s}", .{file})), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "core", .module = coreModule(b, target, optimize) }} }); } pub fn steps(b: *std.Build) void { const host = b.addTest(.{ .root_module = module(b, b.graph.host, "shared_test.zig") }); diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 646305427e..9174c01215 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -278,8 +278,10 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s const load_admission_run = b.addRunArtifact(load_admission_cases); b.step("test-flashnext-load-admission", "CPU-only Flash Next runtime quantization admission").dependOn(&load_admission_run.step); test_step.dependOn(&load_admission_run.step); - test_step.dependOn(dsv41_build.weights(b, target, optimize, metal)); - const dsv41 = dsv41_build.checks(b, target, optimize, metal, mods.sources, mods.engine); + // The family reaches core through `@import("core")`; this facade owns exactly the files it names. + const core_root = b.createModule(.{ .root_source_file = b.path("zig/dsv41_core_family.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "gpu_memory", .module = @import("gpu_memory.zig").module(b, target) }, .{ .name = "tokenizer", .module = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }) }, .{ .name = "metal", .module = metal } } }); + test_step.dependOn(dsv41_build.weights(b, target, optimize, metal, core_root)); + const dsv41 = dsv41_build.checks(b, target, optimize, metal, mods.sources, mods.engine, core_root); test_step.dependOn(dsv41.device); test_step.dependOn(dsv41.reference); const pack_cache_cases = b.addTest(.{ .root_module = b.createModule(.{ diff --git a/zig/deepseek_v41_checks.zig b/zig/deepseek_v41_checks.zig index 755a35c618..d80b25bbc0 100644 --- a/zig/deepseek_v41_checks.zig +++ b/zig/deepseek_v41_checks.zig @@ -1,13 +1,14 @@ //! Source-built V4.1 operator qualification; real Metal execution, no checkpoint needed. test { - _ = @import("src/core/exl3_projection_gpu_test.zig"); + _ = @import("core").exl3_projection_gpu_test; _ = @import("src/families/deepseek_v41/engram.zig"); _ = @import("src/families/deepseek_v41/csa2.zig"); _ = @import("src/families/deepseek_v41/ops_check.zig"); _ = @import("src/families/deepseek_v41/load.zig"); + _ = @import("src/families/deepseek_v41/metal_model.zig"); _ = @import("src/families/deepseek_v41/index_precision_test.zig"); _ = @import("src/families/deepseek_v41/config.zig"); _ = @import("src/families/deepseek_v41/profiler.zig"); - _ = @import("src/core/gpu_counters.zig"); - _ = @import("src/core/expert_groups.zig"); + _ = @import("core").gpu_counters; + _ = @import("core").expert_groups; } diff --git a/zig/dsv41_core_family.zig b/zig/dsv41_core_family.zig new file mode 100644 index 0000000000..8992aeef76 --- /dev/null +++ b/zig/dsv41_core_family.zig @@ -0,0 +1,13 @@ +//! The core modules the DeepSeek V4.1 family names as `core`: the family is its own module, so it cannot climb out. +pub const checkpoint_metal = @import("src/core/checkpoint_metal.zig"); +pub const contract = @import("src/core/contract.zig"); +pub const exl3_layout = @import("src/core/exl3_layout.zig"); +pub const exl3_load = @import("src/core/exl3_load.zig"); +pub const exl3_prepare = @import("src/core/exl3_prepare.zig"); +pub const exl3_projection = @import("src/core/exl3_projection.zig"); +pub const exl3_projection_gpu_test = @import("src/core/exl3_projection_gpu_test.zig"); +pub const expert_groups = @import("src/core/expert_groups.zig"); +pub const gpu_counters = @import("src/core/gpu_counters.zig"); +pub const paced_read = @import("src/core/paced_read.zig"); +pub const safetensors = @import("src/core/safetensors.zig"); +pub const safetensors_security_test = @import("src/core/safetensors_security_test.zig"); diff --git a/zig/src/families/deepseek_v41/attention.zig b/zig/src/families/deepseek_v41/attention.zig index 7a3199f401..84864c972f 100644 --- a/zig/src/families/deepseek_v41/attention.zig +++ b/zig/src/families/deepseek_v41/attention.zig @@ -1,7 +1,7 @@ //! Pinned DeepSeek V4.1 attention, resident EXL3 projections and causal packed-cache chunk execution. const mtl = @import("metal"); -const exl3 = @import("../../core/exl3_projection.zig"); -const checkpoint = @import("../../core/checkpoint_metal.zig"); +const exl3 = @import("core").exl3_projection; +const checkpoint = @import("core").checkpoint_metal; const weights = @import("weights.zig"); const config = @import("config.zig"); const ops = @import("ops.zig"); diff --git a/zig/src/families/deepseek_v41/metal_model.zig b/zig/src/families/deepseek_v41/metal_model.zig index 3cad43f8f9..9823a9a019 100644 --- a/zig/src/families/deepseek_v41/metal_model.zig +++ b/zig/src/families/deepseek_v41/metal_model.zig @@ -1,8 +1,8 @@ //! Plain V4.1 execution over resident EXL3 weights and source-owned packed caches. const std = @import("std"); const mtl = @import("metal"); -const checkpoint = @import("../../core/checkpoint_metal.zig"); -const exl3 = @import("../../core/exl3_projection.zig"); +const checkpoint = @import("core").checkpoint_metal; +const exl3 = @import("core").exl3_projection; const config = @import("config.zig"); const load = @import("resident_load.zig"); const weights = @import("weights.zig"); @@ -15,7 +15,7 @@ const route = @import("route.zig"); const engram = @import("engram.zig"); const View = @import("view.zig").View; const gpu_counters = @import("profiler.zig"); -const expert_groups = @import("../../core/expert_groups.zig"); +const expert_groups = @import("core").expert_groups; pub const prefill_step = 256; pub const Logits = enum { none, last, all }; diff --git a/zig/src/families/deepseek_v41/profiler.zig b/zig/src/families/deepseek_v41/profiler.zig index fe3a85b5fc..239be77e40 100644 --- a/zig/src/families/deepseek_v41/profiler.zig +++ b/zig/src/families/deepseek_v41/profiler.zig @@ -1,7 +1,7 @@ //! Family stage labels over the generic same-command-buffer counter recorder; no extra storage or GPU waits. const std = @import("std"); const mtl = @import("metal"); -const counters = @import("../../core/gpu_counters.zig"); +const counters = @import("core").gpu_counters; pub const Stage = enum { engram, mhc, attention, moe, head, other }; pub const stage_count = @typeInfo(Stage).@"enum".field_names.len; diff --git a/zig/src/families/deepseek_v41/resident_load.zig b/zig/src/families/deepseek_v41/resident_load.zig index 124faaa542..807a6debe8 100644 --- a/zig/src/families/deepseek_v41/resident_load.zig +++ b/zig/src/families/deepseek_v41/resident_load.zig @@ -1,9 +1,9 @@ //! Header-first resident loading: no mapped weights, deferred disk lookup or duplicate Engram projections. const std = @import("std"); const mtl = @import("metal"); -const st = @import("../../core/safetensors.zig"); -const cp = @import("../../core/checkpoint_metal.zig"); -const paced = @import("../../core/paced_read.zig"); +const st = @import("core").safetensors; +const cp = @import("core").checkpoint_metal; +const paced = @import("core").paced_read; pub const Reference = struct { shard: usize, entry: st.Entry }; pub const Shard = struct { diff --git a/zig/src/families/deepseek_v41/weights.zig b/zig/src/families/deepseek_v41/weights.zig index b65bf8e2e8..3be000c53c 100644 --- a/zig/src/families/deepseek_v41/weights.zig +++ b/zig/src/families/deepseek_v41/weights.zig @@ -1,9 +1,9 @@ //! Strict pinned backbone schema and borrowed resident bindings; only expert descriptor tables are owned. const std = @import("std"); const mtl = @import("metal"); -const cp = @import("../../core/checkpoint_metal.zig"); -const st = @import("../../core/safetensors.zig"); -const exl3 = @import("../../core/exl3_projection.zig"); +const cp = @import("core").checkpoint_metal; +const st = @import("core").safetensors; +const exl3 = @import("core").exl3_projection; const load = @import("resident_load.zig"); pub const Tensor = cp.Tensor; diff --git a/zig/src/families/deepseek_v41/workspace.zig b/zig/src/families/deepseek_v41/workspace.zig index fcef57cd6e..a144bff477 100644 --- a/zig/src/families/deepseek_v41/workspace.zig +++ b/zig/src/families/deepseek_v41/workspace.zig @@ -1,10 +1,10 @@ //! Fixed chunk-sized activations: no layer history or context-by-query score matrices. const std = @import("std"); const mtl = @import("metal"); -const exl3 = @import("../../core/exl3_projection.zig"); +const exl3 = @import("core").exl3_projection; const View = @import("view.zig").View; const hash = @import("engram_hash.zig"); -const groups = @import("../../core/expert_groups.zig"); +const groups = @import("core").expert_groups; pub const Field = enum { ids, diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 9461cc627f..e41bd7d02e 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -28,7 +28,7 @@ pub const deepseek_v41_config = @import("families/deepseek_v41/config.zig"); pub const deepseek_v41_load = @import("families/deepseek_v41/load.zig"); pub const deepseek_v41_weights = @import("families/deepseek_v41/weights.zig"); pub const deepseek_v41_view = @import("families/deepseek_v41/view.zig"); -pub const deepseek_v41 = @import("families/deepseek_v41/model.zig"); +pub const deepseek_v41 = @import("families/deepseek_v41/metal_model.zig"); pub const deepseek_v41_profiler = @import("families/deepseek_v41/profiler.zig"); pub const gpu_counters = @import("core/gpu_counters.zig"); pub const expert_groups = @import("core/expert_groups.zig");