From f50abf67e428262c9708c735ea8ca33db86ef54f Mon Sep 17 00:00:00 2001 From: Shinto C V <11365900+cshintov@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:54:23 +0530 Subject: [PATCH 1/3] contract: qwen27 attention_io's pack, append and keep at 24 over 4 and a rank's 12 over 2 Keep checks the staging rows where Scratch sized them, and a device keep (one window, all rows dispatched, the map as tree_commit writes it, half the rows and one kept) runs beside it. On a rank both keeps fail: gather and scatter index 4 heads a staging row in staging sized for 2. --- zig/src/families/qwen27/attention.zig | 3 +- zig/tests/contract/metal_main.zig | 23 ++ zig/tests/contract/metal_q27_attn.zig | 116 ++++++++ zig/tests/contract/metal_q27_attn_io.zig | 362 +++++++++++++++++++++++ 4 files changed, 503 insertions(+), 1 deletion(-) create mode 100644 zig/tests/contract/metal_q27_attn.zig create mode 100644 zig/tests/contract/metal_q27_attn_io.zig diff --git a/zig/src/families/qwen27/attention.zig b/zig/src/families/qwen27/attention.zig index cd9fc461be..e120aff5be 100644 --- a/zig/src/families/qwen27/attention.zig +++ b/zig/src/families/qwen27/attention.zig @@ -8,7 +8,8 @@ pub const planning = @import("attention_plan.zig"); pub const Plan = planning.Plan; pub const Cache = struct { keys: Ref, values: Ref, capacity: u32, key_stride: u32 = 0, value_stride: u32 = 0 }; const CachePtrs = extern struct { keys: u64, values: u64 }; -const PackArgs = extern struct { rows: u32, packed_rows: u32, heads: u32 = 24, kv_heads: u32 = 4, dim: u32 = 256, group: u32 = 6 }; +/// The pack and append kernels' arguments (public for the kernel contract's io check). +pub const PackArgs = extern struct { rows: u32, packed_rows: u32, heads: u32 = 24, kv_heads: u32 = 4, dim: u32 = 256, group: u32 = 6 }; const PromptArgs = extern struct { start: i32, rows: i32, chunks: i32, rp: i32, key_stride: i32, value_stride: i32, scale2: f32, kv_heads: i32 }; pub const Limits = struct { rows: u32 = 128, diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index e74219237d..3188b8cf8a 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -22,6 +22,8 @@ const glm_rows_k = @import("metal_glm_rows.zig"); const pair_k = @import("metal_pair.zig"); const pair_controls = @import("metal_pair_controls.zig"); const minimax_k = @import("metal_minimax.zig"); +const attn_k = @import("metal_q27_attn.zig"); +const io_k = @import("metal_q27_attn_io.zig"); const flash_pool = @import("metal_flashnext_pool.zig"); const flash = @import("metal_flashnext.zig"); @@ -406,6 +408,27 @@ pub fn main(init: std.process.Init) !u8 { var nemotron_unroll: @import("metal_nemotron_unroll.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; defer nemotron_unroll.deinit(); try nemotron_unroll.collect(o.only); + // qwen27 attention_io at the 27B's head counts on one Mac (24 over 4) and on a rank of two (12 over 2) + var ios: std.ArrayList(*io_k.Io) = .empty; + defer { + for (ios.items) |x| { + x.deinit(); + a.destroy(x); + } + ios.deinit(a); + } + var io_regs: std.ArrayList(io_k.Reg) = .empty; + defer io_regs.deinit(a); + for (std.enums.values(attn_k.Heads)) |h| for (std.enums.values(io_k.Which)) |w| try io_regs.append(a, .{ .heads = h, .which = w }); + try io_regs.append(a, .{ .heads = .q27b, .which = .append, .control = true }); + for (io_regs.items) |r| { + if (!wanted(o, io_k.spec(r).name)) continue; + const x = try a.create(io_k.Io); + errdefer a.destroy(x); + x.* = try io_k.Io.init(a, device, queue, r); + try ios.append(a, x); + try kernels.append(a, x.kernel()); + } var minimax: minimax_k.Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; defer minimax.deinit(); try minimax.collect(o.only); diff --git a/zig/tests/contract/metal_q27_attn.zig b/zig/tests/contract/metal_q27_attn.zig new file mode 100644 index 0000000000..7829983109 --- /dev/null +++ b/zig/tests/contract/metal_q27_attn.zig @@ -0,0 +1,116 @@ +//! qwen27 attention's contract helpers: head counts, the engine's source, a round's plan, scratch and binding. +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 q = tf.qwen27; +const attn = q.attention; +const planning = q.attention_plan; + +const D = 256; + +/// The head counts main runs: the 27B's 24 query heads over 4, and its rank of a two-Mac split (12 over 2, #1433). +pub const Heads = enum { q27b, rank }; + +pub fn geometry(h: Heads, shared: bool) planning.Geometry { + return switch (h) { + .q27b => .{ .heads = 24, .kv_heads = 4, .shared_prefix = shared }, + .rank => .{ .heads = 12, .kv_heads = 2, .shared_prefix = shared }, + }; +} + +pub fn label(h: Heads) []const u8 { + return if (h == .q27b) "27B" else "27B rank"; +} + +/// A planted control's one edit to the shipped source; `from` must occur once. +pub const Edit = struct { from: []const u8, to: []const u8 }; + +/// The source as decode_round.zig builds it here (simd_attention.library, the prompt kernels included). +pub fn source(a: std.mem.Allocator, tensor: bool, edit: ?Edit) ![]u8 { + const text = try q.simd_attention.library(a, tensor, true); + const e = edit orelse return text; + defer a.free(text); + if (std.mem.count(u8, text, e.from) != 1) return error.ControlEdit; + return std.mem.replaceOwned(u8, a, text, e.from, e.to); +} + +/// Rounds of the round's rows, query heads and the first stream's cached keys (past 16 rows, streams of 16). +pub fn served(comptime h: Heads) []const c.Shape { + const n: usize = comptime geometry(h, false).heads; + // both head counts are GQA6, so the same rows reach the same tiles: 3 and 6 rows fork at 64 and 65 cached keys + return comptime &.{ .{ .rows = 1, .n = n, .k = 2 }, .{ .rows = 3, .n = n, .k = 64 }, .{ .rows = 6, .n = n, .k = 64 }, .{ .rows = 3, .n = n, .k = 65 }, .{ .rows = 6, .n = n, .k = 65 }, .{ .rows = 16, .n = n, .k = 2046 }, .{ .rows = 16, .n = n, .k = 2047 }, .{ .rows = 128, .n = n, .k = 510 }, .{ .rows = 128, .n = n, .k = 511 }, .{ .rows = 1, .n = n, .k = 8191 }, .{ .rows = 11, .n = n, .k = 8190 }, .{ .rows = 14, .n = n, .k = 8191 } }; +} + +/// The engine's plan, scratch and binding for one round. +pub const State = struct { + plan: attn.Plan, + scratch: attn.Scratch, + binding: attn.Binding, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, r: *const Round, gm: planning.Geometry, caches: []const attn.Cache) !State { + var ws: [planning.max_streams]planning.Window = undefined; + var plan = try attn.Plan.init(a, r.windows(&ws), gm); + errdefer plan.deinit(); + const limits = attn.Limits{ .rows = @intCast(r.rows), .streams = @intCast(r.streams), .keys = @intCast(r.cap), .shared_prefix = gm.shared_prefix, .geometry = gm }; + var scratch = try attn.Scratch.init(device, limits); + errdefer scratch.deinit(); + var binding = try attn.Binding.init(device, limits); + errdefer binding.deinit(); + try scratch.upload(plan); + try binding.upload(plan, caches); + return .{ .plan = plan, .scratch = scratch, .binding = binding }; + } + + pub fn deinit(s: *State) void { + s.binding.deinit(); + s.scratch.deinit(); + s.plan.deinit(); + } +}; + +/// A round's windows: up to 16 rows each, 97 keys apart; at even starts 3 rows or more fork. +pub const Round = struct { + rows: usize = 0, + streams: usize = 0, + cap: usize = 0, + counts: [planning.max_streams]usize = @splat(0), + firsts: [planning.max_streams]usize = @splat(0), + starts: [planning.max_streams]usize = @splat(0), + parents: [planning.max_rows]i32 = @splat(-1), + + pub fn init(shape: c.Shape, one: bool, trees: bool, rng: std.Random) Round { + var r = Round{ .rows = shape.rows, .streams = if (one) 1 else (shape.rows + 15) / 16 }; + var first: usize = 0; + var end: usize = 0; + for (0..r.streams) |s| { + const n = if (one) shape.rows else @min(16, shape.rows - first); + r.counts[s] = n; + r.firsts[s] = first; + r.starts[s] = shape.k + 97 * s; + const tree = trees and r.starts[s] % 2 == 0 and n >= 3; + for (0..n) |i| r.parents[first + i] = if (i == 0) -1 else if (tree and (i <= 2 or rng.uintLessThan(u32, 3) != 0)) @intCast(if (i <= 2) 0 else rng.uintLessThan(usize, i)) else @intCast(i - 1); + first += n; + end = @max(end, r.starts[s] + n); + } + r.cap = end + 64; + return r; + } + + pub fn windows(r: *const Round, ws: *[planning.max_streams]planning.Window) []const planning.Window { + for (0..r.streams) |s| ws[s] = .{ .parents = r.parents[r.firsts[s]..][0..r.counts[s]], .start = @intCast(r.starts[s]), .capacity = @intCast(r.cap) }; + return ws[0..r.streams]; + } + + /// Element 0 of key head `hk`'s cache row `row` in stream `s`'s caches, `kv_heads` heads each. + pub fn at(r: *const Round, kv_heads: usize, s: usize, hk: usize, row: usize) usize { + return ((s * kv_heads + hk) * r.cap + row) * D; + } +}; + +pub fn ref(r: rig_.Ref) q.projection.Ref { + return .{ .buffer = r.buf, .offset = r.off }; +} diff --git a/zig/tests/contract/metal_q27_attn_io.zig b/zig/tests/contract/metal_q27_attn_io.zig new file mode 100644 index 0000000000..4434942b30 --- /dev/null +++ b/zig/tests/contract/metal_q27_attn_io.zig @@ -0,0 +1,362 @@ +//! qwen27's attention_io kernels under the contract on Metal: packing and cache moves copy bf16 bits, checked exactly. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const at_ = @import("metal_q27_attn.zig"); + +const c = contract.spec; +const g = contract.gates; +const attn = tf.qwen27.attention; +const planning = tf.qwen27.attention_plan; + +const D = 256; +const rw = mtl.ResourceUsage.read | mtl.ResourceUsage.write; + +/// pack, append, keep (kept paths moved), device_keep (tree_commit's map: one chain, half the rows and one kept). +pub const Which = enum { pack, append, keep, device_keep }; + +pub const Reg = struct { heads: at_.Heads, which: Which, control: bool = false }; + +const rule = "bits only: NaN and both infinities copy as they are"; + +fn noSums(_: c.Shape) usize { + return 0; +} + +fn nameOf(comptime r: Reg) []const u8 { + if (r.control) return "control: qwen27 attention io append one row late"; + const what = comptime switch (r.which) { + .pack => "pack q27_attn_pack_a q27_attn_pack_b", + .append => "append q27_attn_append", + .keep => "keep q27_attn_gather q27_attn_scatter", + .device_keep => "device keep q27_attn_gather q27_attn_scatter", + }; + return std.fmt.comptimePrint("metal qwen27 attention {s} io {s}", .{ comptime at_.label(r.heads), what }); +} + +/// The device keep's one window: 1 to 64 rows (tree_round.max_rows), short and long caches. +fn deviceServed(comptime h: at_.Heads) []const c.Shape { + const n: usize = comptime at_.geometry(h, false).heads; + return comptime &.{ .{ .rows = 1, .n = n, .k = 2 }, .{ .rows = 3, .n = n, .k = 65 }, .{ .rows = 16, .n = n, .k = 2047 }, .{ .rows = 33, .n = n, .k = 511 }, .{ .rows = 64, .n = n, .k = 8191 } }; +} + +pub fn spec(r: Reg) c.Spec { + return switch (r.heads) { + inline else => |h| switch (r.which) { + inline else => |w| switch (r.control) { + inline else => |k| .{ .name = nameOf(.{ .heads = h, .which = w, .control = k }), .accum = .none, .depth = noSums, .out = .raw, .in_place = w == .keep or w == .device_keep, .nan = .{ .stated = rule }, .tiles = .{ .{ .size = 16 }, .{}, .{ .size = 512 } }, .served = if (w == .device_keep) deviceServed(h) else at_.served(h) }, + }, + }, + }; +} + +pub const Io = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + reg: Reg, + op: attn.Attention, + tensor: bool, + st: ?at_.State = null, + rd: at_.Round = .{}, + host: [3][]u16 = .{ &.{}, &.{}, &.{} }, + paths: [planning.max_streams][]u32 = @splat(&.{}), + path_buf: [planning.max_rows]u32 = undefined, + kept: u32 = 0, + ins: [2]rig_.Ref = undefined, + outs: [2 * planning.max_streams]rig_.Ref = undefined, + caches: [planning.max_streams]attn.Cache = undefined, + msg: [200]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, r: Reg) !Io { + const tensor = device.tensorUnits(); + const edit: ?at_.Edit = if (r.control) .{ .from = "row = uint(meta[base + 4]) + local;", .to = "row = uint(meta[base + 4]) + local + 1;" } else null; + const text = try at_.source(a, tensor, edit); + defer a.free(text); + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + const op = try attn.Attention.init(device, lib); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .reg = r, .op = op, .tensor = tensor }; + } + + pub fn deinit(t: *Io) void { + t.release(); + t.op.deinit(); + t.rig.deinit(); + } + + fn release(t: *Io) void { + if (t.st) |*s| s.deinit(); + t.st = null; + for (&t.host) |*h| { + t.a.free(h.*); + h.* = &.{}; + } + } + + pub fn kernel(t: *Io) g.Kernel { + return .{ .spec = spec(t.reg), .ctx = t, .vt = &vt, .control = if (t.reg.control) .shapes else null }; + } + + const vt = g.hooks(Io); + + fn geo(t: *const Io) planning.Geometry { + return at_.geometry(t.reg.heads, t.tensor); + } + + fn bits(rng: std.Random, xs: []u16) void { + for (xs) |*x| x.* = c.toBf16(rng.float(f32) * 4 - 2); + } + + /// Inputs (pack: queries; append: keys and values; keep: the caches as they start), at gate 1 with NaN and Inf. + pub fn prepare(t: *Io, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const gm = t.geo(); + const heads: usize = gm.heads; + const hk: usize = gm.kv_heads; + if (shape.n != heads or shape.rows == 0 or shape.rows > planning.max_rows) return false; + // the device keep serves one window of at most 64 rows (tree_round.max_rows) + const device = t.reg.which == .device_keep; + if (device and shape.rows > 64) return false; + t.release(); + t.rig.reset(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + t.rd = at_.Round.init(shape, device, !device, rng); + const r = &t.rd; + const per = hk * r.cap * D; + switch (t.reg.which) { + .pack => t.host[0] = try t.a.alloc(u16, r.rows * heads * D), + .append => for (t.host[0..2]) |*h| { + h.* = try t.a.alloc(u16, r.rows * hk * D); + }, + .keep, .device_keep => for (t.host[0..2]) |*h| { + h.* = try t.a.alloc(u16, r.streams * per); + }, + } + for (t.host) |h| bits(rng, h); + if (fill == .nonfinite) for (t.host[0..if (t.reg.which == .pack) 1 else 2], 0..) |h, i| { + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (put, 0..) |v, j| h[(i * 7 + j * 131) % h.len] = c.toBf16(v); + }; + if (t.reg.which == .pack or t.reg.which == .append) for (t.host, 0..) |h, i| if (h.len > 0) { + t.ins[i] = try t.rig.input(std.mem.sliceAsBytes(h)); + }; + // the caches are what append and keep write; pack's binding needs them only to exist + for (0..r.streams) |s| { + var kv: [2]rig_.Ref = undefined; + for (&kv) |*x| x.* = if (t.reg.which == .pack) try t.rig.scratch(per * 2) else try t.rig.output(per * 2); + t.outs[2 * s] = kv[0]; + t.outs[2 * s + 1] = kv[1]; + t.caches[s] = .{ .keys = at_.ref(kv[0]), .values = at_.ref(kv[1]), .capacity = @intCast(r.cap) }; + } + t.st = try at_.State.init(t.a, t.rig.device, r, gm, t.caches[0..r.streams]); + if (t.reg.which == .pack) { + t.outs[0] = try t.rig.output(hk * t.st.?.plan.tiles * 16 * D * 2); + t.outs[1] = try t.rig.output(hk * r.rows * 16 * D * 2); + } + if (t.reg.which == .keep) { + t.pick(rng, &t.path_buf, &t.paths); + const st = &t.st.?; + t.kept = try st.binding.uploadKeep(st.plan, t.caches[0..r.streams], t.paths[0..r.streams]); + } + if (device) { + // tree_commit's map: the accepted chain's first rows/2 + 1 rows, then rows no stream keeps + const half = r.rows / 2 + 1; + for (0..half) |i| t.path_buf[i] = @intCast(i); + t.paths[0] = t.path_buf[0..half]; + const st = &t.st.?; + const map = st.binding.get(.keep_map).buffer.slice(planning.KeepRow, r.rows); + for (map, 0..) |*m, i| m.* = .{ .stream = if (i < half) 0 else 0xffffffff, .source = if (i < half) @intCast(i) else 0, .destination = @intCast(i), .row = @intCast(i) }; + try st.binding.admitDeviceKeep(st.plan, t.caches[0..1], @intCast(r.rows)); + t.kept = @intCast(r.rows); + } + return true; + } + + /// Staging row i (kept rows by stream, then depth) holds its source row's heads at (i * kv_heads + h). + fn staged(t: *Io, paths: []const []u32) !?[]const u8 { + const r = &t.rd; + const hk: usize = t.geo().kv_heads; + const per = hk * r.cap * D; + var rows: usize = 0; + for (paths) |p| rows += p.len; + const st = &t.st.?; + const stage = [2][]const u16{ st.scratch.get(.keep_k).buffer.slice(u16, rows * hk * D), st.scratch.get(.keep_v).buffer.slice(u16, rows * hk * D) }; + var i: usize = 0; + for (paths, 0..) |p, s| for (p) |source| { + for (0..2) |kv| for (0..hk) |h| { + const was = t.host[kv][s * per + (h * r.cap + r.starts[s] + source) * D ..][0..D]; + if (!std.mem.eql(u16, stage[kv][(i * hk + h) * D ..][0..D], was)) return try t.say("staging {s} row {d} head {d} is not stream {d}'s row {d}", .{ if (kv == 0) "keys" else "values", i, h, s, r.starts[s] + source }); + }; + i += 1; + }; + return null; + } + + /// Each stream's kept rows: a random row's path from the root, as a committed draft keeps. + fn pick(t: *const Io, rng: std.Random, buf: *[planning.max_rows]u32, paths: *[planning.max_streams][]u32) void { + const r = &t.rd; + var used: usize = 0; + for (0..r.streams) |s| { + var x: i32 = @intCast(rng.uintLessThan(usize, r.counts[s])); + var n: usize = 0; + var rev: [planning.max_rows]u32 = undefined; + while (x >= 0) : (x = r.parents[r.firsts[s] + @as(usize, @intCast(x))]) { + rev[n] = @intCast(x); + n += 1; + } + for (0..n) |i| buf[used + i] = rev[n - 1 - i]; + paths[s] = buf[used..][0..n]; + used += n; + } + } + + /// The engine's keep of `kept` rows on the caches as they are now. + fn keep(t: *Io, kept: u32) !void { + const st = &t.st.?; + const p = t.rig.begin(); + t.op.keep(p.e, &st.scratch, &st.binding, kept, t.caches[0..t.rd.streams]) catch |err| { + p.e.end(); + p.pool.pop(); + return err; + }; + try t.rig.finish(p); + } + + pub fn run(t: *Io) !void { + const st = &t.st.?; + const r = &t.rd; + const gm = t.geo(); + const per = gm.kv_heads * r.cap * D; + // keep moves rows within the caches it reads, so each run starts from the caches as prepared + if (t.reg.which == .keep or t.reg.which == .device_keep) { + for (0..r.streams) |s| for (0..2) |kv| { + @memcpy(t.rig.bytes(t.outs[2 * s + kv], per * 2), std.mem.sliceAsBytes(t.host[kv][s * per ..][0..per])); + }; + return t.keep(t.kept); + } + const p = t.rig.begin(); + const e = p.e; + const args = attn.PackArgs{ .rows = @intCast(r.rows), .packed_rows = st.plan.tiles * 16, .heads = gm.heads, .kv_heads = gm.kv_heads, .group = gm.group() }; + switch (t.reg.which) { + .pack => { + e.setPipeline(t.op.io[0]); + for ([_]mtl.Buffer{ t.ins[0].buf, st.scratch.get(.q_rows).buffer, t.outs[0].buf }, [_]usize{ t.ins[0].off, 0, t.outs[0].off }, 0..) |b, off, i| e.setBuffer(b, off, i); + e.setValue(args, 3); + e.dispatchThreads(mtl.Size.of(256, args.packed_rows, gm.kv_heads), mtl.Size.of(32, 4, 1)); + e.setPipeline(t.op.io[1]); + e.setBuffer(t.ins[0].buf, t.ins[0].off, 0); + e.setBuffer(t.outs[1].buf, t.outs[1].off, 1); + e.setValue(args, 2); + e.dispatchThreads(mtl.Size.of(256, 16, gm.kv_heads * r.rows), mtl.Size.of(32, 4, 1)); + }, + .append => { + for (st.binding.buffers) |b| e.useResource(b, rw); + for (t.outs[0 .. 2 * r.streams]) |o| e.useResource(o.buf, rw); + e.setPipeline(t.op.io[2]); + for ([_]mtl.Buffer{ t.ins[0].buf, t.ins[1].buf, st.binding.get(.caches).buffer, st.binding.get(.meta).buffer, st.scratch.get(.nodes).buffer }, [_]usize{ t.ins[0].off, t.ins[1].off, 0, 0, 0 }, 0..) |b, off, i| e.setBuffer(b, off, i); + e.setValue(args, 5); + e.dispatchThreads(mtl.Size.of(256, gm.kv_heads, r.rows), mtl.Size.of(32, 4, 1)); + }, + .keep, .device_keep => unreachable, + } + try t.rig.finish(p); + } + + pub fn count(_: *Io) usize { + return 0; + } + + pub fn reference(_: *Io, _: []f64, _: []f64) !void {} + + /// An output's bf16 words (the rig's payloads start 64 bytes in, so they are aligned). + fn words(t: *Io, r: rig_.Ref, n: usize) []const u16 { + return @alignCast(std.mem.bytesAsSlice(u16, t.rig.bytes(r, n * 2))); + } + + fn say(t: *Io, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(&t.msg, fmt, args); + } + + /// Every element against its copy: packed rows (zero past a block's heads), appended rows, kept rows by depth. + pub fn check(t: *Io, _: g.Gate) !?[]const u8 { + const gm = t.geo(); + const heads: usize = gm.heads; + const hk: usize = gm.kv_heads; + const group: usize = gm.group(); + const r = &t.rd; + const per = hk * r.cap * D; + switch (t.reg.which) { + .pack => { + const plan = t.st.?.plan; + const packed_rows = plan.tiles * 16; + const qa = t.words(t.outs[0], hk * packed_rows * D); + const qb = t.words(t.outs[1], hk * r.rows * 16 * D); + for (0..hk) |h| for (0..packed_rows) |y| { + const m: usize = @intCast(plan.q_rows[y]); + const from = t.host[0][((m / group) * heads + h * group + m % group) * D ..][0..D]; + if (!std.mem.eql(u16, qa[(h * packed_rows + y) * D ..][0..D], from)) return try t.say("pack_a key head {d} row {d} differs from query row {d}", .{ h, y, m }); + }; + for (0..hk) |h| for (0..r.rows) |node| for (0..16) |y| { + const got = qb[((h * r.rows + node) * 16 + y) * D ..][0..D]; + const ok = if (y < group) std.mem.eql(u16, got, t.host[0][(node * heads + h * group + y) * D ..][0..D]) else std.mem.allEqual(u16, got, 0); + if (!ok) return try t.say("pack_b key head {d} row {d} block row {d} is not its copy", .{ h, node, y }); + }; + }, + .append => for (0..r.streams) |s| for (0..2) |kv| { + const got = t.words(t.outs[2 * s + kv], per); + for (0..hk) |h| for (0..r.cap) |row| { + const cell = got[(h * r.cap + row) * D ..][0..D]; + const local = row -% r.starts[s]; + const ok = if (local < r.counts[s]) std.mem.eql(u16, cell, t.host[kv][((r.firsts[s] + local) * hk + h) * D ..][0..D]) else std.mem.allEqual(u16, cell, 0xffff); + if (!ok) return try t.say("stream {d} {s} head {d} cache row {d}: {s}", .{ s, if (kv == 0) "keys" else "values", h, row, if (local < r.counts[s]) "not the round's row" else "written outside the round" }); + }; + }, + .device_keep => { + if (try t.moved(.{ t.host[0], t.host[1] }, t.paths[0..1], "")) |why| return why; + return t.staged(t.paths[0..1]); + }, + .keep => { + if (try t.moved(.{ t.host[0], t.host[1] }, t.paths[0..r.streams], "")) |why| return why; + // gather and scatter can agree on a wrong layout: check staging where Scratch sized it + if (try t.staged(t.paths[0..r.streams])) |why| return why; + // a second keep, other paths, on the caches the first left (no restore): keep reads them as they are + var now: [2][]u16 = undefined; + for (&now, 0..) |*x, kv| { + x.* = try t.a.alloc(u16, r.streams * per); + for (0..r.streams) |s| @memcpy(x.*[s * per ..][0..per], t.words(t.outs[2 * s + kv], per)); + } + defer for (now) |x| t.a.free(x); + var prng = std.Random.DefaultPrng.init(t.kept); + var buf: [planning.max_rows]u32 = undefined; + var paths: [planning.max_streams][]u32 = undefined; + t.pick(prng.random(), &buf, &paths); + const st = &t.st.?; + try t.keep(try st.binding.uploadKeep(st.plan, t.caches[0..r.streams], paths[0..r.streams])); + const why = try t.moved(.{ now[0], now[1] }, paths[0..r.streams], "after a first keep, "); + t.kept = try st.binding.uploadKeep(st.plan, t.caches[0..r.streams], t.paths[0..r.streams]); + return why; + }, + } + return null; + } + + /// The caches against `was` with each stream's kept path moved to its window's start, the rest as they were. + fn moved(t: *Io, was_: [2][]const u16, paths: []const []u32, comptime when: []const u8) !?[]const u8 { + const r = &t.rd; + const hk: usize = t.geo().kv_heads; + const per = hk * r.cap * D; + for (0..r.streams) |s| for (0..2) |kv| { + const got = t.words(t.outs[2 * s + kv], per); + const was = was_[kv][s * per ..][0..per]; + const path = paths[s]; + for (0..hk) |h| for (0..r.cap) |row| { + const depth = row -% r.starts[s]; + const from = if (depth < path.len) r.starts[s] + path[depth] else row; + if (!std.mem.eql(u16, got[(h * r.cap + row) * D ..][0..D], was[(h * r.cap + from) * D ..][0..D])) return try t.say(when ++ "stream {d} {s} head {d} cache row {d} is not row {d} as it was", .{ s, if (kv == 0) "keys" else "values", h, row, from }); + }; + }; + return null; + } +}; From 7e9d9c3846da6fc2d8a5c18b905096260569d030 Mon Sep 17 00:00:00 2001 From: Shinto C V <11365900+cshintov@users.noreply.github.com> Date: Sun, 1 Nov 2026 16:02:36 +0530 Subject: [PATCH 2/3] qwen27: keep's gather and scatter stage the binding's KV heads a row, not 4 Scratch sizes keep_k and keep_v at rows * kv_heads, but the kernels indexed row * 4 + head, so on a rank (12 over 2) a keep of more than half the frame's rows wrote past them. The 27B on one Mac is unchanged. The contract's keep control plants the 4 back in both kernels and must fail. --- zig/kernels/metal/qwen27/attention_io.metal | 13 +++++++------ zig/src/families/qwen27/attention.zig | 9 +++++++-- zig/tests/contract/metal_main.zig | 2 +- zig/tests/contract/metal_q27_attn.zig | 6 +++--- zig/tests/contract/metal_q27_attn_io.zig | 6 +++++- 5 files changed, 23 insertions(+), 13 deletions(-) diff --git a/zig/kernels/metal/qwen27/attention_io.metal b/zig/kernels/metal/qwen27/attention_io.metal index 9c6bb169e5..1e5dccb3f7 100644 --- a/zig/kernels/metal/qwen27/attention_io.metal +++ b/zig/kernels/metal/qwen27/attention_io.metal @@ -4,6 +4,7 @@ using namespace metal; struct Q27CachePtrs { device bfloat* keys; device bfloat* values; }; struct Q27PackArgs { uint rows, packed_rows, heads, kv_heads, dim, group; }; struct Q27KeepRow { uint stream, source, destination, row; }; +struct Q27KeepArgs { uint kept, kv_heads; }; // staging rows hold kv_heads heads, as Scratch sizes keep_k and keep_v kernel void q27_attn_pack_a(device const bfloat* Q [[buffer(0)]], device const int* rows [[buffer(1)]], device bfloat* QA [[buffer(2)]], constant Q27PackArgs& a [[buffer(3)]], uint3 at [[thread_position_in_grid]]) { @@ -28,19 +29,19 @@ kernel void q27_attn_append(device const bfloat* K [[buffer(0)]], device const b } kernel void q27_attn_gather(device const Q27CachePtrs* caches [[buffer(0)]], device const int* meta [[buffer(1)]], device const Q27KeepRow* map [[buffer(2)]], device bfloat* K [[buffer(3)]], device bfloat* V [[buffer(4)]], - constant uint& kept [[buffer(5)]], uint3 at [[thread_position_in_grid]]) { - if (at.x >= 256 || at.y >= 4 || at.z >= kept) return; + constant Q27KeepArgs& a [[buffer(5)]], uint3 at [[thread_position_in_grid]]) { + if (at.x >= 256 || at.y >= a.kv_heads || at.z >= a.kept) return; const Q27KeepRow r = map[at.z]; if (r.stream == 0xffffffffu) return; const uint base = 8 + r.stream * 12, source = uint(meta[base + 4]) + r.source; - const ulong destination = (ulong(r.row) * 4 + at.y) * 256 + at.x; + const ulong destination = (ulong(r.row) * a.kv_heads + at.y) * 256 + at.x; K[destination] = caches[r.stream].keys[ulong(at.y) * uint(meta[base + 7]) + source * 256 + at.x]; V[destination] = caches[r.stream].values[ulong(at.y) * uint(meta[base + 8]) + source * 256 + at.x]; } kernel void q27_attn_scatter(device const bfloat* K [[buffer(0)]], device const bfloat* V [[buffer(1)]], device const Q27CachePtrs* caches [[buffer(2)]], device const int* meta [[buffer(3)]], device const Q27KeepRow* map [[buffer(4)]], - constant uint& kept [[buffer(5)]], uint3 at [[thread_position_in_grid]]) { - if (at.x >= 256 || at.y >= 4 || at.z >= kept) return; + constant Q27KeepArgs& a [[buffer(5)]], uint3 at [[thread_position_in_grid]]) { + if (at.x >= 256 || at.y >= a.kv_heads || at.z >= a.kept) return; const Q27KeepRow r = map[at.z]; if (r.stream == 0xffffffffu) return; const uint base = 8 + r.stream * 12, destination = uint(meta[base + 4]) + r.destination; - const ulong source = (ulong(r.row) * 4 + at.y) * 256 + at.x; + const ulong source = (ulong(r.row) * a.kv_heads + at.y) * 256 + at.x; caches[r.stream].keys[ulong(at.y) * uint(meta[base + 7]) + destination * 256 + at.x] = K[source]; caches[r.stream].values[ulong(at.y) * uint(meta[base + 8]) + destination * 256 + at.x] = V[source]; } diff --git a/zig/src/families/qwen27/attention.zig b/zig/src/families/qwen27/attention.zig index e120aff5be..b139e90455 100644 --- a/zig/src/families/qwen27/attention.zig +++ b/zig/src/families/qwen27/attention.zig @@ -10,6 +10,8 @@ pub const Cache = struct { keys: Ref, values: Ref, capacity: u32, key_stride: u3 const CachePtrs = extern struct { keys: u64, values: u64 }; /// The pack and append kernels' arguments (public for the kernel contract's io check). pub const PackArgs = extern struct { rows: u32, packed_rows: u32, heads: u32 = 24, kv_heads: u32 = 4, dim: u32 = 256, group: u32 = 6 }; +/// Keep's gather and scatter: rows kept, and the KV heads a staging row holds. +const KeepArgs = extern struct { kept: u32, kv_heads: u32 }; const PromptArgs = extern struct { start: i32, rows: i32, chunks: i32, rp: i32, key_stride: i32, value_stride: i32, scale2: f32, kv_heads: i32 }; pub const Limits = struct { rows: u32 = 128, @@ -414,6 +416,9 @@ pub const Attention = struct { if (kept > scratch.limits.rows or binding.kept == null or binding.kept.? != kept) return error.AttentionKeep; try binding.checkCaches(caches); const g = binding.limits.geometry; + // the staging rows (scratch's keep_k, keep_v) hold scratch's KV heads; the caches' must be the same count + if (scratch.limits.geometry.kv_heads != g.kv_heads) return error.AttentionKeep; + const args = KeepArgs{ .kept = kept, .kv_heads = g.kv_heads }; for (scratch.buffers) |buffer| e.useResource(buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); for (binding.buffers) |buffer| e.useResource(buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); if (kept == 0) return; @@ -423,12 +428,12 @@ pub const Attention = struct { } e.setPipeline(a.io[5]); bind(e, .{ binding.get(.caches), binding.get(.meta), binding.get(.keep_map), scratch.get(.keep_k), scratch.get(.keep_v) }); - e.setValue(kept, 5); + e.setValue(args, 5); e.dispatchThreads(mtl.Size.of(g.dim, g.kv_heads, kept), mtl.Size.of(32, 4, 1)); e.barrier(); e.setPipeline(a.io[6]); bind(e, .{ scratch.get(.keep_k), scratch.get(.keep_v), binding.get(.caches), binding.get(.meta), binding.get(.keep_map) }); - e.setValue(kept, 5); + e.setValue(args, 5); e.dispatchThreads(mtl.Size.of(g.dim, g.kv_heads, kept), mtl.Size.of(32, 4, 1)); e.barrier(); } diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 3188b8cf8a..d8bae2a208 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -420,7 +420,7 @@ pub fn main(init: std.process.Init) !u8 { var io_regs: std.ArrayList(io_k.Reg) = .empty; defer io_regs.deinit(a); for (std.enums.values(attn_k.Heads)) |h| for (std.enums.values(io_k.Which)) |w| try io_regs.append(a, .{ .heads = h, .which = w }); - try io_regs.append(a, .{ .heads = .q27b, .which = .append, .control = true }); + try io_regs.appendSlice(a, &.{ .{ .heads = .q27b, .which = .append, .control = true }, .{ .heads = .rank, .which = .device_keep, .control = true } }); for (io_regs.items) |r| { if (!wanted(o, io_k.spec(r).name)) continue; const x = try a.create(io_k.Io); diff --git a/zig/tests/contract/metal_q27_attn.zig b/zig/tests/contract/metal_q27_attn.zig index 7829983109..fad56aae29 100644 --- a/zig/tests/contract/metal_q27_attn.zig +++ b/zig/tests/contract/metal_q27_attn.zig @@ -26,15 +26,15 @@ pub fn label(h: Heads) []const u8 { return if (h == .q27b) "27B" else "27B rank"; } -/// A planted control's one edit to the shipped source; `from` must occur once. -pub const Edit = struct { from: []const u8, to: []const u8 }; +/// A planted control's edit to the shipped source; `from` must occur exactly `count` times. +pub const Edit = struct { from: []const u8, to: []const u8, count: usize = 1 }; /// The source as decode_round.zig builds it here (simd_attention.library, the prompt kernels included). pub fn source(a: std.mem.Allocator, tensor: bool, edit: ?Edit) ![]u8 { const text = try q.simd_attention.library(a, tensor, true); const e = edit orelse return text; defer a.free(text); - if (std.mem.count(u8, text, e.from) != 1) return error.ControlEdit; + if (std.mem.count(u8, text, e.from) != e.count) return error.ControlEdit; return std.mem.replaceOwned(u8, a, text, e.from, e.to); } diff --git a/zig/tests/contract/metal_q27_attn_io.zig b/zig/tests/contract/metal_q27_attn_io.zig index 4434942b30..be8a023832 100644 --- a/zig/tests/contract/metal_q27_attn_io.zig +++ b/zig/tests/contract/metal_q27_attn_io.zig @@ -26,6 +26,7 @@ fn noSums(_: c.Shape) usize { } fn nameOf(comptime r: Reg) []const u8 { + if (r.control and r.which == .device_keep) return "control: qwen27 attention io keep staging 4 heads a row on a rank"; if (r.control) return "control: qwen27 attention io append one row late"; const what = comptime switch (r.which) { .pack => "pack q27_attn_pack_a q27_attn_pack_b", @@ -71,7 +72,10 @@ pub const Io = struct { pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, r: Reg) !Io { const tensor = device.tensorUnits(); - const edit: ?at_.Edit = if (r.control) .{ .from = "row = uint(meta[base + 4]) + local;", .to = "row = uint(meta[base + 4]) + local + 1;" } else null; + // controls: an append one row late; keep's gather and scatter both indexing 4 heads a staging row again + const late: at_.Edit = .{ .from = "row = uint(meta[base + 4]) + local;", .to = "row = uint(meta[base + 4]) + local + 1;" }; + const four: at_.Edit = .{ .from = "(ulong(r.row) * a.kv_heads + at.y)", .to = "(ulong(r.row) * 4 + at.y)", .count = 2 }; + const edit: ?at_.Edit = if (!r.control) null else if (r.which == .device_keep) four else late; const text = try at_.source(a, tensor, edit); defer a.free(text); const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); From a3efe3ce58da16fbd1c33cbfd86da9faccc3bfcc Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Tue, 3 Nov 2026 10:06:55 +0000 Subject: [PATCH 3/3] changelog: #2013 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 42c9fcbdf7..2c64ffb3a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Metal: Qwen3.8-27B's keep gathers and scatters the binding's KV heads a row, so a keep stages what the layout holds (#2013). Thanks to @cshintov. + - `TF_LOAD_PREFETCH`: Gemma 4 and Qwen 3.5 on CUDA read each layer ahead while the current one uploads, so the 27B loads about a third faster on a GB10. Replies are unchanged, and `TF_LOAD_PREFETCH=0` is the old path (#1928). Thanks to @jschmied. - Flash Next on CUDA loads its experts about 12% faster with the direct load: the next chunk is read while the current one packs, and the store switches shards without reopening them. Replies are unchanged (#1960). Thanks to @jschmied.