qwen27: keep's gather and scatter stage the binding's KV heads a row, not 4

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.

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];
 }

zig/src/families/qwen27/attention.zig

@@ -8,7 +8,10 @@ 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 };
+/// 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,
@@ -413,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;
@@ -422,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();
     }

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

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 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) != e.count) 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 };
+}