diff --git a/zig/deepseek_v41_run.zig b/zig/deepseek_v41_run.zig index e6ca81e783..0a73c0856b 100644 --- a/zig/deepseek_v41_run.zig +++ b/zig/deepseek_v41_run.zig @@ -176,6 +176,7 @@ pub fn main(init: std.process.Init) !void { const resident_gate_passed = after.disk_read == before.disk_read; const receipt = try std.json.Stringify.valueAlloc(a, .{ .mode = mode, + .attention_parallel = model.pipelines.attention.sparse.parallel_enabled, .evidence_origin = "current_invocation_not_historical_qualification", .context_capacity = model.cache.plan.capacity, .consumed_tokens = model.cache.position, diff --git a/zig/kernels/metal/deepseek_v41/csa2.metal b/zig/kernels/metal/deepseek_v41/csa2.metal index 9d16cfcce5..8083646fae 100644 --- a/zig/kernels/metal/deepseek_v41/csa2.metal +++ b/zig/kernels/metal/deepseek_v41/csa2.metal @@ -100,3 +100,55 @@ kernel void dsv41_sparse_attention(device const ushort *q [[buffer(0)]], device 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); } + +// Each SIMD group scores different keys; every lane retains the predecessor's maxima and value-sum order. +[[max_total_threads_per_threadgroup(256)]] +kernel void dsv41_sparse_attention_parallel(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_simdgroup]], + uint group [[simdgroup_index_in_threadgroup]]) { + uint head = g.x, row = g.y; + if (head >= p.heads || row >= p.rows) return; + const uint groups = min(8u, p.width / 32); + threadgroup float lane_scores[64 * 32]; + float query[16], acc[2] = {0, 0}; + for (uint j = 0; j < p.width / 32; ++j) + query[j] = dsv41_bf(q[row * p.q_stride + head * p.width + lane + j * 32]); + 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) { + for (uint j = group; j < 64; j += groups) { + 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)); + lane_scores[j * 32 + lane] = key ? dsv41_dot_reduce(dot, dot_error) * p.scale : -INFINITY; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float next_max = mx; + for (uint j = 0; j < 64; ++j) next_max = max(next_max, lane_scores[j * 32 + lane]); + float correction = exp(mx - next_max); + float tile_sum = 0; + for (uint d = group; d < p.width / 32; d += groups) acc[d / groups] *= correction; + for (uint j = 0; j < 64; ++j) { + float probability = exp(lane_scores[j * 32] - 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 = group; d < p.width / 32; d += groups) + acc[d / groups] += 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); + } + float denom = sum + exp(sink[head] - mx); + for (uint d = group; d < p.width / 32; d += groups) + out[row * p.q_stride + head * p.width + lane + d * 32] = dsv41_round(acc[d / groups] / denom); +} diff --git a/zig/src/core/contract/dsv41_metal_records.zig b/zig/src/core/contract/dsv41_metal_records.zig new file mode 100644 index 0000000000..a7f7a2a6f0 --- /dev/null +++ b/zig/src/core/contract/dsv41_metal_records.zig @@ -0,0 +1,6 @@ +//! DeepSeek V4.1's sparse attention pipelines: declared geometry, M1 and M2 limits, the parallel launch's reservation. +const limits = @import("limits.zig"); +pub const records = [_]limits.MetalLaunch{ + .{ .name = "deepseek_v41 sparse dsv41_sparse_attention", .threads = 32, .tg_memory = 256, .m1 = 768, .m2 = 768 }, + .{ .name = "deepseek_v41 sparse dsv41_sparse_attention_parallel", .threads = 256, .tg_memory = 8192, .reserved = 256 }, +}; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..f727d81515 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -576,7 +576,7 @@ test { _ = paging; } /// The Metal collector keeps measured slots and the restored vision declarations. -pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records; +pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("dsv41_metal_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; diff --git a/zig/src/families/deepseek_v41/metal/csa2.zig b/zig/src/families/deepseek_v41/metal/csa2.zig index 395de9107a..0ae27d6d07 100644 --- a/zig/src/families/deepseek_v41/metal/csa2.zig +++ b/zig/src/families/deepseek_v41/metal/csa2.zig @@ -42,15 +42,22 @@ pub const Params = extern struct { }; pub const Pipelines = struct { attention: mtl.Pipeline, + attention_parallel: mtl.Pipeline, + parallel_enabled: bool = false, 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) }; + const attention = try mtl.Pipeline.init(device, lib, "dsv41_sparse_attention", false); + errdefer attention.deinit(); + const parallel = try mtl.Pipeline.init(device, lib, "dsv41_sparse_attention_parallel", false); + const mode = std.c.getenv("TF_DSV41_PARALLEL_ATTENTION"); + return .{ .attention = attention, .attention_parallel = parallel, .parallel_enabled = if (mode) |v| std.mem.eql(u8, std.mem.span(v), "1") else false }; } pub fn deinit(p: *Pipelines) void { p.attention.deinit(); + p.attention_parallel.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 { @@ -64,14 +71,18 @@ pub const Pipelines = struct { 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); + const groups: u32 = if (p.parallel_enabled) @min(8, args.width / 32) else 1; + const pipeline = if (p.parallel_enabled) p.attention_parallel else p.attention; + if (pipeline.maxThreads() < groups * 32) return error.InvalidSparseAttentionThreads; + enc.setPipeline(pipeline); 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)); + enc.dispatchGroups(mtl.Size.of(args.heads, args.rows, 1), mtl.Size.of(32 * groups, 1, 1)); } }; test { _ = @import("kv_test.zig"); + _ = @import("csa2_parallel_test.zig"); } diff --git a/zig/src/families/deepseek_v41/metal/csa2_parallel_test.zig b/zig/src/families/deepseek_v41/metal/csa2_parallel_test.zig new file mode 100644 index 0000000000..1fb16a6051 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/csa2_parallel_test.zig @@ -0,0 +1,160 @@ +//! The parallel key schedule against the original attention bytes, including per-lane compensated maxima. +const std = @import("std"); +const mtl = @import("metal"); +const csa = @import("csa2.zig"); +const View = @import("view.zig").View; + +fn bf(x: f32) u16 { + const b: u32 = @bitCast(x); + return @truncate((b +% (0x7fff + ((b >> 16) & 1))) >> 16); +} + +const Case = struct { width: u32, rows: u32, heads: u32, position: u32, ratio: u32, topk: u32, nonfinite: bool = false }; +const Fixture = struct { + buffers: [8]mtl.Buffer, + params: csa.Params, + bytes: usize, + + fn init(device: mtl.Device, spec: Case) !Fixture { + const width = spec.width; + const stride = spec.heads * width + 32; + const cache_stride = width + width / 32; + const main_stride = width / 2 + width / 16; + const index_stride = @max(1, spec.topk); + const bytes = @as(usize, spec.rows) * stride * 2; + const sizes = [_]usize{ bytes, 128 * cache_stride, spec.rows * cache_stride, 513 * main_stride, spec.rows * index_stride * 4, spec.heads * 4, bytes + 32, bytes + 32 }; + var buffers: [8]mtl.Buffer = undefined; + var made: usize = 0; + errdefer for (buffers[0..made]) |b| b.deinit(); + for (sizes, &buffers) |size, *b| { + b.* = try device.buffer(size, mtl.ResourceOptions.shared); + made += 1; + @memset(b.contents()[0..size], 0); + } + for (buffers[0].slice(u16, bytes / 2), 0..) |*v, i| + v.* = bf(@as(f32, @floatFromInt(@as(i32, @intCast((i * 37) % 67)) - 33)) / 32); + if (spec.nonfinite) { + buffers[0].slice(u16, bytes / 2)[0] = 0x7fc0; + buffers[0].slice(u16, bytes / 2)[1] = 0x7f80; + buffers[0].slice(u16, bytes / 2)[2] = 0xff80; + } else { + buffers[0].slice(u16, bytes / 2)[0] = bf(0x1p24); + buffers[0].slice(u16, bytes / 2)[1] = bf(-0x1p24); + buffers[0].slice(u16, bytes / 2)[32 % width] = bf(0x1p-6); + } + const fp8 = [_]u8{ 0x00, 0x38, 0xb8, 0x3c, 0xbc, 0x20, 0xa0, 0x48 }; + for ([_]usize{ 1, 2 }) |which| { + const count: usize = if (which == 1) 128 else spec.rows; + for (0..count) |row| { + for (0..width) |d| buffers[which].contents()[row * cache_stride + d] = fp8[(row * 11 + d * 3) % fp8.len]; + @memset(buffers[which].contents()[row * cache_stride + width ..][0 .. width / 32], 127); + } + } + for (0..513) |row| { + for (0..width / 2) |d| buffers[3].contents()[row * main_stride + d] = @truncate(row * 31 + d * 17); + @memset(buffers[3].contents()[row * main_stride + width / 2 ..][0 .. width / 16], 0x38); + } + for (buffers[4].slice(i32, spec.rows * index_stride), 0..) |*v, i| + v.* = if (i % 7 == 0) -1 else if (i % 11 == 0) 512 else @intCast((i * 73 + 13) % 513); + for (buffers[5].slice(f32, spec.heads), 0..) |*v, i| v.* = @as(f32, @floatFromInt(i % 5)) - 2; + for (buffers[6..]) |b| @memset(b.contents()[0 .. bytes + 32], 0xa7); + return .{ .buffers = buffers, .bytes = bytes, .params = .{ .rows = spec.rows, .heads = spec.heads, .width = width, .q_stride = stride, .position = spec.position, .chunk_rows = spec.rows, .ring_stride = cache_stride, .chunk_stride = cache_stride, .main_stride = main_stride, .ratio = spec.ratio, .compressed_rows = 513, .topk = spec.topk, .index_stride = index_stride } }; + } + + fn deinit(f: *Fixture) void { + for (f.buffers) |b| b.deinit(); + } + + fn inputHash(f: *const Fixture) [32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for (f.buffers[0..6]) |b| h.update(b.contents()[0..b.length()]); + return h.finalResult(); + } + + fn run(f: *Fixture, p: *csa.Pipelines, queue: mtl.Queue, parallel: bool) !f64 { + p.parallel_enabled = parallel; + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + var views: [7]View = undefined; + for (0..6) |i| views[i] = .{ .buffer = f.buffers[i] }; + views[6] = .{ .buffer = f.buffers[if (parallel) 7 else 6], .offset = 16 }; + try p.encode(e, views[0], views[1], views[2], views[3], views[4], views[5], views[6], f.params); + e.end(); + cb.commit(); + cb.wait(); + try std.testing.expect(cb.failure() == null); + return cb.gpuSeconds(); + } + + fn assertEqual(f: *const Fixture) !void { + try std.testing.expectEqualSlices(u8, f.buffers[6].contents()[0 .. f.bytes + 32], f.buffers[7].contents()[0 .. f.bytes + 32]); + for (f.buffers[6..]) |b| { + try std.testing.expect(std.mem.allEqual(u8, b.contents()[0..16], 0xa7)); + try std.testing.expect(std.mem.allEqual(u8, b.contents()[16 + f.bytes ..][0..16], 0xa7)); + for (0..f.params.rows) |row| { + const pad = 16 + (row * f.params.q_stride + f.params.heads * f.params.width) * 2; + try std.testing.expect(std.mem.allEqual(u8, b.contents()[pad..][0..64], 0xa7)); + } + } + } +}; + +test "parallel sparse attention retains predecessor bits, immutable inputs and guards at tile edges" { + 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 csa.Pipelines.load(device); + defer p.deinit(); + var count: usize = 0; + for ([_]u32{ 32, 96, 128, 224, 256, 288, 512 }) |width| { + for ([_]u32{ 0, 1, 63, 127, 128, 129, 8191 }) |position| { + const rows: u32 = if (position % 2 == 0) 3 else 1; + const topk = ([_]u32{ 0, 1, 63, 64, 65, 511, 512 })[count % 7]; + var f = try Fixture.init(device, .{ .width = width, .rows = rows, .heads = if (count % 3 == 0) 3 else 1, .position = position, .ratio = @intCast(count % 3), .topk = topk }); + defer f.deinit(); + const hash = f.inputHash(); + _ = try f.run(&p, queue, false); + for (0..5) |_| { + _ = try f.run(&p, queue, true); + try f.assertEqual(); + try std.testing.expectEqual(hash, f.inputHash()); + } + count += 1; + } + } + for ([_]bool{ false, true }) |nonfinite| { + var f = try Fixture.init(device, .{ .width = 512, .rows = 1, .heads = 64, .position = 1048575, .ratio = 2, .topk = 512, .nonfinite = nonfinite }); + defer f.deinit(); + _ = try f.run(&p, queue, false); + _ = try f.run(&p, queue, true); + try f.assertEqual(); + count += 1; + } + std.debug.print("ATTENTION_PARALLEL_BITS {d} shapes: predecessor bytes and guards equal\n", .{count}); +} + +test "parallel sparse attention reports isolated GPU timings after warmup" { + 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 csa.Pipelines.load(device); + defer p.deinit(); + for ([_]u32{ 128, 8191 }) |position| { + var f = try Fixture.init(device, .{ .width = 512, .rows = 1, .heads = 64, .position = position, .ratio = 2, .topk = 512 }); + defer f.deinit(); + for ([_]bool{ false, true }) |parallel| { + _ = try f.run(&p, queue, parallel); + var times: [32]f64 = undefined; + for (×) |*time| time.* = try f.run(&p, queue, parallel); + std.mem.sort(f64, ×, {}, std.sort.asc(f64)); + std.debug.print("ATTENTION_GPU {{\"parallel\":{s},\"position\":{d},\"samples\":32,\"median_seconds\":{d}}}\n", .{ if (parallel) "true" else "false", position, (times[15] + times[16]) / 2 }); + } + try f.assertEqual(); + } +} diff --git a/zig/src/families/deepseek_v41/metal/model.zig b/zig/src/families/deepseek_v41/metal/model.zig index 49243f0a80..f414e86ef0 100644 --- a/zig/src/families/deepseek_v41/metal/model.zig +++ b/zig/src/families/deepseek_v41/metal/model.zig @@ -8,6 +8,7 @@ const load = @import("load.zig"); const weights = @import("weights.zig"); const work = @import("workspace.zig"); pub const attention = @import("attention.zig"); +pub const csa2 = @import("csa2.zig"); const kv = @import("kv.zig"); pub const ops = @import("ops.zig"); pub const mhc = @import("mhc.zig"); diff --git a/zig/tests/contract/metal_dsv41_attention.zig b/zig/tests/contract/metal_dsv41_attention.zig new file mode 100644 index 0000000000..73f744c79a --- /dev/null +++ b/zig/tests/contract/metal_dsv41_attention.zig @@ -0,0 +1,195 @@ +//! The parallel sparse schedule: predecessor bytes, packed-cache FP64 formula, and uniform-score FP32 stress. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const c = contract.spec; +const g = contract.gates; +const csa = tf.deepseek_v41.csa2; +const View = tf.deepseek_v41_view.View; +const width = 512; +const ring_stride = 528; +const main_stride = 288; + +fn depth(s: c.Shape) usize { + return s.k + 32; +} +pub fn spec() c.Spec { + return .{ .name = "deepseek_v41 sparse dsv41_sparse_attention_parallel", .replaces = "dsv41_sparse_attention", .accum = .f32, .depth = depth, .ops = 32, .mids = 1, .out = .bf16, .tiles = .{ .{ .size = 16 }, .{}, .{ .size = 192, .min = 128 } }, .served = &.{ .{ .rows = 8, .n = width, .k = 640 }, .{ .rows = 256, .n = width, .k = 128 }, .{ .rows = 1, .n = width, .k = 191 }, .{ .rows = 1, .n = width, .k = 192 }, .{ .rows = 1, .n = width, .k = 193 } } }; +} + +pub const Probe = struct { + const vt = g.hooks(Probe); + a: std.mem.Allocator, + rig: rig_.Rig, + p: csa.Pipelines, + refs: [8]rig_.Ref = undefined, + shape: c.Shape = undefined, + params: csa.Params = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !Probe { + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .p = try csa.Pipelines.load(device) }; + } + pub fn deinit(self: *Probe) void { + self.p.deinit(); + self.rig.deinit(); + } + pub fn kernel(self: *Probe) g.Kernel { + return .{ .spec = spec(), .ctx = self, .vt = &vt }; + } + fn bytes(self: *Probe, i: usize) []u8 { + const sizes = [_]usize{ self.shape.rows * width * 2, 128 * ring_stride, self.shape.rows * ring_stride, 513 * main_stride, self.shape.rows * @max(1, self.shape.k - 128) * 4, 4, self.count() * 2, self.count() * 2 }; + return self.rig.bytes(self.refs[i], sizes[i]); + } + pub fn prepare(self: *Probe, s: c.Shape, fill: g.Fill, seed: u64) !bool { + if (s.n != width or s.k < 128 or s.k > 640 or s.rows == 0 or s.rows > 256) return false; + self.rig.reset(); + self.shape = s; + var arena = std.heap.ArenaAllocator.init(self.a); + defer arena.deinit(); + const a = arena.allocator(); + const sizes = [_]usize{ s.rows * width * 2, 128 * ring_stride, s.rows * ring_stride, 513 * main_stride, s.rows * @max(1, s.k - 128) * 4, 4 }; + var rng = std.Random.DefaultPrng.init(seed); + for (sizes, 0..) |size, i| { + const b = try a.alloc(u8, size); + @memset(b, 0); + if (i == 0) { + for (0..s.rows * width) |j| std.mem.writeInt(u16, b[j * 2 ..][0..2], c.toBf16(if (fill == .stress) 0 else (rng.random().float(f32) * 2 - 1)), .little); + if (fill == .nonfinite) { + const vals = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32) }; + for (vals, 0..) |v, row| std.mem.writeInt(u16, b[(row * width + 7) * 2 ..][0..2], c.toBf16(v), .little); + std.mem.writeInt(u16, b[(3 * width + 9) * 2 ..][0..2], c.toBf16(-std.math.inf(f32)), .little); + } + } else if (i == 1 or i == 2) { + const count_ = if (i == 1) 128 else s.rows; + for (0..count_) |row| { + const physical = if (i == 1) row else (8191 + row) % 128; + const code: u8 = if (fill == .stress) (if (physical < 56) @as(u8, @intCast(0x40 + physical % 8)) else if (physical < 64) 0x38 else if (physical < 120) @as(u8, @intCast(0xc0 + (physical - 64) % 8)) else 0) else @as(u8, @intCast(0x30 + rng.random().uintLessThan(u8, 16))); + @memset(b[row * ring_stride ..][0..width], code); + const exponent: u8 = if (fill == .stress and (physical < 56 or (physical >= 64 and physical < 120))) 139 else 127; + @memset(b[row * ring_stride + width ..][0..16], exponent); + } + } else if (i == 3 and fill != .stress) { + for (0..513) |row| { + for (0..256) |j| b[row * main_stride + j] = rng.random().int(u8); + @memset(b[row * main_stride + 256 ..][0..32], 0x38); + } + } else if (i == 4) { + for (0..size / 4) |j| std.mem.writeInt(i32, b[j * 4 ..][0..4], if (j % 7 == 0) -1 else @as(i32, @intCast((j * 73 + 13) % 513)), .little); + } + self.refs[i] = try self.rig.input(b); + } + self.refs[6] = try self.rig.output(self.count() * 2); + self.refs[7] = try self.rig.scratch(self.count() * 2); + self.params = .{ .rows = @intCast(s.rows), .heads = 1, .width = width, .q_stride = width, .position = 8191, .chunk_rows = @intCast(s.rows), .ratio = 2, .compressed_rows = 513, .topk = @intCast(s.k - 128), .index_stride = @intCast(@max(1, s.k - 128)) }; + return true; + } + fn launch(self: *Probe, parallel: bool, out: usize) !void { + self.p.parallel_enabled = parallel; + const pass = self.rig.begin(); + var v: [7]View = undefined; + for (0..6) |i| v[i] = .{ .buffer = self.refs[i].buf, .offset = self.refs[i].off }; + v[6] = .{ .buffer = self.refs[out].buf, .offset = self.refs[out].off }; + try self.p.encode(pass.e, v[0], v[1], v[2], v[3], v[4], v[5], v[6], self.params); + try self.rig.finish(pass); + } + pub fn run(self: *Probe) !void { + try self.launch(true, 6); + } + pub fn replaced(self: *Probe) ![]const u8 { + try self.launch(false, 7); + return self.bytes(7); + } + pub fn count(self: *Probe) usize { + return self.shape.rows * width; + } + + fn fp8(code: u8) f64 { + const sign: f64 = if (code & 128 != 0) -1 else 1; + const e = (code >> 3) & 15; + const m = code & 7; + if (e == 0) return sign * @as(f64, @floatFromInt(m)) * 0x1p-9; + if (e == 15 and m == 7) return std.math.nan(f64); + return sign * (1 + @as(f64, @floatFromInt(m)) / 8) * std.math.pow(f64, 2, @as(f64, @floatFromInt(e)) - 7); + } + fn key(self: *Probe, row: usize, slot: usize, dim: usize) ?f64 { + if (slot < 128) { + const position = 8191 + row + 1 - 128 + slot; + const which: usize = if (position >= 8191) 2 else 1; + const at = (if (which == 2) position - 8191 else position % 128) * ring_stride; + const b = self.bytes(which); + return fp8(b[at + dim]) * std.math.pow(f64, 2, @as(f64, @floatFromInt(b[at + width + dim / 32])) - 127); + } + const b = self.bytes(4); + const off = (row * @max(1, self.shape.k - 128) + slot - 128) * 4; + const index = std.mem.readInt(i32, b[off..][0..4], .little); + if (index < 0 or index >= 513) return null; + const at = @as(usize, @intCast(index)) * main_stride; + const main = self.bytes(3); + const code = (main[at + dim / 2] >> @as(u3, @intCast((dim % 2) * 4))) & 15; + const lut = [_]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + return (if (code & 8 != 0) -lut[code & 7] else lut[code & 7]) * fp8(main[at + 256 + dim / 16]); + } + + pub fn reference(self: *Probe, want: []f64, mag: []f64) !void { + @setFloatMode(.strict); + const query = self.bytes(0); + for (0..self.shape.rows) |row| { + var acc: [width]f64 = @splat(0); + var magnitude: [width]f64 = @splat(0); + var high: f64 = -1e30; + var denominator: f64 = 0; + var base: usize = 0; + while (base < self.shape.k) : (base += 64) { + var scores: [64]f64 = @splat(-std.math.inf(f64)); + var next = high; + for (0..@min(64, self.shape.k - base)) |j| { + if (self.key(row, base + j, 0) == null) continue; + var dot: f64 = 0; + for (0..width) |d| { + const q = c.fromBf16(std.mem.readInt(u16, query[(row * width + d) * 2 ..][0..2], .little)); + dot += @as(f64, q) * self.key(row, base + j, d).?; + } + scores[j] = dot / @sqrt(@as(f64, width)); + if (!std.math.isNan(scores[j])) next = @max(next, scores[j]); + } + const correction = @exp(high - next); + for (&acc, &magnitude) |*v, *m| { + v.* *= correction; + m.* *= correction; + } + denominator *= correction; + for (scores, 0..) |score, j| { + const probability = @exp(score - next); + denominator += probability; + if (base + j >= self.shape.k or self.key(row, base + j, 0) == null) continue; + const rounded: f64 = c.fromBf16(c.toBf16(@floatCast(probability))); + for (0..width) |d| { + const term = rounded * self.key(row, base + j, d).?; + acc[d] += term; + magnitude[d] += @abs(term); + } + } + high = next; + } + denominator += @exp(-high); + for (0..width) |d| { + want[row * width + d] = acc[d] / denominator; + mag[row * width + d] = magnitude[d] / denominator; + } + } + } + pub fn simulate(self: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void { + var terms: [640]f64 = undefined; + for (0..self.shape.rows) |row| for (0..width) |d| { + var live: usize = 0; + for (terms[0..self.shape.k], 0..) |*v, j| { + const value = self.key(row, j, d); + v.* = value orelse 0; + live += @intFromBool(value != null); + } + out[row * width + d] = contract.rows.sumIn(acc, order, terms[0..self.shape.k]) / @as(f64, @floatFromInt(live + 1)); + }; + } +}; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index 821404d2f0..ac1cae9950 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -11,7 +11,7 @@ const limits = contract.limits; const registry = contract.registry; /// A registered pipeline as built here: its declaration, function, this GPU's source and the pre-M5 source. -pub const Launch = struct { decl: limits.MetalLaunch, function: []const u8, here: []const u8, pre_m5: []const u8 }; +pub const Launch = struct { decl: limits.MetalLaunch, function: []const u8, here: []const u8, pre_m5: []const u8, precise: bool = false }; /// The pre-M5 text of a source that includes nax.h (frags.source's branch for GPUs without tensor units). fn preM5(a: std.mem.Allocator, head: []const u8, text: []const u8) ![]u8 { @@ -36,6 +36,10 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La text = tf.flashnext_vision_contract.source(w); return .{ .decl = d, .function = fun, .here = text, .pre_m5 = text }; } + if (std.mem.startsWith(u8, d.name, "deepseek_v41 sparse ")) { + const body = sources.deepseek_v41_kv ++ sources.deepseek_v41_csa2; + return .{ .decl = d, .function = fun, .here = body, .pre_m5 = body, .precise = true }; + } if (std.mem.startsWith(u8, d.name, "drama3 ")) { const text_d3 = try @import("metal_drama3.zig").source(a, d.name); return .{ .decl = d, .function = fun, .here = text_d3, .pre_m5 = text_d3 }; @@ -88,9 +92,18 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La return .{ .decl = d, .function = fun, .here = here, .pre_m5 = pre }; } +fn pipeline(device: mtl.Device, l: Launch) !mtl.Pipeline { + if (!l.precise) return rig_.pipeline(device, l.here, l.function); + var options = mtl.CompileOptions.mlx(); + options.math_functions = .precise; + const lib = try mtl.Library.fromSource(device, l.here, options); + defer lib.deinit(); + return mtl.Pipeline.init(device, lib, l.function, false); +} + /// One launch on this GPU: the pipeline's own thread limit and threadgroup bytes against its declaration. pub fn onThisGpu(a: std.mem.Allocator, device: mtl.Device, l: Launch, control: bool) !usize { - const p = rig_.pipeline(device, l.here, l.function) catch |err| { + const p = pipeline(device, l) catch |err| { std.debug.print("{s} {s}: no pipeline on this GPU ({t})\n", .{ if (control) "RED " else "FAIL", l.decl.name, err }); return 1; }; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 0415713ac1..f08d3b4590 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -11,6 +11,7 @@ const limits_k = @import("metal_limits.zig"); const h3_k = @import("metal_h3.zig"); const h3_topk = @import("metal_h3_topk.zig"); const drama3_k = @import("metal_drama3.zig"); +const dsv41_attention = @import("metal_dsv41_attention.zig"); const flash_pool = @import("metal_flashnext_pool.zig"); const flash = @import("metal_flashnext.zig"); @@ -180,6 +181,12 @@ pub fn main(init: std.process.Init) !u8 { try kernels.append(a, flash_list[flash_made - 1].kernel()); } }; + var dsv41: ?dsv41_attention.Probe = null; + defer if (dsv41) |*probe| probe.deinit(); + if (wanted(o, dsv41_attention.spec().name)) { + dsv41 = try dsv41_attention.Probe.init(a, device, queue); + try kernels.append(a, dsv41.?.kernel()); + } var d3norm: ?drama3_k.Norm = null; defer if (d3norm) |*n| n.deinit(); if (wanted(o, drama3_k.spec().name)) { diff --git a/zig/tests/contract/metal_offline.zig b/zig/tests/contract/metal_offline.zig index 0bf2a1854f..3be65920dc 100644 --- a/zig/tests/contract/metal_offline.zig +++ b/zig/tests/contract/metal_offline.zig @@ -28,16 +28,17 @@ fn exec(x: Ctx, argv: []const []const u8) !void { } /// The source compiled as MLX compiles it at run time, linked into a metallib named by its hash. -fn library(x: Ctx, source: []const u8) ![]const u8 { +fn library(x: Ctx, source: []const u8, precise: bool) ![]const u8 { var h = std.hash.Wyhash.init(0); h.update(source); + h.update(if (precise) "precise" else "fast"); const stem = try std.fmt.allocPrint(x.a, "{s}/src_{x}", .{ x.dir, h.final() }); const lib = try std.fmt.allocPrint(x.a, "{s}.metallib", .{stem}); if (std.Io.Dir.cwd().access(x.io, lib, .{})) |_| return lib else |_| {} const src = try std.fmt.allocPrint(x.a, "{s}.metal", .{stem}); const air = try std.fmt.allocPrint(x.a, "{s}.air", .{stem}); try std.Io.Dir.cwd().writeFile(x.io, .{ .sub_path = src, .data = source }); - try exec(x, &.{ "xcrun", "-sdk", "macosx", "metal", std_flag(), "-fmetal-math-mode=safe", "-fmetal-math-fp32-functions=fast", "-c", src, "-o", air }); + try exec(x, &.{ "xcrun", "-sdk", "macosx", "metal", std_flag(), "-fmetal-math-mode=safe", if (precise) "-fmetal-math-fp32-functions=precise" else "-fmetal-math-fp32-functions=fast", "-c", src, "-o", air }); try exec(x, &.{ "xcrun", "-sdk", "macosx", "metallib", air, "-o", lib }); return lib; } @@ -112,7 +113,7 @@ pub fn run(a: std.mem.Allocator, io: std.Io, dir: []const u8, list: []const limi const control = std.mem.startsWith(u8, d.name, "control"); const safe = d.adapts or d.reserved >= d.threads or (!d.offline and d.threads <= 256); const got = blk: { - const lib = library(x, l.pre_m5) catch |err| break :blk err; + const lib = library(x, l.pre_m5, l.precise) catch |err| break :blk err; const m1 = familyLimit(x, lib, l.function, &m1_archs) catch |err| break :blk err; const m2 = familyLimit(x, lib, l.function, &m2_archs) catch |err| break :blk err; break :blk [2]u32{ m1, m2 };