DeepSeek V4.1: the indexer's merge and finish on a threadgroup a row, the same picks; a decode row's indexer 13-34x faster

docs/kernel-contract.md

@@ -102,8 +102,9 @@ MLX's formula keeps the infinity, because its codes enter the MMAs as q + 128 an
 Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), affine_mm (4 and 8 bits), simd_qmm
 (4 bits) and simd_qmm_bits (3, 6 and 8 bits), and the lane projection (6-bit groups of 32 on every chip; 4- and 8-bit
 formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's
-prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of
-tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10.
+prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles), and DeepSeek V4.1's index
+merge (its picks against a host oracle as sets) and finish (bit-equal to the serial finish). Gate 5 also holds the
+rest of tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10.
 CUDA: the RMS norm rows (16-byte form and the fp32-out tree), Nemotron's router logits (and their 128-row tiles,
 bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and
 shared add over bf16 share sums (bit-equal to `tf_super_combine` and `tf_super_moe_out` over the fp32 shares), the

zig/deepseek_v41_checks.zig

@@ -6,6 +6,7 @@ test {
     _ = @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/index_select_test.zig");
     _ = @import("src/families/deepseek_v41/config.zig");
     _ = @import("src/families/deepseek_v41/profiler.zig");
     _ = @import("src/core/gpu_counters.zig");

zig/kernels/metal/deepseek_v41/index.metal

@@ -60,7 +60,8 @@ kernel void dsv41_index_clear(device Dsv41Pick *top [[buffer(0)]], device Dsv41P
     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)]],
+// The serial merge and finish (a thread a row): the parallel ones' baseline in the tests and the kernel contract.
+kernel void dsv41_index_merge_serial(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;
@@ -93,11 +94,161 @@ inline void dsv41_sort_positions(device Dsv41Pick *p, uint size) {
         p[j] = value;
     }
 }
-kernel void dsv41_index_finish(device Dsv41Pick *top [[buffer(0)]], device Dsv41Pick *candidate [[buffer(1)]],
+kernel void dsv41_index_finish_serial(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);
 }
