EXL3: share decoded weights across prefill rows

RUNBOOK.md

@@ -52,6 +52,12 @@ candidate output bit against the arithmetic decoder, but the independent FP64
 oracle samples at most two rows and three 128-column blocks. Its worst-column
 metrics are explicitly labelled `worst_sampled`; they are not whole-matrix bounds.
 
+Plain prefill windows of at least 32 rows share decoded weight tiles across four
+SIMDgroups when the pipeline supports 128 threads. Smaller windows, filtered
+routes and indexed experts keep the eight-row kernel. If the prefill pipeline
+is unavailable, the operator uses that existing path. The benchmark's arithmetic
+reference disables the optional prefill pipeline to keep its comparison valid.
+
 The borrowed-buffer API in `zig/src/core/exl3_projection.zig` encodes BF16
 inputs, FP16-rounded rotated inputs, compensated FP32 projection dots and BF16
 or F32 outputs. Scratch holds row intermediates, never a dequantized matrix.

zig/exl3_projection_bench.zig

@@ -319,6 +319,9 @@ pub fn main(init: std.process.Init) !void {
     const baseline = try mtl.Pipeline.init(device, lib, "tf_exl3_matmul", false);
     p[0].matmul.deinit();
     p[0].matmul = baseline;
+    // Keep the arithmetic reference independent of the optional prefill kernel.
+    if (p[0].prefill) |pf| pf.deinit();
+    p[0].prefill = null;
     try checkCodebooks(device, queue);
     for ([_]ex.format.Codebook{ .inst3, .mcg, .mul1 }) |book| {
         for ([_]u16{ 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16 }) |width| {

zig/kernels/metal/core/exl3_projection.metal

@@ -48,9 +48,9 @@ inline ushort tf_exl3_state(const device ushort* tile, uint halves, uint p) {
 inline ushort tf_exl3_mapped_state(const device ushort* tile, uint map) {
   const uint wi = map & 63u, off = (map >> 6) & 31u;
   const uint lo = tf_exl3_word(tile, wi);
-  if (off <= 16) return ushort(lo >> (16 - off));
   const uint hi = tf_exl3_word(tile, map >> 11);
-  return ushort((lo << (off - 16)) | (hi >> (48 - off)));
+  const uint carry = (hi >> ((32 - off) & 31u)) & (0u - uint(off != 0));
+  return ushort(((lo << off) | carry) >> 16);
 }
 
 // Both operations are exact in FP32 before the single FP16 round; do not split the mul1 FMA.
@@ -228,3 +228,75 @@ kernel void tf_exl3_rotate_out(const device float* sums [[buffer(0)]],
     else reinterpret_cast<device bfloat*>(y)[at] = bfloat(out);
   }
 }
+
+// Qualification only: compare mapped windows with independent CPU 64-bit concatenation.
+
+kernel void tf_exl3_state_window_check(const device uint4* inputs [[buffer(0)]],
+                              device ushort* outputs [[buffer(1)]],
+                              constant uint& count [[buffer(2)]],
+                              uint tid [[thread_position_in_grid]]) {
+  if (tid >= count) return;
+  const device ushort* tile = reinterpret_cast<const device ushort*>(inputs + tid);
+  const uint off = inputs[tid].z;
+  outputs[tid] = tf_exl3_mapped_state(tile, (off << 6) | (1u << 11));
+}
+
+// Four SIMDgroups share decoded weights while each retains its original eight-row compensated dots.
+kernel void tf_exl3_matmul_prefill(const device half* x [[buffer(0)]],
+                                  const device uchar* trellis [[buffer(1)]],
+                                  device float* sums [[buffer(2)]],
+                                  constant TfExl3Args& a [[buffer(3)]],
+                                  const device uint* routes [[buffer(4)]],
+                                  const device TfExl3Expert* descriptors [[buffer(5)]],
+                                  uint lane [[thread_index_in_simdgroup]],
+                                  uint sg [[simdgroup_index_in_threadgroup]],
+                                  uint3 tg [[threadgroup_position_in_grid]]) {
+  threadgroup float wx[256], xx[512];
+  const uint rb = tg.y * 32, nb = tg.x * 16;
+  if (a.indexed || a.routed || rb >= a.rows) return;
+  uint maps[8];
+  if (sg == 0) {
+    for (uint j = 0; j < 8; ++j) {
+      const uint p = lane + 32 * j, ring = 128 * a.halves;
+      const uint end = ((p + 1) * a.halves - ((a.halves & 1u) ? ((p + 1) & 1u) : 0u)) / 2;
+      const uint start = (end + ring - 16) % ring;
+      const uint wi = start / 32, off = start % 32, next = wi + 1 == 4 * a.halves ? 0 : wi + 1;
+      maps[j] = wi | (off << 6) | (next << 11);
+    }
+  }
+  float high[4] = {0,0,0,0}, low[4] = {0,0,0,0};
+  for (uint kb = 0; kb < a.k; kb += 16) {
+    if (sg == 0) {
+      const device ushort* tile = reinterpret_cast<const device ushort*>(
+          trellis + ulong(kb / 16) * a.tile_k_stride + ulong(nb / 16) * a.tile_n_stride);
+      for (uint j = 0; j < 8; ++j) {
+        const uint p = lane + 32 * j, l = p / 8, e = p % 8;
+        const uint r = 2 * (l % 4) + (e & 1u) + 8 * ((e >> 1) & 1u);
+        const uint c = l / 4 + 8 * (e >> 2);
+        wx[r * 16 + c] = float(tf_exl3_value(tf_exl3_mapped_state(tile, maps[j]), a.codebook));
+      }
+    }
+    for (uint j = lane; j < 128; j += 32) {
+      const uint row = rb + sg * 8 + j / 16;
+      xx[sg * 128 + j] = tf_exl3_active(routes, a, row)
+          ? float(x[ulong(row) * a.k + kb + j % 16]) : 0.0f;
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    for (uint i = 0; i < 4; ++i) {
+      const uint at = lane + 32 * i, row = rb + sg * 8 + at / 16, col = at % 16;
+      if (!tf_exl3_active(routes, a, row)) continue;
+      for (uint k = 0; k < 16; ++k) {
+        const float product = xx[sg * 128 + (at / 16) * 16 + k] * wx[k * 16 + col];
+        const float next = high[i] + product;
+        const float borrowed = next - high[i];
+        low[i] += (high[i] - (next - borrowed)) + (product - borrowed);
+        high[i] = next;
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+  for (uint i = 0; i < 4; ++i) {
+    const uint at = lane + 32 * i, row = rb + sg * 8 + at / 16;
+    if (tf_exl3_active(routes, a, row)) sums[ulong(row) * a.n + nb + at % 16] = high[i] + low[i];
+  }
+}

zig/src/core/exl3_projection.zig

@@ -136,6 +136,7 @@ pub const Pipelines = struct {
     rotate_in: mtl.Pipeline,
     matmul: mtl.Pipeline,
     rotate_out: mtl.Pipeline,
+    prefill: ?mtl.Pipeline = null,
 
     pub fn load(device: mtl.Device) !Pipelines {
         const pool = mtl.objc.Pool.push();
@@ -150,12 +151,20 @@ pub const Pipelines = struct {
         errdefer ro.deinit();
         for ([_]mtl.Pipeline{ ri, mm, ro }) |p|
             if (p.simdWidth() != 32 or p.maxThreads() < 32) return error.UnsupportedExl3Simd;
-        return .{ .rotate_in = ri, .matmul = mm, .rotate_out = ro };
+        var prefill = mtl.Pipeline.init(device, lib, "tf_exl3_matmul_prefill", false) catch null;
+        if (prefill) |pf| {
+            if (pf.simdWidth() != 32 or pf.maxThreads() < 128) {
+                pf.deinit();
+                prefill = null;
+            }
+        }
+        return .{ .rotate_in = ri, .matmul = mm, .rotate_out = ro, .prefill = prefill };
     }
     pub fn deinit(p: *Pipelines) void {
         p.rotate_in.deinit();
         p.matmul.deinit();
         p.rotate_out.deinit();
+        if (p.prefill) |pf| pf.deinit();
     }
 };
 
@@ -316,14 +325,16 @@ fn stages(p: *const Pipelines, e: mtl.ComputeEncoder, a: Args, trellis: View, su
     bind(e, descriptors, 6);
     e.dispatchGroups(.{ .width = a.k / 128, .height = a.rows }, .{ .width = 32 });
     e.barrier();
-    e.setPipeline(p.matmul);
+    const shared = p.prefill != null and a.rows >= 32 and a.indexed == 0 and a.routed == 0;
+    e.setPipeline(if (shared) p.prefill.? else p.matmul);
     bind(e, scratch, 0);
     bind(e, trellis, 1);
     bind(e, sums, 2);
     e.setValue(a, 3);
     bind(e, routes, 4);
     bind(e, descriptors, 5);
-    e.dispatchGroups(.{ .width = a.n / 16, .height = if (a.indexed == 1) a.rows else (@as(usize, a.rows) + 7) / 8 }, .{ .width = 32 });
+    const group_rows: usize = if (shared) 32 else 8;
+    e.dispatchGroups(.{ .width = a.n / 16, .height = if (a.indexed == 1) a.rows else (@as(usize, a.rows) + group_rows - 1) / group_rows }, .{ .width = if (shared) @as(usize, 128) else 32 });
     e.barrier();
     e.setPipeline(p.rotate_out);
     bind(e, sums, 0);

zig/src/core/exl3_projection_gpu_test.zig

@@ -1,6 +1,7 @@
 //! Explicit GPU qualification tests; run through the native test module on a Metal SIMDgroup-matrix device.
 const std = @import("std");
 const mtl = @import("metal");
+const sources = @import("kernel_sources");
 const ex = @import("exl3_projection.zig");
 const decoder = @import("exl3_projection_decode.zig");
 const k = 256;
@@ -386,3 +387,186 @@ test "caller-prepared grouped expert tiles preserve scatter slots and ignore pad
     try complete(cb, e);
     try std.testing.expectEqualSlices(u32, plain.slice(u32, 16 * n), grouped.slice(u32, 16 * n));
 }
+
+const StateWindowInput = extern struct { lo: u32, hi: u32, off: u32, reserved: u32 = 0 };
+
+test "Metal EXL3 mapped windows match independent 64-bit extraction at every offset" {
+    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();
+    const lib = try mtl.Library.fromSource(device, sources.core_exl3_projection, mtl.CompileOptions.mlx());
+    defer lib.deinit();
+    const pipeline = try mtl.Pipeline.init(device, lib, "tf_exl3_state_window_check", false);
+    defer pipeline.deinit();
+    if (pipeline.maxThreads() < 256) return error.UnsupportedProbeThreads;
+    var patterns: [1095][2]u32 = undefined;
+    for (0..32) |bit| {
+        patterns[bit] = .{ @as(u32, 1) << @intCast(bit), 0 };
+        patterns[32 + bit] = .{ 0, @as(u32, 1) << @intCast(bit) };
+    }
+    const special = [_][2]u32{
+        .{ 0, 0 },                   .{ 0xffffffff, 0 },          .{ 0, 0xffffffff },          .{ 0xffffffff, 0xffffffff },
+        .{ 0xaaaaaaaa, 0x55555555 }, .{ 0x55555555, 0xaaaaaaaa }, .{ 0x80000001, 0x80000001 },
+    };
+    for (special, 0..) |p, i| patterns[64 + i] = p;
+    var prng = std.Random.DefaultPrng.init(0x7472656c6c6973);
+    for (patterns[71..]) |*p| p.* = .{ prng.random().int(u32), prng.random().int(u32) };
+    const count = patterns.len * 32;
+    const state_inputs = try device.buffer((count + 2) * @sizeOf(StateWindowInput), mtl.ResourceOptions.shared);
+    defer state_inputs.deinit();
+    const state_outputs = try device.buffer((count + 16) * @sizeOf(u16), mtl.ResourceOptions.shared);
+    defer state_outputs.deinit();
+    @memset(state_outputs.slice(u16, count + 16), 0xa55a);
+    const expected = try std.testing.allocator.alloc(u16, count);
+    defer std.testing.allocator.free(expected);
+    const records = state_inputs.slice(StateWindowInput, count + 2)[1..][0..count];
+    for (0..32) |off| for (patterns, 0..) |p, i| {
+        const at = off * patterns.len + i;
+        records[at] = .{ .lo = p[0], .hi = p[1], .off = @intCast(off) };
+        const joined: u64 = (@as(u64, p[0]) << 32) | p[1];
+        expected[at] = @truncate(joined >> @intCast(48 - off));
+    };
+    const cb = queue.commandBuffer();
+    const e = cb.compute(.concurrent);
+    e.setPipeline(pipeline);
+    e.setBuffer(state_inputs, 16, 0);
+    e.setBuffer(state_outputs, 16, 1);
+    e.setValue(@as(u32, count), 2);
+    e.dispatchThreads(.{ .width = count }, .{ .width = 256 });
+    e.end();
+    cb.commit();
+    cb.wait();
+    if (cb.failure() != null) return error.ProbeGpuFailure;
+    for (expected, state_outputs.slice(u16, count + 16)[8..][0..count], 0..) |want, got, i| {
+        if (want != got) {
+            std.debug.print("state mismatch off={d} index={d} expected={d} actual={d}\n", .{ records[i].off, i, want, got });
+            return error.MappedStateMismatch;
+        }
+    }
+    for (state_outputs.slice(u16, count + 16)[0..8]) |guard| try std.testing.expectEqual(@as(u16, 0xa55a), guard);
+    for (state_outputs.slice(u16, count + 16)[count + 8 ..]) |guard| try std.testing.expectEqual(@as(u16, 0xa55a), guard);
+    std.debug.print("mapped_state_cases={d} offsets=32 mismatches=0 guards=unchanged\n", .{count});
+}
+
+// Larger and resumed windows must retain the original per-row arithmetic and untouched padding.
+test "shared prefill retains exact row bits and guards across tile boundaries and fallback" {
+    const device = try mtl.Device.init();
+    defer device.deinit();
+    var p = try ex.Pipelines.load(device);
+    defer p.deinit();
+    const prefill = p.prefill orelse return error.SharedPrefillUnavailable;
+    try std.testing.expect(prefill.simdWidth() == 32 and prefill.maxThreads() >= 128);
+    var legacy = p;
+    legacy.prefill = null;
+    const queue = try device.queue();
+    defer queue.deinit();
+    const w = try Owned.init(device, 5, .mul1, 0x7368617265);
+    defer w.deinit();
+    for ([_]usize{ 31, 32, 33, 63, 64, 65 }) |count| {
+        const x = try input(device, count);
+        defer x.deinit();
+        const ids = try device.buffer(count * 4, mtl.ResourceOptions.shared);
+        defer ids.deinit();
+        for (ids.slice(u32, count), 0..) |*id, i| id.* = @intCast(i % 2);
+        for ([_]ex.Output{ .bf16, .f32 }) |kind| {
+            const size: usize = if (kind == .bf16) 2 else 4;
+            const bytes = count * (n + 8) * size;
+            const actual = try device.buffer(bytes + 64, mtl.ResourceOptions.shared);
+            defer actual.deinit();
+            const original = try device.buffer(bytes + 64, mtl.ResourceOptions.shared);
+            defer original.deinit();
+            const scratch_bytes = try ex.scratchBytes(k, n, count);
+            const scratch = try device.buffer(scratch_bytes + 32, mtl.ResourceOptions.shared);
+            defer scratch.deinit();
+            for ([_]bool{ false, true }) |filtered| {
+                @memset(actual.slice(u8, bytes + 64), 0xa5);
+                @memset(original.slice(u8, bytes + 64), 0xa5);
+                @memset(scratch.slice(u8, scratch_bytes + 32), 0xcd);
+                const r: ex.Rows = .{ .count = count, .x_stride = k + 16, .y_stride = n + 8, .output = kind, .routes = if (filtered) .{ .ids = .{ .buffer = ids } } else null };
+                const cb = queue.commandBuffer();
+                const e = cb.compute(.concurrent);
+                try ex.encode(&p, e, w.matrix, .{ .buffer = x, .offset = 32 }, .{ .buffer = actual, .offset = 32 }, .{ .buffer = scratch, .offset = 16 }, r);
+                e.barrier();
+                try ex.encode(&legacy, e, w.matrix, .{ .buffer = x, .offset = 32 }, .{ .buffer = original, .offset = 32 }, .{ .buffer = scratch, .offset = 16 }, r);
+                try complete(cb, e);
+                try std.testing.expectEqualSlices(u8, original.slice(u8, bytes + 64), actual.slice(u8, bytes + 64));
+                if (!filtered) {
+                    @memset(original.slice(u8, bytes + 64), 0xa5);
+                    const resumed = queue.commandBuffer();
+                    const re = resumed.compute(.concurrent);
+                    var first = r;
+                    first.count = 31;
+                    try ex.encode(&p, re, w.matrix, .{ .buffer = x, .offset = 32 }, .{ .buffer = original, .offset = 32 }, .{ .buffer = scratch, .offset = 16 }, first);
+                    if (count > 31) {
+                        re.barrier();
+                        var tail = r;
+                        tail.count = count - 31;
+                        try ex.encode(&p, re, w.matrix, .{ .buffer = x, .offset = 32 + 31 * (k + 16) * 2 }, .{ .buffer = original, .offset = 32 + 31 * (n + 8) * size }, .{ .buffer = scratch, .offset = 16 }, tail);
+                    }
+                    try complete(resumed, re);
+                    try std.testing.expectEqualSlices(u8, original.slice(u8, bytes + 64), actual.slice(u8, bytes + 64));
+                }
+                for (actual.slice(u8, bytes + 64)[0..32]) |v| try std.testing.expectEqual(@as(u8, 0xa5), v);
+                for (actual.slice(u8, bytes + 64)[bytes + 32 ..]) |v| try std.testing.expectEqual(@as(u8, 0xa5), v);
+                for (0..count) |row| {
+                    const pad = 32 + (row * (n + 8) + n) * size;
+                    for (actual.slice(u8, bytes + 64)[pad .. pad + 8 * size]) |v| try std.testing.expectEqual(@as(u8, 0xa5), v);
+                    if (filtered and row % 2 != 0) {
+                        const at = 32 + row * (n + 8) * size;
+                        for (actual.slice(u8, bytes + 64)[at .. at + n * size]) |v| try std.testing.expectEqual(@as(u8, 0xa5), v);
+                    }
+                }
+                for (scratch.slice(u8, scratch_bytes + 32)[0..16]) |v| try std.testing.expectEqual(@as(u8, 0xcd), v);
+                for (scratch.slice(u8, scratch_bytes + 32)[scratch_bytes + 16 ..]) |v| try std.testing.expectEqual(@as(u8, 0xcd), v);
+            }
+        }
+    }
+}
+
+test "shared prefill preserves gather and scatter maps across a partial tile" {
+    const device = try mtl.Device.init();
+    defer device.deinit();
+    var p = try ex.Pipelines.load(device);
+    defer p.deinit();
+    if (p.prefill == null) return error.SharedPrefillUnavailable;
+    var legacy = p;
+    legacy.prefill = null;
+    const queue = try device.queue();
+    defer queue.deinit();
+    const w = try Owned.init(device, 3, .mul1, 0x676174686572);
+    defer w.deinit();
+    const x = try input(device, 3);
+    defer x.deinit();
+    const count = 33;
+    const gather = try device.buffer(count * 4 + 16, mtl.ResourceOptions.shared);
+    defer gather.deinit();
+    const scatter = try device.buffer(count * 4 + 16, mtl.ResourceOptions.shared);
+    defer scatter.deinit();
+    for (0..count) |i| {
+        gather.slice(u32, count + 4)[4 + i] = @intCast((i * 5) % 3);
+        scatter.slice(u32, count + 4)[4 + i] = @intCast(count - 1 - i);
+    }
+    for ([_]ex.Output{ .bf16, .f32 }) |kind| {
+        const size: usize = if (kind == .bf16) 2 else 4;
+        const bytes = count * (n + 8) * size + 64;
+        const actual = try device.buffer(bytes, mtl.ResourceOptions.shared);
+        defer actual.deinit();
+        const original = try device.buffer(bytes, mtl.ResourceOptions.shared);
+        defer original.deinit();
+        @memset(actual.slice(u8, bytes), 0xa5);
+        @memset(original.slice(u8, bytes), 0xa5);
+        const scratch = try device.buffer(try ex.scratchBytes(k, n, count), mtl.ResourceOptions.shared);
+        defer scratch.deinit();
+        const r: ex.Rows = .{ .count = count, .x_stride = k + 16, .y_stride = n + 8, .output = kind, .input_ids = .{ .buffer = gather, .offset = 16 }, .output_ids = .{ .buffer = scatter, .offset = 16 }, .input_rows = 3, .output_rows = count };
+        const cb = queue.commandBuffer();
+        const e = cb.compute(.concurrent);
+        try ex.encode(&p, e, w.matrix, .{ .buffer = x, .offset = 32 }, .{ .buffer = actual, .offset = 32 }, .{ .buffer = scratch }, r);
+        e.barrier();
+        try ex.encode(&legacy, e, w.matrix, .{ .buffer = x, .offset = 32 }, .{ .buffer = original, .offset = 32 }, .{ .buffer = scratch }, r);
+        try complete(cb, e);
+        try std.testing.expectEqualSlices(u8, original.slice(u8, bytes), actual.slice(u8, bytes));
+    }
+}