From 3f5b87dd957e353ead3a999be4d79f1fe63182a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:21:28 +0200 Subject: [PATCH 01/11] cuda qwen3_5: attention caches readable through a page table Seg gains pt (an int page table), ps (elements a page) and ptok (positions a page); kv_row() gives every key/value row address: with pt 0 the old expression (cap rows a head), else the page holding the position, [layer][key, value][head][ptok][256] inside it. rope's store, tf_q35_attn_partial and keep's tf_q35k_compact read and write through it. Only addresses change, never the order a sum takes, so paged results are byte-equal to contiguous at any page size. Cache.paged maps a cache's rows to pages; TF_QWEN_KV_PAGE_TEST=n gives each cache's own buffer a reversed table of n-position pages (checks only). Prompt-cache saves and restores copy page runs into the same host layout, so kept states move between paged and contiguous caches. --kv-cache fp8 keeps contiguous rows and refuses a page table. tf-cuda-test qwen-pages (also in kernel-contract): six shapes, cached lengths to 9000, pages of 16, 48, 128, 256 and 512 positions, two streams shuffled through one arena: rotated queries, attention rows, rope's stores and keep's moves equal contiguous, three repeats, a float64 bound, guards and inputs intact; red: a swapped table over placed pages, a planted row. GB10, Qwen3.8-27B MLX-4bit: contiguous == main (3/3 chat, 3/3 8k prefill), 13.8 tok/s; page test 16 and 256 == contiguous, 13.8-14.0 tok/s; with DFlash2 at --parallel 8, drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh. --- zig/kernels/cuda/qwen35.cu | 14 +- zig/kernels/cuda/qwen35.cuh | 11 + zig/kernels/cuda/qwen35_keep.cu | 7 +- zig/src/families/qwen3_5/cuda_cache.zig | 83 ++++++- zig/src/families/qwen3_5/cuda_reuse.zig | 44 +++- zig/tests/cuda/contract_qwen_pages.zig | 296 ++++++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 1 + zig/tests/cuda/main.zig | 2 + 8 files changed, 436 insertions(+), 22 deletions(-) create mode 100644 zig/tests/cuda/contract_qwen_pages.zig diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 103b766320..c377f976a7 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -394,10 +394,9 @@ extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfl struct StoreBf16 { int layer, kvh; __device__ __forceinline__ void operator()(const Seg& sg, int h, int slot, int d, __nv_bfloat16 k, __nv_bfloat16 v) const { - __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * kvh * sg.cap * HD; - const size_t dst = ((size_t)h * sg.cap + slot) * HD + d; - keys[dst] = k; - keys[(size_t)kvh * sg.cap * HD + dst] = v; + __nv_bfloat16* kv = reinterpret_cast<__nv_bfloat16*>(sg.kv); + kv[kv_row(sg, kvh, layer, 0, h, slot) + d] = k; + kv[kv_row(sg, kvh, layer, 1, h, slot) + d] = v; } }; @@ -434,8 +433,7 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r } return; } - const __nv_bfloat16* K = reinterpret_cast(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; - const __nv_bfloat16* V = K + (size_t)KVH * sg.cap * HD; + const __nv_bfloat16* KV = reinterpret_cast(sg.kv); float q[DPL], o[DPL]; #pragma unroll for (int i = 0; i < DPL; i++) { @@ -445,7 +443,7 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r float mx = -INFINITY, l = 0.0f; for (int t = k0; t < k1; t++) { const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); - const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const uint4 kw = *reinterpret_cast(KV + kv_row(sg, KVH, layer, 0, h, pos) + lane * DPL); const __nv_bfloat16* kb = reinterpret_cast(&kw); float dot = 0.0f; #pragma unroll @@ -454,7 +452,7 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r const float mn = fmaxf(mx, s); const float a = __expf(mx - mn), b = __expf(s - mn); l = l * a + b; - const uint4 vw = *reinterpret_cast(V + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const uint4 vw = *reinterpret_cast(KV + kv_row(sg, KVH, layer, 1, h, pos) + lane * DPL); const __nv_bfloat16* vb = reinterpret_cast(&vw); #pragma unroll for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh index 728614bd49..2fe83705f5 100644 --- a/zig/kernels/cuda/qwen35.cuh +++ b/zig/kernels/cuda/qwen35.cuh @@ -34,8 +34,19 @@ struct Seg { unsigned long long before; // exact pre-round state for a one-row in-place update int cap, len, row0, rows, tap_rows, taps; int delta; // an image prompt's rotary shift: rows past it rotate at position + delta + unsigned long long pt; // the stream's KV page table (int a page; 0: kv holds cap contiguous rows) + unsigned long long ps; // elements a KV page holds: every attention layer's keys and values of ptok positions + int ptok; // positions a KV page holds }; +// Seg.kv's element of layer a's key (s 0) or value (s 1) at pos, head h: cap rows a head, or pos's page's row. +__device__ __forceinline__ size_t kv_row(const Seg& sg, int kvh, int a, int s, int h, int pos) { + const size_t blk = ((size_t)a * 2 + s) * kvh + h; + if (sg.pt == 0) return (blk * sg.cap + pos) * HD; + const int* pt = reinterpret_cast(sg.pt); + return (size_t)pt[pos / sg.ptok] * sg.ps + (blk * sg.ptok + pos % sg.ptok) * HD; +} + // A window row's record of one linear layer (Seg.log). struct LogRow { float* g; // [nv] decay diff --git a/zig/kernels/cuda/qwen35_keep.cu b/zig/kernels/cuda/qwen35_keep.cu index 20d00fc212..6bca4e89e0 100644 --- a/zig/kernels/cuda/qwen35_keep.cu +++ b/zig/kernels/cuda/qwen35_keep.cu @@ -44,14 +44,13 @@ extern "C" __global__ void tf_q35k_conv(Geo geo, Seg sg, int last) { extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact(Geo geo, Seg sg, Path path) { const int KVH = geo.kvh; const int layer = blockIdx.x, h = blockIdx.y; - __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + ((size_t)layer * 2 * KVH + h) * sg.cap * HD; - __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + __nv_bfloat16* kv = reinterpret_cast<__nv_bfloat16*>(sg.kv); const int d = threadIdx.x; for (int j = 0; j < path.n; j++) { const int from = sg.len + path.rows[j], to = sg.len + j; if (from != to) { - keys[(size_t)to * HD + d] = keys[(size_t)from * HD + d]; - values[(size_t)to * HD + d] = values[(size_t)from * HD + d]; + kv[kv_row(sg, KVH, layer, 0, h, to) + d] = kv[kv_row(sg, KVH, layer, 0, h, from) + d]; + kv[kv_row(sg, KVH, layer, 1, h, to) + d] = kv[kv_row(sg, KVH, layer, 1, h, from) + d]; } __syncthreads(); } diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig index 925d8183bb..359db822f9 100644 --- a/zig/src/families/qwen3_5/cuda_cache.zig +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -11,11 +11,39 @@ const DeviceBuffer = cuda.DeviceBuffer; pub const tap_rows = 2048 + engine.max_rows; /// qwen35.cuh's Seg and Path. -pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, before: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int, delta: c_int }; +pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, before: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int, delta: c_int, pt: u64 = 0, ps: u64 = 0, ptok: c_int = 0 }; pub const Path = extern struct { n: c_int, rows: [engine.most_window]c_int }; +/// A cache's keys and values as pages (qwen35.cuh kv_row): page i of the stream's positions at `base` + ids[i] pages. +pub const Paged = struct { + /// Where page 0 of the arena starts: the core pool's, or the cache's own buffer under TF_QWEN_KV_PAGE_TEST. + base: u64, + /// The stream's page table on the device, an int a page, `ids` uploaded. + table: DeviceBuffer, + ids: []u32, + /// Positions a page holds. + tokens: usize, + /// The cache owns `ids` (TF_QWEN_KV_PAGE_TEST); a pool's lane maps its lease's. + own_ids: bool = true, + + /// Elements (bf16) a page holds: every attention layer's keys and values of `tokens` positions. + pub fn elems(p: Paged, g: geometry.Geometry) usize { + return g.attentions() * 2 * g.kv_heads * p.tokens * geometry.head_dim; + } +}; + +/// TF_QWEN_KV_PAGE_TEST=n: each cache's own buffer read through a table of n-position pages in reverse order (checks). +pub fn pageTest() !usize { + const v = std.c.getenv("TF_QWEN_KV_PAGE_TEST") orelse return 0; + const n = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadQwenPageTest; + if (n == 0) return error.BadQwenPageTest; + return n; +} + pub const Cache = struct { untaught: bool = false, + /// Keys and values through a page table (null: cap contiguous rows a head). + paged: ?Paged = null, delta: i32 = 0, state: DeviceBuffer, /// The exact state before a recorded one-row round advances its stable slot in place. @@ -34,14 +62,44 @@ pub const Cache = struct { last: ?struct { start: usize, rows: usize } = null, /// An image prompt's rotary shift (get_rope_index's rope_delta): rows after its prompt rotate at position + delta. pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { + const test_pages = try pageTest(); + if (test_pages > 0 and e.kv8) return error.PagedKvNeedsBf16; + const rows = if (test_pages > 0) (cap + test_pages - 1) / test_pages * test_pages else cap; + var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(rows, e.kv8)); + errdefer kv.free(); + try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) + var paged: ?Paged = null; + errdefer if (paged) |*p| freePaged(gpa, p); + if (test_pages > 0) { + const n = rows / test_pages; + const ids = try gpa.alloc(u32, n); + errdefer gpa.free(ids); + for (ids, 0..) |*id, i| id.* = @intCast(n - 1 - i); // a table no contiguous layout matches + var table = try DeviceBuffer.alloc(e.d, n * 4); + errdefer table.free(); + try table.upload(0, std.mem.sliceAsBytes(ids)); + paged = .{ .base = kv.ptr, .table = table, .ids = ids, .tokens = test_pages }; + } + return fixed(gpa, e, cap, kv, paged); + } + + /// A pool lane's cache: its own state, log, tap ring and table of `pages` pages at `base`, mapped at open. + pub fn initPaged(gpa: std.mem.Allocator, e: *const Engine, pages: usize, tokens: usize, base: u64) !Cache { + if (e.kv8) return error.PagedKvNeedsBf16; + var table = try DeviceBuffer.alloc(e.d, pages * 4); + errdefer table.free(); + try table.fill8(0, null); + const none = try DeviceBuffer.alloc(e.d, 0); + return fixed(gpa, e, 0, none, .{ .base = base, .table = table, .ids = &.{}, .tokens = tokens, .own_ids = false }); + } + + /// The parts every cache holds itself beside `kv` and `paged` (the caller's to free on failure). + fn fixed(gpa: std.mem.Allocator, e: *const Engine, cap: usize, kv: DeviceBuffer, paged: ?Paged) !Cache { var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); errdefer state.free(); try state.fill8(0, null); var before = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); errdefer before.free(); - var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap, e.kv8)); - errdefer kv.free(); - try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) var log = try DeviceBuffer.alloc(e.d, e.g.logBytes(engine.most_window)); errdefer log.free(); try log.fill8(0, null); @@ -51,7 +109,7 @@ pub const Cache = struct { tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); try tap.?.fill8(0, null); } - return .{ .state = state, .before = before, .kv = kv, .tap = tap, .log = log, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + return .{ .state = state, .before = before, .kv = kv, .paged = paged, .tap = tap, .log = log, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; } /// Publish a completed forward's metadata, including one-row slots that already advanced in place. @@ -75,6 +133,7 @@ pub const Cache = struct { self.state.free(); if (self.before) |*b| b.free(); self.kv.free(); + if (self.paged) |*p| freePaged(gpa, p); self.log.free(); if (self.tap) |*t| t.free(); gpa.free(self.logits); @@ -83,7 +142,14 @@ pub const Cache = struct { /// The stream's window of `rows` rows from forward row `row0`, after `len` cached positions. pub fn seg(self: *const Cache, e: *const Engine, len: usize, row0: usize, rows: usize) Seg { const g = e.g; - return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .before = if (self.before) |b| b.ptr else 0, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps), .delta = self.delta }; + var sg: Seg = .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .before = if (self.before) |b| b.ptr else 0, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps), .delta = self.delta }; + if (self.paged) |p| { + sg.kv = p.base; + sg.pt = p.table.ptr; + sg.ps = p.elems(g); + sg.ptok = @intCast(p.tokens); + } + return sg; } /// Keeps `path` (window rows, ascending from 0): the state as that chain would leave it, a tree's keys moved down. @@ -125,3 +191,8 @@ pub const Cache = struct { try self.keep(e, all[0..last.rows]); } }; + +fn freePaged(gpa: std.mem.Allocator, p: *Paged) void { + p.table.free(); + if (p.own_ids) gpa.free(p.ids); +} diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig index 8b7eb120ee..63210d3428 100644 --- a/zig/src/families/qwen3_5/cuda_reuse.zig +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -4,6 +4,7 @@ const cuda = @import("cuda"); const lanes = @import("lanes"); const geometry = @import("cuda_geometry.zig"); const engine = @import("cuda_engine.zig"); +const caches = @import("cuda_cache.zig"); const Engine = engine.Engine; const Cache = engine.Cache; @@ -90,6 +91,41 @@ const Layout = struct { } }; +/// A piece's device spans: one, or with a page table a run a page (the host copy keeps the contiguous order). +const Runs = struct { + base: u64, + pages: ?caches.Paged, + page_bytes: usize, + in_page: usize, + len: usize, + page: usize = 0, + done: usize = 0, + + const Run = struct { dev: u64, at: usize, len: usize }; + + fn of(e: *const Engine, cache: *const Cache, l: Layout, i: usize, p: Piece) Runs { + if (!p.kv or cache.paged == null) return .{ .base = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off, .pages = null, .page_bytes = 0, .in_page = 0, .len = p.len }; + const pg = cache.paged.?; + const row = head_dim * 2; // bf16 only: a paged cache refuses fp8 at init + std.debug.assert(l.format == .bf16); + return .{ .base = pg.base, .pages = pg, .page_bytes = pg.elems(e.g) * 2, .in_page = (i - 1) * pg.tokens * row, .len = p.len }; + } + + fn next(r: *Runs) ?Run { + if (r.done >= r.len) return null; + const pg = r.pages orelse { + r.done = r.len; + return .{ .dev = r.base, .at = 0, .len = r.len }; + }; + const span = pg.tokens * head_dim * 2; + const n = @min(span, r.len - r.done); + const run: Run = .{ .dev = r.base + pg.ids[r.page] * r.page_bytes + r.in_page, .at = r.done, .len = n }; + r.page += 1; + r.done += n; + return run; + } +}; + /// Whether a request's images match a kept state's: the same identities at the same rows, for every image before `at`. pub fn sameImages(kept: []const Seen, request: []const Seen, at: usize) bool { var n: usize = 0; @@ -135,8 +171,8 @@ pub fn saveCache(e: *Engine, cache: *const Cache, at: usize, seen: []const Seen) var off: usize = 0; for (0..l.count()) |i| { const p = l.piece(i); - const dev = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off; - if (p.len > 0) try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off, dev, p.len, e.stream.handle), "cuMemcpyDtoHAsync"); + var it = Runs.of(e, cache, l, i, p); + while (it.next()) |r| try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off + r.at, r.dev, r.len, e.stream.handle), "cuMemcpyDtoHAsync"); off += p.len; } try e.stream.synchronize(); @@ -174,8 +210,8 @@ pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque, seen: []const var off: usize = 0; for (0..l.count()) |i| { const p = l.piece(i); - const dev = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off; - if (p.len > 0) try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(dev, copy.host.bytes.ptr + off, p.len, e.stream.handle), "cuMemcpyHtoDAsync"); + var it = Runs.of(e, cache, l, i, p); + while (it.next()) |r| try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(r.dev, copy.host.bytes.ptr + off + r.at, r.len, e.stream.handle), "cuMemcpyHtoDAsync"); off += p.len; } try e.stream.synchronize(); // the host copy may be dropped once this returns diff --git a/zig/tests/cuda/contract_qwen_pages.zig b/zig/tests/cuda/contract_qwen_pages.zig new file mode 100644 index 0000000000..82c7418dc8 --- /dev/null +++ b/zig/tests/cuda/contract_qwen_pages.zig @@ -0,0 +1,296 @@ +//! qwen35.cu's attention cache through a page table: rope stores, attention reads, keep moves, paged == contiguous. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const hd = 256; +const max_chunks = 64; +const attentions = 2; // the fixture's attention layers: the kernels run layer 1, so a page's layer offset counts +const q_heads = 24; // the 27B's heads +const kv_heads = 4; + +// qwen35.cuh's Geo, Seg, Ancestor and Path (src/families/qwen3_5/cuda_cache.zig mirrors them for the engine) +const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int, nl: c_int }; +const Seg = extern struct { rec: u64 = 0, conv: u64 = 0, kv: u64, tap: u64 = 0, log: u64 = 0, before: u64 = 0, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int = 0, taps: c_int = 0, delta: c_int = 0, pt: u64 = 0, ps: u64 = 0, ptok: c_int = 0 }; +const Ancestor = extern struct { bits: u64, depth: u32, chain: u32 }; +const Path = extern struct { n: c_int, rows: [64]c_int }; + +/// Two streams' cached positions, the window rows each adds, positions a page. +const Case = struct { lens: [2]usize, rows: usize, ptok: usize }; +const cases = [_]Case{ + .{ .lens = .{ 0, 3 }, .rows = 1, .ptok = 16 }, + .{ .lens = .{ 127, 128 }, .rows = 4, .ptok = 16 }, + .{ .lens = .{ 129, 1000 }, .rows = 8, .ptok = 128 }, + .{ .lens = .{ 700, 33 }, .rows = 2, .ptok = 48 }, // a page no chunk width divides + .{ .lens = .{ 4095, 5000 }, .rows = 16, .ptok = 256 }, + .{ .lens = .{ 9000, 300 }, .rows = 3, .ptok = 512 }, +}; + +/// One case's host data: the logical caches (cap rows a head), the window's projections, the page tables. +const Fixture = struct { + c: Case, + cap: usize, + kv: [2][]u16, // [attention][key, value][head][cap][hd] a stream + q: []u16, // [rows][q_heads][hd] + k: []u16, // [rows][kv_heads][hd] + v: []u16, + ids: [2][]u32, // each stream's pages in the shared arena, shuffled across both + pages: usize, // the arena's pages (two spare) + + fn blockElems(f: Fixture) usize { + return attentions * 2 * kv_heads * f.cap * hd; + } + + fn pageElems(f: Fixture) usize { + return attentions * 2 * kv_heads * f.c.ptok * hd; + } + + /// The element where logical (attention a, key/value s, head h, position pos) of stream `st` sits in the arena. + fn paged(f: Fixture, st: usize, a: usize, s: usize, h: usize, pos: usize) usize { + const blk = (a * 2 + s) * kv_heads + h; + return f.ids[st][pos / f.c.ptok] * f.pageElems() + (blk * f.c.ptok + pos % f.c.ptok) * hd; + } + + fn logical(f: Fixture, a: usize, s: usize, h: usize, pos: usize) usize { + return (((a * 2 + s) * kv_heads + h) * f.cap + pos) * hd; + } + + fn make(a: std.mem.Allocator, c: Case, seed: u64) !Fixture { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const most = @max(c.lens[0], c.lens[1]) + c.rows; + const cap = (most + c.ptok - 1) / c.ptok * c.ptok; + var f: Fixture = .{ .c = c, .cap = cap, .kv = undefined, .q = undefined, .k = undefined, .v = undefined, .ids = undefined, .pages = 2 * (cap / c.ptok) + 2 }; + for (&f.kv) |*x| x.* = try randoms(a, rng, f.blockElems()); + f.q = try randoms(a, rng, c.rows * 2 * q_heads * hd); + f.k = try randoms(a, rng, c.rows * 2 * kv_heads * hd); + f.v = try randoms(a, rng, c.rows * 2 * kv_heads * hd); + const order = try a.alloc(u32, f.pages); + for (order, 0..) |*o, i| o.* = @intCast(i); + rng.shuffle(u32, order); + const n = cap / c.ptok; + f.ids = .{ order[0..n], order[n .. 2 * n] }; + return f; + } + + /// The arena holding both streams' logical caches at their pages; spare pages random. + fn arena(f: Fixture, a: std.mem.Allocator, rng: std.Random) ![]u16 { + const out = try randoms(a, rng, f.pages * f.pageElems()); + for (0..2) |st| for (0..attentions) |at| for (0..2) |s| for (0..kv_heads) |h| for (0..f.cap) |pos| { + @memcpy(out[f.paged(st, at, s, h, pos)..][0..hd], f.kv[st][f.logical(at, s, h, pos)..][0..hd]); + }; + return out; + } + + /// The arena's pages read back as stream `st`'s logical cache. + fn unpage(f: Fixture, a: std.mem.Allocator, arena_: []const u16, st: usize) ![]u16 { + const out = try a.alloc(u16, f.blockElems()); + for (0..attentions) |at| for (0..2) |s| for (0..kv_heads) |h| for (0..f.cap) |pos| { + @memcpy(out[f.logical(at, s, h, pos)..][0..hd], arena_[f.paged(st, at, s, h, pos)..][0..hd]); + }; + return out; + } +}; + +/// bf16 in [-1, 1): finite, so a sum's bytes are its order's alone. +fn randoms(a: std.mem.Allocator, rng: std.Random, n: usize) ![]u16 { + const out = try a.alloc(u16, n); + for (out) |*x| x.* = @truncate(@as(u32, @bitCast(rng.float(f32) * 2 - 1)) >> 16); + return out; +} + +fn bf(x: u16) f64 { + return @as(f32, @bitCast(@as(u32, x) << 16)); +} + +/// One layout's results: the rotated queries, the attention rows, each stream's cache after rope and after keep. +const Result = struct { q: []u8, out: []u8, after_rope: [2][]u16, after_keep: [2][]u16 }; + +const Fns = struct { rope: cuda.Function, partial: cuda.Function, merge: cuda.Function, compact: cuda.Function }; + +/// rope, attention, keep on layer 1 of both streams: contiguous, or pages placed by `tables`, read through `given`. +fn launch(gpu: check.Gpu, s: cuda.Stream, fns: Fns, f: Fixture, tables: ?[2][]const u32, given: ?[2][]const u32, a: std.mem.Allocator, seed: u64) !Result { + var rig = rig_.Rig{ .gpu = gpu, .s = s }; + defer rig.deinit(); + var prng = std.Random.DefaultPrng.init(seed); + const m = 2 * f.c.rows; + const geo: Geo = .{ .hid = 0, .nk = 0, .nv = 0, .qh = q_heads, .kvh = kv_heads, .nl = 0 }; + var segs: [2]Seg = undefined; + var kv: [2]u64 = undefined; // a stream's buffer, or the arena twice + var arena_host: []u16 = &.{}; + if (tables) |t| { + var g = f; + g.ids = .{ @constCast(t[0]), @constCast(t[1]) }; + arena_host = try g.arena(a, prng.random()); + const ptr = try rig.scratch(arena_host.len * 2); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = ptr, .len = arena_host.len * 2 }).upload(0, std.mem.sliceAsBytes(arena_host)); + kv = .{ ptr, ptr }; + } else for (0..2) |st| { + kv[st] = try rig.scratch(f.blockElems() * 2); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = kv[st], .len = f.blockElems() * 2 }).upload(0, std.mem.sliceAsBytes(f.kv[st])); + } + for (0..2) |st| { + segs[st] = .{ .kv = kv[st], .cap = @intCast(f.cap), .len = @intCast(f.c.lens[st]), .row0 = @intCast(st * f.c.rows), .rows = @intCast(f.c.rows) }; + if (given orelse tables) |t| { + segs[st].pt = try rig.input(std.mem.sliceAsBytes(t[st])); + segs[st].ps = f.pageElems(); + segs[st].ptok = @intCast(f.c.ptok); + } + } + const rowseg = try a.alloc(i32, m); + const rowanc = try a.alloc(Ancestor, m); + for (0..m) |r| { + rowseg[r] = @intCast(r / f.c.rows); + rowanc[r] = .{ .bits = 0, .depth = @intCast(r % f.c.rows), .chain = 1 }; + } + const segs_d = try rig.input(std.mem.sliceAsBytes(&segs)); + const rowseg_d = try rig.input(std.mem.sliceAsBytes(rowseg)); + const rowanc_d = try rig.input(std.mem.sliceAsBytes(rowanc)); + const qn = try rig.scratch(f.q.len * 2); // rotated in place + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = qn, .len = f.q.len * 2 }).upload(0, std.mem.sliceAsBytes(f.q)); + const kn = try rig.input(std.mem.sliceAsBytes(f.k)); + const vr = try rig.input(std.mem.sliceAsBytes(f.v)); + const slots = m * q_heads * max_chunks; + const pm = try rig.scratch(slots * 4); + const pl = try rig.scratch(slots * 4); + const po = try rig.scratch(slots * hd * 4); + const out = try rig.output(m * q_heads * hd * 2); + const before = try a.dupe(u8, try rig.inputBytes()); + const layer: c_int = 1; + const none: u64 = 0; + var args: cuda.Args = .{}; + inline for (.{ geo, qn, kn, vr, segs_d, rowseg_d, rowanc_d, layer, none, none }) |x| args.add(x); + try cuda.launch.launch(fns.rope, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = hd } }, s, &args); + args = .{}; + inline for (.{ geo, qn, segs_d, rowseg_d, rowanc_d, layer, none, @as(f32, 1.0 / 16.0), pm, pl, po }) |x| args.add(x); + try cuda.launch.launch(fns.partial, .{ .grid = .{ .x = max_chunks, .y = kv_heads, .z = @intCast(m) }, .block = .{ .x = 32 * (q_heads / kv_heads) } }, s, &args); + args = .{}; + inline for (.{ geo, pm, pl, po, segs_d, rowseg_d, rowanc_d, none, out }) |x| args.add(x); + try cuda.launch.launch(fns.merge, .{ .grid = .{ .x = q_heads, .y = @intCast(m) }, .block = .{ .x = 32 } }, s, &args); + var r: Result = undefined; + r.q = try a.dupe(u8, try rig.bytes(qn, f.q.len * 2)); + r.out = try a.dupe(u8, try rig.bytes(out, m * q_heads * hd * 2)); + try readCaches(&rig, a, f, tables, kv, &r.after_rope); + // keep: stream 0's window rows 0, 2 and the last (a tree's path), moved down after its cache + if (f.c.rows >= 3) { + var path: Path = .{ .n = 3, .rows = @splat(0) }; + path.rows[1] = 2; + path.rows[2] = @intCast(f.c.rows - 1); + args = .{}; + inline for (.{ geo, segs[0], path }) |x| args.add(x); + try cuda.launch.launch(fns.compact, .{ .grid = .{ .x = attentions, .y = kv_heads }, .block = .{ .x = hd } }, s, &args); + } + try readCaches(&rig, a, f, tables, kv, &r.after_keep); + try check.expect(try rig.intact() and !rig.overlap() and std.mem.eql(u8, before, try rig.inputBytes()), "qwen35 pages: a launch changed an input or guard", .{}); + return r; +} + +fn readCaches(rig: *rig_.Rig, a: std.mem.Allocator, f: Fixture, tables: ?[2][]const u32, kv: [2]u64, out: *[2][]u16) !void { + if (tables) |t| { + var g = f; + g.ids = .{ @constCast(t[0]), @constCast(t[1]) }; + const all = try a.alloc(u16, f.pages * f.pageElems()); + @memcpy(std.mem.sliceAsBytes(all), try rig.bytes(kv[0], all.len * 2)); + for (0..2) |st| out[st] = try g.unpage(a, all, st); + } else for (0..2) |st| { + out[st] = try a.alloc(u16, f.blockElems()); + @memcpy(std.mem.sliceAsBytes(out[st]), try rig.bytes(kv[st], f.blockElems() * 2)); + } +} + +/// Each attention row against float64 softmax(q k / 16) v: a bound (fp32 and fast exp), bf16 rounding of its size. +fn agrees(a: std.mem.Allocator, f: Fixture, r: Result) !bool { + const q: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, r.q)); + const out: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, r.out)); + const keys_most = @max(f.c.lens[0], f.c.lens[1]) + f.c.rows; + const p = try a.alloc(f64, keys_most); + defer a.free(p); + for (0..2 * f.c.rows) |row| { + const st = row / f.c.rows; + const keys = f.c.lens[st] + row % f.c.rows + 1; // the chain's rows are positions after the cache + const cache = r.after_rope[st]; + for (0..q_heads) |qh| { + const h = qh / (q_heads / kv_heads); + const qv = q[(row * q_heads + qh) * hd ..][0..hd]; + var top: f64 = -std.math.inf(f64); + for (0..keys) |t| { + var dot: f64 = 0; + const kr = cache[f.logical(1, 0, h, t)..][0..hd]; + for (0..hd) |d| dot += bf(qv[d]) * bf(kr[d]); + p[t] = dot / 16.0; + top = @max(top, p[t]); + } + var sum: f64 = 0; + for (p[0..keys]) |*x| { + x.* = @exp(x.* - top); + sum += x.*; + } + for (0..hd) |d| { + var o: f64 = 0; + var mag: f64 = 0; + for (0..keys) |t| { + const v = bf(cache[f.logical(1, 1, h, t) + d]); + o += p[t] * v; + mag += p[t] * @abs(v); + } + o /= sum; + mag /= sum; + const got = bf(out[(row * q_heads + qh) * hd + d]); + if (!(@abs(got - o) <= std.math.ldexp(@as(f64, 1), -8) * (@abs(o) + mag))) return false; + } + } + } + return true; +} + +fn eqCaches(x: [2][]u16, y: [2][]u16) bool { + return std.mem.eql(u16, x[0], y[0]) and std.mem.eql(u16, x[1], y[1]); +} + +pub fn run(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.qwen35); + defer mod.unload(); + var keep = try cuda.Module.load(gpu.d, cuda.kernels.qwen35_keep); + defer keep.unload(); + const fns: Fns = .{ .rope = try mod.function("tf_q35_rope"), .partial = try mod.function("tf_q35_attn_partial"), .merge = try mod.function("tf_q35_attn_merge"), .compact = try keep.function("tf_q35k_compact") }; + var red_table = false; + var red_bound = false; + for (cases, 0..) |c, i| { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const f = try Fixture.make(a, c, 0x9a6e + i); + const flat = try launch(gpu, s, fns, f, null, null, a, 1); + try check.expect(try agrees(a, f, flat), "qwen35 pages: contiguous attention off float64 at lens {any}, rows {d}", .{ c.lens, c.rows }); + for (0..3) |rep| { + const pg = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, null, a, 2 + rep); + try check.sameBytes("qwen35 pages: rotated queries", pg.q, flat.q); + try check.sameBytes("qwen35 pages: attention rows, paged against contiguous", pg.out, flat.out); + try check.expect(eqCaches(pg.after_rope, flat.after_rope), "qwen35 pages: rope's stores differ at lens {any}, page {d}", .{ c.lens, c.ptok }); + try check.expect(eqCaches(pg.after_keep, flat.after_keep), "qwen35 pages: keep's moves differ at lens {any}, page {d}", .{ c.lens, c.ptok }); + } + // red controls: a table with two of stream 1's pages swapped reads other keys; a planted row breaks the bound + const n1 = (c.lens[1] + c.rows + c.ptok - 1) / c.ptok; + if (n1 >= 2) { + const swapped = try a.dupe(u32, f.ids[1]); + std.mem.swap(u32, &swapped[0], &swapped[n1 - 1]); + const bad = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, .{ f.ids[0], swapped }, a, 9); + const moved = !std.mem.eql(u8, bad.out, flat.out); + red_table = red_table or moved; + try check.expect(moved, "qwen35 pages: a swapped page table still gave contiguous bytes (control missed) at lens {any}", .{c.lens}); + } + const o = try a.dupe(u8, flat.out); + std.mem.writeInt(u16, o[2 * (q_heads * hd + 5) ..][0..2], 0x4040, .little); // 3.0 where |o| <= 1 + var bad_r = flat; + bad_r.out = o; + const caught = !(try agrees(a, f, bad_r)); + red_bound = red_bound or caught; + try check.expect(caught, "qwen35 pages: a planted attention row passed float64 (control missed)", .{}); + std.debug.print("PASS qwen35 pages: lens {any}, {d} rows a stream, {d}-position pages: paged == contiguous (queries, rows, stores, keep), three repeats, float64 bound, guards\n", .{ c.lens, c.rows, c.ptok }); + } + try check.expect(red_table and red_bound, "qwen35 pages: a control never fired", .{}); + std.debug.print("PASS qwen35 pages: controls RED as planted (swapped table, planted row)\n", .{}); +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 86aa6eb966..b929e4be03 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -281,6 +281,7 @@ pub fn run(gpu: check.Gpu) !void { std.debug.print("PASS gradient component suites (FP64 and finite-difference checks); resource receipts still required\n", .{}); const t1 = check.now(gpu.io); try @import("contract_qwen_served.zig").run(gpu); + try @import("contract_qwen_pages.zig").run(gpu); try @import("contract_qwen_learning.zig").numerical(gpu); const lim = try live(gpu); const t2 = check.now(gpu.io); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 837f1bb070..8380cef677 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -188,6 +188,7 @@ const usage = \\ exl3-tiles [DIR] EXL3's tile candidate: codebooks and tiles bit-equal, FP64, rows through qlinear and Dense \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) + \\ qwen-pages qwen35.cu's attention cache through a page table: paged == contiguous, float64, red controls \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) \\ glm-vision MODEL F32 BF16 IMG GLM's image tower and preparation against vision_ref.py's torch runs @@ -357,6 +358,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "q27-gemm")) return q27_gemm_tests.run(gpu); if (std.mem.eql(u8, cmd, "gdn-back-27b")) return gdn_back_tests.runShape(gpu, 5120, 64); if (std.mem.eql(u8, cmd, "kernel-contract")) return contract_tests.run(gpu); + if (std.mem.eql(u8, cmd, "qwen-pages")) return @import("contract_qwen_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "tree")) return tree_tests.run(gpu); if (std.mem.eql(u8, cmd, "block")) return block_tests.run(gpu); From d5189b3c17347fc51f560c6e287b3fe02bb67306 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:21:28 +0200 Subject: [PATCH 02/11] cuda qwen3_5: the streams beside the first take KV pages from the core pool (TF_QWEN_KV_PAGES) TF_QWEN_KV_PAGES=n (off by default): the first stream keeps its reserved window; every other stream takes n-position pages of one device arena from core kv_pool, as many as its prompt, reply and one more forward need (cuda_pages.zig: layout, Pages, lanes of Cache.initPaged). Admission goes through fits; a request with no lane or pages free waits (LaneMemory). A closed stream's pages go back once the engine's stream drains. qwen3_5_moe serves the same pool. TF_QWEN_KV_PAGES_ALL=1 puts every stream on pages (checks). --kv-cache fp8 and --slide refuse it. The native host's plan: cuda_kv_pool.served() adds a wrapping drafter's per-stream buffers to each stream's fixed part, and a family with no free stream pays its first window before pages are planned. Nemotron (a free stream, no wrapper) plans as before. GB10, Qwen3.8-27B MLX-4bit, n=256, --parallel 4 at 32768: 3 more streams on 387 pages (6.05 GiB) instead of three 2.5 GiB windows; solo bytes == contiguous (3/3 chat, 3/3 8k prefill), 13.9-14.0 tok/s, c=4 48.6 (code) / 48.2 (chat) against 48.4 / 46.9 contiguous. With DFlash2 at --parallel 8: drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh, outputs equal the page test's. Pages only (8 requests on 7 lanes, one waiting): the same bytes, a kept state restored into a lane. --- zig/src/families/qwen3_5/cuda.zig | 1 + zig/src/families/qwen3_5/cuda_native.zig | 50 +++++- zig/src/families/qwen3_5/cuda_native_moe.zig | 3 + zig/src/families/qwen3_5/cuda_pages.zig | 157 +++++++++++++++++++ zig/src/native/cuda.zig | 5 +- zig/src/native/cuda_kv_pool.zig | 18 +++ 6 files changed, 224 insertions(+), 10 deletions(-) create mode 100644 zig/src/families/qwen3_5/cuda_pages.zig diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 0dd79a8acd..2fae132c6d 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -26,6 +26,7 @@ test { _ = @import("cuda_vision_prompt.zig"); _ = reuse; _ = native; + _ = @import("cuda_pages.zig"); } pub const slide = @import("cuda_slide.zig"); diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index d783449cb0..f24554e0fc 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -13,6 +13,8 @@ const Draft = @import("cuda_draft.zig").Draft; const vision_prompt = @import("cuda_vision_prompt.zig"); const reuse = @import("cuda_reuse.zig"); const Mtp = @import("cuda_mtp.zig").Mtp; +const pages_mod = @import("cuda_pages.zig"); +const kv_pool = @import("core").kv_pool; const be = lanes.backend; pub const vision = @import("cuda_vision.zig"); pub const slide = @import("cuda_slide.zig"); @@ -25,6 +27,8 @@ pub const max_segments: u32 = 1; pub const prompt_rows: u32 = E.max_rows; /// --drafter loads the DFlash2 draft model (openDrafter); the checkpoints have no head of their own. pub const takes_drafter = true; +/// A request's pages hold one more forward past its prompt and reply (TF_QWEN_KV_PAGES). +pub const window_rows: u32 = pages_mod.window_rows; /// Kept prompt states for the shared cache (cuda_reuse.zig), in pinned host memory, on the zero-anchored chunk grid. pub const snap_bytes = reuse.bytes; @@ -58,11 +62,30 @@ pub const Loaded = struct { vision: ?*vision.Vision = null, /// What the prompt cache's saves copy: the engine and the streams' caches. target: *reuse.Target, + /// TF_QWEN_KV_PAGES: the streams beside the first take pages of one pool (null: a window each). + layout: ?kv_pool.Layout = null, }; +/// The pool beside the first stream's window: `streams` lanes and `pages` pages, made now. +pub fn attach(p: *anyopaque, streams: u32, pages: usize) !void { + const l: *Lanes = @ptrCast(@alignCast(p)); + l.pages = try pages_mod.Pages.init(l.gpa, l.e, l.capacity, l.page_tokens, streams, pages); + l.paged_only = pages_mod.pagedOnly() and streams > 0; // checks: every stream on pages, the reserved window idle +} + +/// The lane host's admission: the first stream's window is free, or a lane and the request's pages are. +pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { + const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.pool.items.len > 0 and !l.paged_only) return true; + const pg = l.pages orelse return true; + return pg.fits(prompt, max_tokens); +} + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.QwenHasNoCapturedKernels; if (o.slide and o.kv8) return error.UnsupportedLearningFormat; + const page_tokens = try pages_mod.paging(); + if (page_tokens > 0 and (o.kv8 or o.slide)) return error.QwenPagesNeedBf16; const e = try E.Engine.init(gpa, io, ctx, dir, o.kv8); errdefer e.deinit(); if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary @@ -71,7 +94,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ const v: ?*vision.Vision = if (o.vision) |lim| try vision.Vision.open(gpa, io, ctx.d, e.stream, dir, e.g.hidden, lim) else null; errdefer if (v) |x| x.deinit(); const l = try gpa.create(Lanes); - l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .cost_context = ctx, .startup = e.startup, .vision = v, .target = undefined }; + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context, .page_tokens = page_tokens, .cost_context = ctx, .startup = e.startup, .vision = v, .target = undefined }; l.target = .{ .e = e, .caches = &l.caches, .seen = &l.seen }; return .{ .learns = if (o.slide) e else null, @@ -85,6 +108,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .reserve = reserve, .vision = v, .target = &l.target, + .layout = if (page_tokens > 0) pages_mod.layout(e, o.context, page_tokens) else null, }; } @@ -100,7 +124,8 @@ pub fn laneFacts(window: u32, verify: u32) lanes.Model { fn reserve(p: *anyopaque, streams: u32) anyerror!void { const l: *Lanes = @ptrCast(@alignCast(p)); - try reservation.target(E.Cache, E.Part, l, streams); + // with a pool the first stream's window is the one reserved; `attach` made the other streams' lanes + try reservation.target(E.Cache, E.Part, l, if (l.pages != null) @min(streams, 1) else streams); } /// The DFlash2 drafter in `dir` over engine `p`, 4-bit only, and the target's windows timed for the round loop. pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void, prepare: struct { ctx: *anyopaque, run: *const fn (*anyopaque) anyerror!lanes.drafter.Facts, bytes: usize } } { @@ -168,6 +193,8 @@ fn closeDrafter(p: *anyopaque) void { pub fn explain(ctx: ?*anyopaque, err: anyerror) ?[]const u8 { if (ctx == null) return switch (err) { error.NoQwenVision => "--vision: this checkpoint has no image tower (no vision_config): serve it without --vision", + error.QwenPagesNeedBf16 => "TF_QWEN_KV_PAGES: paged caches hold bf16 rows and serve no learner yet; drop --kv-cache fp8 and --slide, or unset it", + error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0 or unset: a window each)", error.UnsupportedQwenVision => "--vision: this checkpoint's image tower or text M-RoPE is not one the CUDA engine reads", else => null, }; @@ -224,6 +251,10 @@ const Lanes = struct { io: std.Io, e: *E.Engine, capacity: usize, + /// TF_QWEN_KV_PAGES's positions a page (0: no pool) and the pool `attach` made. + page_tokens: usize = 0, + pages: ?*pages_mod.Pages = null, + paged_only: bool = false, cost_context: ?*const cuda.Context = null, cost_identity: ?[32]u8 = null, cost_kind: cost_adapter.Kind = .dflash, @@ -241,13 +272,14 @@ const Lanes = struct { fn deinit(self: *Lanes) void { var it = self.caches.valueIterator(); - while (it.next()) |cache| self.drop(cache.*); + while (it.next()) |cache| self.give(cache.*); // a pool lane's pages back, a window to the reserve self.caches.deinit(self.gpa); var seen = self.seen.valueIterator(); while (seen.next()) |v| self.gpa.free(v.*); self.seen.deinit(self.gpa); for (self.pool.items) |cache| self.drop(cache); self.pool.deinit(self.gpa); + if (self.pages) |p| p.deinit(); if (self.vision) |v| v.deinit(); self.e.deinit(); self.gpa.destroy(self); @@ -258,12 +290,13 @@ const Lanes = struct { self.gpa.destroy(cache); } - /// A cache for a new stream: a pooled one reset, else one more (a window beyond the reserve). - fn take(self: *Lanes) !*E.Cache { - if (self.pool.pop()) |cache| { + /// A cache for a new stream: a pooled one reset, a pool lane with its pages, else one more window. + fn take(self: *Lanes, prompt: usize, max_new: usize) !*E.Cache { + if (!self.paged_only) if (self.pool.pop()) |cache| { try cache.reset(self.e); return cache; - } + }; + if (self.pages) |p| return p.open(prompt, max_new); // LaneMemory: the lane host waits for a stream to end const cache = try self.gpa.create(E.Cache); errdefer self.gpa.destroy(cache); cache.* = try E.Cache.init(self.gpa, self.e, self.capacity); @@ -271,6 +304,7 @@ const Lanes = struct { } fn give(self: *Lanes, cache: *E.Cache) void { + if (self.pages) |p| if (p.close(cache)) return; self.pool.append(self.gpa, cache) catch self.drop(cache); } @@ -309,7 +343,7 @@ const Lanes = struct { return err; }; errdefer if (self.seen.fetchRemove(stream)) |kv| self.gpa.free(kv.value); - const cache = try self.take(); + const cache = try self.take(ids.len, stream.max_new); errdefer self.give(cache); // the stream's cache is found by a save at a mark (cuda_reuse.zig) while its pass runs try self.caches.put(self.gpa, stream, cache); diff --git a/zig/src/families/qwen3_5/cuda_native_moe.zig b/zig/src/families/qwen3_5/cuda_native_moe.zig index 5d0f4b0c82..c7ede8d5e5 100644 --- a/zig/src/families/qwen3_5/cuda_native_moe.zig +++ b/zig/src/families/qwen3_5/cuda_native_moe.zig @@ -10,6 +10,9 @@ pub const prompt_rows = native.prompt_rows; pub const takes_drafter = true; pub const Options = native.Options; pub const open = native.open; +pub const attach = native.attach; // TF_QWEN_KV_PAGES: the same pool as the dense models' +pub const fits = native.fits; +pub const window_rows = native.window_rows; pub const openDrafter = native.openMtp; pub const explain = native.explain; /// The dense family's kept prompt states, its open refusals' words and its image tower: the MLP is the only change. diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig new file mode 100644 index 0000000000..a325c4655d --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -0,0 +1,157 @@ +//! Qwen3.5's attention caches from the core pool: the streams beside the first take pages of one device arena. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const engine = @import("cuda_engine.zig"); +const caches = @import("cuda_cache.zig"); + +const Engine = engine.Engine; +const Cache = caches.Cache; + +/// Rows past a request's prompt and reply its pages hold: one more forward's (a verify window or a prompt chunk). +pub const window_rows: u32 = engine.max_rows; + +/// TF_QWEN_KV_PAGES=n: streams beside the first take n-position pages from one pool (unset or 0: a window each). +pub fn paging() !usize { + const v = std.c.getenv("TF_QWEN_KV_PAGES") orelse return 0; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch error.BadQwenKvPages; +} + +/// TF_QWEN_KV_PAGES_ALL=1: every stream on pages, the first stream's reserved window left idle (checks). +pub fn pagedOnly() bool { + const v = std.c.getenv("TF_QWEN_KV_PAGES_ALL") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// Pages a stream maps at most: a full window and one more forward. +fn mostPages(capacity: usize, tokens: usize) usize { + return (capacity + window_rows + tokens - 1) / tokens; +} + +/// A stream's fixed part (state, log, tap ring, table), its keys and values a token, a kept state's recurrences. +pub fn layout(e: *const Engine, capacity: usize, tokens: usize) kv.Layout { + return .{ .fixed_bytes = e.streamBytes(0) + mostPages(capacity, tokens) * 4, .token_bytes = e.g.kvBytes(1, false), .page_tokens = @intCast(tokens), .snap_bytes = e.g.stateBytes() }; +} + +/// One pool lane: a cache whose table maps its stream's pages, and that stream's lease. +const Lane = struct { cache: *Cache, lease: kv.Lease = .{}, used: bool = false }; + +pub const Pages = struct { + gpa: std.mem.Allocator, + e: *Engine, + layout: kv.Layout, + capacity: usize, + pool: kv.Pool, + arena: cuda.DeviceBuffer, + lanes: []Lane, + + /// `streams` lanes and `pages` pages, made and touched now. + pub fn init(gpa: std.mem.Allocator, e: *Engine, capacity: usize, tokens: usize, streams: usize, pages: usize) !*Pages { + const l = layout(e, capacity, tokens); + const p = try gpa.create(Pages); + errdefer gpa.destroy(p); + var arena = try cuda.DeviceBuffer.alloc(e.d, pages * l.pageBytes()); + errdefer arena.free(); + try arena.fill8(0, null); // committed now (the driver commits on first touch) + p.* = .{ .gpa = gpa, .e = e, .layout = l, .capacity = capacity, .pool = try kv.Pool.init(gpa, pages, l.pageBytes()), .arena = arena, .lanes = &.{} }; + errdefer p.pool.deinit(); + p.lanes = try gpa.alloc(Lane, streams); + errdefer gpa.free(p.lanes); + var made: usize = 0; + errdefer for (p.lanes[0..made]) |x| { + x.cache.deinit(gpa); + gpa.destroy(x.cache); + }; + for (p.lanes) |*x| { + const c = try gpa.create(Cache); + errdefer gpa.destroy(c); + c.* = try Cache.initPaged(gpa, e, mostPages(capacity, tokens), tokens, arena.ptr); + x.* = .{ .cache = c }; + made += 1; + } + return p; + } + + pub fn deinit(p: *Pages) void { + p.e.stream.synchronize() catch std.log.warn("qwen3_5: the pool's stream did not drain at close", .{}); + for (p.lanes) |*x| { + if (x.used) p.pool.finish(&x.lease); + x.cache.deinit(p.gpa); + p.gpa.destroy(x.cache); + } + p.gpa.free(p.lanes); + p.arena.free(); + p.pool.deinit(); + p.gpa.destroy(p); + } + + fn need(p: *const Pages, prompt: usize, max_tokens: usize) usize { + return p.layout.need(prompt, max_tokens, window_rows); + } + + /// Whether a lane is free and the pool can promise a request's pages now. + pub fn fits(p: *const Pages, prompt: usize, max_tokens: usize) bool { + for (p.lanes) |x| if (!x.used) return p.pool.fits(p.need(prompt, max_tokens)); + return false; + } + + /// A free lane's cache, reset, holding the pages of `prompt` + `max_tokens` rows and a window; LaneMemory: not now. + pub fn open(p: *Pages, prompt: usize, max_tokens: usize) !*Cache { + const x = for (p.lanes) |*x| { + if (!x.used) break x; + } else return error.LaneMemory; + const n = p.need(prompt, max_tokens); + if (n > mostPages(p.capacity, p.layout.page_tokens)) return error.InvalidQwenPrompt; + try p.pool.promise(&x.lease, n); + errdefer p.pool.finish(&x.lease); + var up: Upload = .{ .p = p, .cache = x.cache }; + try p.pool.cover(&x.lease, n, .{ .ctx = &up, .write = Upload.write }); + try x.cache.reset(p.e); + x.cache.paged.?.ids = x.lease.ids.items; + x.cache.cap = n * p.layout.page_tokens; + x.used = true; + return x.cache; + } + + /// The stream on `cache` ended: its pages go back once the forwards that read them are done. False: not a lane's. + pub fn close(p: *Pages, cache: *Cache) bool { + const x = for (p.lanes) |*x| { + if (x.cache == cache) break x; + } else return false; + if (!x.used) return true; + // a page another stream takes next may still be read by queued work on the engine's stream + p.e.stream.synchronize() catch std.log.warn("qwen3_5: a closed stream's forwards did not drain; its pages go back", .{}); + p.pool.finish(&x.lease); + cache.paged.?.ids = &.{}; + cache.cap = 0; + x.used = false; + return true; + } + + /// Whether `cache` is one of the pool's lanes. + pub fn owns(p: *const Pages, cache: *const Cache) bool { + for (p.lanes) |x| if (x.cache == cache) return true; + return false; + } + + /// A lane's new table entries uploaded to its device table. + const Upload = struct { + p: *Pages, + cache: *Cache, + + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const u: *Upload = @ptrCast(@alignCast(ctx)); + const t = u.cache.paged.?.table; + try t.uploadAsync(from * 4, std.mem.sliceAsBytes(ids[from..]), u.p.e.stream.handle); + try u.p.e.stream.synchronize(); // the ids slice may move as the lease grows + } + }; +}; + +test "a stream's pages: its prompt, reply and one more forward, no more than a full window's" { + const l: kv.Layout = .{ .fixed_bytes = 1, .token_bytes = 64 << 10, .page_tokens = 256, .snap_bytes = 1 }; + try std.testing.expectEqual(@as(usize, 1), l.need(1, 0, window_rows)); + try std.testing.expectEqual(@as(usize, 2), l.need(100, 29, window_rows)); + try std.testing.expectEqual(@as(usize, 129), mostPages(32768, 256)); + try std.testing.expect(l.need(32768 - 1, 1, window_rows) <= mostPages(32768, 256)); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index a8463051a5..00873d7d13 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -346,10 +346,11 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; const free = loaded.free_streams; const grow = budget.growRoom(room, !@hasField(F.Options, "lanes")); - const layout = kv_pages.layout(loaded); + const wrapped = kv_pages.served(kv_pages.layout(loaded), free, loaded.stream_bytes, served.stream_bytes, room); + const layout = wrapped.layout; var paged: ?kv_pages.Plan = null; const streams = if (layout) |l| blk: { - paged = try kv_pages.plan(F, a, l, room, o, @intCast(window), problem, host_cache) orelse return null; + paged = try kv_pages.plan(F, a, l, wrapped.room, o, @intCast(window), problem, host_cache) orelse return null; break :blk paged.?.lanes; } else budget.admit(grow, served.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { diff --git a/zig/src/native/cuda_kv_pool.zig b/zig/src/native/cuda_kv_pool.zig index 754602de1d..355e0e541b 100644 --- a/zig/src/native/cuda_kv_pool.zig +++ b/zig/src/native/cuda_kv_pool.zig @@ -12,6 +12,24 @@ pub fn layout(loaded: anytype) ?kv_pool.Layout { return if (comptime @hasField(@TypeOf(loaded), "layout")) loaded.layout else null; } +/// The plan's view of a wrapped family: drafter buffers join the fixed part; no free stream pays its window first. +pub fn served(l: ?kv_pool.Layout, free_streams: u32, family_bytes: u64, stream_bytes: u64, room: u64) struct { layout: ?kv_pool.Layout, room: u64 } { + var out = l orelse return .{ .layout = null, .room = room }; + out.fixed_bytes += stream_bytes -| family_bytes; + return .{ .layout = out, .room = if (free_streams == 0) room -| stream_bytes else room }; +} + +test "a wrapped family's plan: drafter buffers a stream, a first window without a free stream" { + const l: kv_pool.Layout = .{ .fixed_bytes = 10, .token_bytes = 1, .page_tokens = 4, .snap_bytes = 1 }; + const nemotron = served(l, 1, 500, 500, 1000); // its own sequence is in the model's bytes: nothing moves + try std.testing.expectEqual(@as(u64, 10), nemotron.layout.?.fixed_bytes); + try std.testing.expectEqual(@as(u64, 1000), nemotron.room); + const qwen = served(l, 0, 500, 530, 1000); // a drafter's 30 bytes a stream, the first stream's 530 first + try std.testing.expectEqual(@as(u64, 40), qwen.layout.?.fixed_bytes); + try std.testing.expectEqual(@as(u64, 470), qwen.room); + try std.testing.expectEqual(@as(?kv_pool.Layout, null), served(null, 0, 1, 2, 3).layout); +} + /// The real native startup room before page planning and cache allocation. pub fn startupRoom(comptime F: type, a: Allocator, loaded: anytype, o: api.Open, available: u64, host_budget: u64, problem: *[]const u8) !?u64 { const memory = @import("cuda_family_memory.zig"); From 6037621afb367960f19876d177fc9069e7cb79c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:37:50 +0200 Subject: [PATCH 03/11] cuda qwen3_5: with a KV pool, warm the shared shapes on lanes opened for one row Reservation in pool mode reserves the first stream's window and warms the 2..N-segment shapes on pool lanes opened for one row and closed again, so the first concurrent round runs no shape cold. GB10, Qwen3.8-27B MLX-4bit, --parallel 8, a fresh server's first round of 8 requests: TTFT max 0.97/0.97 s warmed, 0.98/0.98 unwarmed, 0.99 contiguous (qwen captures no graphs, so a cold shape cost nothing visible); the bytes of every arm equal. --- zig/src/families/qwen3_5/cuda_native.zig | 3 +- zig/src/families/qwen3_5/cuda_reserve.zig | 48 ++++++++++++++++------- 2 files changed, 35 insertions(+), 16 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index f24554e0fc..6eb97a7c88 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -124,8 +124,7 @@ pub fn laneFacts(window: u32, verify: u32) lanes.Model { fn reserve(p: *anyopaque, streams: u32) anyerror!void { const l: *Lanes = @ptrCast(@alignCast(p)); - // with a pool the first stream's window is the one reserved; `attach` made the other streams' lanes - try reservation.target(E.Cache, E.Part, l, if (l.pages != null) @min(streams, 1) else streams); + try reservation.target(E.Cache, E.Part, l, streams); // with a pool: the first window, the shapes on its lanes } /// The DFlash2 drafter in `dir` over engine `p`, 4-bit only, and the target's windows timed for the round loop. pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void, prepare: struct { ctx: *anyopaque, run: *const fn (*anyopaque) anyerror!lanes.drafter.Facts, bytes: usize } } { diff --git a/zig/src/families/qwen3_5/cuda_reserve.zig b/zig/src/families/qwen3_5/cuda_reserve.zig index c6fa09e916..b92f7430bd 100644 --- a/zig/src/families/qwen3_5/cuda_reserve.zig +++ b/zig/src/families/qwen3_5/cuda_reserve.zig @@ -2,11 +2,16 @@ const warm = @import("cuda_warm.zig"); const E = @import("cuda_engine.zig"); +/// `streams` windows reserved and warmed; with a page pool (TF_QWEN_KV_PAGES) the first stream's window only, the +/// shared shapes warmed on lanes opened for one row and closed again. pub fn target(comptime Cache: type, comptime Part: type, l: anytype, streams: u32) !void { const started = warm.profile.now(l.io); var warm_ns: i96 = 0; - try l.pool.ensureTotalCapacity(l.gpa, streams); - while (l.pool.items.len < streams) { + const paged = comptime @hasField(@TypeOf(l.*), "pages"); + const pooled = if (paged) l.pages != null else false; + const windows: u32 = if (pooled) @min(streams, 1) else streams; + try l.pool.ensureTotalCapacity(l.gpa, windows); + while (l.pool.items.len < windows) { const cache = try l.gpa.create(Cache); errdefer l.gpa.destroy(cache); cache.* = try Cache.init(l.gpa, l.e, l.capacity); @@ -23,18 +28,17 @@ pub fn target(comptime Cache: type, comptime Part: type, l: anytype, streams: u3 var parts: [E.max_segs]Part = undefined; var ids: [@max(E.max_segs, E.most_window)]u32 = @splat(0); var picks: [@max(E.max_segs, E.most_window)]E.Draw = undefined; - for (0..@min(streams, E.max_segs)) |n| { - parts[n] = .{ .cache = l.pool.items[n], .rows = 1 }; - picks[n] = .{ .rule = null, .position = 1 }; - const shape = E.Engine.Shape.init(n + 1, n + 1, true, l.e.n_taps > 0, 1); - if (l.warmed.contains(shape)) continue; - const warm_start = warm.profile.now(l.io); - try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); - try l.e.stream.synchronize(); - warm_ns += warm.profile.elapsed(l.io, warm_start); - for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); - l.warmed.mark(shape); - } + var warming: [E.max_segs]*Cache = undefined; + var opened: usize = 0; + const most = @min(streams, E.max_segs); + for (l.pool.items[0..@min(windows, most)], 0..) |cache, i| warming[i] = cache; + if (comptime paged) if (l.pages) |p| { + while (windows + opened < most) : (opened += 1) warming[windows + opened] = p.open(1, 0) catch break; + }; + defer if (comptime paged) if (l.pages) |p| for (warming[windows..][0..opened]) |cache| { + _ = p.close(cache); + }; + try warmShapes(Part, l, warming[0 .. @min(windows, most) + opened], &parts, &ids, &picks, &warm_ns); // Pending calibration warms on its capacity-safe cache during synchronous open. if (l.e.n_taps > 0 and l.calibration == null) for (2..l.e.window_rows + 1) |rows| { const shape = E.Engine.Shape.init(rows, 1, true, true, 1); @@ -51,4 +55,20 @@ pub fn target(comptime Cache: type, comptime Part: type, l: anytype, streams: u3 l.startup.reserve += warm.profile.elapsed(l.io, started) - warm_ns; } +/// Each shape of 1..caches.len one-row segments run once on `caches`, then reset. +fn warmShapes(comptime Part: type, l: anytype, caches: anytype, parts: []Part, ids: []u32, picks: []E.Draw, warm_ns: *i96) !void { + for (caches, 0..) |cache, n| { + parts[n] = .{ .cache = cache, .rows = 1 }; + picks[n] = .{ .rule = null, .position = 1 }; + const shape = E.Engine.Shape.init(n + 1, n + 1, true, l.e.n_taps > 0, 1); + if (l.warmed.contains(shape)) continue; + const warm_start = warm.profile.now(l.io); + try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); + try l.e.stream.synchronize(); + warm_ns.* += warm.profile.elapsed(l.io, warm_start); + for (caches[0 .. n + 1]) |c| try c.reset(l.e); + l.warmed.mark(shape); + } +} + pub const Test = if (@import("builtin").is_test) @import("cuda_reserve_test.zig") else struct {}; From 5c09685eb628c125d51f9b6288e7e1ae4e750b7c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:32:58 +0200 Subject: [PATCH 04/11] cuda qwen3_5: --kv-cache fp8 through the page table too qwen35_kv8.cu's store, tf_q35_attn_partial8 and tf_q35k_compact8 take a row's e4m3 codes and fp32 scale through row8(): cap rows a head as before, or with a page table the row's page, laid out as a layer's block at ptok rows (Seg.ps in bytes for fp8). Prompt-cache copies run a page at a time for either format (the piece's layout at cap = ptok inside each page), the pool's token bytes follow the format, and TF_QWEN_KV_PAGES no longer refuses --kv-cache fp8 (only --slide, whose learner reads whole windows). tf-cuda-test qwen-pages runs the six shapes and both controls for bf16 and fp8 (the fp8 reference dequantizes e4m3 x scale in float64): paged == contiguous in every shape, controls red. GB10, Qwen3.8-27B MLX-4bit, --kv-cache fp8: contiguous == the previous kernels (3/3 chat, 3/3 8k prefill), 14.0/13.7 tok/s against 14.05/13.6, c=4 48.9/47.5 against 49.2/47.7; page test 16 and pool 256 == contiguous (pool: 3 streams on 387 pages, 3.07 GiB); with DFlash2 at --parallel 8, contiguous and pages-only both drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh, and equal to each other. --- zig/kernels/cuda/qwen35.cuh | 2 +- zig/kernels/cuda/qwen35_kv8.cu | 62 +++--- zig/src/families/qwen3_5/cuda_cache.zig | 10 +- zig/src/families/qwen3_5/cuda_native.zig | 4 +- zig/src/families/qwen3_5/cuda_pages.zig | 2 +- zig/src/families/qwen3_5/cuda_reuse.zig | 17 +- zig/tests/cuda/contract_qwen_pages.zig | 228 ++++++++++++++--------- 7 files changed, 194 insertions(+), 131 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh index 2fe83705f5..0615636caa 100644 --- a/zig/kernels/cuda/qwen35.cuh +++ b/zig/kernels/cuda/qwen35.cuh @@ -35,7 +35,7 @@ struct Seg { int cap, len, row0, rows, tap_rows, taps; int delta; // an image prompt's rotary shift: rows past it rotate at position + delta unsigned long long pt; // the stream's KV page table (int a page; 0: kv holds cap contiguous rows) - unsigned long long ps; // elements a KV page holds: every attention layer's keys and values of ptok positions + unsigned long long ps; // elements a KV page holds (bf16, or bytes for fp8): every attention layer's ptok rows int ptok; // positions a KV page holds }; diff --git a/zig/kernels/cuda/qwen35_kv8.cu b/zig/kernels/cuda/qwen35_kv8.cu index ed7ccdfab1..32e2cffa56 100644 --- a/zig/kernels/cuda/qwen35_kv8.cu +++ b/zig/kernels/cuda/qwen35_kv8.cu @@ -6,9 +6,27 @@ using namespace tf_q35; -// A layer's block in Seg.kv: kvh * cap rows of HD bytes for keys then values, then kvh * cap fp32 scales for each. -__device__ __forceinline__ unsigned char* layer_at(const Seg& sg, int layer, int kvh) { - return reinterpret_cast(sg.kv) + (size_t)layer * 2 * kvh * sg.cap * (HD + 4); +// A key's and value's codes and scales. +struct Row8 { + unsigned char *k, *v; + float *ks, *vs; +}; + +// Layer `layer`'s rows at pos, head h: a layer's block is kvh * n rows of HD bytes for keys then values, then kvh * n +// fp32 scales for each; n is cap, or with a page table ptok inside pos's page (Seg.ps bytes a page). +__device__ __forceinline__ Row8 row8(const Seg& sg, int layer, int kvh, int h, int pos) { + unsigned char* base = reinterpret_cast(sg.kv); + int n = sg.cap, p = pos; + if (sg.pt != 0) { + base += (size_t)reinterpret_cast(sg.pt)[pos / sg.ptok] * sg.ps; + n = sg.ptok; + p = pos % sg.ptok; + } + unsigned char* keys = base + (size_t)layer * 2 * kvh * n * (HD + 4); + unsigned char* values = keys + (size_t)kvh * n * HD; + float* kscale = reinterpret_cast(values + (size_t)kvh * n * HD); + const size_t row = (size_t)h * n + p; + return {keys + row * HD, values + row * HD, kscale + row, kscale + (size_t)kvh * n + row}; } // The block's 256 threads' row as e4m3 into q[d], its scale (1 for a row of zeros) to *s; tf_kv_quant's bits. @@ -32,13 +50,9 @@ struct StoreFp8 { int layer, kvh; float* red; __device__ __forceinline__ void operator()(const Seg& sg, int h, int slot, int d, __nv_bfloat16 k, __nv_bfloat16 v) const { - unsigned char* keys = layer_at(sg, layer, kvh); - unsigned char* values = keys + (size_t)kvh * sg.cap * HD; - float* kscale = reinterpret_cast(values + (size_t)kvh * sg.cap * HD); - float* vscale = kscale + (size_t)kvh * sg.cap; - const size_t row = (size_t)h * sg.cap + slot; - quant_row(bf(k), d, keys + row * HD, kscale + row, red); - quant_row(bf(v), d, values + row * HD, vscale + row, red); + const Row8 r = row8(sg, layer, kvh, h, slot); + quant_row(bf(k), d, r.k, r.ks, red); + quant_row(bf(v), d, r.v, r.vs, red); } }; @@ -86,10 +100,6 @@ extern "C" __global__ void tf_q35_attn_partial8(Geo geo, const __nv_bfloat16* __ } return; } - const unsigned char* K = layer_at(sg, layer, KVH); - const unsigned char* V = K + (size_t)KVH * sg.cap * HD; - const float* KS = reinterpret_cast(V + (size_t)KVH * sg.cap * HD); - const float* VS = KS + (size_t)KVH * sg.cap; float q[DPL], o[DPL]; #pragma unroll for (int i = 0; i < DPL; i++) { @@ -99,18 +109,18 @@ extern "C" __global__ void tf_q35_attn_partial8(Geo geo, const __nv_bfloat16* __ float mx = -INFINITY, l = 0.0f; for (int t = k0; t < k1; t++) { const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); - const size_t row = (size_t)h * sg.cap + pos; + const Row8 r = row8(sg, layer, KVH, h, pos); float kf[DPL], vf[DPL]; - widen8(*reinterpret_cast(K + row * HD + lane * DPL), kf); + widen8(*reinterpret_cast(r.k + lane * DPL), kf); float dot = 0.0f; #pragma unroll for (int i = 0; i < DPL; i++) dot = fmaf(q[i], kf[i], dot); - const float s = warp_sum(dot) * KS[row] * scale; + const float s = warp_sum(dot) * *r.ks * scale; const float mn = fmaxf(mx, s); const float a = __expf(mx - mn), b = __expf(s - mn); l = l * a + b; - widen8(*reinterpret_cast(V + row * HD + lane * DPL), vf); - const float bv = b * VS[row]; + widen8(*reinterpret_cast(r.v + lane * DPL), vf); + const float bv = b * *r.vs; #pragma unroll for (int i = 0; i < DPL; i++) o[i] = fmaf(bv, vf[i], o[i] * a); mx = mn; @@ -127,20 +137,16 @@ extern "C" __global__ void tf_q35_attn_partial8(Geo geo, const __nv_bfloat16* __ extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact8(Geo geo, Seg sg, Path path) { const int KVH = geo.kvh; const int layer = blockIdx.x, h = blockIdx.y; - unsigned char* keys = layer_at(sg, layer, KVH); - unsigned char* values = keys + (size_t)KVH * sg.cap * HD; - float* kscale = reinterpret_cast(values + (size_t)KVH * sg.cap * HD); - float* vscale = kscale + (size_t)KVH * sg.cap; - const size_t base = (size_t)h * sg.cap; const int d = threadIdx.x; for (int j = 0; j < path.n; j++) { const int from = sg.len + path.rows[j], to = sg.len + j; if (from != to) { - keys[(base + to) * HD + d] = keys[(base + from) * HD + d]; - values[(base + to) * HD + d] = values[(base + from) * HD + d]; + const Row8 a = row8(sg, layer, KVH, h, from), b = row8(sg, layer, KVH, h, to); + b.k[d] = a.k[d]; + b.v[d] = a.v[d]; if (d == 0) { - kscale[base + to] = kscale[base + from]; - vscale[base + to] = vscale[base + from]; + *b.ks = *a.ks; + *b.vs = *a.vs; } } __syncthreads(); diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig index 359db822f9..1332c59cf3 100644 --- a/zig/src/families/qwen3_5/cuda_cache.zig +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -26,9 +26,9 @@ pub const Paged = struct { /// The cache owns `ids` (TF_QWEN_KV_PAGE_TEST); a pool's lane maps its lease's. own_ids: bool = true, - /// Elements (bf16) a page holds: every attention layer's keys and values of `tokens` positions. - pub fn elems(p: Paged, g: geometry.Geometry) usize { - return g.attentions() * 2 * g.kv_heads * p.tokens * geometry.head_dim; + /// Elements a page holds (bf16, or bytes for fp8 rows and their scales): every attention layer's of `tokens` rows. + pub fn elems(p: Paged, g: geometry.Geometry, kv8: bool) usize { + return if (kv8) g.kvBytes(p.tokens, true) else g.kvBytes(p.tokens, false) / 2; } }; @@ -63,7 +63,6 @@ pub const Cache = struct { /// An image prompt's rotary shift (get_rope_index's rope_delta): rows after its prompt rotate at position + delta. pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { const test_pages = try pageTest(); - if (test_pages > 0 and e.kv8) return error.PagedKvNeedsBf16; const rows = if (test_pages > 0) (cap + test_pages - 1) / test_pages * test_pages else cap; var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(rows, e.kv8)); errdefer kv.free(); @@ -85,7 +84,6 @@ pub const Cache = struct { /// A pool lane's cache: its own state, log, tap ring and table of `pages` pages at `base`, mapped at open. pub fn initPaged(gpa: std.mem.Allocator, e: *const Engine, pages: usize, tokens: usize, base: u64) !Cache { - if (e.kv8) return error.PagedKvNeedsBf16; var table = try DeviceBuffer.alloc(e.d, pages * 4); errdefer table.free(); try table.fill8(0, null); @@ -146,7 +144,7 @@ pub const Cache = struct { if (self.paged) |p| { sg.kv = p.base; sg.pt = p.table.ptr; - sg.ps = p.elems(g); + sg.ps = p.elems(g, e.kv8); sg.ptok = @intCast(p.tokens); } return sg; diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 6eb97a7c88..4b1858252b 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -85,7 +85,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ if (kernels != null) return error.QwenHasNoCapturedKernels; if (o.slide and o.kv8) return error.UnsupportedLearningFormat; const page_tokens = try pages_mod.paging(); - if (page_tokens > 0 and (o.kv8 or o.slide)) return error.QwenPagesNeedBf16; + if (page_tokens > 0 and o.slide) return error.QwenPagesNoLearner; const e = try E.Engine.init(gpa, io, ctx, dir, o.kv8); errdefer e.deinit(); if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary @@ -192,7 +192,7 @@ fn closeDrafter(p: *anyopaque) void { pub fn explain(ctx: ?*anyopaque, err: anyerror) ?[]const u8 { if (ctx == null) return switch (err) { error.NoQwenVision => "--vision: this checkpoint has no image tower (no vision_config): serve it without --vision", - error.QwenPagesNeedBf16 => "TF_QWEN_KV_PAGES: paged caches hold bf16 rows and serve no learner yet; drop --kv-cache fp8 and --slide, or unset it", + error.QwenPagesNoLearner => "TF_QWEN_KV_PAGES: the learner reads whole windows, not pages; drop --slide or unset it", error.BadQwenKvPages => "TF_QWEN_KV_PAGES takes the positions a page holds (a whole number; 0 or unset: a window each)", error.UnsupportedQwenVision => "--vision: this checkpoint's image tower or text M-RoPE is not one the CUDA engine reads", else => null, diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig index a325c4655d..2b26e22992 100644 --- a/zig/src/families/qwen3_5/cuda_pages.zig +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -30,7 +30,7 @@ fn mostPages(capacity: usize, tokens: usize) usize { /// A stream's fixed part (state, log, tap ring, table), its keys and values a token, a kept state's recurrences. pub fn layout(e: *const Engine, capacity: usize, tokens: usize) kv.Layout { - return .{ .fixed_bytes = e.streamBytes(0) + mostPages(capacity, tokens) * 4, .token_bytes = e.g.kvBytes(1, false), .page_tokens = @intCast(tokens), .snap_bytes = e.g.stateBytes() }; + return .{ .fixed_bytes = e.streamBytes(0) + mostPages(capacity, tokens) * 4, .token_bytes = e.g.kvBytes(1, e.kv8), .page_tokens = @intCast(tokens), .snap_bytes = e.g.stateBytes() }; } /// One pool lane: a cache whose table maps its stream's pages, and that stream's lease. diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig index 63210d3428..c396c91597 100644 --- a/zig/src/families/qwen3_5/cuda_reuse.zig +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -97,18 +97,20 @@ const Runs = struct { pages: ?caches.Paged, page_bytes: usize, in_page: usize, + span: usize = 0, // a full page's run of the piece len: usize, page: usize = 0, done: usize = 0, const Run = struct { dev: u64, at: usize, len: usize }; - fn of(e: *const Engine, cache: *const Cache, l: Layout, i: usize, p: Piece) Runs { + fn of(cache: *const Cache, l: Layout, i: usize, p: Piece) Runs { if (!p.kv or cache.paged == null) return .{ .base = (if (p.kv) cache.kv.ptr else cache.state.ptr) + p.off, .pages = null, .page_bytes = 0, .in_page = 0, .len = p.len }; const pg = cache.paged.?; - const row = head_dim * 2; // bf16 only: a paged cache refuses fp8 at init - std.debug.assert(l.format == .bf16); - return .{ .base = pg.base, .pages = pg, .page_bytes = pg.elems(e.g) * 2, .in_page = (i - 1) * pg.tokens * row, .len = p.len }; + // the piece's offset inside a page: its layout at cap = tokens (Layout.piece), a run of tokens rows a page + const in_page: Layout = .{ .g = l.g, .format = l.format, .cap = pg.tokens, .at = pg.tokens }; + const q = in_page.piece(i); + return .{ .base = pg.base, .pages = pg, .page_bytes = in_page.kvBytes(), .in_page = q.off, .span = q.len, .len = p.len }; } fn next(r: *Runs) ?Run { @@ -117,8 +119,7 @@ const Runs = struct { r.done = r.len; return .{ .dev = r.base, .at = 0, .len = r.len }; }; - const span = pg.tokens * head_dim * 2; - const n = @min(span, r.len - r.done); + const n = @min(r.span, r.len - r.done); const run: Run = .{ .dev = r.base + pg.ids[r.page] * r.page_bytes + r.in_page, .at = r.done, .len = n }; r.page += 1; r.done += n; @@ -171,7 +172,7 @@ pub fn saveCache(e: *Engine, cache: *const Cache, at: usize, seen: []const Seen) var off: usize = 0; for (0..l.count()) |i| { const p = l.piece(i); - var it = Runs.of(e, cache, l, i, p); + var it = Runs.of(cache, l, i, p); while (it.next()) |r| try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(copy.host.bytes.ptr + off + r.at, r.dev, r.len, e.stream.handle), "cuMemcpyDtoHAsync"); off += p.len; } @@ -210,7 +211,7 @@ pub fn restoreCache(e: *Engine, cache: *Cache, saved: *anyopaque, seen: []const var off: usize = 0; for (0..l.count()) |i| { const p = l.piece(i); - var it = Runs.of(e, cache, l, i, p); + var it = Runs.of(cache, l, i, p); while (it.next()) |r| try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(r.dev, copy.host.bytes.ptr + off + r.at, r.len, e.stream.handle), "cuMemcpyHtoDAsync"); off += p.len; } diff --git a/zig/tests/cuda/contract_qwen_pages.zig b/zig/tests/cuda/contract_qwen_pages.zig index 82c7418dc8..ddae593d81 100644 --- a/zig/tests/cuda/contract_qwen_pages.zig +++ b/zig/tests/cuda/contract_qwen_pages.zig @@ -1,6 +1,7 @@ -//! qwen35.cu's attention cache through a page table: rope stores, attention reads, keep moves, paged == contiguous. +//! qwen35.cu's and qwen35_kv8.cu's attention caches through a page table: rope, attention, keep, paged == contiguous. const std = @import("std"); const cuda = @import("cuda"); +const core = @import("core"); const check = @import("check.zig"); const rig_ = @import("contract_rig.zig"); @@ -9,6 +10,7 @@ const max_chunks = 64; const attentions = 2; // the fixture's attention layers: the kernels run layer 1, so a page's layer offset counts const q_heads = 24; // the 27B's heads const kv_heads = 4; +const e4m3 = core.quant.fp8.e4m3; // qwen35.cuh's Geo, Seg, Ancestor and Path (src/families/qwen3_5/cuda_cache.zig mirrors them for the engine) const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int, nl: c_int }; @@ -16,6 +18,9 @@ const Seg = extern struct { rec: u64 = 0, conv: u64 = 0, kv: u64, tap: u64 = 0, const Ancestor = extern struct { bits: u64, depth: u32, chain: u32 }; const Path = extern struct { n: c_int, rows: [64]c_int }; +/// bf16 rows (qwen35.cu), or e4m3 rows with an fp32 scale each (qwen35_kv8.cu, --kv-cache fp8). +const Format = enum { bf16, fp8 }; + /// Two streams' cached positions, the window rows each adds, positions a page. const Case = struct { lens: [2]usize, rows: usize, ptok: usize }; const cases = [_]Case{ @@ -27,42 +32,64 @@ const cases = [_]Case{ .{ .lens = .{ 9000, 300 }, .rows = 3, .ptok = 512 }, }; +const Span = struct { off: usize, len: usize }; + +/// Bytes of a layer's (or a whole block's, `layers`) rows, `n` positions a head. +fn blockBytes(f: Format, n: usize, layers: usize) usize { + return layers * 2 * kv_heads * n * @as(usize, if (f == .fp8) hd + 4 else hd * 2); +} + +/// Layer a's key (s 0) or value (s 1) at pos, head h, in a block of n positions a head: its row, and fp8's scale. +fn spans(f: Format, n: usize, a: usize, s: usize, h: usize, pos: usize) [2]Span { + const blk = s * kv_heads + h; + const layer = blockBytes(f, n, a); + return switch (f) { + .bf16 => .{ .{ .off = layer + (blk * n + pos) * hd * 2, .len = hd * 2 }, .{ .off = 0, .len = 0 } }, + .fp8 => .{ .{ .off = layer + (blk * n + pos) * hd, .len = hd }, .{ .off = layer + 2 * kv_heads * n * hd + (blk * n + pos) * 4, .len = 4 } }, + }; +} + /// One case's host data: the logical caches (cap rows a head), the window's projections, the page tables. const Fixture = struct { c: Case, + format: Format, cap: usize, - kv: [2][]u16, // [attention][key, value][head][cap][hd] a stream + kv: [2][]u8, q: []u16, // [rows][q_heads][hd] k: []u16, // [rows][kv_heads][hd] v: []u16, ids: [2][]u32, // each stream's pages in the shared arena, shuffled across both pages: usize, // the arena's pages (two spare) - fn blockElems(f: Fixture) usize { - return attentions * 2 * kv_heads * f.cap * hd; + fn streamBytes(f: Fixture) usize { + return blockBytes(f.format, f.cap, attentions); } - fn pageElems(f: Fixture) usize { - return attentions * 2 * kv_heads * f.c.ptok * hd; + fn pageBytes(f: Fixture) usize { + return blockBytes(f.format, f.c.ptok, attentions); } - /// The element where logical (attention a, key/value s, head h, position pos) of stream `st` sits in the arena. - fn paged(f: Fixture, st: usize, a: usize, s: usize, h: usize, pos: usize) usize { - const blk = (a * 2 + s) * kv_heads + h; - return f.ids[st][pos / f.c.ptok] * f.pageElems() + (blk * f.c.ptok + pos % f.c.ptok) * hd; + fn logical(f: Fixture, a: usize, s: usize, h: usize, pos: usize) [2]Span { + return spans(f.format, f.cap, a, s, h, pos); } - fn logical(f: Fixture, a: usize, s: usize, h: usize, pos: usize) usize { - return (((a * 2 + s) * kv_heads + h) * f.cap + pos) * hd; + /// Where logical (a, s, h, pos) of stream `st` sits in the arena under `ids`. + fn paged(f: Fixture, ids: [2][]const u32, st: usize, a: usize, s: usize, h: usize, pos: usize) [2]Span { + var r = spans(f.format, f.c.ptok, a, s, h, pos % f.c.ptok); + for (&r) |*x| x.off += ids[st][pos / f.c.ptok] * f.pageBytes(); + return r; } - fn make(a: std.mem.Allocator, c: Case, seed: u64) !Fixture { + fn make(a: std.mem.Allocator, c: Case, format: Format, seed: u64) !Fixture { var prng = std.Random.DefaultPrng.init(seed); const rng = prng.random(); const most = @max(c.lens[0], c.lens[1]) + c.rows; const cap = (most + c.ptok - 1) / c.ptok * c.ptok; - var f: Fixture = .{ .c = c, .cap = cap, .kv = undefined, .q = undefined, .k = undefined, .v = undefined, .ids = undefined, .pages = 2 * (cap / c.ptok) + 2 }; - for (&f.kv) |*x| x.* = try randoms(a, rng, f.blockElems()); + var f: Fixture = .{ .c = c, .format = format, .cap = cap, .kv = undefined, .q = undefined, .k = undefined, .v = undefined, .ids = undefined, .pages = 2 * (cap / c.ptok) + 2 }; + for (&f.kv) |*x| { + x.* = try a.alloc(u8, f.streamBytes()); + for (0..attentions) |at| for (0..2) |s| for (0..kv_heads) |h| for (0..cap) |pos| fill(f, rng, x.*, f.logical(at, s, h, pos)); + } f.q = try randoms(a, rng, c.rows * 2 * q_heads * hd); f.k = try randoms(a, rng, c.rows * 2 * kv_heads * hd); f.v = try randoms(a, rng, c.rows * 2 * kv_heads * hd); @@ -74,23 +101,54 @@ const Fixture = struct { return f; } - /// The arena holding both streams' logical caches at their pages; spare pages random. - fn arena(f: Fixture, a: std.mem.Allocator, rng: std.Random) ![]u16 { - const out = try randoms(a, rng, f.pages * f.pageElems()); + /// A row's bytes: bf16 in [-1, 1); e4m3 codes but NaN's, a scale in [1/1792, 1/448) (values within 1). + fn fill(f: Fixture, rng: std.Random, out: []u8, sp: [2]Span) void { + switch (f.format) { + .bf16 => for (std.mem.bytesAsSlice(u16, out[sp[0].off..][0..sp[0].len])) |*x| { + x.* = @truncate(@as(u32, @bitCast(rng.float(f32) * 2 - 1)) >> 16); + }, + .fp8 => { + for (out[sp[0].off..][0..sp[0].len]) |*b| { + b.* = rng.int(u8); + if (b.* & 0x7f == 0x7f) b.* ^= 1; + } + const sc: f32 = (1 + 3 * rng.float(f32)) / 1792; + @memcpy(out[sp[1].off..][0..4], std.mem.asBytes(&sc)); + }, + } + } + + /// The arena holding both streams' logical caches at their pages under `ids`; spare pages random. + fn arena(f: Fixture, a: std.mem.Allocator, rng: std.Random, ids: [2][]const u32) ![]u8 { + const out = try a.alloc(u8, f.pages * f.pageBytes()); + rng.bytes(out); for (0..2) |st| for (0..attentions) |at| for (0..2) |s| for (0..kv_heads) |h| for (0..f.cap) |pos| { - @memcpy(out[f.paged(st, at, s, h, pos)..][0..hd], f.kv[st][f.logical(at, s, h, pos)..][0..hd]); + const from = f.logical(at, s, h, pos); + const to = f.paged(ids, st, at, s, h, pos); + for (from, to) |x, y| @memcpy(out[y.off..][0..y.len], f.kv[st][x.off..][0..x.len]); }; return out; } /// The arena's pages read back as stream `st`'s logical cache. - fn unpage(f: Fixture, a: std.mem.Allocator, arena_: []const u16, st: usize) ![]u16 { - const out = try a.alloc(u16, f.blockElems()); + fn unpage(f: Fixture, a: std.mem.Allocator, arena_: []const u8, ids: [2][]const u32, st: usize) ![]u8 { + const out = try a.alloc(u8, f.streamBytes()); for (0..attentions) |at| for (0..2) |s| for (0..kv_heads) |h| for (0..f.cap) |pos| { - @memcpy(out[f.logical(at, s, h, pos)..][0..hd], arena_[f.paged(st, at, s, h, pos)..][0..hd]); + const to = f.logical(at, s, h, pos); + const from = f.paged(ids, st, at, s, h, pos); + for (from, to) |x, y| @memcpy(out[y.off..][0..y.len], arena_[x.off..][0..x.len]); }; return out; } + + /// Logical value d of (a, s, h, pos) in stream cache `cache`. + fn value(f: Fixture, cache: []const u8, a: usize, s: usize, h: usize, pos: usize, d: usize) f64 { + const sp = f.logical(a, s, h, pos); + return switch (f.format) { + .bf16 => bf(std.mem.readInt(u16, cache[sp[0].off + 2 * d ..][0..2], .little)), + .fp8 => e4m3(cache[sp[0].off + d]) * @as(f64, @as(f32, @bitCast(std.mem.readInt(u32, cache[sp[1].off..][0..4], .little)))), + }; + } }; /// bf16 in [-1, 1): finite, so a sum's bytes are its order's alone. @@ -105,7 +163,7 @@ fn bf(x: u16) f64 { } /// One layout's results: the rotated queries, the attention rows, each stream's cache after rope and after keep. -const Result = struct { q: []u8, out: []u8, after_rope: [2][]u16, after_keep: [2][]u16 }; +const Result = struct { q: []u8, out: []u8, after_rope: [2][]u8, after_keep: [2][]u8 }; const Fns = struct { rope: cuda.Function, partial: cuda.Function, merge: cuda.Function, compact: cuda.Function }; @@ -118,23 +176,20 @@ fn launch(gpu: check.Gpu, s: cuda.Stream, fns: Fns, f: Fixture, tables: ?[2][]co const geo: Geo = .{ .hid = 0, .nk = 0, .nv = 0, .qh = q_heads, .kvh = kv_heads, .nl = 0 }; var segs: [2]Seg = undefined; var kv: [2]u64 = undefined; // a stream's buffer, or the arena twice - var arena_host: []u16 = &.{}; if (tables) |t| { - var g = f; - g.ids = .{ @constCast(t[0]), @constCast(t[1]) }; - arena_host = try g.arena(a, prng.random()); - const ptr = try rig.scratch(arena_host.len * 2); - try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = ptr, .len = arena_host.len * 2 }).upload(0, std.mem.sliceAsBytes(arena_host)); + const host = try f.arena(a, prng.random(), t); + const ptr = try rig.scratch(host.len); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = ptr, .len = host.len }).upload(0, host); kv = .{ ptr, ptr }; } else for (0..2) |st| { - kv[st] = try rig.scratch(f.blockElems() * 2); - try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = kv[st], .len = f.blockElems() * 2 }).upload(0, std.mem.sliceAsBytes(f.kv[st])); + kv[st] = try rig.scratch(f.streamBytes()); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = kv[st], .len = f.streamBytes() }).upload(0, f.kv[st]); } for (0..2) |st| { segs[st] = .{ .kv = kv[st], .cap = @intCast(f.cap), .len = @intCast(f.c.lens[st]), .row0 = @intCast(st * f.c.rows), .rows = @intCast(f.c.rows) }; if (given orelse tables) |t| { segs[st].pt = try rig.input(std.mem.sliceAsBytes(t[st])); - segs[st].ps = f.pageElems(); + segs[st].ps = if (f.format == .fp8) f.pageBytes() else f.pageBytes() / 2; segs[st].ptok = @intCast(f.c.ptok); } } @@ -186,17 +241,11 @@ fn launch(gpu: check.Gpu, s: cuda.Stream, fns: Fns, f: Fixture, tables: ?[2][]co return r; } -fn readCaches(rig: *rig_.Rig, a: std.mem.Allocator, f: Fixture, tables: ?[2][]const u32, kv: [2]u64, out: *[2][]u16) !void { +fn readCaches(rig: *rig_.Rig, a: std.mem.Allocator, f: Fixture, tables: ?[2][]const u32, kv: [2]u64, out: *[2][]u8) !void { if (tables) |t| { - var g = f; - g.ids = .{ @constCast(t[0]), @constCast(t[1]) }; - const all = try a.alloc(u16, f.pages * f.pageElems()); - @memcpy(std.mem.sliceAsBytes(all), try rig.bytes(kv[0], all.len * 2)); - for (0..2) |st| out[st] = try g.unpage(a, all, st); - } else for (0..2) |st| { - out[st] = try a.alloc(u16, f.blockElems()); - @memcpy(std.mem.sliceAsBytes(out[st]), try rig.bytes(kv[st], f.blockElems() * 2)); - } + const all = try a.dupe(u8, try rig.bytes(kv[0], f.pages * f.pageBytes())); + for (0..2) |st| out[st] = try f.unpage(a, all, t, st); + } else for (0..2) |st| out[st] = try a.dupe(u8, try rig.bytes(kv[st], f.streamBytes())); } /// Each attention row against float64 softmax(q k / 16) v: a bound (fp32 and fast exp), bf16 rounding of its size. @@ -206,6 +255,8 @@ fn agrees(a: std.mem.Allocator, f: Fixture, r: Result) !bool { const keys_most = @max(f.c.lens[0], f.c.lens[1]) + f.c.rows; const p = try a.alloc(f64, keys_most); defer a.free(p); + const vals = try a.alloc(f64, keys_most * hd); + defer a.free(vals); for (0..2 * f.c.rows) |row| { const st = row / f.c.rows; const keys = f.c.lens[st] + row % f.c.rows + 1; // the chain's rows are positions after the cache @@ -216,10 +267,10 @@ fn agrees(a: std.mem.Allocator, f: Fixture, r: Result) !bool { var top: f64 = -std.math.inf(f64); for (0..keys) |t| { var dot: f64 = 0; - const kr = cache[f.logical(1, 0, h, t)..][0..hd]; - for (0..hd) |d| dot += bf(qv[d]) * bf(kr[d]); + for (0..hd) |d| dot += bf(qv[d]) * f.value(cache, 1, 0, h, t, d); p[t] = dot / 16.0; top = @max(top, p[t]); + for (0..hd) |d| vals[t * hd + d] = f.value(cache, 1, 1, h, t, d); } var sum: f64 = 0; for (p[0..keys]) |*x| { @@ -230,9 +281,8 @@ fn agrees(a: std.mem.Allocator, f: Fixture, r: Result) !bool { var o: f64 = 0; var mag: f64 = 0; for (0..keys) |t| { - const v = bf(cache[f.logical(1, 1, h, t) + d]); - o += p[t] * v; - mag += p[t] * @abs(v); + o += p[t] * vals[t * hd + d]; + mag += p[t] * @abs(vals[t * hd + d]); } o /= sum; mag /= sum; @@ -244,8 +294,8 @@ fn agrees(a: std.mem.Allocator, f: Fixture, r: Result) !bool { return true; } -fn eqCaches(x: [2][]u16, y: [2][]u16) bool { - return std.mem.eql(u16, x[0], y[0]) and std.mem.eql(u16, x[1], y[1]); +fn eqCaches(x: [2][]u8, y: [2][]u8) bool { + return std.mem.eql(u8, x[0], y[0]) and std.mem.eql(u8, x[1], y[1]); } pub fn run(gpu: check.Gpu) !void { @@ -255,42 +305,50 @@ pub fn run(gpu: check.Gpu) !void { defer mod.unload(); var keep = try cuda.Module.load(gpu.d, cuda.kernels.qwen35_keep); defer keep.unload(); - const fns: Fns = .{ .rope = try mod.function("tf_q35_rope"), .partial = try mod.function("tf_q35_attn_partial"), .merge = try mod.function("tf_q35_attn_merge"), .compact = try keep.function("tf_q35k_compact") }; - var red_table = false; - var red_bound = false; - for (cases, 0..) |c, i| { - var arena = std.heap.ArenaAllocator.init(gpu.gpa); - defer arena.deinit(); - const a = arena.allocator(); - const f = try Fixture.make(a, c, 0x9a6e + i); - const flat = try launch(gpu, s, fns, f, null, null, a, 1); - try check.expect(try agrees(a, f, flat), "qwen35 pages: contiguous attention off float64 at lens {any}, rows {d}", .{ c.lens, c.rows }); - for (0..3) |rep| { - const pg = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, null, a, 2 + rep); - try check.sameBytes("qwen35 pages: rotated queries", pg.q, flat.q); - try check.sameBytes("qwen35 pages: attention rows, paged against contiguous", pg.out, flat.out); - try check.expect(eqCaches(pg.after_rope, flat.after_rope), "qwen35 pages: rope's stores differ at lens {any}, page {d}", .{ c.lens, c.ptok }); - try check.expect(eqCaches(pg.after_keep, flat.after_keep), "qwen35 pages: keep's moves differ at lens {any}, page {d}", .{ c.lens, c.ptok }); - } - // red controls: a table with two of stream 1's pages swapped reads other keys; a planted row breaks the bound - const n1 = (c.lens[1] + c.rows + c.ptok - 1) / c.ptok; - if (n1 >= 2) { - const swapped = try a.dupe(u32, f.ids[1]); - std.mem.swap(u32, &swapped[0], &swapped[n1 - 1]); - const bad = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, .{ f.ids[0], swapped }, a, 9); - const moved = !std.mem.eql(u8, bad.out, flat.out); - red_table = red_table or moved; - try check.expect(moved, "qwen35 pages: a swapped page table still gave contiguous bytes (control missed) at lens {any}", .{c.lens}); + var kv8 = try cuda.Module.load(gpu.d, cuda.kernels.qwen35_kv8); + defer kv8.unload(); + const merge = try mod.function("tf_q35_attn_merge"); + for (std.enums.values(Format)) |format| { + const fns: Fns = switch (format) { + .bf16 => .{ .rope = try mod.function("tf_q35_rope"), .partial = try mod.function("tf_q35_attn_partial"), .merge = merge, .compact = try keep.function("tf_q35k_compact") }, + .fp8 => .{ .rope = try kv8.function("tf_q35_rope8"), .partial = try kv8.function("tf_q35_attn_partial8"), .merge = merge, .compact = try kv8.function("tf_q35k_compact8") }, + }; + var red_table = false; + var red_bound = false; + for (cases, 0..) |c, i| { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const f = try Fixture.make(a, c, format, 0x9a6e + i); + const flat = try launch(gpu, s, fns, f, null, null, a, 1); + try check.expect(try agrees(a, f, flat), "qwen35 pages {t}: contiguous attention off float64 at lens {any}, rows {d}", .{ format, c.lens, c.rows }); + for (0..3) |rep| { + const pg = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, null, a, 2 + rep); + try check.sameBytes("qwen35 pages: rotated queries", pg.q, flat.q); + try check.sameBytes("qwen35 pages: attention rows, paged against contiguous", pg.out, flat.out); + try check.expect(eqCaches(pg.after_rope, flat.after_rope), "qwen35 pages {t}: rope's stores differ at lens {any}, page {d}", .{ format, c.lens, c.ptok }); + try check.expect(eqCaches(pg.after_keep, flat.after_keep), "qwen35 pages {t}: keep's moves differ at lens {any}, page {d}", .{ format, c.lens, c.ptok }); + } + // red controls: two of stream 1's pages swapped in the table only (the pages stay); a planted row + const n1 = (c.lens[1] + c.rows + c.ptok - 1) / c.ptok; + if (n1 >= 2) { + const swapped = try a.dupe(u32, f.ids[1]); + std.mem.swap(u32, &swapped[0], &swapped[n1 - 1]); + const bad = try launch(gpu, s, fns, f, .{ f.ids[0], f.ids[1] }, .{ f.ids[0], swapped }, a, 9); + const moved = !std.mem.eql(u8, bad.out, flat.out); + red_table = red_table or moved; + try check.expect(moved, "qwen35 pages {t}: a swapped page table still gave contiguous bytes (control missed) at lens {any}", .{ format, c.lens }); + } + const o = try a.dupe(u8, flat.out); + std.mem.writeInt(u16, o[2 * (q_heads * hd + 5) ..][0..2], 0x4040, .little); // 3.0 where |o| <= 1 + var bad_r = flat; + bad_r.out = o; + const caught = !(try agrees(a, f, bad_r)); + red_bound = red_bound or caught; + try check.expect(caught, "qwen35 pages {t}: a planted attention row passed float64 (control missed)", .{format}); + std.debug.print("PASS qwen35 pages {t}: lens {any}, {d} rows a stream, {d}-position pages: paged == contiguous (queries, rows, stores, keep), three repeats, float64 bound, guards\n", .{ format, c.lens, c.rows, c.ptok }); } - const o = try a.dupe(u8, flat.out); - std.mem.writeInt(u16, o[2 * (q_heads * hd + 5) ..][0..2], 0x4040, .little); // 3.0 where |o| <= 1 - var bad_r = flat; - bad_r.out = o; - const caught = !(try agrees(a, f, bad_r)); - red_bound = red_bound or caught; - try check.expect(caught, "qwen35 pages: a planted attention row passed float64 (control missed)", .{}); - std.debug.print("PASS qwen35 pages: lens {any}, {d} rows a stream, {d}-position pages: paged == contiguous (queries, rows, stores, keep), three repeats, float64 bound, guards\n", .{ c.lens, c.rows, c.ptok }); + try check.expect(red_table and red_bound, "qwen35 pages {t}: a control never fired", .{format}); + std.debug.print("PASS qwen35 pages {t}: controls RED as planted (swapped table, planted row)\n", .{format}); } - try check.expect(red_table and red_bound, "qwen35 pages: a control never fired", .{}); - std.debug.print("PASS qwen35 pages: controls RED as planted (swapped table, planted row)\n", .{}); } From 9a99d1f853b04bb003c99c8c2e966b61fca80a93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:33:22 +0200 Subject: [PATCH 05/11] cuda qwen3_5: one-line comments on row8 and the pool's reservation --- zig/kernels/cuda/qwen35_kv8.cu | 3 +-- zig/src/families/qwen3_5/cuda_reserve.zig | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/zig/kernels/cuda/qwen35_kv8.cu b/zig/kernels/cuda/qwen35_kv8.cu index 32e2cffa56..5c43e5cb2c 100644 --- a/zig/kernels/cuda/qwen35_kv8.cu +++ b/zig/kernels/cuda/qwen35_kv8.cu @@ -12,8 +12,7 @@ struct Row8 { float *ks, *vs; }; -// Layer `layer`'s rows at pos, head h: a layer's block is kvh * n rows of HD bytes for keys then values, then kvh * n -// fp32 scales for each; n is cap, or with a page table ptok inside pos's page (Seg.ps bytes a page). +// Codes and scales at (layer, h, pos): a layer's block of n = cap rows a head, or of ptok rows in pos's page. __device__ __forceinline__ Row8 row8(const Seg& sg, int layer, int kvh, int h, int pos) { unsigned char* base = reinterpret_cast(sg.kv); int n = sg.cap, p = pos; diff --git a/zig/src/families/qwen3_5/cuda_reserve.zig b/zig/src/families/qwen3_5/cuda_reserve.zig index b92f7430bd..71221f8f45 100644 --- a/zig/src/families/qwen3_5/cuda_reserve.zig +++ b/zig/src/families/qwen3_5/cuda_reserve.zig @@ -2,8 +2,7 @@ const warm = @import("cuda_warm.zig"); const E = @import("cuda_engine.zig"); -/// `streams` windows reserved and warmed; with a page pool (TF_QWEN_KV_PAGES) the first stream's window only, the -/// shared shapes warmed on lanes opened for one row and closed again. +/// `streams` windows reserved and warmed; with a pool the first only, the shapes warmed on one-row lanes. pub fn target(comptime Cache: type, comptime Part: type, l: anytype, streams: u32) !void { const started = warm.profile.now(l.io); var warm_ns: i96 = 0; From 1ac2075387ce4c0270ccd2f832428ffe9ca74c06 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:54:26 +0200 Subject: [PATCH 06/11] cuda qwen3_5: a pool lane's recurrent state is a run of pool pages while its stream runs A pool lane held its GDN recurrences and convolution tails, the one-row rollback copy, the verify log and the drafter's tap ring from load on, whether a stream ran on it or not (558 MiB a lane on the 27B). Now they are one run of consecutive pages in the KV arena, taken when a stream opens the lane and given back when it closes: the lane keeps only its page table. Only the buffers' addresses move; no kernel changes. core kv_pool: Layout.run_bytes (0 by default, so Nemotron's plan is unchanged) adds runPages() to a request's need, to each stream's minimum and to the window bound; Pool.findRun / claimRun / releaseRun take k consecutive free pages from the top (the free stack hands out KV pages from the bottom) out of the stream's promise, held as leased so the store never reclaims them. A pool too fragmented for a run makes the request wait (LaneMemory), as one without pages does. GB10, Qwen3.8-27B MLX-4bit, TF_QWEN_KV_PAGES=256, --parallel 4 at 32768: "3 more taking 492 pages of 256 rows (7.69 GiB), 0 MiB fixed each" (was 387 pages and 558 MiB fixed each); solo bytes == contiguous 3/3 chat, 3/3 8k prefill, also pages only; 13.9/14.1 tok/s, c=4 48.8/48.3. With DFlash2 at --parallel 8, pages only, bf16 and fp8: drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh (into a lane), outputs equal the previous runs'; 8 requests on 7 lanes, the waiting one 1.7 s. --- zig/src/core/kv_pool.zig | 62 ++++++++++++++++++++++--- zig/src/core/kv_pool_test.zig | 47 +++++++++++++++++++ zig/src/families/qwen3_5/cuda_cache.zig | 40 +++++++++++++++- zig/src/families/qwen3_5/cuda_pages.zig | 29 ++++++++---- 4 files changed, 161 insertions(+), 17 deletions(-) diff --git a/zig/src/core/kv_pool.zig b/zig/src/core/kv_pool.zig index 8e9d55c923..8f64bcac05 100644 --- a/zig/src/core/kv_pool.zig +++ b/zig/src/core/kv_pool.zig @@ -12,6 +12,8 @@ pub const Layout = struct { page_tokens: u32, /// Bytes of one kept snapshot's slot, outside the pages. snap_bytes: u64, + /// Bytes a running stream takes from the pages beside its keys and values, as one run of whole pages (0: none). + run_bytes: u64 = 0, pub fn pageBytes(l: Layout) u64 { return l.token_bytes * l.page_tokens; @@ -22,9 +24,14 @@ pub const Layout = struct { return @intCast((tokens + l.page_tokens - 1) / l.page_tokens); } - /// Pages a request reserves at admission: its prompt, its reply and one last verify window. + /// Pages a request reserves at admission: its prompt, its reply and one last verify window, and its run. pub fn need(l: Layout, prompt: u64, max_tokens: u64, window: u64) usize { - return l.pages(prompt + max_tokens + window); + return l.pages(prompt + max_tokens + window) + l.runPages(); + } + + /// Consecutive pages a running stream's `run_bytes` take. + pub fn runPages(l: Layout) usize { + return @intCast((l.run_bytes + l.pageBytes() - 1) / l.pageBytes()); } }; @@ -63,13 +70,13 @@ pub fn plan(l: Layout, b: Budget) error{ NoStream, TooMany }!Plan { if (b.engine > b.room or b.lanes == 0) return error.NoStream; var left = b.room - b.engine; const extra: u64 = b.lanes - 1; - const each = l.fixed_bytes + l.pageBytes(); + const each = l.fixed_bytes + (1 + l.runPages()) * l.pageBytes(); const fit = @min(extra, left / each); if (fit < extra and b.fixed) return error.TooMany; left -= fit * l.fixed_bytes; - const c = split(@min(b.cache, left - fit * l.pageBytes()), l.snap_bytes, b.snaps); + const c = split(@min(b.cache, left - fit * (1 + l.runPages()) * l.pageBytes()), l.snap_bytes, b.snaps); left -= c.snaps * l.snap_bytes + c.copied; - const most: u64 = if (b.window > 0) fit * l.pages(b.window) else std.math.maxInt(u64); + const most: u64 = if (b.window > 0) fit * (l.pages(b.window) + l.runPages()) else std.math.maxInt(u64); return .{ .lanes = @intCast(1 + fit), .snaps = c.snaps, .copied = c.copied, .pages = @intCast(@min(left / l.pageBytes(), most)) }; } @@ -77,7 +84,7 @@ pub fn plan(l: Layout, b: Budget) error{ NoStream, TooMany }!Plan { pub fn holds(l: Layout, b: Budget) bool { if (b.engine > b.room or b.lanes == 0) return false; const left = b.room - b.engine; - const each = l.fixed_bytes + l.pageBytes(); + const each = l.fixed_bytes + (1 + l.runPages()) * l.pageBytes(); return left - @min(@as(u64, b.lanes - 1), left / each) * each >= b.cache; } @@ -263,6 +270,49 @@ pub const Pool = struct { } } + /// The first id of `k` consecutive free pages, highest first (the free stack hands out low ids), or null. + pub fn findRun(p: *const Pool, k: usize) ?u32 { + if (k == 0 or k > p.refs.len) return null; + var end = p.refs.len; + while (end >= k) { + const start = end - k; + const busy = for (p.refs[start..end], start..) |r, i| { + if (r != 0) break i; + } else null; + if (busy) |i| end = i else return @intCast(start); + } + return null; + } + + /// `k` consecutive pages from `l`'s promise, held as one run until releaseRun; null: no such run is free now. + pub fn claimRun(p: *Pool, l: *Lease, k: usize) ?u32 { + if (l.promised < k) return null; + const start = p.findRun(k) orelse return null; + var n: usize = 0; + for (p.free[0..p.nfree]) |id| if (id < start or id >= start + k) { + p.free[n] = id; + n += 1; + }; + std.debug.assert(n + k == p.nfree); + p.nfree = n; + @memset(p.refs[start..][0..k], 1); + @memset(p.leased[start..][0..k], 1); // held as a lease holds its pages: never the store's to reclaim + l.promised -= k; + p.promised -= k; + return start; + } + + /// A run claimed at `start` goes back to the free pages. + pub fn releaseRun(p: *Pool, start: u32, k: usize) void { + for (start..start + k) |id| { + if (p.refs[id] != 1 or p.leased[id] != 1) std.debug.panic("kv pool: run page {d} released while not a run's", .{id}); + p.refs[id] = 0; + p.leased[id] = 0; + p.free[p.nfree] = @intCast(id); + p.nfree += 1; + } + } + /// The stream ended (finished, cancelled or failed): its promise and every page it holds go back. pub fn finish(p: *Pool, l: *Lease) void { p.unpromise(l); diff --git a/zig/src/core/kv_pool_test.zig b/zig/src/core/kv_pool_test.zig index 47196e1856..f6ef4e4fe5 100644 --- a/zig/src/core/kv_pool_test.zig +++ b/zig/src/core/kv_pool_test.zig @@ -216,3 +216,50 @@ test "a promise that does not fit asks the store to reclaim first, and fails onl p.finish(&b); try expect(p.idle()); } + +test "a run: consecutive pages from the top, from the stream's promise, back whole; KV pages fill from the bottom" { + const a = std.testing.allocator; + var p = try Pool.init(a, 10, 1); + defer p.deinit(); + const l: kv.Layout = .{ .fixed_bytes = 0, .token_bytes = 1, .page_tokens = 4, .snap_bytes = 0, .run_bytes = 9 }; + try expectEqual(@as(usize, 3), l.runPages()); + try expectEqual(@as(usize, 1 + 3), l.need(1, 0, 0)); + var x: Lease = .{}; + try p.promise(&x, l.need(5, 0, 0)); // two KV pages and a run of three + try expectEqual(@as(?u32, 7), p.claimRun(&x, 3)); + try expectEqual(@as(usize, 2), x.promised); + try p.cover(&x, 2, null); + try expectEqual(@as(u32, 0), x.ids.items[0]); // KV from the bottom + try expectEqual(@as(u32, 1), x.ids.items[1]); + try expectEqual(@as(usize, 5), p.available()); + try expectEqual(@as(usize, 0), p.reclaimable()); // a run is never the store's + // a second stream: the next run below the first, the promise checked first + var y: Lease = .{}; + try expectEqual(@as(?u32, null), p.claimRun(&y, 3)); + try p.promise(&y, 4); + try expectEqual(@as(?u32, 4), p.claimRun(&y, 3)); + // fragmentation: a run that is not free waits, even with pages to spare + var z: Lease = .{}; + try p.promise(&z, 1); + try p.cover(&z, 1, null); // page 2 + try expectEqual(@as(?u32, null), p.findRun(2)); // only page 3 is free + p.releaseRun(7, 3); + try expectEqual(@as(?u32, 7), p.findRun(3)); + p.releaseRun(4, 3); + p.finish(&x); + p.finish(&y); + p.finish(&z); + try expect(p.idle()); + try expectEqual(@as(usize, 10), p.available()); +} + +test "the plan with a run: each stream pays its run's pages beside its first page, the window bound adds them" { + const g: u64 = 1 << 30; + var l = nemotron; + l.run_bytes = 10 * l.pageBytes() - 1; + const p = try kv.plan(l, .{ .room = 40 * g, .engine = g, .lanes = 4, .window = 16384 + 16 }); + try expectEqual(@as(usize, 3 * (33 + 10)), p.pages); + const small = 2 * g + 2 * (64 * mib + 11 * l.pageBytes()) + mib; + try expectEqual(@as(u32, 3), (try kv.plan(l, .{ .room = small, .engine = 2 * g, .lanes = 8 })).lanes); + try expectEqual(nemotron.need(2048, 4096, 16) + 10, l.need(2048, 4096, 16)); +} diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig index 1332c59cf3..7bb6f04eba 100644 --- a/zig/src/families/qwen3_5/cuda_cache.zig +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -82,13 +82,49 @@ pub const Cache = struct { return fixed(gpa, e, cap, kv, paged); } - /// A pool lane's cache: its own state, log, tap ring and table of `pages` pages at `base`, mapped at open. + /// A pool lane's cache: a table of `pages` pages at `base`; its state, log and tap ring a run mapped at open. pub fn initPaged(gpa: std.mem.Allocator, e: *const Engine, pages: usize, tokens: usize, base: u64) !Cache { var table = try DeviceBuffer.alloc(e.d, pages * 4); errdefer table.free(); try table.fill8(0, null); const none = try DeviceBuffer.alloc(e.d, 0); - return fixed(gpa, e, 0, none, .{ .base = base, .table = table, .ids = &.{}, .tokens = tokens, .own_ids = false }); + return .{ .state = none, .before = none, .kv = none, .paged = .{ .base = base, .table = table, .ids = &.{}, .tokens = tokens, .own_ids = false }, .log = none, .cap = 0, .logits = try gpa.alloc(u16, e.g.vocab) }; + } + + /// Bytes of the parts a pool lane maps from its run: state, before, log, tap ring, each 256-byte aligned. + pub fn runBytes(e: *const Engine) usize { + var n: usize = 0; + for (runParts(e)) |len| n += std.mem.alignForward(usize, len, 256); + return n; + } + + fn runParts(e: *const Engine) [4]usize { + return .{ e.g.stateBytes(), e.g.stateBytes(), e.g.logBytes(engine.most_window), e.tapBytes() }; + } + + /// A pool lane's state, before, log and tap ring as views of the run at `base`, zeroed but `before`. + pub fn mapRun(self: *Cache, e: *const Engine, base: u64) !void { + var at = base; + var views: [4]DeviceBuffer = undefined; + for (runParts(e), &views) |len, *v| { + v.* = .{ .d = e.d, .ptr = at, .len = len, .borrowed = true }; + at += std.mem.alignForward(usize, len, 256); + } + self.state = views[0]; + self.before = views[1]; + self.log = views[2]; + self.tap = if (e.n_taps > 0) views[3] else null; + try self.log.fill8(0, e.stream.handle); + if (self.tap) |t| try t.fill8(0, e.stream.handle); + } + + /// A closed lane's views dropped: its run went back to the pool. + pub fn unmapRun(self: *Cache, e: *const Engine) void { + const none: DeviceBuffer = .{ .d = e.d, .ptr = 0, .len = 0, .borrowed = true }; + self.state = none; + self.before = none; + self.log = none; + self.tap = null; } /// The parts every cache holds itself beside `kv` and `paged` (the caller's to free on failure). diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig index 2b26e22992..1bf308a412 100644 --- a/zig/src/families/qwen3_5/cuda_pages.zig +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -28,13 +28,13 @@ fn mostPages(capacity: usize, tokens: usize) usize { return (capacity + window_rows + tokens - 1) / tokens; } -/// A stream's fixed part (state, log, tap ring, table), its keys and values a token, a kept state's recurrences. +/// A lane's table; its keys and values a token; its state, rollback copy, log and tap ring a run; a kept state. pub fn layout(e: *const Engine, capacity: usize, tokens: usize) kv.Layout { - return .{ .fixed_bytes = e.streamBytes(0) + mostPages(capacity, tokens) * 4, .token_bytes = e.g.kvBytes(1, e.kv8), .page_tokens = @intCast(tokens), .snap_bytes = e.g.stateBytes() }; + return .{ .fixed_bytes = mostPages(capacity, tokens) * 4, .token_bytes = e.g.kvBytes(1, e.kv8), .page_tokens = @intCast(tokens), .snap_bytes = e.g.stateBytes(), .run_bytes = Cache.runBytes(e) }; } -/// One pool lane: a cache whose table maps its stream's pages, and that stream's lease. -const Lane = struct { cache: *Cache, lease: kv.Lease = .{}, used: bool = false }; +/// One pool lane: a cache whose table maps its stream's pages, that stream's lease and its run's first page. +const Lane = struct { cache: *Cache, lease: kv.Lease = .{}, used: bool = false, run: u32 = 0 }; pub const Pages = struct { gpa: std.mem.Allocator, @@ -75,7 +75,10 @@ pub const Pages = struct { pub fn deinit(p: *Pages) void { p.e.stream.synchronize() catch std.log.warn("qwen3_5: the pool's stream did not drain at close", .{}); for (p.lanes) |*x| { - if (x.used) p.pool.finish(&x.lease); + if (x.used) { + p.pool.releaseRun(x.run, p.layout.runPages()); + p.pool.finish(&x.lease); + } x.cache.deinit(p.gpa); p.gpa.destroy(x.cache); } @@ -89,9 +92,9 @@ pub const Pages = struct { return p.layout.need(prompt, max_tokens, window_rows); } - /// Whether a lane is free and the pool can promise a request's pages now. + /// Whether a lane is free, the pool can promise a request's pages now and a run of them is free. pub fn fits(p: *const Pages, prompt: usize, max_tokens: usize) bool { - for (p.lanes) |x| if (!x.used) return p.pool.fits(p.need(prompt, max_tokens)); + for (p.lanes) |x| if (!x.used) return p.pool.fits(p.need(prompt, max_tokens)) and p.pool.findRun(p.layout.runPages()) != null; return false; } @@ -100,15 +103,21 @@ pub const Pages = struct { const x = for (p.lanes) |*x| { if (!x.used) break x; } else return error.LaneMemory; - const n = p.need(prompt, max_tokens); + const k = p.layout.runPages(); + const n = p.need(prompt, max_tokens) - k; // keys and values if (n > mostPages(p.capacity, p.layout.page_tokens)) return error.InvalidQwenPrompt; - try p.pool.promise(&x.lease, n); + try p.pool.promise(&x.lease, n + k); errdefer p.pool.finish(&x.lease); + // the run first: KV pages fill from the bottom, runs from the top; a fragmented pool waits (LaneMemory) + const run = p.pool.claimRun(&x.lease, k) orelse return error.LaneMemory; + errdefer p.pool.releaseRun(run, k); var up: Upload = .{ .p = p, .cache = x.cache }; try p.pool.cover(&x.lease, n, .{ .ctx = &up, .write = Upload.write }); + try x.cache.mapRun(p.e, p.arena.ptr + run * p.layout.pageBytes()); try x.cache.reset(p.e); x.cache.paged.?.ids = x.lease.ids.items; x.cache.cap = n * p.layout.page_tokens; + x.run = run; x.used = true; return x.cache; } @@ -121,7 +130,9 @@ pub const Pages = struct { if (!x.used) return true; // a page another stream takes next may still be read by queued work on the engine's stream p.e.stream.synchronize() catch std.log.warn("qwen3_5: a closed stream's forwards did not drain; its pages go back", .{}); + p.pool.releaseRun(x.run, p.layout.runPages()); p.pool.finish(&x.lease); + cache.unmapRun(p.e); cache.paged.?.ids = &.{}; cache.cap = 0; x.used = false; From 51f835402bbf034757e65528c9a088fef1c922bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:41:35 +0200 Subject: [PATCH 07/11] cuda qwen3_5: a DFlash2 stream's ring is part of its pool lane's run; the pool planned after the drafter The pool's layout was planned when the family opened, before a drafter set the target's taps, so a lane's run left out the tap ring that mapRun then mapped after the state: with --drafter, the tap ring ran past the run (into the run above, or past the arena for the top one). The DFlash2 gates of 1ac20753 passed regardless; that was luck, not a guarantee. Now the native host asks a family that places a drafter's buffers itself (pagedLayout) for its layout once the drafter is attached: qwen3_5 puts the target's tap ring and the DFlash2 ring in each lane's run and keeps the rest of the drafter's bytes a stream (its frame scratch's share, an MTP head's buffers) fixed. mapRun refuses a run shorter than its parts (QwenRunTooShort). The drafter borrows a lane's ring for a stream on a pool lane (a context with a borrowed buffer) and keeps its own rings only for the first stream. GB10, Qwen3.8-27B MLX-4bit, TF_QWEN_KV_PAGES=256, DFlash2, --parallel 4 at 32768: "3 more taking 522 pages (8.16 GiB), 23 MiB fixed each" (was 169 MiB fixed each); bytes == contiguous (3/3 chat, 3/3 8k prefill, pages only); DFlash2 decode code/chat 59.9/44.4 tok/s pool, 60.3/44.5 pages only (--parallel 5), 60.5/44.6 contiguous; c=4 161.5/123.1, 164.2/123.4 against 164.2/123.6. Pages only at --parallel 8, bf16 and fp8: drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh, equal to the earlier runs; no run refused. --- zig/src/families/qwen3_5/cuda_cache.zig | 30 +++++++++-------- zig/src/families/qwen3_5/cuda_draft.zig | 20 +++++++++-- zig/src/families/qwen3_5/cuda_native.zig | 35 ++++++++++++++++++-- zig/src/families/qwen3_5/cuda_native_moe.zig | 3 +- zig/src/families/qwen3_5/cuda_pages.zig | 17 +++++----- zig/src/native/cuda.zig | 2 +- zig/src/native/cuda_kv_pool.zig | 8 +++++ 7 files changed, 87 insertions(+), 28 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig index 7bb6f04eba..0655054135 100644 --- a/zig/src/families/qwen3_5/cuda_cache.zig +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -44,6 +44,8 @@ pub const Cache = struct { untaught: bool = false, /// Keys and values through a page table (null: cap contiguous rows a head). paged: ?Paged = null, + /// A pool lane's drafter ring in its run (the drafter borrows it for the lane's stream). + drafter_ring: ?u64 = null, delta: i32 = 0, state: DeviceBuffer, /// The exact state before a recorded one-row round advances its stable slot in place. @@ -91,31 +93,32 @@ pub const Cache = struct { return .{ .state = none, .before = none, .kv = none, .paged = .{ .base = base, .table = table, .ids = &.{}, .tokens = tokens, .own_ids = false }, .log = none, .cap = 0, .logits = try gpa.alloc(u16, e.g.vocab) }; } - /// Bytes of the parts a pool lane maps from its run: state, before, log, tap ring, each 256-byte aligned. - pub fn runBytes(e: *const Engine) usize { + /// Bytes of a pool lane's run: state, before, log, tap ring and a drafter's ring (`ring`), each 256-aligned. + pub fn runBytes(e: *const Engine, ring: usize) usize { var n: usize = 0; - for (runParts(e)) |len| n += std.mem.alignForward(usize, len, 256); + for (runParts(e, ring)) |len| n += std.mem.alignForward(usize, len, 256); return n; } - fn runParts(e: *const Engine) [4]usize { - return .{ e.g.stateBytes(), e.g.stateBytes(), e.g.logBytes(engine.most_window), e.tapBytes() }; + fn runParts(e: *const Engine, ring: usize) [5]usize { + return .{ e.g.stateBytes(), e.g.stateBytes(), e.g.logBytes(engine.most_window), e.tapBytes(), ring }; } - /// A pool lane's state, before, log and tap ring as views of the run at `base`, zeroed but `before`. - pub fn mapRun(self: *Cache, e: *const Engine, base: u64) !void { + /// A lane's state, before, log, taps and drafter ring as views of the run at `base`; a short run is refused. + pub fn mapRun(self: *Cache, e: *const Engine, base: u64, len: usize, ring: usize) !void { + if (runBytes(e, ring) > len) return error.QwenRunTooShort; var at = base; - var views: [4]DeviceBuffer = undefined; - for (runParts(e), &views) |len, *v| { - v.* = .{ .d = e.d, .ptr = at, .len = len, .borrowed = true }; - at += std.mem.alignForward(usize, len, 256); + var views: [5]DeviceBuffer = undefined; + for (runParts(e, ring), &views) |n, *v| { + v.* = .{ .d = e.d, .ptr = at, .len = n, .borrowed = true }; + at += std.mem.alignForward(usize, n, 256); } self.state = views[0]; self.before = views[1]; self.log = views[2]; self.tap = if (e.n_taps > 0) views[3] else null; - try self.log.fill8(0, e.stream.handle); - if (self.tap) |t| try t.fill8(0, e.stream.handle); + self.drafter_ring = if (ring > 0) views[4].ptr else null; + for ([_]DeviceBuffer{ views[2], views[3], views[4] }) |v| if (v.len > 0) try v.fill8(0, e.stream.handle); } /// A closed lane's views dropped: its run went back to the pool. @@ -125,6 +128,7 @@ pub const Cache = struct { self.before = none; self.log = none; self.tap = null; + self.drafter_ring = null; } /// The parts every cache holds itself beside `kv` and `paged` (the caller's to free on failure). diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index 6883c19fc7..f8124174c9 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -100,6 +100,10 @@ pub const Draft = struct { frame_streams: usize = 1, scratch_one: usize = 0, target_taps: usize = 0, + /// TF_QWEN_KV_PAGES: a stream's ring is a view of its target lane's run (null: the contexts own theirs). + ring_of: ?RingOf = null, + + pub const RingOf = struct { ctx: *anyopaque, ring: *const fn (*anyopaque, *Stream) ?u64 }; const words_at = frame.words_at; const hp_at = frame.hp_at; @@ -235,8 +239,9 @@ pub const Draft = struct { const started = profile.now(d.io); const previous_times = d.startup; try memory.resize(d, n); - try d.pool.ensureTotalCapacity(d.gpa, n); - while (d.pool.items.len < n) d.pool.appendAssumeCapacity(try d.create()); + const own = if (d.ring_of != null) @min(n, 1) else n; // with a pool the other streams' rings are lane runs + try d.pool.ensureTotalCapacity(d.gpa, own); + while (d.pool.items.len < own) d.pool.appendAssumeCapacity(try d.create()); d.startup.reserve += profile.elapsed(d.io, started) - (d.startup.warmups - previous_times.warmups) - (d.startup.reserve - previous_times.reserve); } @@ -376,6 +381,12 @@ pub const Draft = struct { fn openFn(p: *anyopaque, s: *Stream) anyerror!void { const d = cast(p); if (d.streams.contains(s)) return error.StreamOpen; + if (d.ring_of) |r| if (r.ring(r.ctx, s)) |ptr| { // the lane's run holds it: a context borrowing it + const ctx = try d.gpa.create(Ctx); + errdefer d.gpa.destroy(ctx); + ctx.* = .{ .ring = .{ .d = d.e.d, .ptr = ptr, .len = d.c.ringBytes(), .borrowed = true } }; + return d.streams.put(d.gpa, s, ctx); + }; const ctx = d.pool.pop() orelse try d.create(); errdefer d.pool.append(d.gpa, ctx) catch d.drop(ctx); ctx.* = .{ .ring = ctx.ring }; @@ -444,7 +455,10 @@ pub const Draft = struct { fn releaseFn(p: *anyopaque, s: *Stream) void { const d = cast(p); - if (d.streams.fetchRemove(s)) |kv| d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + if (d.streams.fetchRemove(s)) |kv| { + if (kv.value.ring.borrowed) return d.gpa.destroy(kv.value); // its lane's run goes back with the lane + d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + } } }; diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 4b1858252b..fc0c4d668e 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -66,10 +66,24 @@ pub const Loaded = struct { layout: ?kv_pool.Layout = null, }; +/// The pool's layout with a drafter (`extra` bytes a stream): DFlash2's ring and taps in the run, the rest fixed. +pub fn pagedLayout(p: *anyopaque, extra: u64) ?kv_pool.Layout { + const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.page_tokens == 0) return null; + const ring = l.ringBytes(); + const in_run = ring + if (l.draft != null) l.e.tapBytes() else 0; + l.extra = @intCast(extra -| in_run); + return pages_mod.layout(l.e, l.capacity, l.page_tokens, ring, l.extra); +} + /// The pool beside the first stream's window: `streams` lanes and `pages` pages, made now. pub fn attach(p: *anyopaque, streams: u32, pages: usize) !void { const l: *Lanes = @ptrCast(@alignCast(p)); - l.pages = try pages_mod.Pages.init(l.gpa, l.e, l.capacity, l.page_tokens, streams, pages); + const ring = l.ringBytes(); + l.pages = try pages_mod.Pages.init(l.gpa, l.e, pages_mod.layout(l.e, l.capacity, l.page_tokens, ring, l.extra), l.capacity, ring, streams, pages); + if (l.draft) |d| if (streams > 0) { + d.ring_of = .{ .ctx = l, .ring = Lanes.ringOf }; + }; l.paged_only = pages_mod.pagedOnly() and streams > 0; // checks: every stream on pages, the reserved window idle } @@ -108,7 +122,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .reserve = reserve, .vision = v, .target = &l.target, - .layout = if (page_tokens > 0) pages_mod.layout(e, o.context, page_tokens) else null, + .layout = if (page_tokens > 0) pages_mod.layout(e, o.context, page_tokens, 0, 0) else null, }; } @@ -137,6 +151,7 @@ pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []con d.target_taps = l.e.tapBytes(); cost_adapter.identify(l, dir, .dflash); l.calibration = d; + l.draft = d; return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter, .prepare = .{ .ctx = l, .run = calibrateDrafter, .bytes = if (l.capacity < l.e.window_rows) l.e.streamBytes(l.e.window_rows) else 0 } }; } @@ -254,6 +269,10 @@ const Lanes = struct { page_tokens: usize = 0, pages: ?*pages_mod.Pages = null, paged_only: bool = false, + /// The DFlash2 drafter, when one is attached: its rings go into the pool lanes' runs. + draft: ?*Draft = null, + /// A drafter's bytes a stream outside the lanes' runs (pagedLayout). + extra: usize = 0, cost_context: ?*const cuda.Context = null, cost_identity: ?[32]u8 = null, cost_kind: cost_adapter.Kind = .dflash, @@ -316,6 +335,18 @@ const Lanes = struct { return .{ .ptr = self, .first_before_draft = true, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .prepare_features = prepareFeaturesFn, .features = featuresFn, .release = releaseFn } }; } + /// A drafter ring's bytes in each lane's run (0: no DFlash2 drafter). + fn ringBytes(self: *const Lanes) usize { + return if (self.draft) |d| d.c.ringBytes() else 0; + } + + /// The ring in the run of the lane `stream`'s cache holds; null: its cache is a reserved window's. + fn ringOf(ptr: *anyopaque, stream: *lanes.Stream) ?u64 { + const self = cast(ptr); + const cache = self.caches.get(stream) orelse return null; + return cache.drafter_ring; + } + fn cast(ptr: *anyopaque) *Lanes { return @ptrCast(@alignCast(ptr)); } diff --git a/zig/src/families/qwen3_5/cuda_native_moe.zig b/zig/src/families/qwen3_5/cuda_native_moe.zig index c7ede8d5e5..3bc6155bd5 100644 --- a/zig/src/families/qwen3_5/cuda_native_moe.zig +++ b/zig/src/families/qwen3_5/cuda_native_moe.zig @@ -10,7 +10,8 @@ pub const prompt_rows = native.prompt_rows; pub const takes_drafter = true; pub const Options = native.Options; pub const open = native.open; -pub const attach = native.attach; // TF_QWEN_KV_PAGES: the same pool as the dense models' +pub const attach = native.attach; +pub const pagedLayout = native.pagedLayout; // TF_QWEN_KV_PAGES: the same pool as the dense models' pub const fits = native.fits; pub const window_rows = native.window_rows; pub const openDrafter = native.openMtp; diff --git a/zig/src/families/qwen3_5/cuda_pages.zig b/zig/src/families/qwen3_5/cuda_pages.zig index 1bf308a412..e4436ae022 100644 --- a/zig/src/families/qwen3_5/cuda_pages.zig +++ b/zig/src/families/qwen3_5/cuda_pages.zig @@ -28,9 +28,9 @@ fn mostPages(capacity: usize, tokens: usize) usize { return (capacity + window_rows + tokens - 1) / tokens; } -/// A lane's table; its keys and values a token; its state, rollback copy, log and tap ring a run; a kept state. -pub fn layout(e: *const Engine, capacity: usize, tokens: usize) kv.Layout { - return .{ .fixed_bytes = mostPages(capacity, tokens) * 4, .token_bytes = e.g.kvBytes(1, e.kv8), .page_tokens = @intCast(tokens), .snap_bytes = e.g.stateBytes(), .run_bytes = Cache.runBytes(e) }; +/// A lane's table and `fixed` bytes; its keys and values a token; a run of its state, log, taps and a `ring`-byte ring. +pub fn layout(e: *const Engine, capacity: usize, tokens: usize, ring: usize, fixed: usize) kv.Layout { + return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + fixed, .token_bytes = e.g.kvBytes(1, e.kv8), .page_tokens = @intCast(tokens), .snap_bytes = e.g.stateBytes(), .run_bytes = Cache.runBytes(e, ring) }; } /// One pool lane: a cache whose table maps its stream's pages, that stream's lease and its run's first page. @@ -41,19 +41,20 @@ pub const Pages = struct { e: *Engine, layout: kv.Layout, capacity: usize, + ring: usize, pool: kv.Pool, arena: cuda.DeviceBuffer, lanes: []Lane, - /// `streams` lanes and `pages` pages, made and touched now. - pub fn init(gpa: std.mem.Allocator, e: *Engine, capacity: usize, tokens: usize, streams: usize, pages: usize) !*Pages { - const l = layout(e, capacity, tokens); + /// `streams` lanes and `pages` pages of layout `l`, made and touched now; each run holds a `ring`-byte ring. + pub fn init(gpa: std.mem.Allocator, e: *Engine, l: kv.Layout, capacity: usize, ring: usize, streams: usize, pages: usize) !*Pages { + const tokens: usize = l.page_tokens; const p = try gpa.create(Pages); errdefer gpa.destroy(p); var arena = try cuda.DeviceBuffer.alloc(e.d, pages * l.pageBytes()); errdefer arena.free(); try arena.fill8(0, null); // committed now (the driver commits on first touch) - p.* = .{ .gpa = gpa, .e = e, .layout = l, .capacity = capacity, .pool = try kv.Pool.init(gpa, pages, l.pageBytes()), .arena = arena, .lanes = &.{} }; + p.* = .{ .gpa = gpa, .e = e, .layout = l, .capacity = capacity, .ring = ring, .pool = try kv.Pool.init(gpa, pages, l.pageBytes()), .arena = arena, .lanes = &.{} }; errdefer p.pool.deinit(); p.lanes = try gpa.alloc(Lane, streams); errdefer gpa.free(p.lanes); @@ -113,7 +114,7 @@ pub const Pages = struct { errdefer p.pool.releaseRun(run, k); var up: Upload = .{ .p = p, .cache = x.cache }; try p.pool.cover(&x.lease, n, .{ .ctx = &up, .write = Upload.write }); - try x.cache.mapRun(p.e, p.arena.ptr + run * p.layout.pageBytes()); + try x.cache.mapRun(p.e, p.arena.ptr + run * p.layout.pageBytes(), k * p.layout.pageBytes(), p.ring); try x.cache.reset(p.e); x.cache.paged.?.ids = x.lease.ids.items; x.cache.cap = n * p.layout.page_tokens; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 00873d7d13..860fcf65c4 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -346,7 +346,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; const free = loaded.free_streams; const grow = budget.growRoom(room, !@hasField(F.Options, "lanes")); - const wrapped = kv_pages.served(kv_pages.layout(loaded), free, loaded.stream_bytes, served.stream_bytes, room); + const wrapped = kv_pages.wrapped(F, loaded, free, served.stream_bytes, room); const layout = wrapped.layout; var paged: ?kv_pages.Plan = null; const streams = if (layout) |l| blk: { diff --git a/zig/src/native/cuda_kv_pool.zig b/zig/src/native/cuda_kv_pool.zig index 355e0e541b..17f52e2bc2 100644 --- a/zig/src/native/cuda_kv_pool.zig +++ b/zig/src/native/cuda_kv_pool.zig @@ -30,6 +30,14 @@ test "a wrapped family's plan: drafter buffers a stream, a first window without try std.testing.expectEqual(@as(?kv_pool.Layout, null), served(null, 0, 1, 2, 3).layout); } +/// served() on the family's layout; a family that places a drafter's buffers itself (pagedLayout) plans it again. +pub fn wrapped(comptime F: type, loaded: anytype, free: u32, stream_bytes: u64, room: u64) @TypeOf(served(null, 0, 0, 0, 0)) { + const own = comptime @hasDecl(F, "pagedLayout"); + const l = layout(loaded); + const planned = if (own) (if (l != null) F.pagedLayout(loaded.ctx, stream_bytes - loaded.stream_bytes) else null) else l; + return served(planned, free, if (own) stream_bytes else loaded.stream_bytes, stream_bytes, room); +} + /// The real native startup room before page planning and cache allocation. pub fn startupRoom(comptime F: type, a: Allocator, loaded: anytype, o: api.Open, available: u64, host_budget: u64, problem: *[]const u8) !?u64 { const memory = @import("cuda_family_memory.zig"); From ce5b78d7c27444b6f152ab2c0118f918333e5868 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:32:52 +0200 Subject: [PATCH 08/11] cuda qwen3_5: the attention partials test the page table once a block, not once a key On Qwen3.6-35B-A3B (decode at 84 tok/s, attention a larger share of a step than on the 27B) the contiguous path paid kv_row's page-table test inside the key loop: an ABAB against main measured -2.7..-3.6 % at one stream and -1.5..-2.1 % at four. tf_q35_attn_partial and tf_q35_attn_partial8 now branch once: without a table they take main's K and V (and scale) bases and step a row at a time, with one they read each row through kv_row / row8. The addresses are the same, so the bytes. GB10, Qwen3.6-35B-A3B MLX-4bit (MTP), --parallel 4 at 32768, ABAB against main, code/chat tok/s: main 84.2/75.0 and 83.3/75.1, this 82.7/74.5 and 84.0/75.2 (c=4 163.7/140.4, 163.2/143.8 against 161.0/142.1, 163.5/143.7); fp8 84.0/75.1 main, 84.4/75.3 this. Bytes == main (3/3 chat, 3/3 8k prefill) for the 35B in bf16 and fp8 and for the 27B; pages only == main; tf-cuda-test qwen-pages passes in both formats with its controls red. --- zig/kernels/cuda/qwen35.cu | 42 +++++++++++++++++++++------------- zig/kernels/cuda/qwen35_kv8.cu | 39 +++++++++++++++++++------------ 2 files changed, 50 insertions(+), 31 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index c377f976a7..2b657e642c 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -441,22 +441,32 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (int t = k0; t < k1; t++) { - const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); - const uint4 kw = *reinterpret_cast(KV + kv_row(sg, KVH, layer, 0, h, pos) + lane * DPL); - const __nv_bfloat16* kb = reinterpret_cast(&kw); - float dot = 0.0f; -#pragma unroll - for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(kb[i]), dot); - const float s = warp_sum(dot) * scale; - const float mn = fmaxf(mx, s); - const float a = __expf(mx - mn), b = __expf(s - mn); - l = l * a + b; - const uint4 vw = *reinterpret_cast(KV + kv_row(sg, KVH, layer, 1, h, pos) + lane * DPL); - const __nv_bfloat16* vb = reinterpret_cast(&vw); -#pragma unroll - for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); - mx = mn; + // the keys in order; a key's and value's rows by `row` (one branch here, not one a key) + const auto scan = [&](auto row) { + for (int t = k0; t < k1; t++) { + const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); + const uint4 kw = *reinterpret_cast(row(0, pos) + lane * DPL); + const __nv_bfloat16* kb = reinterpret_cast(&kw); + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(kb[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const uint4 vw = *reinterpret_cast(row(1, pos) + lane * DPL); + const __nv_bfloat16* vb = reinterpret_cast(&vw); +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); + mx = mn; + } + }; + if (sg.pt == 0) { + const __nv_bfloat16* K = KV + kv_row(sg, KVH, layer, 0, h, 0); + const __nv_bfloat16* V = KV + kv_row(sg, KVH, layer, 1, h, 0); + scan([&](int s, int pos) { return (s == 0 ? K : V) + (size_t)pos * HD; }); + } else { + scan([&](int s, int pos) { return KV + kv_row(sg, KVH, layer, s, h, pos); }); } if (lane == 0) { PM[slot] = mx; diff --git a/zig/kernels/cuda/qwen35_kv8.cu b/zig/kernels/cuda/qwen35_kv8.cu index 5c43e5cb2c..3069f149b1 100644 --- a/zig/kernels/cuda/qwen35_kv8.cu +++ b/zig/kernels/cuda/qwen35_kv8.cu @@ -106,23 +106,32 @@ extern "C" __global__ void tf_q35_attn_partial8(Geo geo, const __nv_bfloat16* __ o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (int t = k0; t < k1; t++) { - const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); - const Row8 r = row8(sg, layer, KVH, h, pos); - float kf[DPL], vf[DPL]; - widen8(*reinterpret_cast(r.k + lane * DPL), kf); - float dot = 0.0f; + // the keys in order; a row's codes and scales by `at` (one branch here, not one a key) + const auto scan = [&](auto at) { + for (int t = k0; t < k1; t++) { + const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); + const Row8 r = at(pos); + float kf[DPL], vf[DPL]; + widen8(*reinterpret_cast(r.k + lane * DPL), kf); + float dot = 0.0f; #pragma unroll - for (int i = 0; i < DPL; i++) dot = fmaf(q[i], kf[i], dot); - const float s = warp_sum(dot) * *r.ks * scale; - const float mn = fmaxf(mx, s); - const float a = __expf(mx - mn), b = __expf(s - mn); - l = l * a + b; - widen8(*reinterpret_cast(r.v + lane * DPL), vf); - const float bv = b * *r.vs; + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], kf[i], dot); + const float s = warp_sum(dot) * *r.ks * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + widen8(*reinterpret_cast(r.v + lane * DPL), vf); + const float bv = b * *r.vs; #pragma unroll - for (int i = 0; i < DPL; i++) o[i] = fmaf(bv, vf[i], o[i] * a); - mx = mn; + for (int i = 0; i < DPL; i++) o[i] = fmaf(bv, vf[i], o[i] * a); + mx = mn; + } + }; + if (sg.pt == 0) { + const Row8 r0 = row8(sg, layer, KVH, h, 0); // row p is row 0's plus p rows + scan([&](int pos) { return Row8{r0.k + (size_t)pos * HD, r0.v + (size_t)pos * HD, r0.ks + pos, r0.vs + pos}; }); + } else { + scan([&](int pos) { return row8(sg, layer, KVH, h, pos); }); } if (lane == 0) { PM[slot] = mx; From 85aeb4bf9085e69bc7695dbdf4e8863efc0933a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:40:43 +0200 Subject: [PATCH 09/11] cuda flashnext: attention caches readable through a page table flashnext_attn.cu's cache kernels take a Pg (an int page table, the elements a page holds, the positions a page holds): fn_attn_prep(_media)'s key, value and indexer-key stores, fn_qsa_pool(_media)'s reads and pooled writes, fn_qsa_scores' pooled reads and fn_attn_chunks(_ahead)'s key and value tiles find a row with pg_at(): the old expression when the table is null, else the page holding the position plus its row there. fn_attn_chunks looks a key's row up once a tile (key_at, beside key_ids), the pool once a block of four. A page of n positions holds, for each attention layer, its keys, values and indexer keys of n positions and the pooled rows of their n / 4 blocks (cuda_pages.zig); a cache pointer is its layer's part of page 0. Only addresses change, never a sum's order, so paged caches give contiguous bytes at any page size that holds whole blocks. Without our glue (a captured Triton set) a page table is refused. TF_FLASHNEXT_KV_PAGE_TEST=n gives each stream's caches one buffer read through a reversed table of n-position pages (checks only). Prompt-cache saves and restores copy page runs into the same layout, so kept states move between paged and contiguous caches. tf-cuda-test flashnext-pages (also in kernel-contract): six cases, cached lengths to 9000, pages of 16, 48, 128, 256 and 512 positions, two streams shuffled through one arena: queries, indexer queries, block scores, dense and sparse attention rows, the stores and pooled rows equal contiguous over three repeats; attention rows within a float64 bound; guards intact. Red: a swapped table over placed pages, a planted row. Records: the sm_80..121 registers and static bytes of the kernels that moved (nvcc 13.0). GB10, nvidia/Qwen3.8-Flash-Next-FP4-Local, --parallel 4 --context 40960, --no-drafts, ABAB against main: c=1 23.51/23.49 and 23.55/23.44 tok/s against 23.46/23.40 and 23.38/23.48; 8k 5.97-6.91 s against 6.03-7.00 s, 31k 25.64/25.74 s against 25.67/26.01 s; replies equal (3 chats, three ~8k, one 31k). Page test 16: the same bytes, 23.56/23.38 tok/s, 31k 26.08 s. With MTP drafts: drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh (7,147 cached tokens restored into pages), equal to contiguous. --- zig/kernels/cuda/flashnext_attn.cu | 71 ++-- zig/src/core/contract/registry.zig | 10 +- zig/src/families/flashnext/cuda_attn.zig | 4 +- zig/src/families/flashnext/cuda_batch.zig | 9 +- .../families/flashnext/cuda_batch_state.zig | 2 + zig/src/families/flashnext/cuda_glue.zig | 15 +- zig/src/families/flashnext/cuda_launch.zig | 22 +- zig/src/families/flashnext/cuda_media.zig | 7 +- .../families/flashnext/cuda_media_test.zig | 9 +- zig/src/families/flashnext/cuda_pages.zig | 132 +++++++ zig/src/families/flashnext/cuda_prompt.zig | 23 +- .../families/flashnext/cuda_prompt_step.zig | 6 +- zig/src/families/flashnext/cuda_rows.zig | 5 +- zig/src/families/flashnext/cuda_snap.zig | 20 +- zig/tests/cuda/contract_fn_pages.zig | 344 ++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 1 + zig/tests/cuda/flashnext_attn_tests.zig | 8 +- zig/tests/cuda/main.zig | 2 + 18 files changed, 621 insertions(+), 69 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_pages.zig create mode 100644 zig/tests/cuda/contract_fn_pages.zig diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu index 2f66798379..2ee410ebdc 100644 --- a/zig/kernels/cuda/flashnext_attn.cu +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -53,6 +53,19 @@ __device__ __forceinline__ int partner(int d, int half) { return d < half ? d + half : (d < 2 * half ? d - half : d); } +// A stream's cache pages (cuda_pages.zig's layout): pt[i] holds positions [i * ptok, (i + 1) * ptok); null: contiguous. +struct Pg { + const int* pt; + long long ps; + int ptok, pad; +}; + +// Element offset of item i (a position, or a block with per = ptok / RATIO) of `row` elements, contiguous or paged. +__device__ __forceinline__ long long pg_at(const Pg& pg, int i, int per, long long row) { + if (pg.pt == nullptr) return static_cast(i) * row; + return static_cast(pg.pt[i / per]) * pg.ps + static_cast(i % per) * row; +} + } // namespace fn_attn using namespace fn_attn; @@ -70,13 +83,13 @@ __device__ __forceinline__ void fn_attn_prep_body(const __nv_bfloat16* __restric __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, - int NI, int IHD, int HALF, const int* rope3, int n, int delta) { + int NI, int IHD, int HALF, const int* rope3, int n, int delta, Pg pg) { __shared__ float scratch[32]; const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; const int pos = POS0[0] + r; const long long row = static_cast(r) * PW; if (head >= NQ + NKV + NI) { - if (d < IHD) IKC[static_cast(pos) * IHD + d] = P[row + NQ * 2 * HD + 2 * NKV * HD + NI * IHD + d]; + if (d < IHD) IKC[pg_at(pg, pos, pg.ptok, IHD) + d] = P[row + NQ * 2 * HD + 2 * NKV * HD + NI * IHD + d]; return; } const bool index = head >= NQ + NKV; @@ -94,9 +107,9 @@ __device__ __forceinline__ void fn_attn_prep_body(const __nv_bfloat16* __restric if (head < NQ) { Q[(static_cast(r) * NQ + head) * HD + d] = out; } else if (!index) { - const long long slot = static_cast(pos) * NKV + head - NQ; - KC[slot * HD + d] = out; - VC[slot * HD + d] = P[row + NQ * 2 * HD + NKV * HD + (head - NQ) * HD + d]; + const long long slot = pg_at(pg, pos, pg.ptok, static_cast(NKV) * HD) + (head - NQ) * HD; + KC[slot + d] = out; + VC[slot + d] = P[row + NQ * 2 * HD + NKV * HD + (head - NQ) * HD + d]; } else { IQ[(static_cast(r) * NI + head - NQ - NKV) * IHD + d] = out; } @@ -108,8 +121,8 @@ extern "C" __global__ void __launch_bounds__(256) fn_attn_prep(const __nv_bfloat __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, - int NI, int IHD, int HALF) { - fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, nullptr, 0, 0); + int NI, int IHD, int HALF, Pg pg) { + fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, nullptr, 0, 0, pg); } extern "C" __global__ void __launch_bounds__(256) fn_attn_prep_media(const __nv_bfloat16* __restrict__ P, const int* __restrict__ POS0, @@ -118,8 +131,8 @@ extern "C" __global__ void __launch_bounds__(256) fn_attn_prep_media(const __nv_ __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, - int NI, int IHD, int HALF, const int* rope3, int n, int delta) { - fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, rope3, n, delta); + int NI, int IHD, int HALF, const int* rope3, int n, int delta, Pg pg) { + fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, rope3, n, delta, pg); } // glue._attn_gate, program (row, head) of HD threads: bf16(o / (1 + e^-gate)) and its 32-group sums. @@ -138,41 +151,42 @@ extern "C" __global__ void __launch_bounds__(256) fn_attn_gate(const __nv_bfloat __device__ __forceinline__ void fn_qsa_pool_body(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, float eps, int R, int DI, int HALF, - int RATIO, const int* rope3, int n, int delta) { + int RATIO, const int* rope3, int n, int delta, Pg pg) { __shared__ float scratch[32]; __shared__ float xs[256]; const int p0 = POS0[0]; const int b = p0 / RATIO + blockIdx.x, d = threadIdx.x; if (RATIO * b + RATIO > p0 + R) return; - float acc = bf(IKC[static_cast(RATIO) * b * DI + d]); - for (int k = 1; k < RATIO; ++k) acc = __fadd_rn(acc, bf(IKC[(static_cast(RATIO) * b + k) * DI + d])); + const long long at = pg_at(pg, RATIO * b, pg.ptok, DI) + d; // a block's RATIO rows share a page + float acc = bf(IKC[at]); + for (int k = 1; k < RATIO; ++k) acc = __fadd_rn(acc, bf(IKC[at + static_cast(k) * DI])); const float x = rbf(__fdiv_rn(acc, static_cast(RATIO))); xs[d] = x; const float rinv = inv_rms(block_sum(__fmul_rn(x, x), scratch), DI, eps); const int p = partner(d, HALF); const float xn = rbf(__fmul_rn(__fmul_rn(x, rinv), W[d])); const float xpn = rbf(__fmul_rn(__fmul_rn(xs[p], rinv), W[p])); - POOLED[static_cast(b) * DI + d] = __float2bfloat16_rn(rope(xn, xpn, d, HALF, media_pos(RATIO * b, d, HALF, rope3, n, delta), INV)); + POOLED[pg_at(pg, b, pg.ptok / RATIO, DI) + d] = __float2bfloat16_rn(rope(xn, xpn, d, HALF, media_pos(RATIO * b, d, HALF, rope3, n, delta), INV)); } extern "C" __global__ void __launch_bounds__(128) fn_qsa_pool(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, float eps, int R, int DI, int HALF, - int RATIO) { - fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, nullptr, 0, 0); + int RATIO, Pg pg) { + fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, nullptr, 0, 0, pg); } extern "C" __global__ void __launch_bounds__(128) fn_qsa_pool_media(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, float eps, int R, int DI, int HALF, - int RATIO, const int* rope3, int n, int delta) { - fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, rope3, n, delta); + int RATIO, const int* rope3, int n, int delta, Pg pg) { + fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, rope3, n, delta, pg); } // attention._scores, program (row, 64 blocks): the heads' sum of relu(q . k) / sqrt(DI), rows past TOP (DI 128). extern "C" __global__ void __launch_bounds__(256) fn_qsa_scores(const __nv_bfloat16* __restrict__ IQ, const __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, float* __restrict__ SC, int NB, int HI, - int RATIO, int TOP) { + int RATIO, int TOP, Pg pg) { constexpr int DI = 128; __shared__ float qs[8][DI]; const int r = blockIdx.x, j = blockIdx.y; @@ -186,7 +200,7 @@ extern "C" __global__ void __launch_bounds__(256) fn_qsa_scores(const __nv_bfloa #pragma unroll for (int i = 0; i < 8; ++i) { const int b = j * 64 + warp * 8 + i; - raw[i] = b < complete ? *reinterpret_cast(POOLED + static_cast(b) * DI + 4 * lane) : make_uint2(0, 0); + raw[i] = b < complete ? *reinterpret_cast(POOLED + pg_at(pg, b, pg.ptok / RATIO, DI) + 4 * lane) : make_uint2(0, 0); } #pragma unroll for (int i = 0; i < 8; ++i) { @@ -210,7 +224,7 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, - int H, int HK, int G, int CH, int NCH, float SCALE, int IDW, int QSA) { + int H, int HK, int G, int CH, int NCH, float SCALE, int IDW, int QSA, Pg pg) { extern __shared__ __align__(16) unsigned char smem_raw[]; typedef __nv_bfloat16 Row[HD + 8]; Row* qs = reinterpret_cast(smem_raw); // [QR] the row's query heads @@ -218,6 +232,7 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, __nv_bfloat16 (*ps)[TK + 8] = reinterpret_cast<__nv_bfloat16 (*)[TK + 8]>(tiles_kv + NB * 2 * TK); // [QR] bf16 p __shared__ float red_m[4][QR], red_l[4][QR]; __shared__ int key_ids[NB][TK]; + __shared__ long long key_at[NB][TK]; // each key's row offset, its page looked up once a tile const int r = blockIdx.x, hk = blockIdx.y, c = blockIdx.z; const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; int n = POS0[0] + r + 1; @@ -236,15 +251,17 @@ __device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, auto load = [&](int t0, int buf) { if (threadIdx.x < TK) { const int ki = start + t0 + threadIdx.x; - key_ids[buf][threadIdx.x] = ki < n ? (sparse ? IDS[static_cast(r) * IDW + ki] : ki) : -1; + const int key = ki < n ? (sparse ? IDS[static_cast(r) * IDW + ki] : ki) : -1; + key_ids[buf][threadIdx.x] = key; + key_at[buf][threadIdx.x] = pg_at(pg, max(key, 0), pg.ptok, static_cast(HK) * HD) + hk * HD; } __syncthreads(); for (int pass = 0; pass < 2; ++pass) { const __nv_bfloat16* src = pass == 0 ? KC : VC; Row* dst = tiles_kv + (buf * 2 + pass) * TK; for (int e = threadIdx.x; e < TK * HD / 8; e += blockDim.x) { - const int k = e / (HD / 8), v = e % (HD / 8), key = key_ids[buf][k]; - qmm_frag::cp16z(&dst[k][8 * v], src + (static_cast(max(key, 0)) * HK + hk) * HD + 8 * v, key >= 0); + const int k = e / (HD / 8), v = e % (HD / 8); + qmm_frag::cp16z(&dst[k][8 * v], src + key_at[buf][k] + 8 * v, key_ids[buf][k] >= 0); } qmm_frag::commit(); } @@ -368,8 +385,8 @@ extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks(const __nv_bflo float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, - float SCALE, int IDW, int QSA) { - attn_chunks<1>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); + float SCALE, int IDW, int QSA, Pg pg) { + attn_chunks<1>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA, pg); } // attn_chunks with the next tile's loads in flight (a decode round's few rows): the same sums. @@ -378,8 +395,8 @@ extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks_ahead(const __n float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, - float SCALE, int IDW, int QSA) { - attn_chunks<2>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); + float SCALE, int IDW, int QSA, Pg pg) { + attn_chunks<2>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA, pg); } // attention._merge, program (row, head) of HD threads: the head's chunk partials merged in chunk order, o / l. diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..f1c855e47f 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -20,11 +20,11 @@ const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "flashnext_gdn flashnext_gdn_chain", .symbol = "flashnext_gdn_chain", .image = "flashnext_gdn", .threads = 1024 }, .usage = .{ .{ .regs = 51, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 62, .static = 2060, .bound = 1024 }, .{ .regs = 62, .static = 2060, .bound = 1024 } } }, .{ .launch = .{ .name = "flashnext_hc fn_add_streams", .symbol = "fn_add_streams", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = .{ .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 }, .{ .regs = 72, .static = 640, .bound = 128 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 71, .static = 768, .bound = 128 }, .{ .regs = 71, .static = 768, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = .{ .{ .regs = 71, .static = 896, .bound = 128 }, .{ .regs = 71, .static = 896, .bound = 128 }, .{ .regs = 71, .static = 896, .bound = 128 }, .{ .regs = 73, .static = 896, .bound = 128 }, .{ .regs = 73, .static = 896, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 1280, .bound = 128 }, .{ .regs = 72, .static = 1280, .bound = 128 }, .{ .regs = 72, .static = 1280, .bound = 128 }, .{ .regs = 71, .static = 1280, .bound = 128 }, .{ .regs = 71, .static = 1280, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_gate", .symbol = "fn_attn_gate", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_merge", .symbol = "fn_attn_merge", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_batch fn_gdn_front", .symbol = "fn_gdn_front", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 32, .static = 1024, .bound = 256 }, .{ .regs = 32, .static = 1024, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_batch fn_gdn_norm", .symbol = "fn_gdn_norm", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 24, .static = 516, .bound = 128 }, .{ .regs = 26, .static = 516, .bound = 128 }, .{ .regs = 26, .static = 516, .bound = 128 }, .{ .regs = 24, .static = 516, .bound = 128 }, .{ .regs = 24, .static = 516, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows", .symbol = "fn_gdn_rows", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, @@ -39,7 +39,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_slices", .symbol = "fn_hc_writeback_slices", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 256, .bound = 256 }, .{ .regs = 38, .static = 256, .bound = 256 }, .{ .regs = 38, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_hc fn_ple_conv", .symbol = "fn_ple_conv", .image = "flashnext_hc", .threads = 512 }, .usage = .{ .{ .regs = 23, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 } } }, .{ .launch = .{ .name = "flashnext_hc fn_ple_gate", .symbol = "fn_ple_gate", .image = "flashnext_hc", .threads = 512 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 35, .static = 128, .bound = 512 }, .{ .regs = 35, .static = 128, .bound = 512 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool", .symbol = "fn_qsa_pool", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool", .symbol = "fn_qsa_pool", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_attn fn_qsa_scores", .symbol = "fn_qsa_scores", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 46, .static = 4096, .bound = 256 }, .{ .regs = 46, .static = 4096, .bound = 256 }, .{ .regs = 48, .static = 4096, .bound = 256 }, .{ .regs = 48, .static = 4096, .bound = 256 } } }, .{ .launch = .{ .name = "flashnext_batch fn_qsa_select", .symbol = "fn_qsa_select", .image = "flashnext_batch", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 44, .static = 1160, .bound = 1024 }, .{ .regs = 40, .static = 1160, .bound = 1024 } } }, .{ .launch = .{ .name = "flashnext_mtp fn_row_best", .symbol = "fn_row_best", .image = "flashnext_mtp", .threads = 1024 }, .usage = .{ .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 } } }, @@ -129,7 +129,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "qwen_vision tf_qv_bias_cast", .symbol = "tf_qv_bias_cast", .image = "qwen_vision", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "qwen_vision tf_qv_splice", .symbol = "tf_qv_splice", .image = "qwen_vision", .threads = 256 }, .usage = .{ .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_prep_media", .symbol = "fn_attn_prep_media", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 } } }, - .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool_media", .symbol = "fn_qsa_pool_media", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool_media", .symbol = "fn_qsa_pool_media", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 } } }, // EXL3: the linear's instantiations at their most rows' shared bytes, the routed experts, Dense's tile candidate .{ .launch = .{ .name = "exl3_linear tf_exl3_rot_in", .symbol = "tf_exl3_rot_in", .image = "exl3_linear", .threads = 128 }, .usage = .{ .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 22, .static = 0, .bound = 128 }, .{ .regs = 22, .static = 0, .bound = 128 }, .{ .regs = 22, .static = 0, .bound = 128 }, .{ .regs = 22, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "exl3_linear tf_exl3_linear_2_2_4", .symbol = "tf_exl3_linear_2_2_4", .image = "exl3_linear", .threads = 128, .dynamic = 16384 }, .usage = .{ .{ .regs = 168, .static = 16, .bound = 128 }, .{ .regs = 168, .static = 16, .bound = 128 }, .{ .regs = 168, .static = 16, .bound = 128 }, .{ .regs = 152, .static = 16, .bound = 128 }, .{ .regs = 152, .static = 16, .bound = 128 } } }, diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index e93e1a2e30..965e5a7bcf 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -258,7 +258,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, defer ik_b.free(); var dummy = try cuda.DeviceBuffer.alloc(driver, 256); defer dummy.free(); - try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0); + try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0, .{}); var po_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); defer po_b.free(); @@ -268,7 +268,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, defer pl_b.free(); var o_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); defer o_b.free(); - try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1, false); + try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1, false, .{}); try tri.attnMerge(po_b.ptr, pm_b.ptr, pl_b.ptr, pos_b.ptr, o_b.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, false); var gated_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 5e79eee832..ae28ffa519 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -382,9 +382,10 @@ pub fn Pass(comptime Tri: type) type { const vc = self.seq.vc[ai].ptr; const ikc = self.seq.ikc[ai].ptr; const pooled = self.seq.pooled[ai].ptr; - if (self.seq.media) |m| try m.attn(self.stream, pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, eps, n) else try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0); + const pg = self.seq.pg; + if (self.seq.media) |m| try m.attn(self.stream, pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, eps, n, pg) else try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0, pg); // attention.qsa_pool for the blocks these rows complete, then each block of rows' scores and attention - if (self.seq.media) |m| try m.pooled(self.stream, ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n) else try self.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n); + if (self.seq.media) |m| try m.pooled(self.stream, ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n, pg) else try self.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n, pg); const nb = scoreStride(self.seq.cap); var r0: usize = 0; var b: usize = 0; @@ -398,7 +399,7 @@ pub fn Pass(comptime Tri: type) type { const pos = at.pos + b * 16; const iq = at.iq + r0 * index_heads * index_dim * 2; // rows past `top` complete blocks score theirs; _scores leaves the rest alone, so skip it when none is - if (ends / ratio > top) try self.tri.qsaScores(iq, pooled, pos, self.seq.scores, nb, m, @min(nb, cdiv(ends, ratio))); + if (ends / ratio > top) try self.tri.qsaScores(iq, pooled, pos, self.seq.scores, nb, m, @min(nb, cdiv(ends, ratio)), pg); var a: cuda.Args = .{}; a.add(self.seq.scores); a.add(pos); @@ -411,7 +412,7 @@ pub fn Pass(comptime Tri: type) type { a.add(cint(ratio)); try cuda.launch.launch(self.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, self.stream, &a); const chunks = cdiv(@min(ends, wt.dense_keys), 512); - try self.tri.attnChunks(at.q + r0 * q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true); + try self.tri.attnChunks(at.q + r0 * q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true, pg); try self.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * out_k * 2, s.qnk, s.qspr, m, kv_heads, true); } } diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index dc1749c657..d92c8d276c 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -497,6 +497,8 @@ pub const SeqState = struct { cap: usize, /// Block scores for an attention block's rows [min(rows, att_rows), scoreStride(cap)] fp32 (the engine's). scores: u64, + /// The attention caches' pages (pt 0: kc, vc, ikc and pooled are contiguous rows). + pg: @import("cuda_glue.zig").Pg = .{}, }; /// The block scores' row stride for a cache of `cap` tokens: its blocks rounded up to 16 (the captured _scores'). diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index e000c8f5df..cce104ba8d 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -22,6 +22,9 @@ fn cdiv(a: usize, b: usize) usize { return (a + b - 1) / b; } +/// flashnext_attn.cu's Pg: an int table a page, elements and positions a page (pt 0: contiguous rows). +pub const Pg = extern struct { pt: u64 = 0, ps: i64 = 0, ptok: c_int = 0, pad: c_int = 0 }; + /// The native kernels, by their extern "C" names. const Fns = struct { wb0: cuda.Function, @@ -177,17 +180,18 @@ pub const Glue = struct { } /// Text positions, a bf16 cache: the model's heads (cuda_prompt_wt.zig) normed, roped and stored. - pub fn attnPrep(g: *const Glue, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize) !void { + pub fn attnPrep(g: *const Glue, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, pg: Pg) !void { if (wt.head_dim != tile_head_dim) return error.RowShape; var a: cuda.Args = .{}; ptrs(&a, &.{ proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc }); a.add(eps); ints(&a, &.{ wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half }); + a.add(pg); try go(g.f.prep, s, .{ rows, wt.q_heads + wt.kv_heads + wt.index_heads + 1, 1 }, tile_head_dim, 0, &a); } /// Attention over 512-key chunks; `qsa` the sparse form (NCH 5, rows' key lists), else dense (NCH 3). - pub fn attnChunks(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { + pub fn attnChunks(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool, pg: Pg) !void { const group = wt.q_heads / kv; if (group > tile_heads) return error.RowShape; var a: cuda.Args = .{}; @@ -195,6 +199,7 @@ pub const Glue = struct { ints(&a, &.{ wt.q_heads, kv, group, 512, if (qsa) wt.nch else wt.dense_nch }); a.add(@as(f32, 0.0625)); ints(&a, &.{ wt.index_idw, @intFromBool(qsa) }); + a.add(pg); const ahead = rows <= ahead_rows; try go(if (ahead) g.f.chunks_ahead else g.f.chunks, s, .{ rows, kv, chunks }, 128, chunkSmem(if (ahead) 2 else 1), &a); } @@ -206,19 +211,21 @@ pub const Glue = struct { try go(g.f.merge, s, .{ rows, wt.q_heads, 1 }, tile_head_dim, 0, &a); } - pub fn qsaPool(g: *const Glue, s: cuda.Stream, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { + pub fn qsaPool(g: *const Glue, s: cuda.Stream, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { var a: cuda.Args = .{}; ptrs(&a, &.{ ikc, pooled, pos, k_scale, inv }); a.add(eps); ints(&a, &.{ rows, wt.index_dim, wt.half, wt.index_ratio }); + a.add(pg); try go(g.f.pool, s, .{ rows / wt.index_ratio + 2, 1, 1 }, wt.index_dim, 0, &a); } - pub fn qsaScores(g: *const Glue, s: cuda.Stream, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { + pub fn qsaScores(g: *const Glue, s: cuda.Stream, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize, pg: Pg) !void { if (wt.index_dim != 128 or wt.index_heads > 8) return error.RowShape; var a: cuda.Args = .{}; ptrs(&a, &.{ iq, pooled, pos, sc }); ints(&a, &.{ nb, wt.index_heads, wt.index_ratio, wt.index_top }); + a.add(pg); try go(g.f.scores, s, .{ rows, cdiv(blocks, 64), 1 }, 256, 0, &a); } diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 99e25c0c12..2c460da902 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -33,6 +33,7 @@ pub const native = @import("cuda_lanes.zig").Namespace(Tri); /// Our own glue kernels, which every launch takes when no captured set is loaded pub const Glue = glue_.Glue; const glue_ = @import("cuda_glue.zig"); +pub const Pg = glue_.Pg; const p = aot.ptr; @@ -131,8 +132,9 @@ pub const Tri = struct { } /// glue.attn_prep for text positions, bf16 KV (BITS 0, the scale pointers unread); `length` unused at MODE 0. - pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { - if (t.glue("_attn_prep")) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows); + pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize, pg: Pg) !void { + if (t.glue("_attn_prep")) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows, pg); + if (pg.pt != 0) return error.PagedKvNeedsGlue; try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ p("P", "*bf16", proj), p("POS0", "*i32", pos), @@ -155,8 +157,9 @@ pub const Tri = struct { } /// attention._chunks over 512-key chunks, bf16 KV; `qsa` the sparse form (NCH 5, QSA 1), else dense (NCH 3, QSA 0). - pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { - if (t.glue("_chunks")) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa); + pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool, pg: Pg) !void { + if (t.glue("_chunks")) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa, pg); + if (pg.pt != 0) return error.PagedKvNeedsGlue; try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ p("Q", "*bf16", q), p("KC", "*bf16", kc), @@ -188,8 +191,9 @@ pub const Tri = struct { } /// attention.qsa_pool for text rows (MODE 0): every complete block of 4 indexer keys pooled, normed and roped. - pub fn qsaPool(t: Tri, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { - if (t.glue("_pool")) |g| return g.qsaPool(t.s, ikc, pooled, pos, k_scale, inv, eps, rows); + pub fn qsaPool(t: Tri, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { + if (t.glue("_pool")) |g| return g.qsaPool(t.s, ikc, pooled, pos, k_scale, inv, eps, rows, pg); + if (pg.pt != 0) return error.PagedKvNeedsGlue; try t.run("_pool", .{ u(rows / 4 + 2), 1, 1 }, &.{ p("IKC", "*bf16", ikc), p("POOLED", "*bf16", pooled), @@ -203,8 +207,9 @@ pub const Tri = struct { } /// attention._scores: rows' complete blocks scored; `nb` is the row stride (a multiple of 16), `blocks` the grid. - pub fn qsaScores(t: Tri, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { - if (t.glue("_scores")) |g| return g.qsaScores(t.s, iq, pooled, pos, sc, nb, rows, blocks); + pub fn qsaScores(t: Tri, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize, pg: Pg) !void { + if (t.glue("_scores")) |g| return g.qsaScores(t.s, iq, pooled, pos, sc, nb, rows, blocks, pg); + if (pg.pt != 0) return error.PagedKvNeedsGlue; try t.run("_scores", .{ u(rows), cdiv(blocks, 64), 1 }, &.{ p("IQ", "*bf16", iq), p("POOLED", "*bf16", pooled), @@ -384,6 +389,7 @@ test { _ = mtp; _ = native; _ = @import("cuda_snap.zig"); + _ = @import("cuda_pages.zig"); _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); diff --git a/zig/src/families/flashnext/cuda_media.zig b/zig/src/families/flashnext/cuda_media.zig index 6e66da20a0..014df50784 100644 --- a/zig/src/families/flashnext/cuda_media.zig +++ b/zig/src/families/flashnext/cuda_media.zig @@ -2,6 +2,7 @@ const cuda = @import("cuda"); const shared = @import("qwen_vision"); const wt = @import("cuda_prompt_wt.zig"); +const Pg = @import("cuda_glue.zig").Pg; pub const State = struct { image: shared.media.Media, prep: cuda.Function, @@ -17,7 +18,7 @@ pub const View = struct { image: shared.media.View, prep: cuda.Function, pool: cuda.Function, - pub fn attn(v: View, stream: cuda.Stream, proj: u64, pos: u64, qw: u64, kw: u64, iw: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize) !void { + pub fn attn(v: View, stream: cuda.Stream, proj: u64, pos: u64, qw: u64, kw: u64, iw: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, pg: Pg) !void { var a: cuda.Args = .{}; for ([_]u64{ proj, pos, qw, kw, iw, inv, q, kc, vc, iq, ikc }) |x| a.add(x); a.add(eps); @@ -25,9 +26,10 @@ pub const View = struct { a.add(v.image.rope); a.add(@as(i32, @intCast(v.image.n))); a.add(v.image.delta); + a.add(pg); try cuda.launch.launch(v.prep, .{ .grid = .{ .x = @intCast(rows), .y = wt.q_heads + wt.kv_heads + wt.index_heads + 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, stream, &a); } - pub fn pooled(v: View, stream: cuda.Stream, ikc: u64, out: u64, pos: u64, w: u64, inv: u64, eps: f32, rows: usize) !void { + pub fn pooled(v: View, stream: cuda.Stream, ikc: u64, out: u64, pos: u64, w: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { var a: cuda.Args = .{}; for ([_]u64{ ikc, out, pos, w, inv }) |x| a.add(x); a.add(eps); @@ -35,6 +37,7 @@ pub const View = struct { a.add(v.image.rope); a.add(@as(i32, @intCast(v.image.n))); a.add(v.image.delta); + a.add(pg); try cuda.launch.launch(v.pool, .{ .grid = .{ .x = @intCast(rows / wt.index_ratio + 2), .y = 1, .z = 1 }, .block = .{ .x = wt.index_dim, .y = 1, .z = 1 } }, stream, &a); } }; diff --git a/zig/src/families/flashnext/cuda_media_test.zig b/zig/src/families/flashnext/cuda_media_test.zig index 0f58af79ba..302a670334 100644 --- a/zig/src/families/flashnext/cuda_media_test.zig +++ b/zig/src/families/flashnext/cuda_media_test.zig @@ -3,10 +3,12 @@ const std = @import("std"); const cuda = @import("cuda"); const media = @import("cuda_media.zig"); const wt = @import("cuda_prompt_wt.zig"); +const Pg = @import("cuda_glue.zig").Pg; const Capture = struct { var calls: usize = 0; var values: [8][21]u64 = undefined; var configs: [8][7]u32 = undefined; + var pages: [8]Pg = undefined; fn launch(f: cuda.abi.Function, gx: c_uint, gy: c_uint, gz: c_uint, bx: c_uint, by: c_uint, bz: c_uint, sm: c_uint, s: cuda.abi.Stream, args: ?[*]?*anyopaque, extra: ?[*]?*anyopaque) callconv(.c) c_int { std.debug.assert(extra == null and s == @as(cuda.abi.Stream, @ptrFromInt(99))); const id = @intFromPtr(f); @@ -15,6 +17,7 @@ const Capture = struct { const wide = if (id == 1) i < 3 else if (id == 2) i < 11 or i == 18 else i < 5 or i == 10; values[calls][i] = if (wide) @as(*const u64, @ptrCast(@alignCast(args.?[i].?))).* else @as(*const u32, @ptrCast(@alignCast(args.?[i].?))).*; } + if (id != 1) pages[calls] = @as(*const Pg, @ptrCast(@alignCast(args.?[count].?))).*; configs[calls] = .{ gx, gy, gz, bx, by, bz, sm }; calls += 1; return 0; @@ -32,14 +35,16 @@ test "the splice, attention and pooled indexer launches get every pointer, scala try m.image.inject(s, 4000, 2, 3, 4); try std.testing.expectEqualSlices(u64, &.{ 1000, 2000, 4000, 2, 20, 2560, 4 }, Capture.values[0][0..7]); try std.testing.expectEqual([7]u32{ 3, 40, 1, 256, 1, 1, 0 }, Capture.configs[0]); - try m.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 3); + try m.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 3, .{ .pt = 12, .ps = 13, .ptok = 16 }); const eps: u32 = @bitCast(@as(f32, 1e-6)); const delta: u32 = @bitCast(@as(i32, -2)); try std.testing.expectEqualSlices(u64, &.{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, eps, wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half, 3000, 20, delta }, &Capture.values[1]); try std.testing.expectEqual([7]u32{ 3, wt.q_heads + wt.kv_heads + wt.index_heads + 1, 1, 256, 1, 1, 0 }, Capture.configs[1]); - try m.pooled(s, 1, 2, 3, 4, 5, 1e-6, 8); + try std.testing.expectEqual(Pg{ .pt = 12, .ps = 13, .ptok = 16 }, Capture.pages[1]); + try m.pooled(s, 1, 2, 3, 4, 5, 1e-6, 8, .{}); try std.testing.expectEqualSlices(u64, &.{ 1, 2, 3, 4, 5, eps, 8, wt.index_dim, wt.half, wt.index_ratio, 3000, 20, delta }, Capture.values[2][0..13]); try std.testing.expectEqual([7]u32{ 8 / wt.index_ratio + 2, 1, 1, 128, 1, 1, 0 }, Capture.configs[2]); + try std.testing.expectEqual(Pg{}, Capture.pages[2]); const Row = struct { seq: struct { media: ?media.View }, pos: usize }; const rows = [_]Row{ .{ .seq = .{ .media = m }, .pos = 2 }, .{ .seq = .{ .media = null }, .pos = 2 }, .{ .seq = .{ .media = view(&d, 1100) }, .pos = 5 } }; try media.inject(null, &rows, s, 9000, 0, 3); diff --git a/zig/src/families/flashnext/cuda_pages.zig b/zig/src/families/flashnext/cuda_pages.zig new file mode 100644 index 0000000000..74ccaffc15 --- /dev/null +++ b/zig/src/families/flashnext/cuda_pages.zig @@ -0,0 +1,132 @@ +//! Flash Next's attention caches as pages: per attention layer, keys, values, indexer keys, then pooled rows (bf16). + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const Pg = @import("cuda_glue.zig").Pg; + +const kv_row = wt.kv_heads * wt.head_dim * 2; +const index_row = wt.index_dim * 2; +const ratio = wt.index_ratio; + +pub const Part = enum { key, value, index, pooled }; + +/// Bytes one row of `part` takes: a position's (a block's for pooled). +pub fn rowBytes(part: Part) usize { + return switch (part) { + .key, .value => kv_row, + .index, .pooled => index_row, + }; +} + +/// Rows of `part` a page of `tokens` positions holds. +pub fn rows(part: Part, tokens: usize) usize { + return if (part == .pooled) tokens / ratio else tokens; +} + +/// Bytes one attention layer's share of a page takes. +pub fn layerBytes(tokens: usize) usize { + return tokens * (2 * kv_row + index_row) + tokens / ratio * index_row; +} + +/// Bytes a page takes: every attention layer's share. +pub fn pageBytes(tokens: usize) usize { + return wt.attn_n * layerBytes(tokens); +} + +/// Where `part` of attention layer `layer` starts inside a page. +pub fn offset(tokens: usize, layer: usize, part: Part) usize { + const at = layer * layerBytes(tokens); + return at + switch (part) { + .key => 0, + .value => tokens * kv_row, + .index => 2 * tokens * kv_row, + .pooled => 2 * tokens * kv_row + tokens * index_row, + }; +} + +/// A page size the kernels take: whole blocks of the indexer's ratio. +pub fn valid(tokens: usize) bool { + return tokens > 0 and tokens % ratio == 0; +} + +/// TF_FLASHNEXT_KV_PAGE_TEST=n: each stream's caches read through reversed n-position pages (checks); 0 when unset. +pub fn pageTest() !usize { + const v = std.c.getenv("TF_FLASHNEXT_KV_PAGE_TEST") orelse return 0; + const n = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadFlashNextPageTest; + if (!valid(n)) return error.BadFlashNextPageTest; + return n; +} + +/// One stream's pages: page i of its positions at `base` + ids[i] pages. +pub const Paged = struct { + /// Where page 0 of the arena starts. + base: u64, + /// The stream's table on the device, an int a page (`ids` uploaded). + table: cuda.DeviceBuffer, + ids: []u32, + /// Positions a page holds. + tokens: usize, + /// The paged state owns `ids` (the page test's); a pool's lane maps its lease's. + own_ids: bool = true, + /// The page test's own arena. + arena: ?cuda.DeviceBuffer = null, + + pub fn pg(p: *const Paged) Pg { + return .{ .pt = p.table.ptr, .ps = @intCast(pageBytes(p.tokens) / 2), .ptok = @intCast(p.tokens) }; + } + + /// `part` of attention layer `layer` in page 0: the cache pointer a kernel takes beside pg(). + pub fn at(p: *const Paged, layer: usize, part: Part) u64 { + return p.base + offset(p.tokens, layer, part); + } + + /// Byte address of row `row` of `part` of `layer` (a position, or a block for pooled). + pub fn rowAt(p: *const Paged, layer: usize, part: Part, row: usize) u64 { + const per = rows(part, p.tokens); + return p.base + @as(u64, p.ids[row / per]) * pageBytes(p.tokens) + offset(p.tokens, layer, part) + (row % per) * rowBytes(part); + } + + /// Borrowed views of each attention layer's `part` in page 0 (what kc, vc, ikc and pooled hold when paged). + pub fn views(p: *const Paged, gpa: std.mem.Allocator, d: *const cuda.Driver, part: Part) ![]cuda.DeviceBuffer { + const out = try gpa.alloc(cuda.DeviceBuffer, wt.attn_n); + for (out, 0..) |*b, layer| b.* = .{ .d = d, .ptr = p.at(layer, part), .len = 0, .borrowed = true }; + return out; + } + + pub fn deinit(p: *Paged, gpa: std.mem.Allocator) void { + p.table.free(); + if (p.own_ids) gpa.free(p.ids); + if (p.arena) |*a| a.free(); + } +}; + +/// The page test's zeroed pages for `cap` positions in one buffer, in reverse order (no contiguous layout matches). +pub fn testPages(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize, tokens: usize) !Paged { + const n = (cap + tokens - 1) / tokens; + var arena = try cuda.DeviceBuffer.alloc(d, n * pageBytes(tokens)); + errdefer arena.free(); + try arena.fill8(0, stream.handle); + const ids = try gpa.alloc(u32, n); + errdefer gpa.free(ids); + for (ids, 0..) |*id, i| id.* = @intCast(n - 1 - i); + var table = try cuda.DeviceBuffer.alloc(d, n * 4); + errdefer table.free(); + try table.upload(0, std.mem.sliceAsBytes(ids)); + return .{ .base = arena.ptr, .table = table, .ids = ids, .tokens = tokens, .arena = arena }; +} + +test "a page's parts tile it in order and pooled rows follow their positions' page" { + const t = 16; + try std.testing.expectEqual(@as(usize, 0), offset(t, 0, .key)); + try std.testing.expectEqual(t * kv_row, offset(t, 0, .value)); + try std.testing.expectEqual(offset(t, 0, .pooled) + t / ratio * index_row, offset(t, 1, .key)); + try std.testing.expectEqual(pageBytes(t), offset(t, wt.attn_n - 1, .pooled) + t / ratio * index_row); + try std.testing.expect(!valid(6) and valid(8) and !valid(0)); + var ids = [_]u32{ 2, 0, 1 }; + const p: Paged = .{ .base = 1 << 20, .table = undefined, .ids = &ids, .tokens = t }; + // position 17 is page 1 (id 0) row 1; block 5 (positions 20..23) is page 1 row 1 of the pooled rows + try std.testing.expectEqual(p.base + offset(t, 3, .value) + kv_row, p.rowAt(3, .value, 17)); + try std.testing.expectEqual(p.base + offset(t, 3, .pooled) + index_row, p.rowAt(3, .pooled, 5)); + try std.testing.expectEqual(p.base + 2 * pageBytes(t) + offset(t, 0, .index), p.rowAt(0, .index, 0)); +} diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 0ff498c22d..f70726260c 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -113,6 +113,7 @@ pub fn pooledBytes(cap: usize) usize { } const round_graph = @import("cuda_round_graph.zig"); +const pages = @import("cuda_pages.zig"); /// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. pub const Engine = struct { @@ -295,13 +296,16 @@ pub const Engine = struct { errdefer freeOwned(e.gpa, conv); const state = try zeros(e.gpa, e.driver, e.stream.*, linear_n, state_n * 4); errdefer freeOwned(e.gpa, state); - const kc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + const test_pages = try pages.pageTest(); + var paged: ?pages.Paged = if (test_pages > 0) try pages.testPages(e.gpa, e.driver, e.stream.*, cap, test_pages) else null; + errdefer if (paged) |*p| p.deinit(e.gpa); + const kc = if (paged) |*p| try p.views(e.gpa, e.driver, .key) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, kc); - const vc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + const vc = if (paged) |*p| try p.views(e.gpa, e.driver, .value) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, vc); - const ikc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); + const ikc = if (paged) |*p| try p.views(e.gpa, e.driver, .index) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); errdefer freeOwned(e.gpa, ikc); - const pooled = try zeros(e.gpa, e.driver, e.stream.*, attn_n, pooledBytes(cap)); + const pooled = if (paged) |*p| try p.views(e.gpa, e.driver, .pooled) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, pooledBytes(cap)); errdefer freeOwned(e.gpa, pooled); var tail = try cuda.DeviceBuffer.alloc(e.driver, ple_tail_n * wide * 2); errdefer tail.free(); @@ -313,6 +317,7 @@ pub const Engine = struct { .vc = vc, .ikc = ikc, .pooled = pooled, + .paged = paged, .tail = tail, .history = .{ e.eos, e.eos }, .fed = 0, @@ -449,6 +454,8 @@ pub const Seq = struct { vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, pooled: []cuda.DeviceBuffer, + /// kc, vc, ikc and pooled as borrowed views of page 0 of these pages (null: each its own contiguous buffer). + paged: ?pages.Paged = null, tail: cuda.DeviceBuffer, history: [2]i64, fed: usize, @@ -459,7 +466,12 @@ pub const Seq = struct { /// The caches a batched pass reads and leaves; `scores` the engine's block scores (Engine.reserveScores). pub fn view(s: *Seq, scores: u64) batch.SeqState { - return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores }; + return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores, .pg = s.pg() }; + } + + /// The attention kernels' page argument (zero: contiguous caches). + pub fn pg(s: *const Seq) @import("cuda_glue.zig").Pg { + return if (s.paged) |*p| p.pg() else .{}; } pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { @@ -476,6 +488,7 @@ pub const Seq = struct { gpa.free(s.vc); gpa.free(s.ikc); gpa.free(s.pooled); + if (s.paged) |*p| p.deinit(gpa); if (s.tail.ptr != 0) s.tail.free(); gpa.free(s.logits); } diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig index 36d364152f..58f76a512f 100644 --- a/zig/src/families/flashnext/cuda_prompt_step.zig +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -235,16 +235,16 @@ pub fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Strea try qmm.partials(driver, stream.*, s.gout.ptr, s.gxs.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, gdn.value_dim, out_sk); } -pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { +pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32, pg: @import("cuda_glue.zig").Pg) !void { try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); // stream-ordered (queued kernels still read s.pos), then waited for: `pos` lives on this frame try s.pos.uploadAsync(0, std.mem.asBytes(&pos), stream.handle); try stream.synchronize(); - try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); + try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0, pg); // attention.attention, dense: chunks of 512 keys up to the row's position (the row-by-row path stays dense) const keys: usize = @as(usize, @intCast(pos)) + 1; if (keys > wt.dense_nch * 512) return error.UnexpectedTensor; - try tri.attnChunks(s.q.ptr, kc, vc, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512, false); + try tri.attnChunks(s.q.ptr, kc, vc, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512, false, pg); try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, false); try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); diff --git a/zig/src/families/flashnext/cuda_rows.zig b/zig/src/families/flashnext/cuda_rows.zig index 89095b53ee..49df3faa45 100644 --- a/zig/src/families/flashnext/cuda_rows.zig +++ b/zig/src/families/flashnext/cuda_rows.zig @@ -75,7 +75,7 @@ pub const Rows = struct { // dense attention only: the reference stops where the indexer would start choosing if (seq.fed + ids.len > weights.dense_nch * 512) return error.UnexpectedTensor; try r.grow(e.stream.*, ids.len); - var pass = Pass{ .gpa = e.gpa, .io = e.io, .driver = e.driver, .stream = e.stream, .model_dir = e.model_dir, .s = &r.scratch, .table = e.table, .mode = r.mode, .history = &seq.history, .conv = seq.conv, .state = seq.state, .kc = seq.kc, .vc = seq.vc, .ikc = seq.ikc, .tail = seq.tail, .mixer = e.mixer, .head_q = e.head_q, .logits = seq.logits, .plan_fn = r.plan_fn, .up_fn = r.up_fn, .down_fn = r.down_fn }; + var pass = Pass{ .gpa = e.gpa, .io = e.io, .driver = e.driver, .stream = e.stream, .model_dir = e.model_dir, .s = &r.scratch, .table = e.table, .mode = r.mode, .history = &seq.history, .conv = seq.conv, .state = seq.state, .kc = seq.kc, .vc = seq.vc, .ikc = seq.ikc, .pg = seq.pg(), .tail = seq.tail, .mixer = e.mixer, .head_q = e.head_q, .logits = seq.logits, .plan_fn = r.plan_fn, .up_fn = r.up_fn, .down_fn = r.down_fn }; @memset(pass.mode, 0); for (ids, 0..) |token, t| try embedRow(Tri, &pass, tri, token, pass.s.h.ptr + @as(u64, t) * wide * 2); for (e.layers[0..e.n_layers], 0..) |*lw, layer| try runLayer(Tri, &pass, tri, lw, layer, ids, seq.fed); @@ -99,6 +99,7 @@ const Pass = struct { kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, + pg: @import("cuda_glue.zig").Pg, tail: cuda.DeviceBuffer, mixer: HcW, head_q: Q4, @@ -137,7 +138,7 @@ fn runLayer(comptime Tri: type, p: *Pass, tri: Tri, lw: *Layer, layer: usize, id try step.gdnStep(p.gpa, p.driver, p.stream, p.s, lw.proj, lw.out, p.conv[li], lw.cw.ptr, lw.a_log.ptr, lw.dt.ptr, lw.norm.ptr, &p.state[li]); } else { const ai = weights.attnIndex(layer); - try step.attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); + try step.attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t), p.pg); } try step.writeback(tri, p.s, h, 4, inj, y, wt); try step.readout(tri, p.s, lw.mlp_hc, h, inj); diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index 0e551809b3..79ccf2b04b 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -6,6 +6,7 @@ const prompt = @import("cuda_prompt.zig"); const wt = @import("cuda_prompt_wt.zig"); const gdn = @import("cuda_gdn.zig"); const Draft = @import("cuda_mtp_draft.zig").Draft; +const pages = @import("cuda_pages.zig"); const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; const state_bytes = wt.state_n * 4; @@ -59,8 +60,25 @@ fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, pieces[7..11].* = attnPieces((&h.kc)[0..1], (&h.vc)[0..1], (&h.ikc)[0..1], (&h.pooled)[0..1], at); count = 11; } - for (pieces[0..count]) |p| for (p.bufs) |b| { + for (pieces[0..count], 0..) |p, pi| for (p.bufs, 0..) |b, layer| { if (p.n == 0) continue; + if (pi >= 2 and pi < 6) if (seq.paged) |*pg| { + // the same bytes in the same order as a contiguous cache's, a run a page + const part: pages.Part = @fromBackingInt(@intCast(pi - 2)); + const row = pages.rowBytes(part); + const per = pages.rows(part, pg.tokens); + const n = p.n / row; + var r: usize = 0; + while (r < n) { + const run = @min(per - r % per, n - r); + const at_row = pg.rowAt(layer, part, r); + const dst, const src = if (way == .save) .{ base + off, at_row } else .{ at_row, base + off }; + try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, run * row, stream.handle), "cuMemcpyDtoDAsync"); + off += run * row; + r += run; + } + continue; + }; const dst, const src = if (way == .save) .{ base + off, b.ptr } else .{ b.ptr, base + off }; try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, p.n, stream.handle), "cuMemcpyDtoDAsync"); off += p.n; diff --git a/zig/tests/cuda/contract_fn_pages.zig b/zig/tests/cuda/contract_fn_pages.zig new file mode 100644 index 0000000000..713670f7b8 --- /dev/null +++ b/zig/tests/cuda/contract_fn_pages.zig @@ -0,0 +1,344 @@ +//! flashnext_attn.cu's caches through a page table: prep, pooling, scores and attention, paged == contiguous, float64. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const q_heads = 24; +const kv_heads = 2; +const hd = 256; +const index_heads = 4; +const index_dim = 128; +const half = 32; +const proj_n = 13952; +const ratio = 4; +const idw = 2052; +const nch = 5; +const top = 512; +const eps: f32 = 1e-6; +const layers = 2; // a page's attention layers: the kernels run layer 1, so the layer offset counts + +const kv_row = kv_heads * hd * 2; +const index_row = index_dim * 2; + +const Part = enum { key, value, index, pooled }; + +fn rowBytes(p: Part) usize { + return if (p == .key or p == .value) kv_row else index_row; +} + +fn perPage(p: Part, t: usize) usize { + return if (p == .pooled) t / ratio else t; +} + +/// A page's layout, written out apart from cuda_pages.zig: per layer keys, values, indexer keys, then pooled rows. +fn layerBytes(t: usize) usize { + return t * (2 * kv_row + index_row) + t / ratio * index_row; +} + +fn partOffset(t: usize, p: Part) usize { + return layerBytes(t) * 1 + switch (p) { + .key => 0, + .value => t * kv_row, + .index => 2 * t * kv_row, + .pooled => 2 * t * kv_row + t * index_row, + }; +} + +/// Two streams' cached positions, the rows each adds, positions a page. +const Case = struct { lens: [2]usize, rows: usize, ptok: usize }; +const cases = [_]Case{ + .{ .lens = .{ 0, 3 }, .rows = 1, .ptok = 16 }, + .{ .lens = .{ 127, 128 }, .rows = 4, .ptok = 16 }, + .{ .lens = .{ 700, 33 }, .rows = 3, .ptok = 48 }, // a page no 512-key chunk or 32-key tile divides + .{ .lens = .{ 2100, 2300 }, .rows = 16, .ptok = 128 }, // block scores and dense attention together + .{ .lens = .{ 4095, 5000 }, .rows = 17, .ptok = 256 }, + .{ .lens = .{ 9000, 300 }, .rows = 2, .ptok = 512 }, +}; + +const Fixture = struct { + c: Case, + cap: usize, + caches: [2][4][]u8, // each stream's logical key, value, indexer key and pooled rows + p: [2][]u16, // projections [rows][proj_n] + qw: []f32, + kw: []f32, + iw: []f32, + ikw: []f32, + inv: []f32, + ids: [2][]i32, // sparse rows' key lists [rows][idw] + nk: [2][]i32, + ids_pages: [2][]const u32, // each stream's pages in the shared arena, shuffled across both + pages: usize, // the arena's pages (two spare) + + fn pageBytes(f: Fixture) usize { + return layers * layerBytes(f.c.ptok); + } + + /// Byte offset of `row` of `part` of stream `st` in the arena under `ids`. + fn paged(f: Fixture, ids: [2][]const u32, st: usize, part: Part, row: usize) usize { + const per = perPage(part, f.c.ptok); + return ids[st][row / per] * f.pageBytes() + partOffset(f.c.ptok, part) + (row % per) * rowBytes(part); + } + + fn rowsOf(f: Fixture, part: Part) usize { + return perPage(part, f.cap); + } + + fn make(a: std.mem.Allocator, c: Case, seed: u64) !Fixture { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const most = @max(c.lens[0], c.lens[1]) + c.rows; + const cap = (most + c.ptok - 1) / c.ptok * c.ptok; + var f: Fixture = undefined; + f.c = c; + f.cap = cap; + for (0..2) |st| { + for (std.enums.values(Part), 0..) |part, pi| { + f.caches[st][pi] = try a.alloc(u8, f.rowsOf(part) * rowBytes(part)); + fillBf(rng, std.mem.bytesAsSlice(u16, f.caches[st][pi]), 1); + } + f.p[st] = try a.alloc(u16, c.rows * proj_n); + fillBf(rng, f.p[st], 2); + f.ids[st] = try a.alloc(i32, c.rows * idw); + f.nk[st] = try a.alloc(i32, c.rows); + for (0..c.rows) |r| { + const keys = c.lens[st] + r + 1; + f.nk[st][r] = @intCast(@max(1, @min(idw, keys) -| r % 3)); + for (0..idw) |i| f.ids[st][r * idw + i] = @intCast((i * 1237 + r * 31 + st * 7) % keys); + } + } + f.qw = try floats(a, rng, hd, 0.5, 1.5); + f.kw = try floats(a, rng, hd, 0.5, 1.5); + f.iw = try floats(a, rng, index_dim, 0.5, 1.5); + f.ikw = try floats(a, rng, index_dim, 0.5, 1.5); + f.inv = try a.alloc(f32, half); + for (f.inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / half)); + f.pages = 2 * (cap / c.ptok) + 2; + const order = try a.alloc(u32, f.pages); + for (order, 0..) |*o, i| o.* = @intCast(i); + rng.shuffle(u32, order); + const n = cap / c.ptok; + f.ids_pages = .{ order[0..n], order[n .. 2 * n] }; + return f; + } + + /// Both streams' logical caches at their pages under `ids`; every other byte (layer 0, spare pages) random. + fn arena(f: Fixture, a: std.mem.Allocator, rng: std.Random, ids: [2][]const u32) ![]u8 { + const out = try a.alloc(u8, f.pages * f.pageBytes()); + rng.bytes(out); + for (0..2) |st| for (std.enums.values(Part), 0..) |part, pi| { + const n = rowBytes(part); + for (0..f.rowsOf(part)) |row| @memcpy(out[f.paged(ids, st, part, row)..][0..n], f.caches[st][pi][row * n ..][0..n]); + }; + return out; + } + + /// Stream `st`'s logical caches read back out of the arena's pages. + fn unpage(f: Fixture, a: std.mem.Allocator, all: []const u8, ids: [2][]const u32, st: usize) ![4][]u8 { + var out: [4][]u8 = undefined; + for (std.enums.values(Part), 0..) |part, pi| { + const n = rowBytes(part); + out[pi] = try a.alloc(u8, f.rowsOf(part) * n); + for (0..f.rowsOf(part)) |row| @memcpy(out[pi][row * n ..][0..n], all[f.paged(ids, st, part, row)..][0..n]); + } + return out; + } +}; + +fn fillBf(rng: std.Random, out: []align(1) u16, scale: f32) void { + for (out) |*x| x.* = @truncate(@as(u32, @bitCast((rng.float(f32) * 2 - 1) * scale)) >> 16); +} + +fn floats(a: std.mem.Allocator, rng: std.Random, n: usize, lo: f32, hi: f32) ![]f32 { + const v = try a.alloc(f32, n); + for (v) |*x| x.* = lo + rng.float(f32) * (hi - lo); + return v; +} + +fn bf(x: u16) f64 { + return @as(f32, @bitCast(@as(u32, x) << 16)); +} + +fn scoreStride(cap: usize) usize { + return ((cap + ratio - 1) / ratio + 15) / 16 * 16; +} + +/// One layout's results a stream: queries, indexer queries, block scores, attention rows and the caches after. +const Result = struct { q: [2][]u8, iq: [2][]u8, sc: [2][]u8, dense: [2]?[]u8, sparse: [2][]u8, caches: [2][4][]u8 }; + +/// Both streams' rows through prep, pool, scores and attention on layer 1: contiguous, or paged by `tables` (`given`). +fn launch(gpu: check.Gpu, s: cuda.Stream, g: *const flash.Glue, f: Fixture, tables: ?[2][]const u32, given: ?[2][]const u32, a: std.mem.Allocator, seed: u64) !Result { + var rig = rig_.Rig{ .gpu = gpu, .s = s }; + defer rig.deinit(); + var prng = std.Random.DefaultPrng.init(seed); + const rows = f.c.rows; + const nb = scoreStride(f.cap); + const qw = try rig.input(std.mem.sliceAsBytes(f.qw)); + const kw = try rig.input(std.mem.sliceAsBytes(f.kw)); + const iw = try rig.input(std.mem.sliceAsBytes(f.iw)); + const ikw = try rig.input(std.mem.sliceAsBytes(f.ikw)); + const inv = try rig.input(std.mem.sliceAsBytes(f.inv)); + const po = try rig.scratch(rows * nch * q_heads * hd * 4); + const pm = try rig.scratch(rows * nch * q_heads * 4); + const pl = try rig.scratch(rows * nch * q_heads * 4); + var arena_ptr: u64 = 0; + if (tables) |t| { + const host = try f.arena(a, prng.random(), t); + arena_ptr = try rig.scratch(host.len); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = arena_ptr, .len = host.len }).upload(0, host); + } + var res: Result = undefined; + var flat: [2][4]u64 = undefined; + var outs: [2][5]u64 = undefined; + for (0..2) |st| { + const pos0: i32 = @intCast(f.c.lens[st]); + const ends = f.c.lens[st] + rows; + const pos = try rig.input(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); + const p = try rig.input(std.mem.sliceAsBytes(f.p[st])); + const ids = try rig.input(std.mem.sliceAsBytes(f.ids[st])); + const nk = try rig.input(std.mem.sliceAsBytes(f.nk[st])); + const ones = try a.alloc(i32, rows); + @memset(ones, 1); + const spr1 = try rig.input(std.mem.sliceAsBytes(ones)); + @memset(ones, 0); + const spr0 = try rig.input(std.mem.sliceAsBytes(ones)); + var pg: flash.Pg = .{}; + var at: [4]u64 = undefined; + if (tables) |t| { + pg = .{ .pt = try rig.input(std.mem.sliceAsBytes((given orelse t)[st])), .ps = @intCast(f.pageBytes() / 2), .ptok = @intCast(f.c.ptok) }; + for (std.enums.values(Part), 0..) |part, pi| at[pi] = arena_ptr + partOffset(f.c.ptok, part); + } else for (0..4) |pi| { + at[pi] = try rig.scratch(f.caches[st][pi].len); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = at[pi], .len = f.caches[st][pi].len }).upload(0, f.caches[st][pi]); + } + flat[st] = at; + outs[st] = .{ try rig.output(rows * q_heads * hd * 2), try rig.output(rows * index_heads * index_dim * 2), try rig.output(rows * nb * 4), try rig.output(rows * q_heads * hd * 2), try rig.output(rows * q_heads * hd * 2) }; + const o = outs[st]; + try g.attnPrep(s, p, pos, qw, kw, iw, inv, o[0], at[0], at[1], o[1], at[2], eps, rows, pg); + try g.qsaPool(s, at[2], at[3], pos, ikw, inv, eps, rows, pg); + if (ends / ratio > top) try g.qsaScores(s, o[1], at[3], pos, o[2], nb, rows, @min(nb, (ends + ratio - 1) / ratio), pg); + if (ends <= nch * 512) { + try g.attnChunks(s, o[0], at[0], at[1], pos, po, pm, pl, ids, nk, spr0, rows, kv_heads, (ends + 511) / 512, true, pg); + try g.attnMerge(s, po, pm, pl, pos, o[3], nk, spr0, rows, kv_heads, true); + } + try g.attnChunks(s, o[0], at[0], at[1], pos, po, pm, pl, ids, nk, spr1, rows, kv_heads, (idw + 511) / 512, true, pg); + try g.attnMerge(s, po, pm, pl, pos, o[4], nk, spr1, rows, kv_heads, true); + } + for (0..2) |st| { + const o = outs[st]; + res.q[st] = try a.dupe(u8, try rig.bytes(o[0], rows * q_heads * hd * 2)); + res.iq[st] = try a.dupe(u8, try rig.bytes(o[1], rows * index_heads * index_dim * 2)); + res.sc[st] = try a.dupe(u8, try rig.bytes(o[2], rows * nb * 4)); + res.dense[st] = if (f.c.lens[st] + rows <= nch * 512) try a.dupe(u8, try rig.bytes(o[3], rows * q_heads * hd * 2)) else null; + res.sparse[st] = try a.dupe(u8, try rig.bytes(o[4], rows * q_heads * hd * 2)); + } + if (tables) |t| { + const all = try a.dupe(u8, try rig.bytes(arena_ptr, f.pages * f.pageBytes())); + for (0..2) |st| res.caches[st] = try f.unpage(a, all, t, st); + } else for (0..2) |st| for (0..4) |pi| { + res.caches[st][pi] = try a.dupe(u8, try rig.bytes(flat[st][pi], f.caches[st][pi].len)); + }; + try check.expect(try rig.intact() and !rig.overlap(), "flashnext pages: a launch wrote past a buffer", .{}); + return res; +} + +/// Each attention row against float64 softmax(q k / 16) v over the caches the pass left (p . v takes p as bf16). +fn agrees(f: Fixture, r: Result, a: std.mem.Allocator) !bool { + const most = @max(f.c.lens[0], f.c.lens[1]) + f.c.rows + idw; + const sc = try a.alloc(f64, most); + for (0..2) |st| for ([_]bool{ false, true }) |sparse| { + const out_bytes = if (sparse) r.sparse[st] else r.dense[st] orelse continue; + const out: []align(1) const u16 = std.mem.bytesAsSlice(u16, out_bytes); + const q: []align(1) const u16 = std.mem.bytesAsSlice(u16, r.q[st]); + const kc: []align(1) const u16 = std.mem.bytesAsSlice(u16, r.caches[st][0]); + const vc: []align(1) const u16 = std.mem.bytesAsSlice(u16, r.caches[st][1]); + for (0..f.c.rows) |row| for (0..q_heads) |h| { + const kvh = h / (q_heads / kv_heads); + const n: usize = if (sparse) @intCast(f.nk[st][row]) else f.c.lens[st] + row + 1; + var hi = -std.math.inf(f64); + for (0..n) |j| { + const key: usize = if (sparse) @intCast(f.ids[st][row * idw + j]) else j; + var dot: f64 = 0; + for (0..hd) |d| dot += bf(q[(row * q_heads + h) * hd + d]) * bf(kc[(key * kv_heads + kvh) * hd + d]); + sc[j] = dot * 0.0625; + hi = @max(hi, sc[j]); + } + var l: f64 = 0; + for (sc[0..n]) |*v| { + v.* = @exp(v.* - hi); + l += v.*; + } + for (0..hd) |d| { + var o: f64 = 0; + var abs: f64 = 0; + for (0..n) |j| { + const key: usize = if (sparse) @intCast(f.ids[st][row * idw + j]) else j; + const v = bf(vc[(key * kv_heads + kvh) * hd + d]); + o += sc[j] * v; + abs += sc[j] * @abs(v); + } + const want = o / l; + const got = bf(out[(row * q_heads + h) * hd + d]); + if (!(@abs(got - want) <= abs / l * 0x1p-8 + 0.5 * @abs(want) * 0x1p-7 + 1e-6)) return false; + } + }; + }; + return true; +} + +fn same(x: []const u8, y: []const u8) bool { + return std.mem.eql(u8, x, y); +} + +fn sameResult(x: Result, y: Result) bool { + for (0..2) |st| { + if (!same(x.q[st], y.q[st]) or !same(x.iq[st], y.iq[st]) or !same(x.sc[st], y.sc[st]) or !same(x.sparse[st], y.sparse[st])) return false; + if ((x.dense[st] == null) != (y.dense[st] == null)) return false; + if (x.dense[st]) |d| if (!same(d, y.dense[st].?)) return false; + for (0..4) |pi| if (!same(x.caches[st][pi], y.caches[st][pi])) return false; + } + return true; +} + +pub fn run(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + var red_table = false; + var red_bound = false; + for (cases, 0..) |c, i| { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const f = try Fixture.make(a, c, 0xf9a6 + i); + const flat = try launch(gpu, s, &g, f, null, null, a, 1); + try check.expect(try agrees(f, flat, a), "flashnext pages: contiguous attention off float64 at lens {any}, rows {d}", .{ c.lens, c.rows }); + for (0..3) |rep| { + const pg = try launch(gpu, s, &g, f, f.ids_pages, null, a, 2 + rep); + try check.expect(sameResult(pg, flat), "flashnext pages: paged bytes differ from contiguous at lens {any}, rows {d}, page {d} (repeat {d})", .{ c.lens, c.rows, c.ptok, rep }); + } + // red controls: stream 1's first and last pages swapped in the table only (the pages stay); a planted row + const n1 = (c.lens[1] + c.rows + c.ptok - 1) / c.ptok; + if (n1 >= 2) { + const swapped = try a.dupe(u32, f.ids_pages[1]); + std.mem.swap(u32, &swapped[0], &swapped[n1 - 1]); + const bad = try launch(gpu, s, &g, f, f.ids_pages, .{ f.ids_pages[0], swapped }, a, 9); + const moved = !same(bad.sparse[1], flat.sparse[1]); + red_table = red_table or moved; + try check.expect(moved, "flashnext pages: a swapped page table still gave contiguous bytes (control missed) at lens {any}", .{c.lens}); + } + var bad_r = flat; + bad_r.sparse[0] = try a.dupe(u8, flat.sparse[0]); + const planted: usize = if (c.rows >= 2) q_heads * hd + 5 else 5; + std.mem.writeInt(u16, bad_r.sparse[0][2 * planted ..][0..2], 0x4040, .little); // 3.0 where |o| <= 1 + const caught = !(try agrees(f, bad_r, a)); + red_bound = red_bound or caught; + try check.expect(caught, "flashnext pages: a planted attention row passed float64 (control missed)", .{}); + check.pass("flashnext pages: lens {any}, {d} rows a stream, {d}-position pages: paged == contiguous (queries, indexer queries, block scores, dense and sparse rows, stores, pooled), three repeats, float64 bound, guards", .{ c.lens, c.rows, c.ptok }); + } + try check.expect(red_table and red_bound, "flashnext pages: a control never fired", .{}); + check.pass("flashnext pages: controls RED as planted (swapped table, planted row)", .{}); +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index b929e4be03..4b772b43ed 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -283,6 +283,7 @@ pub fn run(gpu: check.Gpu) !void { try @import("contract_qwen_served.zig").run(gpu); try @import("contract_qwen_pages.zig").run(gpu); try @import("contract_qwen_learning.zig").numerical(gpu); + try @import("contract_fn_pages.zig").run(gpu); const lim = try live(gpu); const t2 = check.now(gpu.io); std.debug.print("SUMMARY {d} kernels: {d} pass gates 1-4, {d} fail, {d} known findings; {d} controls: {d} red as planted, {d} missed; {d} limit problems; gates {d} ms, limits {d} ms\n", .{ t.kernels, t.passed, t.failed, t.known, t.controls, t.caught, t.missed, lim, @divTrunc(t1 - t0, std.time.ns_per_ms), @divTrunc(t2 - t1, std.time.ns_per_ms) }); diff --git a/zig/tests/cuda/flashnext_attn_tests.zig b/zig/tests/cuda/flashnext_attn_tests.zig index 806a7207bf..744792d380 100644 --- a/zig/tests/cuda/flashnext_attn_tests.zig +++ b/zig/tests/cuda/flashnext_attn_tests.zig @@ -77,7 +77,7 @@ fn prep(r: *Rig, g: *const flash.Glue) !void { kcs[c] = try r.out(cap * kv_heads * head_dim * 2); vcs[c] = try r.out(cap * kv_heads * head_dim * 2); ikcs[c] = try r.out(cap * index_dim * 2); - try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, qs[c], kcs[c], vcs[c], iqs[c], ikcs[c], eps, m); + try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, qs[c], kcs[c], vcs[c], iqs[c], ikcs[c], eps, m, .{}); } const last = counts.len - 1; const q = try r.back(u16, qs[last], most * q_heads * head_dim); @@ -142,7 +142,7 @@ fn indexer(r: *Rig, g: *const flash.Glue) !void { const dikc = try r.dev(std.mem.sliceAsBytes(ikc)); const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); const pooled = try r.out((cap / 4) * index_dim * 2); - try g.qsaPool(r.s, dikc, pooled, dpos, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(inv)), eps, rows); + try g.qsaPool(r.s, dikc, pooled, dpos, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(inv)), eps, rows, .{}); const got = try r.back(u16, pooled, (cap / 4) * index_dim); var b: usize = pos0 / 4; while (4 * b + 4 <= pos0 + rows) : (b += 1) { @@ -171,7 +171,7 @@ fn indexer(r: *Rig, g: *const flash.Glue) !void { var scs: [2]u64 = undefined; for ([_]usize{ 1, rows }, 0..) |m, c| { scs[c] = try r.out(m * nb * 4); - try g.qsaScores(r.s, diq, dkeys, dpos, scs[c], nb, m, (pos0 + rows) / 4); + try g.qsaScores(r.s, diq, dkeys, dpos, scs[c], nb, m, (pos0 + rows) / 4, .{}); } const sc = try r.back(f32, scs[1], rows * nb); for (0..rows) |row| { @@ -227,7 +227,7 @@ fn attention(r: *Rig, g: *const flash.Glue, sparse: bool) !void { for (counts, 0..) |m, c| { outs[c] = try r.out(m * q_heads * head_dim * 2); const chunks = if (sparse) (idw + 511) / 512 else (pos0 + m + 511) / 512; - try g.attnChunks(r.s, dq, dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, m, kv_heads, chunks, true); + try g.attnChunks(r.s, dq, dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, m, kv_heads, chunks, true, .{}); try g.attnMerge(r.s, po, pm, pl, dpos, outs[c], dnk, dspr, m, kv_heads, true); } const got = try r.back(u16, outs[counts.len - 1], rows * q_heads * head_dim); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 8380cef677..d9f648f416 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -106,6 +106,7 @@ const usage = \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) + \\ flashnext-pages the attention glue's caches through a page table: paged == contiguous, float64, controls \\ info driver, device and library versions \\ smoke copies, fills, launches, argument packing, module globals \\ shared-context Context.shared and Stream.borrowed as a library inside a host's process sees them @@ -253,6 +254,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-pages")) return @import("contract_fn_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "info")) return info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); if (std.mem.eql(u8, cmd, "shared-context")) return runtime_tests.sharedContext(gpu); From 336cb3f7937f87867a995129a8242f94d44892c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:05:05 +0200 Subject: [PATCH 10/11] cuda flashnext: the streams beside the first take KV pages from the core pool (TF_FLASHNEXT_KV_PAGES) TF_FLASHNEXT_KV_PAGES=n (off by default, n a multiple of 4): one stream keeps its own caches as before; every other stream takes a pool lane (cuda_kv_lanes.zig). A lane holds only its page table while idle. Its stream's attention caches are n-position pages of one arena, as many as its prompt, reply and a drafted window need; with the MTP head a page holds the head's cache too, as the layer after the model's. Its DeltaNet states, convolution and n-gram tails and the head's buffers (window streams, scores, saved rows, spare states; Draft.initIn) are one run of consecutive pool pages, claimed when the stream opens (Pool.claimRun) and given back when it closes, so only addresses move. The native host plans the pool from Loaded.layout (fixed: the table and image rows; a token's pages; the run) and admits through fits; a request with no lane, pages or free run waits (LaneMemory). A lane's pages and run go back once the engine's stream drains. Prompt-cache copies of the head's cache follow its pages. TF_FLASHNEXT_KV_PAGES_ALL=1 puts every stream on pages (checks). Two ranks, --slide and a captured Triton set refuse it. GB10, nvidia/Qwen3.8-Flash-Next-FP4-Local, n=256, --parallel 4 --context 40960, MTP: "3 more taking 594 pages of 256 rows (4.36 GiB) as they need them, 0 MiB fixed each" (was 369 MiB fixed each, the head's cache and buffers and the states held from load). With MTP: concurrent == alone 4/4, serial 1/1, outputs equal to contiguous; drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh (7,147 tokens), equal to contiguous. Pages only (--parallel 5: 4 lanes, 792 pages, 8 requests, some waiting): the same bytes, a kept state (and its head cache) restored into a lane. Without drafts: solo bytes == contiguous, c=4 65.0/64.5 against 65.4/65.0 (measured before states moved into runs). --- zig/src/families/flashnext/cuda_kv_lanes.zig | 268 ++++++++++++++++++ zig/src/families/flashnext/cuda_lanes.zig | 30 +- .../flashnext/cuda_lanes_admission_test.zig | 1 + .../families/flashnext/cuda_lanes_load.zig | 9 + zig/src/families/flashnext/cuda_launch.zig | 1 + zig/src/families/flashnext/cuda_mtp_draft.zig | 42 +++ zig/src/families/flashnext/cuda_mtp_ops.zig | 2 +- zig/src/families/flashnext/cuda_pages.zig | 13 +- zig/src/families/flashnext/cuda_prompt.zig | 23 +- zig/src/families/flashnext/cuda_snap.zig | 9 +- 10 files changed, 384 insertions(+), 14 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_kv_lanes.zig diff --git a/zig/src/families/flashnext/cuda_kv_lanes.zig b/zig/src/families/flashnext/cuda_kv_lanes.zig new file mode 100644 index 0000000000..d4d988a1e0 --- /dev/null +++ b/zig/src/families/flashnext/cuda_kv_lanes.zig @@ -0,0 +1,268 @@ +//! Flash Next's streams beside the first take attention-cache pages from the core pool (TF_FLASHNEXT_KV_PAGES=n). +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const prompt = @import("cuda_prompt.zig"); +const pages = @import("cuda_pages.zig"); +const mtp = @import("cuda_mtp.zig"); +const snap = @import("cuda_snap.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +/// Rows past a request's prompt and reply its pages hold: one drafted window's. +pub const window_rows: u32 = mtp.max_rows; + +/// TF_FLASHNEXT_KV_PAGES=n: streams beside the first take n-position pages from one pool (unset or 0: none). +pub fn paging() !usize { + const v = std.c.getenv("TF_FLASHNEXT_KV_PAGES") orelse return 0; + const n = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadFlashNextKvPages; + if (n != 0 and !pages.valid(n)) return error.BadFlashNextKvPages; + return n; +} + +/// TF_FLASHNEXT_KV_PAGES_ALL=1: every stream on pages, none on its own contiguous caches (checks). +pub fn pagedOnly() bool { + const v = std.c.getenv("TF_FLASHNEXT_KV_PAGES_ALL") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// Pages a stream maps at most: a full window and one more drafted window. +fn mostPages(capacity: usize, tokens: usize) usize { + return (capacity + window_rows + tokens - 1) / tokens; +} + +fn al(n: usize) usize { + return std.mem.alignForward(usize, n, 256); +} + +const conv_bytes = mtp.conv_bytes; +const state_bytes = wt.state_n * 4; +const tail_bytes = wt.ple_tail_n * wt.wide * 2; + +/// A running stream's own states in its run: each DeltaNet layer's convolution tail and state, the n-gram tail. +fn seqRunBytes() usize { + return wt.linear_n * (al(conv_bytes) + al(state_bytes)) + al(tail_bytes); +} + +/// The MTP head a drafting lane runs: its draft vocabulary and the widest prompt chunk whose streams it keeps. +pub const Head = struct { draft_n: usize, rows: usize }; + +/// Fixed: a lane's table and `extra` (image rows); paged: the attention caches (the head's too); a run: the states. +pub fn layout(capacity: usize, tokens: usize, head: ?Head, extra: usize) kv.Layout { + const layers = wt.attn_n + @intFromBool(head != null); + const run = seqRunBytes() + if (head) |h| mtp.Draft.runBytes(h.draft_n, h.rows) else 0; + return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + extra, .token_bytes = pages.pageBytesOf(tokens, layers) / tokens, .page_tokens = @intCast(tokens), .snap_bytes = snap.bytes(0, head != null), .run_bytes = run }; +} + +/// Where a stream's caches came from: no pool, the one stream on its own caches beside a pool, or a pool lane. +pub const From = enum { none, own, pool }; + +/// The backend's paging: the pool `attach` made, its layout from load, and whether the own stream is taken. +pub const State = struct { + pool: ?*Pool = null, + layout: ?kv.Layout = null, + head: ?Head = null, + paged_only: bool = false, + own_busy: bool = false, + + pub fn choose(s: *const State) From { + if (s.pool == null) return .none; + return if (!s.paged_only and !s.own_busy) .own else .pool; + } + + pub fn opened(s: *State, from: From) void { + if (from == .own) s.own_busy = true; + } + + /// A stream's caches given back: a lane's to the pool (true), else the caller frees them (false). + pub fn close(s: *State, from: From, seq: *prompt.Seq) bool { + switch (from) { + .pool => return s.pool.?.close(seq), + .own => s.own_busy = false, + .none => {}, + } + return false; + } + + /// The lane host's admission: the own stream is free, or a lane and the request's pages are. + pub fn fits(s: *const State, prompt_rows: u32, max_tokens: u32) bool { + const p = s.pool orelse return true; + if (!s.paged_only and !s.own_busy) return true; + return p.fits(@as(usize, prompt_rows) + max_tokens + window_rows); + } +}; + +/// One pool lane: its cache's views (its run's states, its stream's pages), that stream's lease and its run. +const Lane = struct { seq: prompt.Seq, lease: kv.Lease = .{}, used: bool = false, run: u32 = 0 }; + +pub const Pool = struct { + gpa: std.mem.Allocator, + e: *prompt.Engine, + layout: kv.Layout, + capacity: usize, + head: ?Head, + pool: kv.Pool, + arena: cuda.DeviceBuffer, + lanes: []Lane, + + /// `streams` lanes and `count` pages of layout `l`, made and touched now; a lane holds only its table idle. + pub fn init(gpa: std.mem.Allocator, e: *prompt.Engine, l: kv.Layout, capacity: usize, head: ?Head, streams: usize, count: usize) !*Pool { + const p = try gpa.create(Pool); + errdefer gpa.destroy(p); + var arena = try cuda.DeviceBuffer.alloc(e.driver, count * l.pageBytes()); + errdefer arena.free(); + try arena.fill8(0, e.stream.handle); // committed now (the driver commits on first touch) + p.* = .{ .gpa = gpa, .e = e, .layout = l, .capacity = capacity, .head = head, .pool = try kv.Pool.init(gpa, count, l.pageBytes()), .arena = arena, .lanes = &.{} }; + errdefer p.pool.deinit(); + p.lanes = try gpa.alloc(Lane, streams); + errdefer gpa.free(p.lanes); + var made: usize = 0; + errdefer for (p.lanes[0..made]) |*x| x.seq.deinit(gpa); + for (p.lanes) |*x| { + x.* = .{ .seq = try p.laneSeq() }; + made += 1; + } + return p; + } + + /// A lane's cache: views only (states mapped from its run at open, attention caches on the arena) and a table. + fn laneSeq(p: *Pool) !prompt.Seq { + const gpa = p.gpa; + const d = p.e.driver; + var table = try cuda.DeviceBuffer.alloc(d, mostPages(p.capacity, p.layout.page_tokens) * 4); + errdefer table.free(); + try table.fill8(0, p.e.stream.handle); + const paged: pages.Paged = .{ .base = p.arena.ptr, .table = table, .ids = &.{}, .tokens = p.layout.page_tokens, .own_ids = false, .layers = wt.attn_n + @intFromBool(p.head != null) }; + const conv = try gpa.alloc(cuda.DeviceBuffer, wt.linear_n); + errdefer gpa.free(conv); + const state = try gpa.alloc(cuda.DeviceBuffer, wt.linear_n); + errdefer gpa.free(state); + const kc = try paged.views(gpa, d, .key); + errdefer gpa.free(kc); + const vc = try paged.views(gpa, d, .value); + errdefer gpa.free(vc); + const ikc = try paged.views(gpa, d, .index); + errdefer gpa.free(ikc); + const pooled = try paged.views(gpa, d, .pooled); + errdefer gpa.free(pooled); + const none: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0, .borrowed = true }; + @memset(conv, none); + @memset(state, none); + return .{ .conv = conv, .state = state, .kc = kc, .vc = vc, .ikc = ikc, .pooled = pooled, .paged = paged, .tail = none, .history = .{ p.e.eos, p.e.eos }, .fed = 0, .logits = try gpa.alloc(u8, prompt.vocabulary * 2), .cap = 0 }; + } + + pub fn deinit(p: *Pool) void { + p.e.stream.synchronize() catch std.log.warn("flashnext: the pool's stream did not drain at close", .{}); + for (p.lanes) |*x| { + if (x.used) { + p.pool.releaseRun(x.run, p.layout.runPages()); + p.pool.finish(&x.lease); + } + x.seq.deinit(p.gpa); + } + p.gpa.free(p.lanes); + p.arena.free(); + p.pool.deinit(); + p.gpa.destroy(p); + } + + /// Whether a lane is free, the pool can promise a cache of `cap` rows and a run now, and a run is free. + pub fn fits(p: *const Pool, cap: usize) bool { + const k = p.layout.runPages(); + for (p.lanes) |x| if (!x.used) return p.pool.fits(p.layout.pages(cap) + k) and p.pool.findRun(k) != null; + return false; + } + + /// A free lane's cleared cache (a copy) with `cap` rows' pages and its run; LaneMemory: not now. + pub fn open(p: *Pool, cap: usize) !prompt.Seq { + const x = for (p.lanes) |*x| { + if (!x.used) break x; + } else return error.LaneMemory; + const k = p.layout.runPages(); + const n = p.layout.pages(cap); + if (n > mostPages(p.capacity, p.layout.page_tokens)) return error.PromptTooLong; + try p.pool.promise(&x.lease, n + k); + errdefer p.pool.finish(&x.lease); + // the run first: pages fill from the bottom, runs from the top; a fragmented pool waits (LaneMemory) + const run = p.pool.claimRun(&x.lease, k) orelse return error.LaneMemory; + errdefer p.pool.releaseRun(run, k); + var up: Upload = .{ .p = p, .table = x.seq.paged.?.table }; + try p.pool.cover(&x.lease, n, .{ .ctx = &up, .write = Upload.write }); + x.seq.paged.?.ids = x.lease.ids.items; + x.run = run; + p.mapStates(&x.seq); + try p.e.clearSeq(&x.seq, cap); + x.used = true; + return x.seq; + } + + fn runBase(p: *const Pool, x: *const Lane) u64 { + return p.arena.ptr + @as(u64, x.run) * p.layout.pageBytes(); + } + + /// The lane's states as views of its run, laid out afresh (a kept window swaps state buffers with the head's). + fn mapStates(p: *Pool, seq: *prompt.Seq) void { + const x = p.laneOf(seq).?; + var at = p.runBase(x); + for (seq.conv) |*b| { + b.* = .{ .d = p.e.driver, .ptr = at, .len = conv_bytes, .borrowed = true }; + at += al(conv_bytes); + } + for (seq.state) |*b| { + b.* = .{ .d = p.e.driver, .ptr = at, .len = state_bytes, .borrowed = true }; + at += al(state_bytes); + } + seq.tail = .{ .d = p.e.driver, .ptr = at, .len = tail_bytes, .borrowed = true }; + } + + /// The MTP head of the stream on lane cache `seq`: its cache on the stream's pages, its buffers in the run. + pub fn draft(p: *Pool, seq: *const prompt.Seq, cap: usize) !mtp.Draft { + const h = p.head orelse return error.FlashNextRunTooShort; + const x = p.laneOf(seq) orelse return error.UnknownStream; + const run_len = p.layout.runPages() * p.layout.pageBytes(); + return mtp.Draft.initIn(p.gpa, p.e.driver, cap, h.draft_n, h.rows, &x.seq.paged.?, p.runBase(x) + seqRunBytes(), run_len - seqRunBytes()); + } + + fn laneOf(p: *Pool, seq: *const prompt.Seq) ?*Lane { + const t = (seq.paged orelse return null).table.ptr; + for (p.lanes) |*x| if (x.seq.paged.?.table.ptr == t) return x; + return null; + } + + /// A lane's stream ended: its pages and run go back once the forwards reading them are done. False: not a lane's. + pub fn close(p: *Pool, seq: *prompt.Seq) bool { + const x = p.laneOf(seq) orelse return false; + if (!x.used) return true; + // a page another stream takes next may still be read by queued work on the engine's stream + p.e.stream.synchronize() catch std.log.warn("flashnext: a closed stream's forwards did not drain; its pages go back", .{}); + p.pool.releaseRun(x.run, p.layout.runPages()); + p.pool.finish(&x.lease); + if (seq.media) |*m| m.deinit(); // the stream's image state, held by its copy + seq.media = null; + x.seq.paged.?.ids = &.{}; + x.seq.cap = 0; + x.used = false; + return true; + } + + /// A lane's new table entries uploaded to its device table. + const Upload = struct { + p: *Pool, + table: cuda.DeviceBuffer, + + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const u: *Upload = @ptrCast(@alignCast(ctx)); + try u.table.uploadAsync(from * 4, std.mem.sliceAsBytes(ids[from..]), u.p.e.stream.handle); + try u.p.e.stream.synchronize(); // the ids slice may move as the lease grows + } + }; +}; + +test "a stream's pages: its rows and a drafted window, no more than a full window's" { + const l = layout(40960, 256, null, 0); + try std.testing.expectEqual(pages.pageBytes(256), l.pageBytes()); + const h = layout(40960, 256, .{ .draft_n = 32768, .rows = 1536 }, 0); + try std.testing.expectEqual(pages.pageBytesOf(256, wt.attn_n + 1), h.pageBytes()); + try std.testing.expect(h.run_bytes >= seqRunBytes() + mtp.Draft.runBytes(32768, 1536)); + try std.testing.expectEqual(@as(usize, 161), mostPages(40960, 256)); + try std.testing.expect(l.pages(40960 + window_rows) <= mostPages(40960, 256)); +} diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 226c988b55..f940142611 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -13,6 +13,7 @@ const Glue = @import("cuda_glue.zig").Glue; const living = @import("living"); const slide_mod = @import("cuda_slide.zig"); const train = @import("cuda_train.zig"); +const kvl = @import("cuda_kv_lanes.zig"); const be = lanes.backend; const ring = 1024; @@ -40,6 +41,17 @@ pub fn Namespace(comptime Tri: type) type { pub const snap_save = snapSave; pub const snap_restore = snapRestore; pub const snap_drop = snapDrop; + /// TF_FLASHNEXT_KV_PAGES: the pool's lanes, the pages of a request's rows and a drafted window (kvl). + pub const window_rows = kvl.window_rows; + pub fn attach(p: *anyopaque, streams: u32, count: usize) !void { + const own = of(p); + own.kv.pool = try kvl.Pool.init(own.gpa, own.engine, own.kv.layout.?, own.context, own.kv.head, streams, count); + own.kv.paged_only = kvl.pagedOnly() and streams > 0; + } + pub fn fits(p: *anyopaque, rows: u32, max_tokens: u32) bool { + return of(p).kv.fits(rows, max_tokens); + } + /// The pass ignores its cuts, so states keep at any token; the MTP head reads one token on (lookahead 1). pub const cache_rules = struct { pub const lookahead: u32 = 1; @@ -85,6 +97,8 @@ pub fn Namespace(comptime Tri: type) type { /// Whether our own glue runs (no captured set given, or one without this family's entry points). own_glue: bool = true, follow: ?*const fn (*anyopaque) anyerror!void = null, + /// TF_FLASHNEXT_KV_PAGES: the streams beside the first take pages of one pool (null: their own caches). + layout: ?@import("core").kv_pool.Layout = null, }; /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. @@ -128,6 +142,8 @@ pub fn Namespace(comptime Tri: type) type { return switch (err) { error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + error.BadFlashNextKvPages, error.BadFlashNextPageTest => "TF_FLASHNEXT_KV_PAGES and _PAGE_TEST take positions a page, a multiple of 4", + error.FlashNextPagesUnsupported => "TF_FLASHNEXT_KV_PAGES needs one rank, our own glue and no --slide: unset it", error.UnsupportedCheckpoint => "the checkpoint differs from the supported Flash Next CUDA layout; the log names its config field", error.TensorRankDraftSnapshotUnavailable => "tensor-rank MTP snapshots are unavailable: disable the prompt cache for this configuration", error.RankDisconnected, error.RankTimeout, error.RemoteRank, error.BadRankFrame => "a tensor-parallel rank failed; restart both ranks from the same checkpoint before retrying", @@ -145,6 +161,8 @@ pub fn Namespace(comptime Tri: type) type { chain: mtp.Chain = .{}, /// The prompt's next row while its pass runs a chunk a call; null once the prompt is in. filled: ?usize = 0, + /// Where its caches came from: the pool's lanes, or its own (the one a pool leaves beside it). + from: kvl.From = .none, }; const Ops = mtp.Ops(Tri); @@ -169,6 +187,7 @@ pub fn Namespace(comptime Tri: type) type { next: u64 = 0, /// --slide: the pocket's layout, from config.json at load. lay: living.Layout = undefined, + kv: kvl.State = .{}, fn tri(self: *Cuda) Tri { return .{ .set = if (self.set) |*x| x else null, .own = &self.glue, .s = self.stream }; @@ -273,17 +292,19 @@ pub fn Namespace(comptime Tri: type) type { } // a window's rows past the reply's last token (its drafts) still need cache slots const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); - gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else self.engine.newSeq(cap)) catch |err| { + const from = self.kv.choose(); + gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else if (from == .pool) self.kv.pool.?.open(cap) else self.engine.newSeq(cap)) catch |err| { self.lanes.removeByPtr(gop.key_ptr); return err; - } }; + }, .from = from }; + self.kv.opened(from); const l = gop.value_ptr; l.seq.untaught = s.untaught; errdefer { dropLane(self, l); self.lanes.removeByPtr(gop.key_ptr); } - if (drafting) l.draft = try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); + if (drafting) l.draft = if (from == .pool) try self.kv.pool.?.draft(&l.seq, cap) else try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); if (s.images.len != 0) { const v = self.vision orelse return error.VisionNotEnabled; l.seq.media = try v.state(ids, s.images); @@ -370,6 +391,7 @@ pub fn Namespace(comptime Tri: type) type { if (own.head) |h| h.deinit(); if (own.vision) |v| v.deinit(); if (own.ranks) |r| r.deinit(); + if (own.kv.pool) |kp| kp.deinit(); own.engine.deinit(); if (own.set) |*x| x.deinit(); own.glue.deinit(); @@ -379,7 +401,7 @@ pub fn Namespace(comptime Tri: type) type { fn dropLane(self: *Cuda, l: *Lane) void { if (l.draft) |*d| d.deinit(); - if (self.ranks) |r| r.drop(&l.seq) else l.seq.deinit(self.gpa); + if (self.ranks) |r| r.drop(&l.seq) else if (!self.kv.close(l.from, &l.seq)) l.seq.deinit(self.gpa); } fn of(ptr: *anyopaque) *Cuda { diff --git a/zig/src/families/flashnext/cuda_lanes_admission_test.zig b/zig/src/families/flashnext/cuda_lanes_admission_test.zig index 5589587c80..898c0bd9c6 100644 --- a/zig/src/families/flashnext/cuda_lanes_admission_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_admission_test.zig @@ -31,6 +31,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, coun self.ranks = null; self.head = &head; self.lanes = .empty; + self.kv = .{}; defer { var it = self.lanes.valueIterator(); while (it.next()) |lane| drop(&self, lane); diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 69c300ef69..5de4cba5da 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -19,6 +19,8 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; if (o.slide and (mode != .affine or o.tp != 1)) return error.UnsupportedLearningFormat; if (o.slide and captured != null) return error.TrainingOnCapturedGlue; + const page_tokens = try @import("cuda_kv_lanes.zig").paging(); + if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null)) return error.FlashNextPagesUnsupported; const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); @@ -79,7 +81,14 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); if (o.vision) |v| stream_bytes += @as(usize, v.image_tokens) * @import("cuda_prompt_wt.zig").dims * 2 + 16 * o.context; + if (page_tokens > 0) { + // the lane's head runs at the engine's chunk: its cache on the stream's pages, its buffers in the run + own.kv.head = if (own.head) |h| .{ .draft_n = h.ids.len, .rows = own.engine.chunk } else null; + const head_bytes = if (own.head) |h| mtp.Draft.bytes(window, h.ids.len, own.engine.chunk) else 0; + own.kv.layout = @import("cuda_kv_lanes.zig").layout(o.context, page_tokens, own.kv.head, stream_bytes - prompt.seqBytes(window) - head_bytes); + } return .{ + .layout = own.kv.layout, .vision = own.vision, .backend = own.backend(), .facts = own.facts(), diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 2c460da902..5f69440023 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -390,6 +390,7 @@ test { _ = native; _ = @import("cuda_snap.zig"); _ = @import("cuda_pages.zig"); + _ = @import("cuda_kv_lanes.zig"); _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index 1ae18ae2ab..ddabdbd187 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -7,6 +7,7 @@ const batch = @import("cuda_batch.zig"); const prompt = @import("cuda_prompt.zig"); const gdn = @import("cuda_gdn.zig"); const mtp = @import("cuda_mtp.zig"); +const pages = @import("cuda_pages.zig"); const dims = wt.dims; const streams = wt.streams; @@ -51,6 +52,8 @@ pub const Draft = struct { ikc: cuda.DeviceBuffer, /// Its pooled indexer keys [cap / 4, 128] bf16, for sparse attention past 2,048 keys (Python State.mtp_pooled). pooled: cuda.DeviceBuffer, + /// The cache's pages (a pool lane's: its stream's, the layer after the model's), or contiguous. + pg: @import("cuda_glue.zig").Pg = .{}, /// Entries the head's cache holds (Python mtp_len), the last `drafted` of them a chain's. len: usize = 0, drafted: usize = 0, @@ -124,6 +127,12 @@ pub const Draft = struct { errdefer gpa.free(d.host); // zeroed on the compute stream, ahead of the passes that read them for ([_]cuda.DeviceBuffer{ d.kc, d.vc, d.ikc, d.pooled }) |b| try b.fill8(0, stream.handle); + d.point(); + return d; + } + + /// The keep's history and projection rows in `saved`, and the end states into the spares. + fn point(d: *Draft) void { var li: usize = 0; for (&d.keep.layers, 0..) |*lk, layer| { if (!wt.linearLayer(layer)) continue; @@ -132,6 +141,39 @@ pub const Draft = struct { li += 1; } d.aim(); + } + + /// Device bytes `initIn` takes from a run: every buffer but the head's cache, each on a 256-byte boundary. + pub fn runBytes(draft_n: usize, rows: usize) usize { + return bytes(0, draft_n, rows) + (6 + linear_n) * 256; + } + + /// A pool lane's head: its cache on `paged`'s pages after the model's layers, its buffers from `base`'s run. + pub fn initIn(gpa: std.mem.Allocator, driver: *cuda.Driver, cap: usize, draft_n: usize, rows: usize, paged: *const pages.Paged, base: u64, len: usize) !Draft { + if (len < runBytes(draft_n, rows)) return error.FlashNextRunTooShort; + var at = base; + const Carve = struct { + fn take(d: *const cuda.Driver, p: *u64, n: usize) cuda.DeviceBuffer { + const b: cuda.DeviceBuffer = .{ .d = d, .ptr = p.*, .len = n, .borrowed = true }; + p.* += std.mem.alignForward(usize, n, 256); + return b; + } + fn page(d: *const cuda.Driver, ptr: u64) cuda.DeviceBuffer { + return .{ .d = d, .ptr = ptr, .len = 0, .borrowed = true }; + } + }; + const head = wt.attn_n; + var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = Carve.page(driver, paged.at(head, .key)), .vc = Carve.page(driver, paged.at(head, .value)), .ikc = Carve.page(driver, paged.at(head, .index)), .pooled = Carve.page(driver, paged.at(head, .pooled)), .pg = paged.pg(), .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; + d.streams = Carve.take(driver, &at, d.rows_cap * wide * 2); + d.out = Carve.take(driver, &at, wide * 2); + d.logits = Carve.take(driver, &at, max_rows * vocab * 2); + d.dlogits = Carve.take(driver, &at, draft_n * 2); + d.best = Carve.take(driver, &at, max_rows * 16); + d.saved = Carve.take(driver, &at, saved_bytes); + for (&d.spare) |*b| b.* = Carve.take(driver, &at, state_n * 4); + std.debug.assert(at - base <= len); + d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); + d.point(); return d; } diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index cfcd9fdd05..eab7e37b97 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -124,7 +124,7 @@ pub fn Ops(comptime Tri: type) type { .k = &e.batch_k, .s = s, .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram }, - .seq = .{ .media = d.media, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, + .seq = .{ .media = d.media, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .pg = d.pg, .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, .hooks = .{ .logits = .none }, }; var at: usize = 0; diff --git a/zig/src/families/flashnext/cuda_pages.zig b/zig/src/families/flashnext/cuda_pages.zig index 74ccaffc15..a59c02555c 100644 --- a/zig/src/families/flashnext/cuda_pages.zig +++ b/zig/src/families/flashnext/cuda_pages.zig @@ -31,7 +31,12 @@ pub fn layerBytes(tokens: usize) usize { /// Bytes a page takes: every attention layer's share. pub fn pageBytes(tokens: usize) usize { - return wt.attn_n * layerBytes(tokens); + return pageBytesOf(tokens, wt.attn_n); +} + +/// Bytes a page of `layers` attention layers takes (the model's, and the MTP head's cache when it drafts). +pub fn pageBytesOf(tokens: usize, layers: usize) usize { + return layers * layerBytes(tokens); } /// Where `part` of attention layer `layer` starts inside a page. @@ -71,9 +76,11 @@ pub const Paged = struct { own_ids: bool = true, /// The page test's own arena. arena: ?cuda.DeviceBuffer = null, + /// Attention layers a page holds: the model's, and the MTP head's after them in a pool's pages. + layers: usize = wt.attn_n, pub fn pg(p: *const Paged) Pg { - return .{ .pt = p.table.ptr, .ps = @intCast(pageBytes(p.tokens) / 2), .ptok = @intCast(p.tokens) }; + return .{ .pt = p.table.ptr, .ps = @intCast(pageBytesOf(p.tokens, p.layers) / 2), .ptok = @intCast(p.tokens) }; } /// `part` of attention layer `layer` in page 0: the cache pointer a kernel takes beside pg(). @@ -84,7 +91,7 @@ pub const Paged = struct { /// Byte address of row `row` of `part` of `layer` (a position, or a block for pooled). pub fn rowAt(p: *const Paged, layer: usize, part: Part, row: usize) u64 { const per = rows(part, p.tokens); - return p.base + @as(u64, p.ids[row / per]) * pageBytes(p.tokens) + offset(p.tokens, layer, part) + (row % per) * rowBytes(part); + return p.base + @as(u64, p.ids[row / per]) * pageBytesOf(p.tokens, p.layers) + offset(p.tokens, layer, part) + (row % per) * rowBytes(part); } /// Borrowed views of each attention layer's `part` in page 0 (what kc, vc, ikc and pooled hold when paged). diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index f70726260c..883acaf326 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -288,6 +288,14 @@ pub const Engine = struct { /// A fresh cache. `cap` is how many tokens it can hold, prompt plus generated. pub fn newSeq(e: *Engine, cap: usize) !Seq { + const test_pages = try pages.pageTest(); + return e.newSeqOn(cap, if (test_pages > 0) try pages.testPages(e.gpa, e.driver, e.stream.*, cap, test_pages) else null); + } + + /// A fresh cache whose attention caches are `given` pages (owned by the cache from here), else contiguous. + pub fn newSeqOn(e: *Engine, cap: usize, given: ?pages.Paged) !Seq { + var paged = given; + errdefer if (paged) |*p| p.deinit(e.gpa); e.round_graph.clear(); if (cap == 0) return error.UnexpectedTensor; const logits = try e.gpa.alloc(u8, vocab * 2); @@ -296,9 +304,6 @@ pub const Engine = struct { errdefer freeOwned(e.gpa, conv); const state = try zeros(e.gpa, e.driver, e.stream.*, linear_n, state_n * 4); errdefer freeOwned(e.gpa, state); - const test_pages = try pages.pageTest(); - var paged: ?pages.Paged = if (test_pages > 0) try pages.testPages(e.gpa, e.driver, e.stream.*, cap, test_pages) else null; - errdefer if (paged) |*p| p.deinit(e.gpa); const kc = if (paged) |*p| try p.views(e.gpa, e.driver, .key) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, kc); const vc = if (paged) |*p| try p.views(e.gpa, e.driver, .value) else try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); @@ -326,6 +331,18 @@ pub const Engine = struct { }; } + /// A pool lane's cache made fresh for its next stream: zero states and tail, nothing fed, `cap` rows. + pub fn clearSeq(e: *Engine, s: *Seq, cap: usize) !void { + e.round_graph.clear(); + for (s.conv) |b| try b.fill8(0, e.stream.handle); + for (s.state) |b| try b.fill8(0, e.stream.handle); + try s.tail.fill8(0, e.stream.handle); + s.history = .{ e.eos, e.eos }; + s.fed = 0; + s.cap = cap; + s.untaught = false; + } + /// Run `ids` on from `seq.fed` and leave the next token's scores in `seq.logits`. pub fn predict(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32) !void { return e.forward(Tri, tri, seq, ids, .{}); diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index 79ccf2b04b..96871a3cc1 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -62,16 +62,19 @@ fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, } for (pieces[0..count], 0..) |p, pi| for (p.bufs, 0..) |b, layer| { if (p.n == 0) continue; - if (pi >= 2 and pi < 6) if (seq.paged) |*pg| { + // the model's caches on the stream's pages; the head's too in a pool lane (the layer after the model's) + const head = pi >= 7 and d != null and d.?.pg.pt != 0; + if ((pi >= 2 and pi < 6) or head) if (seq.paged) |*pg| { // the same bytes in the same order as a contiguous cache's, a run a page - const part: pages.Part = @fromBackingInt(@intCast(pi - 2)); + const part: pages.Part = @fromBackingInt(@intCast(if (head) pi - 7 else pi - 2)); + const at_layer = if (head) wt.attn_n else layer; const row = pages.rowBytes(part); const per = pages.rows(part, pg.tokens); const n = p.n / row; var r: usize = 0; while (r < n) { const run = @min(per - r % per, n - r); - const at_row = pg.rowAt(layer, part, r); + const at_row = pg.rowAt(at_layer, part, r); const dst, const src = if (way == .save) .{ base + off, at_row } else .{ at_row, base + off }; try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, run * row, stream.handle), "cuMemcpyDtoDAsync"); off += run * row; From e0e69c521d7145538a042baec733f84e2283a765 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 18:41:02 +0200 Subject: [PATCH 11/11] cuda qwen3_5: the reuse tests' caches say they have no page table Two cuda_reuse tests build a Cache as undefined and set only the fields they meant to read. Since the page table (3f5b87dd), restoreCache also reads Cache.paged to decide how to copy the rows, so those tests read undefined memory: on CI's x86_64 Debug run the bytes formed a page table and piece()'s offsets overflowed ("a failed restore's host copy stays through a drop ..."). On aarch64 the same bytes read as null, so the test passed here (with the stack filled with 0xaa it still reads null: the optional's null is a niche value of its bool). Both tests now set paged = null. --- zig/src/families/qwen3_5/cuda_reuse.zig | 2 ++ 1 file changed, 2 insertions(+) diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig index c396c91597..6d715b1d5e 100644 --- a/zig/src/families/qwen3_5/cuda_reuse.zig +++ b/zig/src/families/qwen3_5/cuda_reuse.zig @@ -331,6 +331,7 @@ test "a restore refuses another cache format, another length, a position past th if (comptime @hasField(Engine, "kv8")) e.kv8 = false; var cache: Cache = undefined; cache.cap = 64; + cache.paged = null; // contiguous: restoreCache reads it, and undefined bytes can read as a page table const d: u64 = 9; var images = [_]Seen{.{ .at = 2, .id = d }}; const l: Layout = .{ .g = e.g, .format = .bf16, .cap = 64, .at = 10 }; @@ -389,6 +390,7 @@ test "a failed restore's host copy stays through a drop until a wait, the restor if (comptime @hasField(Engine, "kv8")) e.kv8 = false; var cache: Cache = undefined; cache.cap = 64; + cache.paged = null; // contiguous: restoreCache reads it, and undefined bytes can read as a page table cache.kv = .{ .d = &d, .ptr = 0x10000, .len = 0 }; cache.state = .{ .d = &d, .ptr = 0x20000, .len = 0 }; const l: Layout = .{ .g = e.g, .format = .bf16, .cap = 64, .at = 10 };