+// A row's merge and finish on a threadgroup, the serial kernels' picks and bytes; a key ranks as dsv41_less does.
+inline ulong dsv41_key(float score, int position) {
+    uint u = as_type<uint>(score == 0.0f ? 0.0f : score);
+    u = (u & 0x80000000u) ? ~u : (u | 0x80000000u);
+    return (ulong(u) << 32) | ulong(~uint(position));
+}
+constant uint dsv41_lanes = 1024;
+// The k-th largest of the group's keys (each thread holding N, `live` marking real ones): eight 8-bit passes.
+template <uint N>
+inline ulong dsv41_kth(thread const ulong (&keys)[N], thread const bool (&live)[N], uint k, uint t,
+                       threadgroup atomic_uint *hist, threadgroup uint *pick) {
+    ulong prefix = 0;
+    for (int shift = 56; shift >= 0; shift -= 8) {
+        for (uint b = t; b < 256; b += dsv41_lanes) atomic_store_explicit(&hist[b], 0, memory_order_relaxed);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        const ulong high = shift == 56 ? 0 : prefix >> (shift + 8);
+        for (uint i = 0; i < N; ++i)
+            if (live[i] && (shift == 56 || (keys[i] >> (shift + 8)) == high))
+                atomic_fetch_add_explicit(&hist[uint(keys[i] >> shift) & 255], 1, memory_order_relaxed);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (t == 0) {
+            uint left = k, b = 255;
+            for (;; --b) {
+                const uint c = atomic_load_explicit(&hist[b], memory_order_relaxed);
+                if (c >= left || b == 0) break;
+                left -= c;
+            }
+            pick[0] = b;
+            pick[1] = left;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        prefix |= ulong(pick[0]) << shift;
+        k = pick[1];
+    }
+    return prefix;
+}
+// Keeps the best `k` of the group's live picks in out[0, k) in thread order, the rest of out as {-inf, -1}.
+template <uint N>
+inline void dsv41_keep(thread const ulong (&keys)[N], thread const bool (&live)[N], thread const Dsv41Pick (&picks)[N],
+                       uint k, device Dsv41Pick *out, uint t, threadgroup atomic_uint *hist, threadgroup uint *pick,
+                       threadgroup uint *sums) {
+    for (uint b = t; b < 256; b += dsv41_lanes) atomic_store_explicit(&hist[b], 0, memory_order_relaxed);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    uint mine = 0;
+    for (uint i = 0; i < N; ++i) mine += uint(live[i]);
+    atomic_fetch_add_explicit(&hist[0], mine, memory_order_relaxed);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    const uint count = atomic_load_explicit(&hist[0], memory_order_relaxed);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    const ulong least = count > k ? dsv41_kth(keys, live, k, t, hist, pick) : 0;
+    // each thread's kept picks after those of the threads before it (out may be one of the inputs: all have read it)
+    uint keep = 0;
+    for (uint i = 0; i < N; ++i) keep += uint(live[i] && keys[i] >= least);
+    const uint lane = t % 32, sg = t / 32;
+    uint at = simd_prefix_exclusive_sum(keep);
+    if (lane == 31) sums[sg] = at + keep;
+    threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup);
+    for (uint g = 0; g < sg; ++g) at += sums[g];
+    for (uint i = 0; i < N; ++i)
+        if (live[i] && keys[i] >= least) out[at++] = picks[i];
+    uint kept = 0;
+    for (uint g = 0; g < dsv41_lanes / 32; ++g) kept += sums[g];
+    for (uint i = kept + t; i < k; i += dsv41_lanes) out[i] = {-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 [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]]) {
+    threadgroup atomic_uint hist[256];
+    threadgroup uint pick[2];
+    threadgroup uint sums[dsv41_lanes / 32];
+    if (row >= p.rows) return;
+    {   // the row's 512 held picks and this tile's 4,096 scores (finite ones only, as the heap takes them)
+        constexpr uint N = (512 + 4096 + dsv41_lanes - 1) / dsv41_lanes;
+        ulong keys[N];
+        bool live[N];
+        Dsv41Pick picks[N];
+        for (uint i = 0; i < N; ++i) {
+            const uint at = t + i * dsv41_lanes;
+            live[i] = false;
+            if (at < p.topk) {
+                picks[i] = top[row * p.topk + at];
+                live[i] = isfinite(picks[i].score);
+            } else if (at - p.topk < p.tile_width) {
+                const uint local = at - p.topk;
+                picks[i] = {scores[row * p.tile_width + local], dsv41_index_position(local, row, candidate, p)};
+                live[i] = isfinite(picks[i].score);
+            }
+            keys[i] = live[i] ? dsv41_key(picks[i].score, picks[i].position) : 0;
+        }
+        dsv41_keep(keys, live, picks, p.topk, top + row * p.topk, t, hist, pick, sums);
+    }
+    if (!p.make_candidates) return;
+    threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup);
+    {   // the row's 2,048 held blocks and this tile's 512 block maxima, the newest visible block pinned
+        constexpr uint N = (2048 + 512 + dsv41_lanes - 1) / dsv41_lanes;
+        const uint visible = min(p.keys, (p.position + row + 1) / p.ratio);
+        ulong keys[N];
+        bool live[N];
+        Dsv41Pick picks[N];
+        for (uint i = 0; i < N; ++i) {
+            const uint at = t + i * dsv41_lanes;
+            live[i] = false;
+            if (at < p.candidate_topk) {
+                picks[i] = candidate[row * p.candidate_topk + at];
+                live[i] = picks[i].score > -INFINITY;
+            } else if ((at - p.candidate_topk) * p.block_size < p.tile_width) {
+                const uint base = (at - p.candidate_topk) * 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]);
+                const uint block = (p.tile_start + base) / p.block_size;
+                if (visible && block == (visible - 1) / p.block_size) score = INFINITY;
+                picks[i] = {score, int(block)};
+                live[i] = score > -INFINITY;
+            }
+            keys[i] = live[i] ? dsv41_key(picks[i].score, picks[i].position) : 0;
+        }
+        dsv41_keep(keys, live, picks, p.candidate_topk, candidate + row * p.candidate_topk, t, hist, pick, sums);
+    }
+}
+// Picks by position ascending, the -1 placeholders last (dsv41_sort_positions' order): a bitonic sort in groups.
+inline void dsv41_sort_group(device Dsv41Pick *p, uint size, uint t, threadgroup Dsv41Pick *tmp) {
+    for (uint i = t; i < size; i += dsv41_lanes) tmp[i] = p[i];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    for (uint k = 2; k <= size; k <<= 1)
+        for (uint j = k >> 1; j > 0; j >>= 1) {
+            for (uint i = t; i < size; i += dsv41_lanes) {
+                const uint l = i ^ j;
+                if (l > i) {
+                    const uint a = uint(tmp[i].position), b = uint(tmp[l].position);  // -1 is the largest
+                    if ((a > b) == ((i & k) == 0)) {
+                        const Dsv41Pick x = tmp[i];
+                        tmp[i] = tmp[l];
+                        tmp[l] = x;
+                    }
+                }
+            }
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+    for (uint i = t; i < size; i += dsv41_lanes) p[i] = tmp[i];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+}
+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 [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]]) {
+    threadgroup Dsv41Pick tmp[2048];
+    if (row >= p.rows) return;
+    dsv41_sort_group(top + row * p.topk, p.topk, t, tmp);
+    for (uint i = t; i < p.topk; i += dsv41_lanes) positions[row * p.topk + i] = top[row * p.topk + i].position;
+    if (p.make_candidates) dsv41_sort_group(candidate + row * p.candidate_topk, p.candidate_topk, t, tmp);
+}

