DeepSeek V4.1: sparse attention on eight SIMD groups, the same bytes

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,

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

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

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;

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

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 (&times) |*time| time.* = try f.run(&p, queue, parallel);
+            std.mem.sort(f64, &times, {}, 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();
+    }
+}

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