From 4dbeca4995214e9e910273ab7e0d7b7b95904c7a Mon Sep 17 00:00:00 2001 From: mrv777 Date: Mon, 12 Oct 2026 01:23:51 -0500 Subject: [PATCH 1/2] EXL3: use branch-free trellis window extraction --- zig/kernels/metal/core/exl3_projection.metal | 16 ++++- zig/src/core/exl3_projection_gpu_test.zig | 64 ++++++++++++++++++++ 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/zig/kernels/metal/core/exl3_projection.metal b/zig/kernels/metal/core/exl3_projection.metal index a76a912ab3..3bd8f71591 100644 --- a/zig/kernels/metal/core/exl3_projection.metal +++ b/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,15 @@ kernel void tf_exl3_rotate_out(const device float* sums [[buffer(0)]], else reinterpret_cast(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(inputs + tid); + const uint off = inputs[tid].z; + outputs[tid] = tf_exl3_mapped_state(tile, (off << 6) | (1u << 11)); +} diff --git a/zig/src/core/exl3_projection_gpu_test.zig b/zig/src/core/exl3_projection_gpu_test.zig index f3eb537e24..55eba0b7d4 100644 --- a/zig/src/core/exl3_projection_gpu_test.zig +++ b/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,66 @@ 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}); +} From 8eaf366208d16786293846c2bdc12ec85ab2a484 Mon Sep 17 00:00:00 2001 From: mrv777 Date: Mon, 12 Oct 2026 07:30:22 -0500 Subject: [PATCH 2/2] EXL3: share decoded weights across prefill rows --- RUNBOOK.md | 6 + zig/exl3_projection_bench.zig | 3 + zig/kernels/metal/core/exl3_projection.metal | 60 ++++++++++ zig/src/core/exl3_projection.zig | 17 ++- zig/src/core/exl3_projection_gpu_test.zig | 120 +++++++++++++++++++ 5 files changed, 203 insertions(+), 3 deletions(-) diff --git a/RUNBOOK.md b/RUNBOOK.md index 09630beb41..7c031bef6a 100644 --- a/RUNBOOK.md +++ b/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. diff --git a/zig/exl3_projection_bench.zig b/zig/exl3_projection_bench.zig index 06c4fb2f9a..5569cee9fd 100644 --- a/zig/exl3_projection_bench.zig +++ b/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| { diff --git a/zig/kernels/metal/core/exl3_projection.metal b/zig/kernels/metal/core/exl3_projection.metal index 3bd8f71591..0eca0d80c5 100644 --- a/zig/kernels/metal/core/exl3_projection.metal +++ b/zig/kernels/metal/core/exl3_projection.metal @@ -240,3 +240,63 @@ kernel void tf_exl3_state_window_check(const device uint4* inputs [[buffer(0)]], 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( + 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]; + } +} diff --git a/zig/src/core/exl3_projection.zig b/zig/src/core/exl3_projection.zig index 91e2e5485c..8cba8106e4 100644 --- a/zig/src/core/exl3_projection.zig +++ b/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); diff --git a/zig/src/core/exl3_projection_gpu_test.zig b/zig/src/core/exl3_projection_gpu_test.zig index 55eba0b7d4..6f6445c433 100644 --- a/zig/src/core/exl3_projection_gpu_test.zig +++ b/zig/src/core/exl3_projection_gpu_test.zig @@ -450,3 +450,123 @@ test "Metal EXL3 mapped windows match independent 64-bit extraction at every off 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)); + } +}