zig/src/core/contract/registry.zig

@@ -397,6 +397,8 @@ pub fn checkCuda(a: std.mem.Allocator, list: []const CudaRecord, quiet: bool) !u
 
 /// Metal pipelines as the engine launches them; tg_memory is the most any chip reflects, m1 and m2 the compiled limits.
 pub const metal = [_]M{
+    .{ .name = "deepseek_v41 index merge dsv41_index_merge", .threads = 1024, .tg_memory = 1168, .m1 = 1024, .m2 = 1024 },
+    .{ .name = "deepseek_v41 index finish dsv41_index_finish", .threads = 1024, .tg_memory = 16384, .m1 = 1024, .m2 = 1024 },
     .{ .name = "softmax tf_softmax_bf16", .threads = 1024, .tg_memory = 256, .m1 = 1024, .m2 = 1024 },
     .{ .name = "rms tf_rms_norm_bf16", .threads = 1024, .tg_memory = 144, .m1 = 1024, .m2 = 1024 },
     .{ .name = "top-k tf_bf16_topk", .threads = 256, .tg_memory = 4096, .m1 = 768, .m2 = 768 },

zig/src/families/deepseek_v41/csa2_index.zig

@@ -5,6 +5,8 @@ const mtl = @import("metal");
 const sources = @import("kernel_sources");
 const View = @import("view.zig").View;
 pub const score_tile_width = 4096;
+/// Threads a row's merge and finish take (index.metal's dsv41_lanes).
+pub const select_threads = 1024;
 pub const Params = extern struct {
     rows: u32,
     heads: u32 = 32,
@@ -126,14 +128,14 @@ pub const Pipelines = struct {
             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.dispatchGroups(mtl.Size.of(args.rows, 1, 1), mtl.Size.of(select_threads, 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.dispatchGroups(mtl.Size.of(args.rows, 1, 1), mtl.Size.of(select_threads, 1, 1));
         enc.barrier();
     }
 };

zig/src/families/deepseek_v41/index_select_test.zig

@@ -0,0 +1,69 @@
+//! The indexer's merges and finish a threadgroup a row against the serial ones: the same heaps, candidates and picks.
+const std = @import("std");
+const mtl = @import("metal");
+const sources = @import("kernel_sources");
+const index = @import("csa2_index.zig");
+const ts = @import("test_support.zig");
+
+/// The indexer's pass twice, parallel and serial, on the same inputs: every selection buffer byte for byte.
+fn same(device: mtl.Device, queue: mtl.Queue, parallel: *const index.Pipelines, serial: *const index.Pipelines, rows: u32, keys: u32, ties: bool, seed: u64) !void {
+    var b: ts.Buffers = .{ .device = device };
+    defer b.deinit();
+    var prng = std.Random.DefaultPrng.init(seed);
+    const r = prng.random();
+    const q = try b.get(u16, @as(usize, rows) * 32 * 128);
+    const w = try b.get(u16, @as(usize, rows) * 32);
+    const k = try b.get(u8, @as(usize, keys) * 68);
+    for (q.slice(u16, @as(usize, rows) * 32 * 128)) |*x| x.* = ts.bf(if (ties) @as(f32, @floatFromInt(r.uintLessThan(u8, 3))) - 1 else (r.float(f32) - 0.5) * 2);
+    for (w.slice(u16, @as(usize, rows) * 32)) |*x| x.* = ts.bf(if (ties) 1 else (r.float(f32) - 0.5) * 2);
+    for (0..keys) |i| {
+        const key = k.contents()[i * 68 ..][0..68];
+        if (ties) @memset(key[0..64], @intCast(0x11 * (i % 3))) else r.bytes(key[0..64]);
+        for (key[64..]) |*e| e.* = 125 + r.uintLessThan(u8, 4);
+    }
+    var scratch: [2]index.Scratch = .{ try index.Scratch.init(device, rows), try index.Scratch.init(device, rows) };
+    defer for (&scratch) |*x| x.deinit();
+    // the source layer's pass makes the candidates, the next one scores them
+    for ([_][2]u32{ .{ 1, 0 }, .{ 0, 1 } }) |mode| {
+        for ([_]*const index.Pipelines{ parallel, serial }, &scratch) |p, sc| {
+            const cb = queue.commandBuffer();
+            const enc = cb.compute(.serial);
+            try p.encode(enc, ts.view(q), ts.view(w), ts.view(k), sc, .{ .rows = rows, .q_stride = 32 * 128, .weight_stride = 32, .position = keys * 2 - rows, .ratio = 2, .keys = keys, .make_candidates = mode[0], .use_candidates = mode[1] });
+            try ts.done(cb, enc);
+        }
+        const n = index.Scratch.bytes(rows);
+        const scores = @as(usize, rows) * index.score_tile_width * 4; // the last tile's raw scores: work, not a result
+        const got = scratch[0].buffer.contents()[scores..n];
+        const want = scratch[1].buffer.contents()[scores..n];
+        if (!std.mem.eql(u8, want, got)) {
+            std.debug.print("rows {d} keys {d} ties {} mode {any}: the selections differ\n", .{ rows, keys, ties, mode });
+            return error.ParallelSelectionDiffers;
+        }
+    }
+}
+
+test "the parallel merges and finish give the serial ones' heaps, candidates and picks: tiles, ties, few keys" {
+    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 parallel = try index.Pipelines.load(device);
+    defer parallel.deinit();
+    var serial = try index.Pipelines.load(device);
+    defer serial.deinit();
+    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();
+    serial.merge.deinit();
+    serial.merge = try mtl.Pipeline.init(device, lib, "dsv41_index_merge_serial", false);
+    serial.finish.deinit();
+    serial.finish = try mtl.Pipeline.init(device, lib, "dsv41_index_finish_serial", false);
+    var seed: u64 = 1;
+    for ([_]u32{ 1, 3, 17 }) |rows| for ([_]u32{ 37, 600, index.score_tile_width + 5, 20000 }) |keys| for ([_]bool{ false, true }) |ties| {
+        try same(device, queue, &parallel, &serial, rows, keys, ties, seed);
+        seed += 1;
+    };
+}