qwen27 metal: keys and values in core pool pages behind TF_QWEN27_METAL_KV_PAGES=1 (#1293)

CHANGELOG.md

@@ -4,6 +4,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- Qwen3.8-27B on Metal, opt-in with `TF_QWEN27_METAL_KV_PAGES=1`: each stream's keys and values live in pages of the core KV pool instead of their own buffers, with greedy replies equal to the default on one and two streams. The paged kernels are derived from the default ones, and the page row, page tensor and source edits they share with Nemotron's and Qwen3.5's now live in one place. It has run on Macs without tensor units only; an M5 refuses the flag by name for now (#2283, part of #1293).
 - HIP: the qwen3_5 engine reads a DFlash2 drafter's `config.json` against its target, refusing one it can't run, and loads its weights as MLX affine-4 (groups of 64), quantized on the host as on CUDA; nothing serves the drafter yet (#2477, part 2 of #1563). Thanks to @AbrahamCete.
 
 - Developers: the Qwen3.8-27B Metal lane check initializes its draft-window costs, and the served draft check compares complete state at the consumed token position, with exact tokens and state bytes (#2489).

zig/build/qwen27.zig

@@ -108,10 +108,12 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     const family: []const std.Build.Module.Import = &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "core", .module = qwen_core }, .{ .name = "metal", .module = metal } };
     _ = program(b, "tf-qwen27-run", "Build the native greedy, chunk and prompt-speed check on real weights", "zig/tests/qwen27/run.zig", target, true, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "core", .module = qwen_core }, .{ .name = "metal", .module = metal }, .{ .name = "core_lane", .module = core_lane }, .{ .name = "qwen_runtime_sources", .module = s.runtime_sources }, .{ .name = "qwen_gdn_source", .module = s.gdn_source } });
     _ = program(b, "tf-qwen27-attention-fused", "Build the fused-chain versus tail-kernel attention byte check on this Mac's tensor units, no models", "zig/tests/qwen27/attention_fused.zig", target, false, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "metal", .module = metal }, .{ .name = "qwen_runtime_sources", .module = s.runtime_sources } });
+    _ = program(b, "tf-qwen27-attention-pages", "Build the KV-page attention kernels against main's, bit for bit through shuffled pages, no models", "zig/tests/qwen27/attention_pages.zig", target, false, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "metal", .module = metal } });
     _ = program(b, "tf-qwen27-attention-precision", "Build decode attention's error against a float64 reference on this Mac, no models", "zig/tests/qwen27/attention_precision.zig", target, false, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "metal", .module = metal }, .{ .name = "qwen_runtime_sources", .module = s.runtime_sources } });
     _ = program(b, "tf-qwen27-projection-cases", "Build real-input production projection precision replay", "zig/tests/qwen27/projection_cases.zig", target, true, &.{ .{ .name = "qwen27", .module = qwen27 }, .{ .name = "metal", .module = metal } });
     _ = program(b, "tf-qwen27-dflash-run", "Build the explicit draft prototype and its plain comparison on real weights", "zig/tests/qwen27/dflash_run.zig", target, true, family);
     const engine_only: []const std.Build.Module.Import = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "metal", .module = metal } };
+    _ = program(b, "tf-qwen27-pages-resume", "Build kept states saved, restored and moved through KV pages against main's, no models", "zig/tests/qwen27/pages_resume.zig", target, false, engine_only);
     _ = program(b, "tf-qwen27-empty-keep", "Build shared empty/full keep state checks", "zig/tests/qwen27/empty_keep.zig", target, false, engine_only);
     _ = program(b, "tf-qwen27-keep-error", "Build the check that a failing keep returns its error, never aborting", "zig/tests/qwen27/keep_error.zig", target, false, engine_only);
     _ = program(b, "tf-qwen27-synthetic", "Build bounded synthetic native hybrid chunk checks, no models", "zig/tests/qwen27/synthetic.zig", target, false, engine_only);

zig/kernels/metal/nemotron/paged_tensor.h → zig/kernels/metal/core/kv_page_tensor.h

@@ -1,8 +1,8 @@
-// A 64-key tile's rows of kv head hk as a tensor over its page: rows to L, at most to the page's end.
+// A tile's rows of kv head hk from kt as an MPP tensor over its head-major page: rows to L, at most to the page's end.
 template <int D>
