diff --git a/CHANGELOG.md b/CHANGELOG.md index b480dda9f0..6e4ce04513 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index 524cf517e6..88f939e798 100644 --- a/zig/build/qwen27.zig +++ b/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); diff --git a/zig/kernels/metal/nemotron/paged_tensor.h b/zig/kernels/metal/core/kv_page_tensor.h similarity index 57% rename from zig/kernels/metal/nemotron/paged_tensor.h rename to zig/kernels/metal/core/kv_page_tensor.h index 66fab520d5..3d98807db8 100644 --- a/zig/kernels/metal/nemotron/paged_tensor.h +++ b/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 -inline tensor, 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, 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, tensor_inline>(p, dextents(D, rows), array({1, D})); } diff --git a/zig/kernels/metal/core/kv_pages.h b/zig/kernels/metal/core/kv_pages.h index d08507f5bf..c12ad07535 100644 --- a/zig/kernels/metal/core/kv_pages.h +++ b/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; +} diff --git a/zig/kernels/metal/core/source_edit.zig b/zig/kernels/metal/core/source_edit.zig new file mode 100644 index 0000000000..7e68a23a2d --- /dev/null +++ b/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); +} diff --git a/zig/kernels/metal/nemotron/paged.h b/zig/kernels/metal/nemotron/paged.h deleted file mode 100644 index e40076c83f..0000000000 --- a/zig/kernels/metal/nemotron/paged.h +++ /dev/null @@ -1,4 +0,0 @@ -// A Nemotron page holds [kv_heads][TF_PAGE_TOKENS][hd] bf16 of one cache (keys or values): head h's rows in a run. -inline ulong nemo_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; -} diff --git a/zig/kernels/metal/nemotron/paged.metal b/zig/kernels/metal/nemotron/paged.metal index 47d47ddd78..f8e194b76e 100644 --- a/zig/kernels/metal/nemotron/paged.metal +++ b/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]; } diff --git a/zig/kernels/metal/nemotron/paged.zig b/zig/kernels/metal/nemotron/paged.zig index 92ede00c7f..b6ec1b5c94 100644 --- a/zig/kernels/metal/nemotron/paged.zig +++ b/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 \nusing namespace metal;\n" ++ kv_pages ++ paged_h ++ @embedFile("paged.metal"); +pub const glue = "#include \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(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(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(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(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" ++ diff --git a/zig/kernels/metal/qwen27/paged.zig b/zig/kernels/metal/qwen27/paged.zig new file mode 100644 index 0000000000..b48517aeb9 --- /dev/null +++ b/zig/kernels/metal/qwen27/paged.zig @@ -0,0 +1,75 @@ +//! The 27B's KV-page twins: main's attention sources, only the K/V address edited; names kept, a library swaps whole. +const se = @import("../core/source_edit.zig"); +const edit = se.edit; +const between = se.between; + +const headers = @embedFile("../core/kv_pages.h") ++ @embedFile("../core/kv_page_tensor.h"); + +/// A prompt step's keys or values as a base that its own `kt * PD` rows land in kt's page (a 32-key step is in one). +const prompt_base = + "inline const device bfloat* q27_page_base(const device bfloat* p, Q27AttnReadPtrs c, int kt, int hk) {\n" ++ + " return p + int64_t(tf_paged_row(c.table, uint(kt), c.page_bytes, uint(hk), 256) / 2) - int64_t(kt) * 256;\n}\n"; + +/// Prefix (each 64-key tile from its page), tail (each gathered row from its page), and the read pointers' struct. +pub const attention_mpp = blk: { + var s: []const u8 = @embedFile("attention_mpp.metal"); + s = edit(s, "using namespace mpp::tensor_ops;\n", "using namespace mpp::tensor_ops;\n" ++ headers, 1); + s = edit(s, "struct Q27AttnReadPtrs { device const bfloat* keys; device const bfloat* values; };\n", "struct Q27AttnReadPtrs { device const bfloat* keys; device const bfloat* values; device const uint* table; ulong page_bytes; };\n" ++ prompt_base, 1); + s = edit(s, " auto bK = tK.slice(0, kt);\n", " auto bK = tf_page_tensor(Kb, caches[st].table, caches[st].page_bytes, kt, L, hk).slice(0, 0);\n", 1); + s = edit(s, " auto bV = tV.slice(hi ? 128 : 0, kt);\n", " auto bV = tf_page_tensor(Vb, caches[st].table, caches[st].page_bytes, kt, L, hk).slice(hi ? 128 : 0, 0);\n", 1); + // the tail's head base and row step are the page's now + s = edit(s, " const device bfloat* kbase = (const device bfloat*)Kb + (int64_t)hk * meta[MG + st * MS + 7];\n" ++ + " const device bfloat* vbase = (const device bfloat*)Vb + (int64_t)hk * meta[MG + st * MS + 8];\n const int64_t kstep = D, vstep = D;\n", "", 1); + s = edit(s, "(kbase + phys * kstep + col)", "((const device bfloat*)((const device uchar*)Kb + tf_paged_row(caches[st].table, uint(phys), caches[st].page_bytes, hk, D)) + col)", 1); + s = edit(s, "(vbase + phys * vstep + col)", "((const device bfloat*)((const device uchar*)Vb + tf_paged_row(caches[st].table, uint(phys), caches[st].page_bytes, hk, D)) + col)", 1); + // the non-finite scan over a key range: main's, each key's row from its page (q27_nonfinite_paged, beside it) + s = edit(s, "\n// SG tiles a threadgroup, a simdgroup pair each", "\n" ++ scan(s) ++ "\n// SG tiles a threadgroup, a simdgroup pair each", 1); + // the masked prefix: its half's base without a head stride, its scan and each masked key's values through the table + s = edit(s, "Vb + (int64_t)hk * meta[MG + st * MS + 8] + (hi ? 128 : 0)", "Vb + (hi ? 128 : 0)", 1); + s = edit(s, "q27_nonfinite(vh, max(seen, kbeg), kend, lane)", "q27_nonfinite_paged(vh, caches[st], hk, max(seen, kbeg), kend, lane)", 1); + s = edit(s, "vh + (int64_t)key * D + fn", "vh + (int64_t)(tf_paged_row(caches[st].table, uint(key), caches[st].page_bytes, hk, D) / 2) + fn", 1); + break :blk s; +}; + +/// main's q27_nonfinite as q27_nonfinite_paged: `v` the half's base, a key's row in its page of `c`, head `hk`. +fn scan(comptime text: []const u8) []const u8 { + var f = between(text, "inline bool q27_nonfinite(", "\n// SG tiles"); + f = edit(f, "q27_nonfinite(const device bfloat* v, int lo,", "q27_nonfinite_paged(const device bfloat* v, Q27AttnReadPtrs c, uint hk, int lo,", 1); + f = edit(f, "(v + (int64_t)key * D + (lane & 15) * 8)", "(v + (int64_t)(tf_paged_row(c.table, uint(key), c.page_bytes, hk, D) / 2) + (lane & 15) * 8)", 1); + return f; +} + +/// Append, gather and scatter: each row through its stream's table. +pub const attention_io = blk: { + var s: []const u8 = @embedFile("attention_io.metal"); + s = edit(s, "struct Q27CachePtrs { device bfloat* keys; device bfloat* values; };\n", "struct Q27CachePtrs { device bfloat* keys; device bfloat* values; device const uint* table; ulong page_bytes; };\n" ++ + "inline device bfloat* q27_paged(device bfloat* base, Q27CachePtrs c, uint hk, uint row, uint dim) {\n" ++ + " return (device bfloat*)((device uchar*)base + tf_paged_row(c.table, row, c.page_bytes, hk, dim));\n}\n", 1); + s = edit(s, "caches[st].keys[ulong(at.y) * uint(meta[base + 7]) + row * a.dim + at.x]", "q27_paged(caches[st].keys, caches[st], at.y, row, a.dim)[at.x]", 1); + s = edit(s, "caches[st].values[ulong(at.y) * uint(meta[base + 8]) + row * a.dim + at.x]", "q27_paged(caches[st].values, caches[st], at.y, row, a.dim)[at.x]", 1); + s = edit(s, "caches[r.stream].keys[ulong(at.y) * uint(meta[base + 7]) + source * 256 + at.x]", "q27_paged(caches[r.stream].keys, caches[r.stream], at.y, source, 256)[at.x]", 1); + s = edit(s, "caches[r.stream].values[ulong(at.y) * uint(meta[base + 8]) + source * 256 + at.x]", "q27_paged(caches[r.stream].values, caches[r.stream], at.y, source, 256)[at.x]", 1); + s = edit(s, "caches[r.stream].keys[ulong(at.y) * uint(meta[base + 7]) + destination * 256 + at.x]", "q27_paged(caches[r.stream].keys, caches[r.stream], at.y, destination, 256)[at.x]", 1); + s = edit(s, "caches[r.stream].values[ulong(at.y) * uint(meta[base + 8]) + destination * 256 + at.x]", "q27_paged(caches[r.stream].values, caches[r.stream], at.y, destination, 256)[at.x]", 1); + break :blk s; +}; + +/// A prompt kernel's keys and values: no head stride (a page's head runs are inside it), each step based at its page. +fn prompt(comptime text: []const u8) []const u8 { + var s = text; + s = edit(s, "caches[0].keys + (int64_t)hk * args.key_stride + h * PH", "caches[0].keys + h * PH", 1); + s = edit(s, "caches[0].values + (int64_t)hk * args.value_stride + h * PH", "caches[0].values + h * PH", 1); + // every step (main's, the masked steps and the edge) based at its page + s = edit(s, "(r, kt, kend, q, k, v, xs", "(r, kt, kend, q, q27_page_base(k, caches[0], kt, hk), q27_page_base(v, caches[0], kt, hk), xs", 4); + s = edit(s, "q27_nonfinite(v, max(seen, kbeg), kend, lane)", "q27_nonfinite_paged(v, caches[0], uint(hk), max(seen, kbeg), kend, lane)", 1); + return s; +} + +/// The prompt kernels and the scan picking between them: each scanned key's values through stream 0's table. +pub const attention_prompt = blk: { + var s = prompt(@embedFile("attention_prompt.metal")); + s = edit(s, "caches[0].values + (int64_t)tg.x * args.value_stride;", "caches[0].values;", 1); + s = edit(s, "(v + (int64_t)key * PD + lane * 8)", "(v + (int64_t)(tf_paged_row(caches[0].table, uint(key), caches[0].page_bytes, tg.x, PD) / 2) + lane * 8)", 1); + break :blk s; +}; +pub const attention_prompt_sg = prompt(@embedFile("attention_prompt_sg.metal")); diff --git a/zig/kernels/metal/qwen3_5/paged.metal b/zig/kernels/metal/qwen3_5/paged.metal index 11f230264d..eb10e4ac3f 100644 --- a/zig/kernels/metal/qwen3_5/paged.metal +++ b/zig/kernels/metal/qwen3_5/paged.metal @@ -1,9 +1,5 @@ // Qwen3.5-2B's keys and attention over core-pool pages: the contiguous kernels with only the K/V address changed. -inline ulong qwen35_page_row(device const uint* table, uint pos, ulong page_bytes, uint h) { - return tf_page_offset(table, pos, page_bytes, 256 * 2) + ulong(h) * TF_PAGE_TOKENS * 256 * 2; -} - kernel void qwen35_keys_paged(const device bfloat* k [[buffer(0)]], const device bfloat* v [[buffer(1)]], device uchar* keys [[buffer(2)]], @@ -14,7 +10,7 @@ kernel void qwen35_keys_paged(const device bfloat* k [[buffer(0)]], uint3 p [[thread_position_in_grid]]) { const uint d = p.x, h = p.y, r = p.z; const uint m = dims.z + r; - const ulong at = qwen35_page_row(table, dims.x + r, page_bytes, h); + const ulong at = tf_paged_row(table, dims.x + r, page_bytes, h, 256); ((device bfloat*)(keys + at))[d] = qwen35_rotate(k + (size_t(m) * 2 + h) * 256, d, dims.x + r); ((device bfloat*)(values + at))[d] = v[(size_t(m) * 2 + h) * 256 + d]; } diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 9663500992..45d4848a42 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -7,6 +7,8 @@ pub const dsv41 = @import("deepseek_v41/sources.zig"); pub const nemotron = @import("nemotron/kernels.zig"); /// Nemotron's KV in pages of the core pool: twins of its seven K/V kernels, each library beside its original. pub const nemotron_paged = @import("nemotron/paged.zig"); +/// The 27B's KV in pages of the core pool: its attention sources with the K/V address through a stream's table. +pub const qwen27_paged = @import("qwen27/paged.zig"); pub const qwen35 = @import("qwen3_5/kernels.zig"); pub const qwen35_layout = @embedFile("qwen3_5/layout.metal"); diff --git a/zig/qwen27_core.zig b/zig/qwen27_core.zig index 58ab61fe50..669a5c74ce 100644 --- a/zig/qwen27_core.zig +++ b/zig/qwen27_core.zig @@ -23,6 +23,9 @@ pub const quant = @import("src/core/quant/quant.zig"); pub const draft_ops = @import("draft_ops"); pub const shared_attention = @import("src/core/shared_attention.zig"); +/// The core pool and its Metal pages, which TF_QWEN27_METAL_KV_PAGES keeps the 27B's keys and values in. +pub const kv_pool = @import("src/core/kv_pool.zig"); +pub const metal_kv_pages = @import("src/core/metal_kv_pages.zig"); pub const affine4 = @import("src/core/affine4.zig"); pub const affine4_lane = @import("src/core/affine4_lane.zig"); diff --git a/zig/src/families/qwen27/attention.zig b/zig/src/families/qwen27/attention.zig index 7ccd2be6b0..4057cb8ec3 100644 --- a/zig/src/families/qwen27/attention.zig +++ b/zig/src/families/qwen27/attention.zig @@ -6,8 +6,20 @@ const Ref = @import("projection.zig").Ref; const glue = @import("glue.zig"); 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 }; +pub const Cache = struct { + keys: Ref, + values: Ref, + capacity: u32, + key_stride: u32 = 0, + value_stride: u32 = 0, + /// TF_QWEN27_METAL_KV_PAGES: keys and values are this layer's part of every page; rows found through the table. + pages: ?Pages = null, +}; +/// A paged stream's table (u32 page ids, 256 rows a page) and the bytes from one page to the next in the arena. +pub const Pages = struct { table: Ref, page_bytes: u64 }; const CachePtrs = extern struct { keys: u64, values: u64 }; +/// The paged kernels' read pointers (kernels/metal/qwen27/paged.zig): main's two and the stream's table and page size. +const PagedPtrs = extern struct { keys: u64, values: u64, table: u64, page_bytes: 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. @@ -30,6 +42,7 @@ pub const Binding = struct { uploaded: ?[32]u8 = null, streams: u32 = 0, pointers: [planning.max_streams]CachePtrs = @splat(.{ .keys = 0, .values = 0 }), + paged: bool = false, strides: [planning.max_streams][2]u32 = @splat(.{ 0, 0 }), kept: ?u32 = null, pub fn init(device: mtl.Device, limits: Limits) !Binding { @@ -37,7 +50,7 @@ pub const Binding = struct { var buffers: [3]mtl.Buffer = undefined; var made: usize = 0; errdefer for (buffers[0..made]) |buffer| buffer.deinit(); - for ([_]usize{ (planning.global_words + planning.max_streams * planning.stream_words) * 4, planning.max_streams * @sizeOf(CachePtrs), limits.rows * 16 }, &buffers) |allocation_bytes, *buffer| { + for ([_]usize{ (planning.global_words + planning.max_streams * planning.stream_words) * 4, planning.max_streams * @sizeOf(PagedPtrs), limits.rows * 16 }, &buffers) |allocation_bytes, *buffer| { buffer.* = try device.buffer(allocation_bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); made += 1; } @@ -53,11 +66,24 @@ pub const Binding = struct { if (plan.rows > b.limits.rows or plan.streams > b.limits.streams or caches.len != plan.streams or !plan.geometry.sameShape(b.limits.geometry)) return error.AttentionLimits; var meta = plan.meta; var table: [planning.max_streams]CachePtrs = @splat(.{ .keys = 0, .values = 0 }); + var paged: [planning.max_streams]PagedPtrs = @splat(.{ .keys = 0, .values = 0, .table = 0, .page_bytes = 0 }); + const pages = caches.len > 0 and caches[0].pages != null; for (caches, 0..) |cache, i| { const at = 8 + i * 12; const start: u32 = @intCast(meta[at + 4]); const rows: u32 = @intCast(meta[at + 2]); if (cache.capacity == 0 or cache.capacity > std.math.maxInt(i32) / b.limits.geometry.dim or @as(u64, start) + rows > cache.capacity) return error.AttentionCapacity; + if ((cache.pages != null) != pages) return error.AttentionMixedPages; + if (cache.pages) |p| { + // a page's head runs are [kv_heads][256][dim]: no stride, and no span past the table's pages + try spanAligned(p.table, pageCount(cache.capacity) * 4, 4); + if (cache.key_stride != 0 or cache.value_stride != 0 or p.page_bytes % 16 != 0) return error.AttentionCacheStride; + table[i] = .{ .keys = cache.keys.buffer.gpuAddress() + cache.keys.offset, .values = cache.values.buffer.gpuAddress() + cache.values.offset }; + paged[i] = .{ .keys = table[i].keys, .values = table[i].values, .table = p.table.buffer.gpuAddress() + p.table.offset, .page_bytes = p.page_bytes }; + meta[at + 7] = @intCast(cache.capacity * b.limits.geometry.dim); + meta[at + 8] = @intCast(cache.capacity * b.limits.geometry.dim); + continue; + } const ks = if (cache.key_stride == 0) cache.capacity * b.limits.geometry.dim else cache.key_stride; const vs = if (cache.value_stride == 0) cache.capacity * b.limits.geometry.dim else cache.value_stride; if (ks < @as(u64, cache.capacity) * b.limits.geometry.dim or vs < @as(u64, cache.capacity) * b.limits.geometry.dim or ks > std.math.maxInt(i32) or vs > std.math.maxInt(i32)) return error.AttentionCacheStride; @@ -71,11 +97,12 @@ pub const Binding = struct { meta[at + 7] = @intCast(ks); meta[at + 8] = @intCast(vs); } - @memcpy(b.get(.caches).buffer.slice(CachePtrs, planning.max_streams), &table); + if (pages) @memcpy(b.get(.caches).buffer.slice(PagedPtrs, planning.max_streams), &paged) else @memcpy(b.get(.caches).buffer.slice(CachePtrs, planning.max_streams), &table); @memcpy(b.get(.meta).buffer.slice(i32, meta.len), &meta); b.uploaded = fingerprint(plan); b.streams = plan.streams; b.pointers = table; + b.paged = pages; for (caches, 0..) |_, i| b.strides[i] = .{ @intCast(meta[8 + i * 12 + 7]), @intCast(meta[8 + i * 12 + 8]) }; b.kept = null; } @@ -91,6 +118,11 @@ pub const Binding = struct { fn checkCaches(b: *const Binding, caches: []const Cache) !void { if (b.uploaded == null or caches.len != b.streams) return error.AttentionNotUploaded; for (caches, 0..) |cache, i| { + if ((cache.pages != null) != b.paged) return error.AttentionNotUploaded; + if (cache.pages != null) { + if (cache.keys.buffer.gpuAddress() + cache.keys.offset != b.pointers[i].keys or cache.values.buffer.gpuAddress() + cache.values.offset != b.pointers[i].values) return error.AttentionNotUploaded; + continue; + } if (cache.capacity == 0 or cache.capacity > std.math.maxInt(i32) / b.limits.geometry.dim or cache.keys.buffer.gpuAddress() + cache.keys.offset != b.pointers[i].keys or cache.values.buffer.gpuAddress() + cache.values.offset != b.pointers[i].values) return error.AttentionNotUploaded; const ks = if (cache.key_stride == 0) cache.capacity * b.limits.geometry.dim else cache.key_stride; const vs = if (cache.value_stride == 0) cache.capacity * b.limits.geometry.dim else cache.value_stride; @@ -190,6 +222,20 @@ fn cacheBytes(cache: Cache, key: bool, geometry: planning.Geometry) usize { const stride: usize = if (requested == 0) @as(usize, cache.capacity) * geometry.dim else requested; return ((@as(usize, geometry.kv_heads) - 1) * stride + @as(usize, cache.capacity) * geometry.dim) * 2; } +/// The caches' buffers, and a paged stream's table, resident for the pass's indirect reads. +fn useCaches(e: mtl.ComputeEncoder, caches: []const Cache) void { + for (caches) |cache| { + e.useResource(cache.keys.buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); + e.useResource(cache.values.buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); + if (cache.pages) |p| e.useResource(p.table.buffer, mtl.ResourceUsage.read); + } +} +/// Pages a stream of `rows` rows takes (metal_kv_pages.page_tokens a page). +pub fn pageCount(rows: u32) usize { + return (@as(usize, rows) + page_rows - 1) / page_rows; +} +/// Rows a page holds: kernels/metal/core/kv_pages.h's TF_PAGE_TOKENS. +pub const page_rows = 256; fn bind(e: mtl.ComputeEncoder, refs: anytype) void { inline for (refs, 0..) |r, i| e.setBuffer(r.buffer, r.offset, i); } @@ -316,15 +362,17 @@ pub const Attention = struct { } const operands = [_]struct { ref: Ref, bytes: usize }{ .{ .ref = q, .bytes = q_bytes }, .{ .ref = k, .bytes = kv_bytes }, .{ .ref = v, .bytes = kv_bytes }, .{ .ref = out, .bytes = q_bytes } }; for (caches) |cache| for (operands) |operand| { + if (cache.pages != null) { + // the arena holds only pages: an operand in it at all is an alias + if (operand.ref.buffer.id == cache.keys.buffer.id or operand.ref.buffer.id == cache.values.buffer.id) return error.AttentionAlias; + continue; + } try disjoint(cache.keys, cacheBytes(cache, true, plan.geometry), operand.ref, operand.bytes); try disjoint(cache.values, cacheBytes(cache, false, plan.geometry), operand.ref, operand.bytes); }; 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); - for (caches) |cache| { - e.useResource(cache.keys.buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); - e.useResource(cache.values.buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); - } + useCaches(e, caches); const args = PackArgs{ .rows = plan.rows, .packed_rows = plan.tiles * 16, .heads = g.heads, .kv_heads = g.kv_heads, .dim = g.dim, .group = g.group() }; e.setPipeline(a.io[2]); bind(e, .{ k, v, binding.get(.caches), binding.get(.meta), scratch.get(.nodes) }); @@ -449,10 +497,7 @@ pub const Attention = struct { 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; - for (caches) |cache| { - e.useResource(cache.keys.buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); - e.useResource(cache.values.buffer, mtl.ResourceUsage.read | mtl.ResourceUsage.write); - } + useCaches(e, caches); 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(args, 5); @@ -467,6 +512,7 @@ pub const Attention = struct { }; test "cache and pack descriptors match native pointer and six-word ABI" { try std.testing.expectEqual(@as(usize, 16), @sizeOf(CachePtrs)); + try std.testing.expectEqual(@as(usize, 32), @sizeOf(PagedPtrs)); try std.testing.expectEqual(@as(usize, 24), @sizeOf(PackArgs)); try std.testing.expectEqual(@as(usize, 16), @sizeOf(planning.KeepRow)); } diff --git a/zig/src/families/qwen27/attention_pick_test.zig b/zig/src/families/qwen27/attention_pick_test.zig index 570f0f5987..9c7535d870 100644 --- a/zig/src/families/qwen27/attention_pick_test.zig +++ b/zig/src/families/qwen27/attention_pick_test.zig @@ -45,7 +45,7 @@ test "prompt pick: one non-finite value at every key, head and kind; masked iff const device = mtl.Device.init() catch return error.SkipZigTest; defer device.deinit(); const a = std.testing.allocator; - const text = try simd_attention.grouped(a, device.tensorUnits(), true, 6, kv_heads); + const text = try simd_attention.grouped(a, device.tensorUnits(), true, 6, kv_heads, false); defer a.free(text); const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); defer lib.deinit(); diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index 283b2f40f7..3238df10b2 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -17,6 +17,11 @@ const Round = @import("round_plan.zig").Round; const Taps = @import("taps.zig").Taps; const image_rows = @import("image_rows.zig"); const mrope = @import("core").qwen_vl.rope; +const kv_pages = @import("kv_pages.zig"); +pub const pagesOn = kv_pages.on; +pub const KvPages = kv_pages.KvPages; +pub const KvSpan = kv_pages.KvSpan; +const kvp = kv_pages.kvp; /// Whether `cb` failed on the GPU; its Metal error (out of memory, a timeout) is logged with `what` it ran. pub fn gpuFailed(cb: mtl.CommandBuffer, what: []const u8) bool { @@ -43,6 +48,7 @@ pub const Runner = struct { slots: u32, ready: u32, // slots [0, ready) have their K/V rows; a later slot takes its own when first used (ensure) capacity: u32, + kv: ?KvPages = null, active: ?[32]u8 = null, failed: bool = false, capture_taps: bool = false, @@ -72,6 +78,15 @@ pub const Runner = struct { /// A runner of `slots` whose first `ready` take their K/V rows now, the rest when first used (ensure). pub fn initReady(a: std.mem.Allocator, m: *Model, slots: u32, capacity: u32, ready: u32) !Runner { + return build(a, m, slots, capacity, ready, false); + } + + /// A runner whose slots all take their K/V rows from one arena of core-pool pages, each promised `capacity` rows. + pub fn initPaged(a: std.mem.Allocator, m: *Model, slots: u32, capacity: u32) !Runner { + return build(a, m, slots, capacity, slots, true); + } + + fn build(a: std.mem.Allocator, m: *Model, slots: u32, capacity: u32, ready: u32, paged: bool) !Runner { if (slots == 0 or slots > attn.planning.max_streams or capacity == 0 or m.frame.capacity % 16 != 0 or ready == 0 or ready > slots) return error.BadRoundCapacity; const geometry = try attn.planning.Geometry.fromConfig(m.config, m.device.tensorUnits()); const linear_index = try a.alloc(u32, m.config.layers); @@ -97,13 +112,13 @@ pub const Runner = struct { // before M5 the tail attention reads its tensor-op results in the simdgroup-matrix layout (simd_attention.zig) if (!units) try @import("simd_attention.zig").check(m.device, m.queue); // both builds take the kernels' KV group as the defines attention_mpp.metal reads - const text = try @import("simd_attention.zig").grouped(a, units, true, geometry.group(), geometry.kv_heads); + const text = try @import("simd_attention.zig").grouped(a, units, true, geometry.group(), geometry.kv_heads, paged); defer a.free(text); // before M5 the prompt kernels are optional: a chip that cannot build them keeps the decode path for prompts const library = mtl.Library.fromSource(m.device, text, mtl.CompileOptions.mlx()) catch |err| blk: { if (units) return err; std.log.info("qwen27: the prompt attention kernels do not build here ({t}); prompt rows take the decode path", .{err}); - const base = try @import("simd_attention.zig").grouped(a, units, false, geometry.group(), geometry.kv_heads); + const base = try @import("simd_attention.zig").grouped(a, units, false, geometry.group(), geometry.kv_heads, paged); defer a.free(base); break :blk try mtl.Library.fromSource(m.device, base, mtl.CompileOptions.mlx()); }; @@ -130,14 +145,33 @@ pub const Runner = struct { for (owned.items) |b| b.deinit(); owned.deinit(a); } - for (0..na) |layer| for (0..ready) |slot| { + const kv: ?KvPages = if (paged) try kv_pages.caches(a, m, na, slots, capacity, caches) else null; + errdefer if (kv) |k| kv_pages.free(k); + if (!paged) for (0..na) |layer| for (0..ready) |slot| { caches[layer * slots + slot] = try slotCache(a, m, &owned, capacity); }; const offsets = try a.alloc(u32, slots); errdefer a.free(offsets); @memset(offsets, 0); const taps = try Taps.init(m.device, m.frame.capacity, @intCast(m.config.hidden)); - return .{ .allocator = a, .model = m, .geometry = geometry, .gdn = pool, .attention = attention, .scratch = scratch, .bindings = bindings, .caches = caches, .linear_index = linear_index, .attention_index = attention_index, .offsets = offsets, .owned = owned, .taps = taps, .slots = slots, .ready = ready, .capacity = capacity }; + return .{ .allocator = a, .model = m, .geometry = geometry, .gdn = pool, .attention = attention, .scratch = scratch, .bindings = bindings, .caches = caches, .linear_index = linear_index, .attention_index = attention_index, .offsets = offsets, .owned = owned, .taps = taps, .slots = slots, .ready = ready, .capacity = capacity, .kv = kv }; + } + + /// Slot `slot` holds the pages of its first `rows` rows before a pass or a copy writes them (paged runners). + pub fn cover(r: *Runner, slot: u32, rows: u32) !void { + const k = r.kv orelse return; + try k.pages.cover(slot, @min(rows, r.capacity)); + } + + /// Slot `slot`'s pages past its committed rows go back to its promise; no GPU work in flight reads them. + fn trimSlot(r: *Runner, slot: u32) void { + const k = r.kv orelse return; + k.pages.trim(slot, r.offsets[slot]); + } + + /// Slot `slot`'s first `at` rows of each attention layer's keys then values, head by head, page by page. + pub fn kvSpans(r: *const Runner, a: std.mem.Allocator, slot: u32, at: u32, list: *std.ArrayList(KvSpan)) !void { + return kv_pages.spans(r, a, slot, at, list); } /// One slot's K/V rows across the attention layers. @@ -163,6 +197,7 @@ pub const Runner = struct { /// The K/V buffers of the ready slots (for a residency set). pub fn readyBuffers(r: *const Runner, a: std.mem.Allocator, out: *std.ArrayList(mtl.Buffer)) !void { + if (r.kv) |k| return out.appendSlice(a, &.{ k.arena, k.tables }); for (0..r.caches.len / r.slots) |layer| for (0..r.ready) |slot| { const c = r.caches[layer * r.slots + slot]; try out.appendSlice(a, &.{ c.keys.buffer, c.values.buffer }); @@ -188,6 +223,7 @@ pub const Runner = struct { r.allocator.free(r.linear_index); r.allocator.free(r.attention_index); r.allocator.free(r.offsets); + if (r.kv) |k| kv_pages.free(k); } pub fn profile(r: *Runner, trace: ?*@import("core").gpu_profile.Trace) void { @@ -222,6 +258,7 @@ pub const Runner = struct { // A rank's head writes only this rank's vocabulary columns; the tail's matcher reads the whole vocabulary. if (r.tail != null and r.reduce != null) return error.TailOnRank; for (round.slots, round.windows) |slot, window| if (slot >= r.ready or r.offsets[slot] != window.start) return error.CachePositionDiffers; + for (round.slots, round.windows) |slot, window| try r.cover(slot, window.start + @as(u32, @intCast(window.parents.len))); const paths = try r.allocator.alloc([]const u32, round.windows.len); defer r.allocator.free(paths); @memset(paths, &.{}); @@ -322,7 +359,10 @@ pub const Runner = struct { r.failed = true; return error.RoundGpuFailure; } - for (round.slots, round.windows, paths) |slot, window, path| r.offsets[slot] = (try window.keep(path)).length; + for (round.slots, round.windows, paths) |slot, window, path| { + r.offsets[slot] = (try window.keep(path)).length; + r.trimSlot(slot); // the keep waited: its dropped rows' pages are read by nothing + } r.active = null; } diff --git a/zig/src/families/qwen27/kv_pages.zig b/zig/src/families/qwen27/kv_pages.zig new file mode 100644 index 0000000000..9d24e88b26 --- /dev/null +++ b/zig/src/families/qwen27/kv_pages.zig @@ -0,0 +1,86 @@ +//! TF_QWEN27_METAL_KV_PAGES: the 27B's keys and values in pages of the core pool, a slot a lane (decode_round.Runner). +const std = @import("std"); +const mtl = @import("metal"); +const attn = @import("attention.zig"); +const Ref = @import("projection.zig").Ref; +const Model = @import("model.zig").Model; +const kv_pool = @import("core").kv_pool; +pub const kvp = @import("core").metal_kv_pages; + +/// TF_QWEN27_METAL_KV_PAGES=1: every slot's keys and values in pages of the core pool (off by default). +pub fn on() bool { + const v = std.c.getenv("TF_QWEN27_METAL_KV_PAGES") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// A page: each attention layer's keys [kv_heads][256][head_dim] then its values, 256 rows of every layer. +pub fn pagedLayout(config: anytype, attention_layers: usize) kv_pool.Layout { + return .{ .fixed_bytes = 0, .token_bytes = attention_layers * 2 * config.kvDim() * 2, .page_tokens = kvp.page_tokens, .snap_bytes = 0 }; +} + +/// The pool's pages, arena and tables a paged runner holds; slot i is lane i, promised its whole capacity. +pub const KvPages = struct { pages: *kvp.Pages, arena: mtl.Buffer, tables: mtl.Buffer }; + +/// A span of a slot's K/V rows in a buffer, as snapshot and digest copy and hash them. +pub const KvSpan = struct { buffer: mtl.Buffer, offset: usize, len: usize }; + +/// Slot `slot`'s first `at` rows of each attention layer's keys then values, head by head, page by page. +pub fn spans(r: anytype, a: std.mem.Allocator, slot: u32, at: u32, list: *std.ArrayList(KvSpan)) !void { + const head_dim = r.model.config.head_dim; + const row_bytes = head_dim * @sizeOf(u16); + for (0..r.caches.len / r.slots) |layer| { + const cache = r.caches[layer * r.slots + slot]; + for ([_]Ref{ cache.keys, cache.values }, [_]u32{ cache.key_stride, cache.value_stride }) |ref, stride| { + if (r.kv) |k| { + const ids = k.pages.ids(slot); + if (ids.len * kvp.page_tokens < at) return error.PagesNotCovered; + for (0..r.model.config.kv_heads) |head| { + var row: usize = 0; + while (row < at) : (row += kvp.page_tokens) { + const rows = @min(kvp.page_tokens, at - row); + const page = ref.offset + @as(usize, @intCast(k.pages.arena.offset(ids[row / kvp.page_tokens]))); + try list.append(a, .{ .buffer = ref.buffer, .offset = page + head * kvp.page_tokens * row_bytes, .len = rows * row_bytes }); + } + } + continue; + } + const head_stride = @as(usize, if (stride == 0) cache.capacity * head_dim else stride) * 2; + for (0..r.model.config.kv_heads) |head| try list.append(a, .{ .buffer = ref.buffer, .offset = ref.offset + head * head_stride, .len = @as(usize, at) * row_bytes }); + } + } +} + +/// The pool over `slots` lanes of `capacity` rows, every lane promised them; each layer's caches its part of a page. +pub fn caches(a: std.mem.Allocator, m: *Model, layers: usize, slots: u32, capacity: u32, out: []attn.Cache) !KvPages { + const l = pagedLayout(m.config, layers); + const pages = @as(usize, slots) * l.pages(capacity); + const sized = kvp.Pages.bytes(l, slots, pages, capacity); + if (sized.arena > m.device.maxBufferLength()) return error.ArenaTooLarge; + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + const arena = try m.device.buffer(@intCast(sized.arena), opts); + errdefer arena.deinit(); + const tables = try m.device.buffer(@intCast(sized.tables), opts); + errdefer tables.deinit(); + const p = try kvp.Pages.initIn(a, l, slots, pages, capacity, arena, tables); + errdefer p.deinit(); + for (0..slots) |slot| try p.admit(slot, l.pages(capacity)); + const part = @as(usize, kvp.page_tokens) * m.config.kvDim() * 2; + for (0..layers) |layer| for (0..slots) |slot| { + const table: Ref = .{ .buffer = tables, .offset = p.tables.offset(slot) }; + out[layer * slots + slot] = .{ .keys = .{ .buffer = arena, .offset = 2 * layer * part }, .values = .{ .buffer = arena, .offset = (2 * layer + 1) * part }, .capacity = capacity, .pages = .{ .table = table, .page_bytes = l.pageBytes() } }; + }; + std.log.info("qwen27: KV in pages of the core pool: {d} slots over {d} pages of 256 rows ({d:.2} GiB, {d:.2} MiB a page)", .{ slots, pages, @as(f64, @floatFromInt(sized.arena)) / (1 << 30), @as(f64, @floatFromInt(l.pageBytes())) / (1 << 20) }); + return .{ .pages = p, .arena = arena, .tables = tables }; +} + +/// Every lane finished, every page must be free again: the census says so, or warns. +pub fn free(k: KvPages) void { + for (0..k.pages.leases.len) |lane| k.pages.finish(lane); + const c = k.pages.pool.census(); + if (k.pages.idle() and c.free == c.total) { + std.log.info("qwen27: KV pages idle, {d} of {d} free", .{ c.free, c.total }); + } else std.log.warn("qwen27: every slot closed, yet {d} of {d} KV pages are not free", .{ c.total - c.free, c.total }); + k.pages.deinit(); + k.arena.deinit(); + k.tables.deinit(); +} diff --git a/zig/src/families/qwen27/session.zig b/zig/src/families/qwen27/session.zig index 220c9b205d..aa0b14de72 100644 --- a/zig/src/families/qwen27/session.zig +++ b/zig/src/families/qwen27/session.zig @@ -61,14 +61,10 @@ pub const Session = struct { hash.update(try r.gdn.committed(layer, 0, true)); hash.update(try r.gdn.committed(layer, 0, false)); } - const head_dim = r.model.config.head_dim; - for (0..r.caches.len / r.slots) |layer| { - const cache = r.caches[layer * r.slots]; - for ([_]@import("projection.zig").Ref{ cache.keys, cache.values }, [_]u32{ cache.key_stride, cache.value_stride }) |ref, stride| { - const head_stride = if (stride == 0) cache.capacity * head_dim else stride; - for (0..r.model.config.kv_heads) |head| hash.update(ref.buffer.contents()[ref.offset + head * @as(usize, head_stride) * 2 ..][0 .. @as(usize, r.offsets[0]) * head_dim * @sizeOf(u16)]); - } - } + var spans: std.ArrayList(@import("decode_round.zig").KvSpan) = .empty; + defer spans.deinit(r.allocator); + try r.kvSpans(r.allocator, 0, r.offsets[0], &spans); // a paged slot's rows page by page: the same bytes in order + for (spans.items) |x| hash.update(x.buffer.contents()[x.offset..][0..x.len]); hash.update(r.model.frame.get(.logits).buffer.contents()[0 .. r.model.config.vocab * 2]); var out: [32]u8 = undefined; hash.final(&out); diff --git a/zig/src/families/qwen27/simd_attention.zig b/zig/src/families/qwen27/simd_attention.zig index 2af0db3a56..0515af2fc6 100644 --- a/zig/src/families/qwen27/simd_attention.zig +++ b/zig/src/families/qwen27/simd_attention.zig @@ -32,7 +32,24 @@ const rewrites = [_][2][]const u8{ /// The attention kernels as this chip builds them (before M5: the tail rewritten, the prompt kernels if `prompt`). pub fn library(a: std.mem.Allocator, units: bool, prompt: bool) ![]u8 { - const source = @import("qwen_runtime_sources"); + return compose(a, units, prompt, @import("qwen_runtime_sources")); +} + +/// As `library`, its K/V kernels reading pages of the core pool (TF_QWEN27_METAL_KV_PAGES), under main's names. +pub fn pagedLibrary(a: std.mem.Allocator, units: bool, prompt: bool) ![]u8 { + return compose(a, units, prompt, Paged); +} + +const Paged = struct { + const twins = @import("kernel_sources").qwen27_paged; + pub const glue = @import("qwen_runtime_sources").glue; + pub const attention_mpp = twins.attention_mpp; + pub const attention_io = twins.attention_io; + pub const attention_prompt = twins.attention_prompt; + pub const attention_prompt_sg = twins.attention_prompt_sg; +}; + +fn compose(a: std.mem.Allocator, units: bool, prompt: bool, comptime source: type) ![]u8 { if (units) return a.dupe(u8, source.glue ++ source.attention_mpp ++ source.attention_io ++ source.attention_prompt); const tail = try rewrite(a, source.glue ++ source.attention_mpp ++ source.attention_io); if (!prompt) return tail; @@ -40,9 +57,9 @@ pub fn library(a: std.mem.Allocator, units: bool, prompt: bool) ![]u8 { return std.mem.concat(a, u8, &.{ tail, source.attention_prompt, source.attention_prompt_sg }); } -/// `library` behind the KV-group defines attention_mpp.metal reads: Q27_G query heads over each of Q27_KVH heads. -pub fn grouped(a: std.mem.Allocator, units: bool, prompt: bool, group: u32, kv_heads: u32) ![]u8 { - const body = try library(a, units, prompt); +/// `library`, or `pagedLibrary`, behind the KV-group defines attention_mpp.metal reads: Q27_G over each of Q27_KVH. +pub fn grouped(a: std.mem.Allocator, units: bool, prompt: bool, group: u32, kv_heads: u32, paged: bool) ![]u8 { + const body = try if (paged) pagedLibrary(a, units, prompt) else library(a, units, prompt); defer a.free(body); var defines: [64]u8 = undefined; return std.mem.concat(a, u8, &.{ try std.fmt.bufPrint(&defines, "#define Q27_G {d}\n#define Q27_KVH {d}\n", .{ group, kv_heads }), body }); @@ -144,7 +161,7 @@ test "both attention builds, with and without the prompt kernels, open with the const a = std.testing.allocator; const defines = "#define Q27_G 4\n#define Q27_KVH 4\n"; for ([_]bool{ false, true }) |units| for ([_]bool{ false, true }) |prompt| { - const text = try grouped(a, units, prompt, 4, 4); + const text = try grouped(a, units, prompt, 4, 4, false); defer a.free(text); const body = try library(a, units, prompt); defer a.free(body); diff --git a/zig/src/families/qwen27/snapshot.zig b/zig/src/families/qwen27/snapshot.zig index d99a2c3da4..660d0bce0a 100644 --- a/zig/src/families/qwen27/snapshot.zig +++ b/zig/src/families/qwen27/snapshot.zig @@ -139,14 +139,10 @@ fn parts(r: *const Runner, slot: u32, ring: ?Ring, at: u32, list: *std.ArrayList }; try list.append(a, .{ .buffer = span.buffer, .offset = span.offset, .len = span.bytes }); }; - const row_bytes = r.model.config.head_dim * @sizeOf(u16); - for (0..r.caches.len / r.slots) |layer| { - const cache = r.caches[layer * r.slots + slot]; - for ([_]@import("projection.zig").Ref{ cache.keys, cache.values }, [_]u32{ cache.key_stride, cache.value_stride }) |ref, stride| { - const head_stride = @as(usize, if (stride == 0) cache.capacity * r.model.config.head_dim else stride) * 2; - for (0..r.model.config.kv_heads) |head| try list.append(a, .{ .buffer = ref.buffer, .offset = ref.offset + head * head_stride, .len = @as(usize, at) * row_bytes }); - } - } + var spans: std.ArrayList(@import("decode_round.zig").KvSpan) = .empty; + defer spans.deinit(a); + try r.kvSpans(a, slot, at, &spans); + for (spans.items) |x| try list.append(a, .{ .buffer = x.buffer, .offset = x.offset, .len = x.len }); if (ring) |g| try list.append(a, .{ .buffer = g.buf, .offset = 0, .len = g.bytes(at) }); } @@ -243,6 +239,7 @@ pub fn restoreSlot(a: std.mem.Allocator, r: *Runner, slot: u32, ring: ?Ring, st: if (st.bytes != slotBytes(r, ring, st.at) or st.at > r.capacity) return error.SnapshotLayout; try st.landed(); try r.reset(slot); + try r.cover(slot, st.at); // a paged slot's pages for the state's rows before the copy writes them try copy(a, r, slot, ring, st, false); r.offsets[slot] = st.at; if (ring) |g| g.end.* = st.draft_end; @@ -255,6 +252,7 @@ pub fn moveSlot(a: std.mem.Allocator, r: *Runner, from: u32, from_ring: ?Ring, t if (from == to or r.failed or r.active != null or from >= r.slots or to >= r.slots or (from_ring == null) != (to_ring == null)) return error.BadSlotMove; if (from_ring) |g| if (g.end.* != at) return error.NotAtMark; try r.reset(to); + try r.cover(to, at); var src: std.ArrayList(Part) = .empty; defer src.deinit(a); var dst: std.ArrayList(Part) = .empty; diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 5a6c1b3ba3..6a8d50c5ab 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -320,6 +320,11 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q if (capacity == 0 or capacity > model.config.max_position) return error.ContextFull; const policy = @import("qwen27_window.zig"); const units = model.device.tensorUnits(); + // the paged kernels have run only before the M5: a Mac with tensor units refuses the flag until they run there + if (q.decode_round.pagesOn() and units) { + why.* = "TF_QWEN27_METAL_KV_PAGES=1 has run only on Macs without tensor units (M1 to M4) so far; unset it on this Mac"; + return error.KvPagesUntested; + } const window = lb.windowRows(units) catch |err| { why.* = policy.problem(err) orelse why.*; return err; @@ -372,7 +377,8 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q h.disk_stage = null; h.vision = o.images; // closed with the host; a failed attach leaves it to the caller h.trace = null; - h.runner = try q.decode_round.Runner.initReady(gpa, model, slots, capacity + window, 1); + // TF_QWEN27_METAL_KV_PAGES=1: every slot's rows in core-pool pages now (Fit charged each its whole window above) + h.runner = try if (q.decode_round.pagesOn()) q.decode_round.Runner.initPaged(gpa, model, slots, capacity + window) else q.decode_round.Runner.initReady(gpa, model, slots, capacity + window, 1); errdefer h.runner.deinit(); if (profileEvery()) |every| { h.trace = .{ .queue = model.queue, .report_every = every }; // enabled once measure() has priced the windows diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index b4929097f9..2fbc0f150c 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -122,6 +122,8 @@ pub const shared_attention = @import("core/shared_attention.zig"); pub const tile_qmm = @import("core/tile_qmm.zig"); pub const gpu_profile = @import("core/gpu_profile.zig"); +/// The core KV pool on Metal: one arena of pages and each stream's table (qwen27's TF_QWEN27_METAL_KV_PAGES). +pub const metal_kv_pages = @import("core/metal_kv_pages.zig"); pub const draft_ops = @import("draft_ops"); diff --git a/zig/tests/contract/metal_q27_attn_rows.zig b/zig/tests/contract/metal_q27_attn_rows.zig index 07e81860cb..a752aec6fb 100644 --- a/zig/tests/contract/metal_q27_attn_rows.zig +++ b/zig/tests/contract/metal_q27_attn_rows.zig @@ -107,7 +107,7 @@ pub const Attn = struct { const tensor = device.tensorUnits(); if (!tensor) try q.simd_attention.check(device, queue); const gm = at_.geometry(r.heads, tensor); - const body = try q.simd_attention.grouped(a, tensor, true, gm.group(), gm.kv_heads); + const body = try q.simd_attention.grouped(a, tensor, true, gm.group(), gm.kv_heads, false); defer a.free(body); const text = if (!r.control) try a.dupe(u8, body) else edit: { const e = controlEdit(r.which, tensor) orelse return error.ControlEdit; diff --git a/zig/tests/qwen27/attention_pages.zig b/zig/tests/qwen27/attention_pages.zig new file mode 100644 index 0000000000..e9fbbab6dc --- /dev/null +++ b/zig/tests/qwen27/attention_pages.zig @@ -0,0 +1,229 @@ +//! TF_QWEN27_METAL_KV_PAGES kernels against main's: the same rounds, outputs and cache rows bit for bit, pages shuffled +const std = @import("std"); +const mtl = @import("metal"); +const q = @import("qwen27"); +const attn = q.attention; +const planning = attn.planning; + +const dim = 256; +const page = attn.page_rows; +const most_rows = 128; +const capacity: u32 = 12288; +const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +const Rng = struct { + s: u64, + fn next(r: *Rng) u64 { + r.s ^= r.s >> 12; + r.s ^= r.s << 25; + r.s ^= r.s >> 27; + return r.s *% 0x2545f4914f6cdd1d; + } + fn bf16(r: *Rng) u16 { + const unit: f32 = @as(f32, @floatFromInt(r.next() >> 40)) / @as(f32, 1 << 24); + return @intCast(@as(u32, @bitCast((unit - 0.5) * 4.0)) >> 16); + } + fn below(r: *Rng, n: usize) usize { + return @intCast(r.next() % n); + } +}; + +/// A round: its streams' windows (start, rows, a tree or a chain), the prompt kernel or the decode path, and a keep. +const Case = struct { starts: []const u32, rows: []const u32, tree: bool = false, prompt: bool = false, planted: bool = false }; +const cases = [_]Case{ + .{ .starts = &.{0}, .rows = &.{1} }, .{ .starts = &.{255}, .rows = &.{1} }, .{ .starts = &.{256}, .rows = &.{1} }, + .{ .starts = &.{511}, .rows = &.{1} }, .{ .starts = &.{8191}, .rows = &.{1} }, .{ .starts = &.{12287}, .rows = &.{1} }, + .{ .starts = &.{250}, .rows = &.{16} }, .{ .starts = &.{4090}, .rows = &.{16}, .tree = true }, .{ .starts = &.{8184}, .rows = &.{16} }, + .{ .starts = &.{ 100, 767, 3000 }, .rows = &.{ 16, 9, 16 }, .tree = true }, .{ .starts = &.{ 0, 255, 256, 9000 }, .rows = &.{ 3, 6, 16, 1 }, .tree = true }, .{ .starts = &.{0}, .rows = &.{128}, .prompt = true }, + .{ .starts = &.{200}, .rows = &.{128}, .prompt = true }, .{ .starts = &.{1000}, .rows = &.{128}, .prompt = true }, .{ .starts = &.{8100}, .rows = &.{128}, .prompt = true }, + .{ .starts = &.{12160}, .rows = &.{128}, .prompt = true }, + // Inf and NaN in a round's own values, which a window's earlier rows must not see: the masked kernels and scan + .{ .starts = &.{250}, .rows = &.{16}, .planted = true }, .{ .starts = &.{4090}, .rows = &.{16}, .tree = true, .planted = true }, + .{ .starts = &.{ 100, 767, 3000 }, .rows = &.{ 16, 9, 16 }, .tree = true, .planted = true }, .{ .starts = &.{200}, .rows = &.{128}, .prompt = true, .planted = true }, .{ .starts = &.{8100}, .rows = &.{128}, .prompt = true, .planted = true }, +}; +const most_streams = 4; + +/// One side: main's contiguous caches, or every stream's layer part in one arena of shuffled pages. +const Side = struct { + op: attn.Attention, + binding: attn.Binding, + caches: [most_streams]attn.Cache, + out: mtl.Buffer, +}; + +pub fn main(init: std.process.Init) !void { + const a = init.arena.allocator(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + const units = device.tensorUnits(); + // --shared-prefix: the M5's plan (prefix kernels over every row's keys) on any chip; the bits check addressing + const args = try init.minimal.args.toSlice(init.arena.allocator()); + const shared = units or (args.len > 1 and std.mem.eql(u8, args[1], "--shared-prefix")); + const geometry = attn.planning.Geometry{ .shared_prefix = shared }; + const heads = geometry.heads; + const kvh = geometry.kv_heads; + if (!units) try q.simd_attention.check(device, queue); + const limits = attn.Limits{ .rows = most_rows, .streams = most_streams, .keys = capacity, .shared_prefix = shared, .geometry = geometry }; + var scratch = try attn.Scratch.init(device, limits); + defer scratch.deinit(); + + // main's caches: a stream's keys [kv_heads][capacity][dim], its values the same + const cache_bytes = @as(usize, capacity) * kvh * dim * 2; + var contiguous: [most_streams][2]mtl.Buffer = undefined; + for (&contiguous) |*s| for (s) |*b| { + b.* = try device.buffer(cache_bytes, options); + }; + // the arena: a page is this layer's keys [kv_heads][256][dim] then its values, ids shuffled across streams + const part = @as(u64, kvh) * page * dim * 2; + const page_bytes = 2 * part + 512; // a stride past the two parts, as when other layers share the page + const per = attn.pageCount(capacity); + const total = most_streams * per + 5; + const arena = try device.buffer(@intCast(total * page_bytes), options); + const tables = try device.buffer(most_streams * per * 4, options); + var rng = Rng{ .s = 0x9e3779b97f4a7c15 }; + const ids = try a.alloc(u32, total); + for (ids, 0..) |*id, i| id.* = @intCast(i); + var i: usize = ids.len; + while (i > 1) : (i -= 1) std.mem.swap(u32, &ids[i - 1], &ids[rng.below(i)]); + @memcpy(tables.slice(u32, most_streams * per), ids[0 .. most_streams * per]); + @memset(arena.slice(u8, total * page_bytes), 0x5a); + + var sides: [2]Side = undefined; + for (&sides, 0..) |*s, paged| { + const text = try q.simd_attention.grouped(a, units, true, geometry.group(), kvh, paged == 1); + const library = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer library.deinit(); + s.op = try attn.Attention.init(device, library); + s.binding = try attn.Binding.init(device, limits); + s.out = try device.buffer(most_rows * heads * dim * 2, options); + for (&s.caches, 0..) |*c, st| c.* = if (paged == 0) + .{ .keys = .{ .buffer = contiguous[st][0] }, .values = .{ .buffer = contiguous[st][1] }, .capacity = capacity } + else + .{ .keys = .{ .buffer = arena }, .values = .{ .buffer = arena, .offset = part }, .capacity = capacity, .pages = .{ .table = .{ .buffer = tables, .offset = st * per * 4 }, .page_bytes = page_bytes } }; + } + std.debug.print("attention_pages on {s}, tensor units {s}, shared prefix {s}: {d} rounds, {d} pages of {d} bytes shuffled\n", .{ std.mem.span(device.name()), if (units) "yes" else "no", if (shared) "yes" else "no", cases.len, total, page_bytes }); + + // the same random rows in both layouts + for (0..most_streams) |st| for (0..2) |kv| for (0..kvh) |hk| for (0..capacity) |row| { + const c = contiguous[st][kv].slice(u16, cache_bytes / 2)[(hk * capacity + row) * dim ..][0..dim]; + for (c) |*x| x.* = rng.bf16(); + @memcpy(pagedRow(arena, tables, st, per, page_bytes, part, kv, hk, row), c); + }; + try same(contiguous, arena, tables, per, page_bytes, part, kvh, 0, "the start"); + + const query = try device.buffer(most_rows * heads * dim * 2, options); + const key = try device.buffer(most_rows * kvh * dim * 2, options); + const val = try device.buffer(most_rows * kvh * dim * 2, options); + var compared: usize = 0; + for (cases, 0..) |cs, n| { + var parents: [most_rows]i32 = undefined; + var ws: [most_streams]planning.Window = undefined; + var rows: usize = 0; + for (cs.starts, cs.rows, 0..) |start, count, st| { + for (0..count) |r| parents[rows + r] = if (r == 0) -1 else if (cs.tree and r > 2) @intCast(rng.below(r)) else if (cs.tree) 0 else @intCast(r - 1); + ws[st] = .{ .parents = parents[rows..][0..count], .start = start, .capacity = capacity }; + rows += count; + } + for (query.slice(u16, rows * heads * dim)) |*x| x.* = rng.bf16(); + for (key.slice(u16, rows * kvh * dim), val.slice(u16, rows * kvh * dim)) |*k, *v| { + k.* = rng.bf16(); + v.* = rng.bf16(); + } + if (cs.planted) plant(val.slice(u16, rows * kvh * dim), cs.rows, kvh); + var plan = try attn.Plan.init(a, ws[0..cs.starts.len], geometry); + defer plan.deinit(); + try scratch.upload(plan); + for (&sides) |*s| { + const caches = s.caches[0..cs.starts.len]; + try s.binding.upload(plan, caches); + @memset(s.out.slice(u16, most_rows * heads * dim), 0xffff); + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + try s.op.encode(e, &scratch, &s.binding, plan, .{ .buffer = query }, .{ .buffer = key }, .{ .buffer = val }, .{ .buffer = s.out }, caches, null, cs.prompt, false); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("FAIL round {d}: the GPU: {s}\n", .{ n, msg }); + return error.GpuFailed; + } + } + const got = [2][]const u16{ sides[0].out.slice(u16, rows * heads * dim), sides[1].out.slice(u16, rows * heads * dim) }; + if (!std.mem.eql(u16, got[0], got[1])) { + std.debug.print("FAIL round {d} (starts {any}, rows {any}): attention outputs differ\n", .{ n, cs.starts, cs.rows }); + return error.Differs; + } + compared += got[0].len; + try same(contiguous, arena, tables, per, page_bytes, part, kvh, n, "the append"); + if (cs.tree) { + // each stream keeps its last row's path, moved to the window's first rows + var paths: [most_streams][]u32 = undefined; + var first: usize = 0; + for (cs.rows, 0..) |count, st| { + var chain: std.ArrayList(u32) = .empty; + var at: i32 = @intCast(count - 1); + while (at >= 0) : (at = parents[first + @as(usize, @intCast(at))]) try chain.append(a, @intCast(at)); + std.mem.reverse(u32, chain.items); + paths[st] = chain.items; + first += count; + } + for (&sides) |*s| { + const caches = s.caches[0..cs.starts.len]; + const kept = try s.binding.uploadKeep(plan, caches, paths[0..cs.starts.len]); + const cb = queue.commandBuffer(); + const e = cb.compute(.concurrent); + try s.op.keep(e, &scratch, &s.binding, kept, caches); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("FAIL round {d}'s keep: the GPU: {s}\n", .{ n, msg }); + return error.GpuFailed; + } + } + try same(contiguous, arena, tables, per, page_bytes, part, kvh, n, "the keep"); + } + } + // no page past the tables' was written: the spare pages still hold the fill + for (ids[most_streams * per ..]) |id| for (arena.slice(u8, total * page_bytes)[id * page_bytes ..][0..page_bytes]) |x| if (x != 0x5a) { + std.debug.print("FAIL a spare page {d} was written\n", .{id}); + return error.Differs; + }; + std.debug.print("PASS {d} rounds: {d} output values and every cache row the same bits, paged and contiguous\n", .{ cases.len, compared }); +} + +/// Each window's later rows: an Inf, a -Inf and a NaN in a few heads' values, which its first row must not see. +fn plant(values: []u16, rows: []const u32, kvh: usize) void { + var first: usize = 0; + for (rows) |count| { + var r: usize = 1; + while (r < count) : (r += 3) { + const row = values[((first + r) * kvh + r % kvh) * dim ..][0..dim]; + row[r % dim] = 0x7f80; + row[(r * 7 + 128) % dim] = 0xff80; + row[(r * 13 + 5) % dim] = 0x7fc1; + } + first += count; + } +} + +/// Stream `st`'s cache `kv` (0 keys, 1 values), head `hk`, row `row` in the arena. +fn pagedRow(arena: mtl.Buffer, tables: mtl.Buffer, st: usize, per: usize, page_bytes: u64, part: u64, kv: usize, hk: usize, row: usize) []u16 { + const id = tables.slice(u32, (st + 1) * per)[st * per + row / page]; + const at = id * page_bytes + kv * part + (hk * page + row % page) * dim * 2; + return @alignCast(std.mem.bytesAsSlice(u16, arena.slice(u8, @intCast(at + dim * 2))[@intCast(at)..])); +} + +/// Every stream's every row the same in both layouts. +fn same(contiguous: [most_streams][2]mtl.Buffer, arena: mtl.Buffer, tables: mtl.Buffer, per: usize, page_bytes: u64, part: u64, kvh: usize, n: usize, what: []const u8) !void { + for (0..most_streams) |st| for (0..2) |kv| for (0..kvh) |hk| for (0..capacity) |row| { + const c = contiguous[st][kv].slice(u16, kvh * capacity * dim)[(hk * capacity + row) * dim ..][0..dim]; + if (!std.mem.eql(u16, c, pagedRow(arena, tables, st, per, page_bytes, part, kv, hk, row))) { + std.debug.print("FAIL round {d}, after {s}: stream {d} {s} head {d} row {d} differs\n", .{ n, what, st, if (kv == 0) "keys" else "values", hk, row }); + return error.Differs; + } + }; +} diff --git a/zig/tests/qwen27/pages_resume.zig b/zig/tests/qwen27/pages_resume.zig new file mode 100644 index 0000000000..e3be3a5236 --- /dev/null +++ b/zig/tests/qwen27/pages_resume.zig @@ -0,0 +1,62 @@ +//! TF_QWEN27_METAL_KV_PAGES's saved, restored and moved slots against main's: the same digests, KV rows included. +const std = @import("std"); +const mtl = @import("metal"); +const q = @import("tensorfold").qwen27; +const graph = @import("synthetic_graph.zig"); + +const Result = struct { resumed: [32]u8, moved: [32]u8, fresh: [32]u8 }; + +fn run(a: std.mem.Allocator, model: *q.model.Model, paged: bool, first: []const u32, other: []const u32) !Result { + var r = try if (paged) q.decode_round.Runner.initPaged(a, model, 2, 1024) else q.decode_round.Runner.init(a, model, 2, 1024); + defer r.deinit(); + if (r.caches.len == 0) return error.NoAttentionLayer; + const s: q.session.Session = .{ .runner = &r }; + try r.reset(0); + try s.prefill(first, 128); + const st = try q.snapshot.save(a, &r, null, @intCast(first.len)); + defer q.snapshot.drop(a, st); + try st.landed(); + var out: Result = undefined; + // the slot is overwritten, then the kept state comes back and decodes on + try r.reset(0); + try s.prefill(other, 128); + try q.snapshot.restore(a, &r, null, st); + for ([_]u32{ 7, 8, 9 }) |t| try s.step(t); + out.resumed = try s.fingerprint(); + // it moves to slot 1 and back past another prompt in slot 0 + try q.snapshot.moveSlot(a, &r, 0, null, 1, null); + try r.reset(0); + try s.prefill(other, 128); + try q.snapshot.moveSlot(a, &r, 1, null, 0, null); + out.moved = try s.fingerprint(); + // the same prompt and steps without a kept state + try r.reset(0); + try s.prefill(first, 128); + for ([_]u32{ 7, 8, 9 }) |t| try s.step(t); + out.fresh = try s.fingerprint(); + return out; +} + +pub fn main(init: std.process.Init) !void { + const a = init.gpa; + const autorelease = mtl.objc.Pool.push(); + defer autorelease.pop(); + const model = try graph.createLayers(a, 5); + defer model.deinit(); + var first: [600]u32 = undefined; // past two pages of 256 rows + var other: [300]u32 = undefined; + for (&first, 0..) |*t, i| t.* = @intCast((i * 37 + 11) % 96); + for (&other, 0..) |*t, i| t.* = @intCast((i * 53 + 5) % 96); + const main_ = try run(a, model, false, &first, &other); + const pages = try run(a, model, true, &first, &other); + for ([_]struct { what: []const u8, x: [32]u8, y: [32]u8 }{ + .{ .what = "restored digest, main and pages", .x = main_.resumed, .y = pages.resumed }, + .{ .what = "moved digest, main and pages", .x = main_.moved, .y = pages.moved }, + .{ .what = "fresh digest, main and pages", .x = main_.fresh, .y = pages.fresh }, + .{ .what = "pages: restored and fresh", .x = pages.resumed, .y = pages.fresh }, + }) |c| if (!std.mem.eql(u8, &c.x, &c.y)) { + std.debug.print("FAIL {s} differ\n", .{c.what}); + return error.Differs; + }; + std.debug.print("PASS kept states through pages: restored equal to fresh, and restored, moved and fresh slots the same digests as main's\n", .{}); +}