From ea1b0c8bf6a26a431c831529b2907a99edf6d1ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= <1188270+cwschroeder@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:49:15 +0200 Subject: [PATCH 1/2] DeepSeek V4.1: sparse attention scores a tile's keys across eight SIMD groups, the same bytes dsv41_sparse_attention runs one SIMD group per head and row: it scores the 64 keys of each tile one after another, then accumulates all sixteen 32-wide output columns. dsv41_sparse_attention_parallel launches eight SIMD groups (256 threads). They score different keys of a tile into threadgroup memory; every lane then reads the tile's scores in key order, and each group accumulates two of the sixteen output columns over the same keys. Each dot product, running maximum, tile sum and value sum keeps the predecessor's order, so the outputs are the same bytes. TF_DSV41_PARALLEL_ATTENTION=1 selects it; the default stays the predecessor. Encoding refuses a pipeline that cannot take the launch. tf-dsv41-run records the choice as attention_parallel. csa2_parallel_test compares both kernels byte for byte over 51 shapes, including tile edges, nonfinite inputs and the guards around each output. Co-Authored-By: Claude Opus 5.5 --- zig/deepseek_v41_run.zig | 1 + zig/kernels/metal/deepseek_v41/csa2.metal | 52 ++++++ zig/src/families/deepseek_v41/metal/csa2.zig | 17 +- .../deepseek_v41/metal/csa2_parallel_test.zig | 160 ++++++++++++++++++ 4 files changed, 227 insertions(+), 3 deletions(-) create mode 100644 zig/src/families/deepseek_v41/metal/csa2_parallel_test.zig 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/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(); + } +} From ec1f47f2afc94f7b702fbdcd85cf6813f503ea0f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= <1188270+cwschroeder@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:49:16 +0200 Subject: [PATCH 2/2] kernel contract: DeepSeek V4.1's sparse attention, the parallel launch bit-equal to the one it replaces metal_dsv41_attention.zig runs dsv41_sparse_attention_parallel through gates 1 to 4: nonfinite classes against float64, five byte-identical repeats with untouched guards, 14 shapes bit-equal to dsv41_sparse_attention, and exact sums at k 640. Its launch declares [[max_total_threads_per_threadgroup(256)]] and is registered with that reservation; the predecessor is registered with its 32 threads and the M1 and M2 limits of Xcode's offline compiler. The records live in dsv41_metal_records.zig and join all_metal, as the DeltaNet slots and Flash Next vision records do, so registry.zig stays at its line count. The family compiles these kernels with precise math functions; the limits check and the offline sweep now build them the same way. Co-Authored-By: Claude Opus 5.5 --- zig/src/core/contract/dsv41_metal_records.zig | 6 + zig/src/core/contract/registry.zig | 2 +- zig/src/families/deepseek_v41/metal/model.zig | 1 + zig/tests/contract/metal_dsv41_attention.zig | 195 ++++++++++++++++++ zig/tests/contract/metal_limits.zig | 17 +- zig/tests/contract/metal_main.zig | 7 + zig/tests/contract/metal_offline.zig | 7 +- 7 files changed, 229 insertions(+), 6 deletions(-) create mode 100644 zig/src/core/contract/dsv41_metal_records.zig create mode 100644 zig/tests/contract/metal_dsv41_attention.zig 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/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 };