-inline tensor<device bfloat, dextents<int32_t, 2>, tensor_inline> nemo_page_tensor(const device bfloat16_t* base, device const uint* table,
-                                                                                   ulong page_bytes, int kt, int L, uint hk) {
+inline tensor<device bfloat, dextents<int32_t, 2>, tensor_inline> tf_page_tensor(const device bfloat* base, device const uint* table,
+                                                                             ulong page_bytes, int kt, int L, uint hk) {
   const int rows = min(TF_PAGE_TOKENS - kt % TF_PAGE_TOKENS, L - kt);
-  device bfloat* p = (device bfloat*)((const device uchar*)base + nemo_paged_row(table, uint(kt), page_bytes, hk, D));
+  device bfloat* p = (device bfloat*)((const device uchar*)base + tf_paged_row(table, uint(kt), page_bytes, hk, D));
   return tensor<device bfloat, dextents<int32_t, 2>, tensor_inline>(p, dextents<int32_t, 2>(D, rows), array<int32_t, 2>({1, D}));
 }

zig/kernels/metal/core/kv_pages.h

@@ -3,3 +3,7 @@
 inline ulong tf_page_offset(device const uint* table, uint pos, ulong page_bytes, ulong row_bytes) {
   return ulong(table[pos / TF_PAGE_TOKENS]) * page_bytes + ulong(pos % TF_PAGE_TOKENS) * row_bytes;
 }
+// A head-major page holds [kv_heads][TF_PAGE_TOKENS][hd] bf16 of one cache (keys or values): head h's rows in a run.
+inline ulong tf_paged_row(device const uint* table, uint pos, ulong page_bytes, uint h, uint hd) {
+  return tf_page_offset(table, pos, page_bytes, ulong(hd) * 2) + ulong(h) * TF_PAGE_TOKENS * hd * 2;
+}

zig/kernels/metal/core/source_edit.zig

@@ -0,0 +1,29 @@
+//! Count-checked comptime edits of a kernel's source text, for paged twins derived from main's kernels.
+const std = @import("std");
+
+/// Every match of `from` in `text` becomes `to`, and there must be `count` of them, or the build stops naming it.
+pub fn edit(comptime text: []const u8, comptime from: []const u8, comptime to: []const u8, comptime count: usize) []const u8 {
+    @setEvalBranchQuota(20_000_000);
+    if (std.mem.count(u8, text, from) != count) @compileError("source edit: expected " ++ std.fmt.comptimePrint("{d}", .{count}) ++ " of: " ++ from);
+    var out: []const u8 = "";
+    var rest = text;
+    while (std.mem.indexOf(u8, rest, from)) |i| {
+        out = out ++ rest[0..i] ++ to;
+        rest = rest[i + from.len ..];
+    }
+    return out ++ rest;
+}
+
+/// The text from `start` up to (not including) the next `end`.
+pub fn between(comptime text: []const u8, comptime start: []const u8, comptime end: []const u8) []const u8 {
+    @setEvalBranchQuota(20_000_000);
+    const a = std.mem.indexOf(u8, text, start) orelse @compileError("source edit: no " ++ start);
+    const b = std.mem.indexOfPos(u8, text, a, end) orelse @compileError("source edit: no " ++ end);
+    return text[a..b];
+}
+
+/// The build stops when `text` no longer holds `what`.
+pub fn present(comptime text: []const u8, comptime what: []const u8) void {
+    @setEvalBranchQuota(20_000_000);
+    if (std.mem.indexOf(u8, text, what) == null) @compileError("source edit: the source no longer holds " ++ what);
+}

zig/kernels/metal/nemotron/paged.metal

@@ -11,7 +11,7 @@ kernel void nemotron_kv_write_paged(const device bfloat* qkv [[buffer(0)]],
                                     uint3 size [[threads_per_grid]]) {
   const uint d = pos.x, h = pos.y, t = pos.z, hd = size.x, kvh = size.y;
   const size_t src = size_t(t) * dims.x + dims.y + h * hd + d;
-  const ulong dst = nemo_paged_row(table, dims.w + t, page_bytes, h, hd) + ulong(d) * 2;
+  const ulong dst = tf_paged_row(table, dims.w + t, page_bytes, h, hd) + ulong(d) * 2;
   *(device bfloat*)(kc + dst) = qkv[src];
   *(device bfloat*)(vc + dst) = qkv[src + kvh * hd];
 }
@@ -28,8 +28,8 @@ kernel void nemotron_kv_compact_paged(device uchar* K [[buffer(0)]],
   for (uint j = 0; j < n; j++) {
     const uint src = len + uint(path[j]), dst = len + j;
     if (src == dst) continue;
-    const ulong to = nemo_paged_row(table, dst, page_bytes, hk, hd) + ulong(d) * 2;
-    const ulong from = nemo_paged_row(table, src, page_bytes, hk, hd) + ulong(d) * 2;
+    const ulong to = tf_paged_row(table, dst, page_bytes, hk, hd) + ulong(d) * 2;
+    const ulong from = tf_paged_row(table, src, page_bytes, hk, hd) + ulong(d) * 2;
     *(device bfloat*)(K + to) = *(device bfloat*)(K + from);
     *(device bfloat*)(V + to) = *(device bfloat*)(V + from);
   }
@@ -47,6 +47,6 @@ kernel void nemotron_kv_put_paged(const device bfloat* src [[buffer(0)]],
   if (d >= D || row >= P[4] + P[0]) {
     return;
   }
-  const ulong to = nemo_paged_row(table, uint(row), page_bytes, uint(kv), uint(D)) + ulong(d) * 2;
+  const ulong to = tf_paged_row(table, uint(row), page_bytes, uint(kv), uint(D)) + ulong(d) * 2;
   *(device bfloat*)(dst + to) = row < P[4] ? old[(long(kv) * P[5] + row) * D + d] : src[(long(row - P[4]) * P[1] + kv) * D + d];
 }

zig/kernels/metal/nemotron/paged.zig

@@ -1,40 +1,17 @@
 //! Nemotron's KV-page twins: main's kernels with only the K/V address changed, the big three by count-checked edits.
 const std = @import("std");
 const generated = @import("kernels.zig");
+const se = @import("../core/source_edit.zig");
+const edit = se.edit;
+const between = se.between;
+const present = se.present;
 
 const kv_pages = @embedFile("../core/kv_pages.h");
-const paged_h = @embedFile("paged.h");
-
-/// One count-checked edit of a source's text: every match of `from` becomes `to`, and there must be `count` of them.
-fn edit(comptime text: []const u8, comptime from: []const u8, comptime to: []const u8, comptime count: usize) []const u8 {
-    @setEvalBranchQuota(20_000_000);
-    if (std.mem.count(u8, text, from) != count) @compileError("nemotron paged: expected " ++ std.fmt.comptimePrint("{d}", .{count}) ++ " of: " ++ from);
-    var out: []const u8 = "";
-    var rest = text;
-    while (std.mem.indexOf(u8, rest, from)) |i| {
-        out = out ++ rest[0..i] ++ to;
-        rest = rest[i + from.len ..];
-    }
-    return out ++ rest;
-}
-
-/// The text from `start` up to (not including) the next `end`.
-fn between(comptime text: []const u8, comptime start: []const u8, comptime end: []const u8) []const u8 {
-    @setEvalBranchQuota(20_000_000);
-    const a = std.mem.indexOf(u8, text, start) orelse @compileError("nemotron paged: no " ++ start);
-    const b = std.mem.indexOfPos(u8, text, a, end) orelse @compileError("nemotron paged: no " ++ end);
-    return text[a..b];
-}
-
-fn present(comptime text: []const u8, comptime what: []const u8) void {
-    @setEvalBranchQuota(20_000_000);
-    if (std.mem.indexOf(u8, text, what) == null) @compileError("nemotron paged: the source no longer holds " ++ what);
-}
 
 const table_args = "  device const uint* table [[buffer({d})]],\n  constant ulong& page_bytes [[buffer({d})]],\n";
 
 /// tf_kv_write, tf_kv_compact and prefill's kv_put through a stream's table (paged.metal), in one library.
-pub const glue = "#include <metal_stdlib>\nusing namespace metal;\n" ++ kv_pages ++ paged_h ++ @embedFile("paged.metal");
+pub const glue = "#include <metal_stdlib>\nusing namespace metal;\n" ++ kv_pages ++ @embedFile("paged.metal");
 pub const glue_names = [_][:0]const u8{ "nemotron_kv_write_paged", "nemotron_kv_compact_paged", "nemotron_kv_put_paged" };
 
 /// The decode attention's partial pass (attn_partial_1..16): each 64-key tile's keys and values from its page.
@@ -50,9 +27,9 @@ pub const partial = blk: {
     body = edit(body, name ++ "(", "nemotron_attn_partial_paged(", 1);
     body = edit(body, "  device float* PL [[buffer(9)]],\n", "  device float* PL [[buffer(9)]],\n" ++ std.fmt.comptimePrint(table_args, .{ 10, 11 }), 1);
     // a tile's slice of the whole cache, as a slice of its page bounded the same way: the rows to L, at most the page's
-    body = edit(body, "    auto bK = tK.slice(0, kt);\n", "    auto bK = nemo_page_tensor<D>(K, table, page_bytes, kt, L, hk).slice(0, 0);\n", 1);
-    body = edit(body, "    auto bVlo = tV.slice(0, kt);\n", "    auto bVlo = nemo_page_tensor<D>(V, table, page_bytes, kt, L, hk).slice(0, 0);\n", 1);
-    var out: []const u8 = source[0..at] ++ kv_pages ++ paged_h ++ @embedFile("paged_tensor.h") ++ body;
+    body = edit(body, "    auto bK = tK.slice(0, kt);\n", "    auto bK = tf_page_tensor<D>(K, table, page_bytes, kt, L, hk).slice(0, 0);\n", 1);
+    body = edit(body, "    auto bVlo = tV.slice(0, kt);\n", "    auto bVlo = tf_page_tensor<D>(V, table, page_bytes, kt, L, hk).slice(0, 0);\n", 1);
+    var out: []const u8 = source[0..at] ++ kv_pages ++ @embedFile("../core/kv_page_tensor.h") ++ body;
     for (1..17) |n| {
         const args = std.fmt.comptimePrint("<16, 128, {d}, 512, 64>", .{n});
         out = out ++ std.fmt.comptimePrint("template [[host_name(\"nemotron_attn_partial_paged_{d}\")]] [[kernel]] ", .{n}) ++
@@ -75,10 +52,10 @@ pub const tree = blk: {
     var body = between(source, "// A node's keys past the shared pass", "// A node's attention output");
     body = edit(body, "[[kernel]] void tf_tree_tail(", "[[kernel]] void nemotron_tree_tail_paged(", 1);
     body = edit(body, "  device float* PL [[buffer(14)]],\n", "  device float* PL [[buffer(14)]],\n" ++ std.fmt.comptimePrint(table_args, .{ 15, 16 }), 1);
-    body = edit(body, "(kbase + phys * kstep + col)", "((const device uchar*)K + nemo_paged_row(table, uint(phys), page_bytes, hk, D) + col * 2)", 1);
-    body = edit(body, "(vbase + phys * vstep + col)", "((const device uchar*)V + nemo_paged_row(table, uint(phys), page_bytes, hk, D) + col * 2)", 1);
+    body = edit(body, "(kbase + phys * kstep + col)", "((const device uchar*)K + tf_paged_row(table, uint(phys), page_bytes, hk, D) + col * 2)", 1);
+    body = edit(body, "(vbase + phys * vstep + col)", "((const device uchar*)V + tf_paged_row(table, uint(phys), page_bytes, hk, D) + col * 2)", 1);
     const args = "<16, 128, 512, 64, 64>";
-    break :blk source ++ kv_pages ++ paged_h ++ body ++ "template [[host_name(\"nemotron_tree_tail_paged\")]] [[kernel]] " ++
+    break :blk source ++ kv_pages ++ body ++ "template [[host_name(\"nemotron_tree_tail_paged\")]] [[kernel]] " ++
         "decltype(nemotron_tree_tail_paged" ++ args ++ ") nemotron_tree_tail_paged" ++ args ++ ";\n";
 };
 pub const tree_name: [:0]const u8 = "nemotron_tree_tail_paged";
@@ -92,12 +69,12 @@ pub const prefill = blk: {
     // a stream's keys start at its table's first page: no batch or head offset, a block's rows found by position
     body = edit(body, "  K += long(tg.z) * P[15] + long(kv) * P[16];\n  V += long(tg.z) * P[18] + long(kv) * P[19];\n", "", 1);
     body = edit(body, "    const bool last_k = kb == P[8];\n", "    const bool last_k = kb == P[8];\n" ++
-        "    const device T* Kb = (const device T*)((const device uchar*)K + nemo_paged_row(table, uint(kb * BK), page_bytes, uint(kv), uint(ks)));\n" ++
-        "    const device T* Vb = (const device T*)((const device uchar*)V + nemo_paged_row(table, uint(kb * BK), page_bytes, uint(kv), uint(vs)));\n", 1);
+        "    const device T* Kb = (const device T*)((const device uchar*)K + tf_paged_row(table, uint(kb * BK), page_bytes, uint(kv), uint(ks)));\n" ++
+        "    const device T* Vb = (const device T*)((const device uchar*)V + tf_paged_row(table, uint(kb * BK), page_bytes, uint(kv), uint(vs)));\n", 1);
     body = edit(body, ", K, ks, ", ", Kb, ks, ", 4);
     body = edit(body, ", V, vs, ", ", Vb, vs, ", 4);
     body = edit(body, "    K += BK * ks;\n    V += BK * vs;\n", "", 1);
-    var out: []const u8 = source ++ kv_pages ++ paged_h ++ "namespace tfp {\n" ++ body ++ "}  // namespace tfp\n";
+    var out: []const u8 = source ++ kv_pages ++ "namespace tfp {\n" ++ body ++ "}  // namespace tfp\n";
     for (prefill_tags) |t| {
         present(source, "custom_kernel_tf_attention_nax_bf16_" ++ t ++ "_t_bfloat16_t");
         out = out ++ "[[max_total_threads_per_threadgroup(128)]]\n[[kernel]] void nemotron_attention_paged_" ++ t ++ "(\n" ++