From 0239d0ebfb27d09108f7e572cc6902154f643bd5 Mon Sep 17 00:00:00 2001 From: "Riel St. Amand" <54451475+CanadaApollo6@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:26:48 -0400 Subject: [PATCH 1/5] qwen27 CUDA: split-K matmuls take lane_gemv 16 rows a launch before sm_90 (no clusters) --- zig/src/families/qwen27/cuda/kernels.zig | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/zig/src/families/qwen27/cuda/kernels.zig b/zig/src/families/qwen27/cuda/kernels.zig index 4b486922f8..f7aa6a4c1c 100644 --- a/zig/src/families/qwen27/cuda/kernels.zig +++ b/zig/src/families/qwen27/cuda/kernels.zig @@ -74,6 +74,7 @@ pub const Kernels = struct { gemv_blocks: usize, // resident lane_gemv CTAs: per SM times SMs gemv_rows: usize, // windows of at most this many rows take lane_gemv (none; TF_QWEN27_GEMV_ROWS) gb10: bool, + clusters: bool, // sm_90 and later: K slices meet in a cluster; older chips take lane_gemv 16 rows at a time discrete: bool, // the card has its own memory: checkpoint bytes reach it through page-locked slots /// Loads every module; `triton_dir` holds the captured aot.json and cubins for this GPU. @@ -122,6 +123,7 @@ pub const Kernels = struct { const major = try ctx.attribute(.compute_capability_major); const minor = try ctx.attribute(.compute_capability_minor); k.gb10 = major == 12 and minor == 1; + k.clusters = major >= 9; k.discrete = try ctx.attribute(.integrated) == 0; k.gemv_rows = 0; // a GB10 runs the 27B's projections faster in Python's clusters, 1 to 16 rows alike if (std.c.getenv("TF_QWEN27_GEMV_ROWS")) |v| k.gemv_rows = std.fmt.parseInt(usize, std.mem.span(v), 10) catch k.gemv_rows; @@ -167,6 +169,7 @@ pub const Ops = struct { /// qmm.matmul on sm_12x: x (rows, k) bf16 with group sums xs -> out (rows, n) bf16, qmm_group's bits. pub fn dense(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize) !void { const sk = splitK(q.n, q.k); + if (sk > 1 and !o.k.clusters) return o.gemvRows(x, xs, q, out, rows, sk); if (rows > 16) return o.wideTile(x, xs, q, out, rows, sk); // Python's clusters, a CTA a K slice; lane_gemv (byte-equal) when TF_QWEN27_GEMV_ROWS asks for it return if (sk > 1 and rows <= o.k.gemv_rows) o.gemv(x, xs, q, out, rows, sk) else o.cluster(x, xs, q, out, rows, sk); @@ -174,6 +177,7 @@ pub const Ops = struct { /// qmm.matmul at a K split the caller names (matmul_rows takes the stacked weight's for each part). pub fn denseSplit(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + if (sk > 1 and !o.k.clusters) return o.gemvRows(x, xs, q, out, rows, sk); if (rows > 16) return o.wideTile(x, xs, q, out, rows, sk); return if (sk > 1 and rows <= o.k.gemv_rows) o.gemv(x, xs, q, out, rows, sk) else o.cluster(x, xs, q, out, rows, sk); } @@ -198,6 +202,14 @@ pub const Ops = struct { try cuda.launch.launch(o.k.gemv, cfg, o.s, &a); } + /// Before sm_90 (no clusters): lane_gemv on each 16 rows in turn; a row's bits never depend on the others. + fn gemvRows(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + var r: usize = 0; + while (r < rows) : (r += 16) { + try o.gemv(x + r * q.k * 2, xs + r * (q.k / 64) * 4, q, out + r * q.n * 2, @min(16, rows - r), sk); + } + } + /// qmm_group's tile 2 with one K slice: a CTA a column tile. fn cluster(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { const tiles = (q.n + 63) / 64; From a638859d8a2e8f7b97213192159a1d39fb3eebdd Mon Sep 17 00:00:00 2001 From: "Riel St. Amand" <54451475+CanadaApollo6@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:26:48 -0400 Subject: [PATCH 2/5] qwen27 serve: 16-byte aligned round parts; aligned launches may use a captured unaligned variant The shared-round staging packed each part back to back, so _gdn_pre's WIN pointer landed off 16 bytes at some window sizes and no captured variant matched. Each part now starts on 16 bytes, and an aligned pointer or int can fall back to a variant captured without the alignment hint (never the reverse). --- zig/src/cuda/aot.zig | 19 ++++++++++++++----- zig/src/families/qwen27/cuda/round.zig | 25 ++++++++++++++++++------- zig/src/families/qwen27/cuda/state.zig | 2 +- 3 files changed, 33 insertions(+), 13 deletions(-) diff --git a/zig/src/cuda/aot.zig b/zig/src/cuda/aot.zig index b8e58d2e8e..fce864b3d1 100644 --- a/zig/src/cuda/aot.zig +++ b/zig/src/cuda/aot.zig @@ -122,7 +122,12 @@ pub const Set = struct { pub fn find(self: *const Set, function: []const u8, args: []const Arg, consts: []const Const) !*const Variant { for (self.variants) |*v| { if (!std.mem.eql(u8, v.spec.@"fn", function)) continue; - if (matches(v.spec, args, consts)) return v; + if (matches(v.spec, args, consts, false)) return v; + } + // else one compiled without assuming 16-byte alignment where these are aligned: the same code, valid for any value + for (self.variants) |*v| { + if (!std.mem.eql(u8, v.spec.@"fn", function)) continue; + if (matches(v.spec, args, consts, true)) return v; } std.log.err("no captured Triton variant of {s} for this launch:", .{function}); for (args) |a| switch (a.value) { @@ -167,7 +172,8 @@ fn lookup(args: []const Arg, name: []const u8) ?Arg { return null; } -fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { +/// `relaxed`: an argument divisible by 16 may also take a variant that doesn't assume it (never the reverse). +fn matches(k: KernelJson, args: []const Arg, consts: []const Const, relaxed: bool) bool { for (consts) |c| { const got = k.consts.map.get(c.name) orelse return false; if (c.int) |x| if (got.int == null or got.int.? != x) return false; @@ -189,11 +195,13 @@ fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { const p = param.?; if (!std.mem.eql(u8, p.type, "i32")) return false; if (!p.nospec and x == 1) return false; - if (p.div16 != (!p.nospec and @mod(x, 16) == 0)) return false; + const div16 = !p.nospec and @mod(x, 16) == 0; + if (p.div16 != div16 and !(relaxed and div16 and !p.div16)) return false; }, .ptr => |x| { const p = param orelse return false; - if (!std.mem.eql(u8, p.type, x.ty) or p.div16 != (x.addr % 16 == 0)) return false; + const div16 = x.addr % 16 == 0; + if (!std.mem.eql(u8, p.type, x.ty) or (p.div16 != div16 and !(relaxed and div16 and !p.div16))) return false; }, .f32 => { const p = param orelse return false; @@ -201,7 +209,8 @@ fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { }, .u64 => |x| { const p = param orelse return false; - if (!std.mem.eql(u8, p.type, "u64") or p.div16 != (x % 16 == 0)) return false; + const div16 = x % 16 == 0; + if (!std.mem.eql(u8, p.type, "u64") or (p.div16 != div16 and !(relaxed and div16 and !p.div16))) return false; }, } runtime += 1; diff --git a/zig/src/families/qwen27/cuda/round.zig b/zig/src/families/qwen27/cuda/round.zig index 2b2d4e2697..33ac29cf26 100644 --- a/zig/src/families/qwen27/cuda/round.zig +++ b/zig/src/families/qwen27/cuda/round.zig @@ -179,24 +179,35 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, try t.paths(s.attn_plan + (W + 4 + 3 * n_items) * 4, s.paths, s.depths, W); return r; } - // multi_tree_forward: positions, stream ids, plan entries, starts, ids, windows, attention's plan in one copy + // multi_tree_forward: positions, stream ids, plan entries and starts, ids, windows, attention's plan in one copy, each + // part from a 16-byte boundary (the captured Triton variants are specialized on their pointers' alignment) + const pad = struct { + fn to16(gpa_: std.mem.Allocator, list: *std.ArrayList(i32)) !usize { + while (list.items.len % 4 != 0) try list.append(gpa_, 0); + return list.items.len; + } + }.to16; try host.appendSlice(gpa, h.pos.items); + const at_sids = try pad(gpa, &host); for (parts, 0..) |p, k| for (p.tokens) |_| try host.append(gpa, @intCast(k)); + const at_plan = try pad(gpa, &host); try host.appendSlice(gpa, h.entries.items); for (r.starts_host[0 .. S + 1]) |x| try host.append(gpa, @intCast(x)); + const at_ids = try pad(gpa, &host); for (parts) |p| for (p.tokens) |tok| try host.append(gpa, @intCast(tok)); + const at_windows = try pad(gpa, &host); try host.appendSlice(gpa, h.windows.items); - const at_attn = host.items.len; + const at_attn = try pad(gpa, &host); const n_items = try attnPlan(gpa, &host, parts, group); try ops.upload(s.round, std.mem.sliceAsBytes(host.items)); try ops.s.synchronize(); const base = s.round; r.pos = base; - r.sids = base + W * 4; - r.plan = base + 2 * W * 4; - r.starts = base + 5 * W * 4; - r.ids = base + (5 * W + S + 1) * 4; - r.windows = base + (6 * W + S + 1) * 4; + r.sids = base + at_sids * 4; + r.plan = base + at_plan * 4; + r.starts = r.plan + 3 * W * 4; + r.ids = base + at_ids * 4; + r.windows = base + at_windows * 4; const ab = base + at_attn * 4; r.attn = .{ .rows = ab, .streams = ab + W * 4, .items = ab + (W + 4 * S) * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; // the state pointer table and cache offsets, one copy each; every stream's conv rows stacked by layer diff --git a/zig/src/families/qwen27/cuda/state.zig b/zig/src/families/qwen27/cuda/state.zig index 4cbe06c07e..47e4d46722 100644 --- a/zig/src/families/qwen27/cuda/state.zig +++ b/zig/src/families/qwen27/cuda/state.zig @@ -276,7 +276,7 @@ pub const Scratch = struct { s.taps = try A.take(R * 5 * g.hidden * bf); s.replay_table = try A.take(g.layers * (4 + max_streams) * 8); s.replay_rows = try A.take(max_streams * (W + 1) * 4); - s.round = try A.take((6 * W + max_streams + 1 + W * c.conv_taps + W + 4 * max_streams + 3 * items * max_streams + W) * 4); + s.round = try A.take((6 * W + max_streams + 1 + W * c.conv_taps + W + 4 * max_streams + 3 * items * max_streams + W + 32) * 4); // + each part's 16-byte padding s.tables = try A.take(g.layers * max_streams * 8); s.conv_cat = try A.take(g.layers * max_streams * (c.conv_taps - 1) * g.convDim() * bf); for (0..g.layers) |i| { From c78ecb499ec114710f34cc769edfc430cfec4451 Mon Sep 17 00:00:00 2001 From: "Riel St. Amand" <54451475+CanadaApollo6@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:26:48 -0400 Subject: [PATCH 3/5] qwen27 CUDA: the scratch sized by --parallel; prompt chunk rows and a host token table by env The server sized the scratch for 16 streams of 16 rows at any --parallel, and the attention partials in it grow with the window: on a 24 GB card no 32K stream fit beside it. The scratch now takes `streams` windows (the lane core's admission and the cost curve follow), so --parallel 1 needs one stream's rows. - TF_QWEN27_PROMPT_ROWS: a prompt chunk's rows (default 4096; 1024 frees ~1 GiB at a ~10% slower prompt). - TF_QWEN27_HOST_EMBED=1: the token table in mapped host memory (0.67 GiB; only prompt rows read it). --- zig/src/cli/qwen27_cuda.zig | 5 +++-- zig/src/families/qwen27/cuda/drafts.zig | 4 ++-- zig/src/families/qwen27/cuda/engine.zig | 12 ++++++++---- zig/src/families/qwen27/cuda/forward.zig | 8 ++++---- zig/src/families/qwen27/cuda/lanes.zig | 9 ++++----- zig/src/families/qwen27/cuda/lone.zig | 2 +- zig/src/families/qwen27/cuda/native.zig | 15 ++++++++++----- zig/src/families/qwen27/cuda/round.zig | 8 ++++---- zig/src/families/qwen27/cuda/state.zig | 22 +++++++++++++++------- zig/src/families/qwen27/cuda/weights.zig | 22 +++++++++++++++++++--- zig/src/native/cuda.zig | 1 + 11 files changed, 71 insertions(+), 37 deletions(-) diff --git a/zig/src/cli/qwen27_cuda.zig b/zig/src/cli/qwen27_cuda.zig index a92416ce7b..aab471404b 100644 --- a/zig/src/cli/qwen27_cuda.zig +++ b/zig/src/cli/qwen27_cuda.zig @@ -66,7 +66,8 @@ pub fn main(init: std.process.Init) !u8 { defer ctx.deinit(); const clock = std.Io.Clock.awake; var t0 = clock.now(init.io).toNanoseconds(); - const e = try qwen.Engine.init(gpa, init.io, &ctx, o.model, o.kernels, .{ .context = o.context }); + const rows_env = if (std.c.getenv("TF_QWEN27_PROMPT_ROWS")) |v| std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadPromptRows else qwen.engine.prompt_rows; + const e = try qwen.Engine.init(gpa, init.io, &ctx, o.model, o.kernels, .{ .context = o.context, .prompt_rows = rows_env }); defer e.deinit(); std.debug.print("loaded in {d:.1}s, weights {d:.2} GiB\n", .{ @as(f64, @floatFromInt(clock.now(init.io).toNanoseconds() - t0)) / 1e9, @as(f64, @floatFromInt(e.w.bytes())) / (1 << 30) }); if (windows) return checkWindows(gpa, arena, init.io, e, prompts, o.dump, o.cases); @@ -435,7 +436,7 @@ fn runMulti(gpa: std.mem.Allocator, a: std.mem.Allocator, io: std.Io, e: *qwen.E m.accepted += len - 1; } try f.commit(parts[0..n], paths[0..n]); - try qwen.drafts.absorb(d, f.ops, e.scratch.taps, e.scratch.copy_items, kept[0..n]); + try qwen.drafts.absorb(d, f.ops, e.scratch.taps, e.scratch.window_rows, e.scratch.copy_items, kept[0..n]); for (live[0..n]) |k| { const m = &ms[k]; if (m.out.items.len >= count or e.w.isEos(m.out.items[m.out.items.len - 1])) m.done = true; diff --git a/zig/src/families/qwen27/cuda/drafts.zig b/zig/src/families/qwen27/cuda/drafts.zig index 02e3659524..08ad3fc9ac 100644 --- a/zig/src/families/qwen27/cuda/drafts.zig +++ b/zig/src/families/qwen27/cuda/drafts.zig @@ -86,12 +86,12 @@ pub fn propose(d: *dr.DFlash2, asks: []const Ask) !void { pub const Kept = struct { x: *Stream, rows: []const u32 }; /// add_taps_streams: every stream's kept rows' taps gathered after the round's rows, then absorbed in one batch. -pub fn absorb(d: *dr.DFlash2, ops: kern.Ops, taps: u64, copy_items: u64, kept: []const Kept) !void { +pub fn absorb(d: *dr.DFlash2, ops: kern.Ops, taps: u64, round_rows: usize, copy_items: u64, kept: []const Kept) !void { var items: std.ArrayList(kern.Copy) = .empty; defer items.deinit(d.gpa); var absorbs: [st.max_streams]dr.Absorb = undefined; const row = taps_width * bf; - const gathered = taps + st.round_rows * row; + const gathered = taps + round_rows * row; // after the round's rows for (kept, 0..) |k, j| { for (k.rows) |r| try items.append(d.gpa, .{ .dst = gathered + items.items.len * row, .src = taps + r * row, .bytes = row }); absorbs[j] = .{ .ctx = &k.x.ctx, .rows = k.rows.len }; diff --git a/zig/src/families/qwen27/cuda/engine.zig b/zig/src/families/qwen27/cuda/engine.zig index 52576d7cd3..b25a872f6e 100644 --- a/zig/src/families/qwen27/cuda/engine.zig +++ b/zig/src/families/qwen27/cuda/engine.zig @@ -19,7 +19,8 @@ pub fn chunkBounds(start: usize, end: usize, size: usize, j: usize) [2]usize { return .{ start + (end - start) * j / n, start + (end - start) * (j + 1) / n }; } -pub const Options = struct { context: usize }; +/// `streams`: a shared round's windows at most (the scratch's window rows follow); `prompt_rows`: a prompt chunk's rows. +pub const Options = struct { context: usize, streams: usize = st.max_streams, prompt_rows: usize = prompt_rows }; pub const Engine = struct { gpa: std.mem.Allocator, @@ -29,6 +30,7 @@ pub const Engine = struct { w: wts.Weights, scratch: st.Scratch, context: usize, + prompt_rows: usize, pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: []const u8, o: Options) !*Engine { const cfg = try c.Config.read(gpa, io, dir); @@ -38,13 +40,15 @@ pub const Engine = struct { e.gpa = gpa; e.ctx = ctx; e.context = o.context; + if (o.prompt_rows < 1) return error.BadPromptRows; + e.prompt_rows = o.prompt_rows; e.s = try cuda.Stream.init(ctx.d, true); errdefer e.s.deinit(); e.k = try kern.Kernels.load(gpa, io, ctx, kernels, g.linear_v_heads); errdefer e.k.deinit(); e.w = try wts.load(gpa, io, e.ops(), dir, cfg); errdefer e.w.deinit(); - e.scratch = try st.Scratch.init(gpa, ctx.d, g, prompt_rows, o.context); + e.scratch = try st.Scratch.init(gpa, ctx.d, g, o.prompt_rows, o.context, o.streams); errdefer e.scratch.deinit(); try e.s.synchronize(); return e; @@ -77,9 +81,9 @@ pub const Engine = struct { if (prompt.len >= seq.capacity) return error.PromptTooLong; try seq.reset(e.ops(), e.w.g); var f = e.forward(); - const n = (prompt.len + prompt_rows - 1) / prompt_rows; + const n = (prompt.len + e.prompt_rows - 1) / e.prompt_rows; for (0..n) |j| { - const b = chunkBounds(0, prompt.len, prompt_rows, j); + const b = chunkBounds(0, prompt.len, e.prompt_rows, j); try f.chunk(prompt[b[0]..b[1]], seq, j == n - 1); } } diff --git a/zig/src/families/qwen27/cuda/forward.zig b/zig/src/families/qwen27/cuda/forward.zig index 4e4964be4e..867ce1ee68 100644 --- a/zig/src/families/qwen27/cuda/forward.zig +++ b/zig/src/families/qwen27/cuda/forward.zig @@ -295,7 +295,7 @@ pub const Forward = struct { const g_ = f.g(); const w = f.w; const S = pieces.len; - if (S == 0 or S > st.max_streams) return error.InvalidQwenWindows; + if (S == 0 or S > s.streams()) return error.InvalidQwenWindows; var starts: [st.max_streams + 1]usize = undefined; starts[0] = 0; for (pieces, 0..) |p, k| { @@ -329,7 +329,7 @@ pub const Forward = struct { defer items.deinit(f.gpa); var lin_n: usize = 0; for (0..g_.layers) |i| if (c.linear(i)) { - const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + const cat = s.conv_cat + lin_n * s.streams() * keep * row; for (pieces, 0..) |p, k| try items.append(f.gpa, .{ .dst = cat + k * keep * row, .src = p.seq.conv[i], .bytes = keep * row }); lin_n += 1; }; @@ -340,7 +340,7 @@ pub const Forward = struct { at_layer[i] = items.items.len; if (c.linear(i)) { // the last `keep` rows of [kept rows; the stream's prompt rows] - const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + const cat = s.conv_cat + lin_n * s.streams() * keep * row; for (pieces, 0..) |p, k| { const n = p.ids.len; const o = starts[k]; @@ -378,7 +378,7 @@ pub const Forward = struct { try f.ops.prefillDense(s.y, d.z, s.z, W); try f.ops.prefillDense(s.y, d.b, s.b, W); try f.ops.prefillDense(s.y, d.a, s.a, W); - const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + const cat = s.conv_cat + lin_n * s.streams() * keep * row; try f.t.gdnPre(s.qkv, cat, d.conv, s.windows, s.sids, s.a, s.b, d.a_log, d.dt_bias, s.gq, s.gk, s.gv, s.gg, s.gbeta, W, f.gdnShape()); for (pieces, 0..) |p, k| { const o = starts[k]; diff --git a/zig/src/families/qwen27/cuda/lanes.zig b/zig/src/families/qwen27/cuda/lanes.zig index 8a5f86fad8..e73796e3d4 100644 --- a/zig/src/families/qwen27/cuda/lanes.zig +++ b/zig/src/families/qwen27/cuda/lanes.zig @@ -7,7 +7,6 @@ const st = @import("state.zig"); const Engine = @import("engine.zig").Engine; const Part = @import("forward.zig").Part; const Forward = @import("forward.zig").Forward; -const prompt_rows = @import("engine.zig").prompt_rows; const dl = @import("draft_load.zig"); const dr = @import("draft.zig"); const drafts = @import("drafts.zig"); @@ -78,7 +77,7 @@ pub const Cuda = struct { /// The facts the round loop reads: 16-row windows in one forward; with DFlash2, host trees with node chances. pub fn facts(x: *const Cuda) lanes.Model { - var m: lanes.Model = .{ .exact_width = st.max_rows, .hidden_rows = true, .batch_rows = st.round_rows, .max_streams = @min(x.streams, st.max_streams), .streams_exact = true }; + var m: lanes.Model = .{ .exact_width = st.max_rows, .hidden_rows = true, .batch_rows = @intCast(x.e.scratch.window_rows), .max_streams = @intCast(@min(x.streams, x.e.scratch.window_rows / st.max_rows)), .streams_exact = true }; if (x.drafter != null) { m.mtp = true; m.speculate = true; @@ -248,14 +247,14 @@ pub const Cuda = struct { continue; }; const ids = s.prompt(); - if (ids.len > prompt_rows) { + if (ids.len > x.e.prompt_rows) { x.prefillOpened(s, seq) catch |e| { x.dropSeq(s); err.* = e; }; continue; } - if (n == st.max_streams or rows + ids.len > prompt_rows) { + if (n == x.e.scratch.streams() or rows + ids.len > x.e.prompt_rows) { x.pass(streams, errs, pieces[0..n], which[0..n]); n = 0; rows = 0; @@ -471,7 +470,7 @@ pub const Cuda = struct { na += 1; } const f = x.e.forward(); - try drafts.absorb(d, f.ops, x.e.scratch.taps, x.e.scratch.copy_items, kept[0..nk]); + try drafts.absorb(d, f.ops, x.e.scratch.taps, x.e.scratch.window_rows, x.e.scratch.copy_items, kept[0..nk]); try drafts.propose(d, asks[0..na]); } diff --git a/zig/src/families/qwen27/cuda/lone.zig b/zig/src/families/qwen27/cuda/lone.zig index 0397ea706f..5d75100bf6 100644 --- a/zig/src/families/qwen27/cuda/lone.zig +++ b/zig/src/families/qwen27/cuda/lone.zig @@ -79,7 +79,7 @@ pub fn prefill(e: *Engine, seq: *st.Seq, ctx: *dr.Context, d: *dr.DFlash2, promp ctx.skip(tap_from); var f = e.forward(); f.taps = e.scratch.taps; - const rows = @import("engine.zig").prompt_rows; + const rows = e.prompt_rows; const chunks = (prompt.len + rows - 1) / rows; for (0..chunks) |j| { const b = @import("engine.zig").chunkBounds(0, prompt.len, rows, j); diff --git a/zig/src/families/qwen27/cuda/native.zig b/zig/src/families/qwen27/cuda/native.zig index 6e3f362fdf..aa576026f7 100644 --- a/zig/src/families/qwen27/cuda/native.zig +++ b/zig/src/families/qwen27/cuda/native.zig @@ -13,9 +13,11 @@ pub const model_type = "qwen3_5"; pub const formats: []const []const u8 = &.{"mlx-q4g64"}; pub const default_context: i64 = 32768; pub const max_segments: u32 = 1; -pub const prompt_rows: u32 = engine.prompt_rows; +/// A prompt chunk's rows: TF_QWEN27_PROMPT_ROWS, else engine.prompt_rows (set when the server opens the model). +pub var prompt_rows: u32 = engine.prompt_rows; -pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, drafter: ?[]const u8 = null }; +/// `streams`: --parallel, the shared round's windows at most (the scratch is sized for them). +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, drafter: ?[]const u8 = null, streams: u32 = most_streams }; /// A lone drafted stream's own driver (none here: one stream decodes in the lane core too, as under --parallel). pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool) anyerror!bool; @@ -67,11 +69,14 @@ const most_streams = 32; pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { // the 27B's glue runs the Python engine's Triton kernels, captured on this chip (no own glue kernels yet) - const e = try engine.Engine.init(gpa, io, ctx, dir, kernels orelse return error.QwenNeedsCapturedKernels, .{ .context = o.context }); + const streams: u32 = @max(1, @min(o.streams, st.max_streams)); + if (getenv("TF_QWEN27_PROMPT_ROWS")) |v| prompt_rows = std.fmt.parseInt(u32, v, 10) catch return error.BadPromptRows; + if (prompt_rows < 16) return error.BadPromptRows; + const e = try engine.Engine.init(gpa, io, ctx, dir, kernels orelse return error.QwenNeedsCapturedKernels, .{ .context = o.context, .streams = streams, .prompt_rows = prompt_rows }); errdefer e.deinit(); const own = try gpa.create(Owned); errdefer gpa.destroy(own); - own.* = .{ .gpa = gpa, .e = e, .lanes = try Lanes.init(gpa, e, most_streams) }; + own.* = .{ .gpa = gpa, .e = e, .lanes = try Lanes.init(gpa, e, streams) }; errdefer own.lanes.deinit(); if (o.drafts) if (try drafterDir(gpa, io, o.drafter)) |draft_dir| { defer gpa.free(draft_dir); @@ -80,7 +85,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ own.lanes.drafter = own.draft; std.log.info("DFlash2 drafts every stream: {s}", .{draft_dir}); // the forward's ms by rows: a shared round's tree widths follow the curve - try own.lanes.calibrate(io, most_streams); + try own.lanes.calibrate(io, streams); }; return .{ .backend = own.lanes.backend(), diff --git a/zig/src/families/qwen27/cuda/round.zig b/zig/src/families/qwen27/cuda/round.zig index 33ac29cf26..e3f6f30621 100644 --- a/zig/src/families/qwen27/cuda/round.zig +++ b/zig/src/families/qwen27/cuda/round.zig @@ -33,7 +33,7 @@ pub const Round = struct { /// The conv rows gdn_pre reads for layer `i` (its DeltaNet index `j`): the stream's own, or the stacked copy. pub fn conv(r: *const Round, parts: []const Part, s: *const st.Scratch, g: c.Geometry, i: usize, j: usize) u64 { if (!r.multi) return parts[0].seq.conv[i]; - return s.conv_cat + j * st.max_streams * (c.conv_taps - 1) * g.convDim() * bf; + return s.conv_cat + j * s.streams() * (c.conv_taps - 1) * g.convDim() * bf; } /// Layer `i`'s attention cache offsets: tree_forward's (one stream, by layer) or the packed (layer, streams) rows. @@ -117,7 +117,7 @@ fn plans(gpa: std.mem.Allocator, parts: []const Part, starts: []const usize, mul /// Uploads the round's inputs and runs attention's _paths; the caller's layer loop reads the returned Round. pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, g: c.Geometry, parts: []const Part) !Round { - if (parts.len == 0 or parts.len > st.max_streams) return error.InvalidQwenWindows; + if (parts.len == 0 or parts.len > s.streams()) return error.InvalidQwenWindows; var W: usize = 0; var r: Round = undefined; r.starts_host[0] = 0; @@ -128,7 +128,7 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, W += p.tokens.len; r.starts_host[k + 1] = W; } - if (W > st.round_rows) return error.WindowTooWide; + if (W > s.window_rows) return error.WindowTooWide; r.rows = W; r.streams = parts.len; r.multi = parts.len > 1; @@ -222,7 +222,7 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, for (parts) |p| try offs.appendSlice(gpa, &.{ offsetOf(s, p.seq.keys[i]), offsetOf(s, p.seq.values[i]) }); continue; } - const cat = s.conv_cat + j * st.max_streams * keep * g.convDim() * bf; + const cat = s.conv_cat + j * s.streams() * keep * g.convDim() * bf; for (parts, 0..) |p, k| { try ptrs.append(gpa, p.seq.rec[i]); try stack.append(gpa, .{ .dst = cat + k * keep * g.convDim() * bf, .src = p.seq.conv[i], .bytes = keep * g.convDim() * bf }); diff --git a/zig/src/families/qwen27/cuda/state.zig b/zig/src/families/qwen27/cuda/state.zig index 47e4d46722..695029a483 100644 --- a/zig/src/families/qwen27/cuda/state.zig +++ b/zig/src/families/qwen27/cuda/state.zig @@ -122,6 +122,7 @@ pub const AttnRecord = struct { k: u64, v: u64 }; pub const Scratch = struct { arena: Arena, rows: usize, + window_rows: usize, // a shared round's rows at most: max_rows a stream // the residual stream and the normed rows with their group sums x: u64, h: u64, @@ -192,9 +193,11 @@ pub const Scratch = struct { attn_records: []AttnRecord, // by layer: the round's rotated keys and values gpa: std.mem.Allocator, - /// Scratch for prompt chunks of `prompt_rows` and windows of `max_rows`, over caches of `capacity` positions. - pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: c.Geometry, prompt_rows: usize, capacity: usize) !Scratch { + /// Scratch for prompt chunks of `prompt_rows` and rounds of `streams` windows of `max_rows`, over caches of `capacity` positions. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: c.Geometry, prompt_rows: usize, capacity: usize, stream_count: usize) !Scratch { + if (stream_count < 1 or stream_count > max_streams) return error.BadStreams; var s: Scratch = undefined; + s.window_rows = max_rows * stream_count; s.gpa = gpa; s.records = try gpa.alloc(GdnRecord, g.layers); errdefer gpa.free(s.records); @@ -208,9 +211,14 @@ pub const Scratch = struct { return s; } + /// The streams a shared round or batched prompt may hold. + pub fn streams(s: *const Scratch) usize { + return s.window_rows / max_rows; + } + fn layout(s: *Scratch, g: c.Geometry, prompt_rows: usize, capacity: usize) !void { - const R = @max(prompt_rows, round_rows); - const W = round_rows; + const W = s.window_rows; + const R = @max(prompt_rows, W); const bf = 2; const chunks = (capacity + W + tri.chunk - 1) / tri.chunk; // attention.plan_host: a query tile's items per committed chunk and folded group @@ -271,14 +279,14 @@ pub const Scratch = struct { s.part_m = try A.take(chunks * W * g.query_heads * 4); s.part_l = try A.take(chunks * W * g.query_heads * 4); s.conv_tmp = try A.take((c.conv_taps - 1 + R) * g.convDim() * bf); - s.conv_old = try A.take(max_streams * g.layers * (c.conv_taps - 1) * g.convDim() * bf); + s.conv_old = try A.take(s.streams() * g.layers * (c.conv_taps - 1) * g.convDim() * bf); s.copy_items = try A.take(2 * copy_items * @sizeOf(kern.Copy)); - s.taps = try A.take(R * 5 * g.hidden * bf); + s.taps = try A.take(@max(R, 2 * W) * 5 * g.hidden * bf); // a round's rows, then its kept rows gathered s.replay_table = try A.take(g.layers * (4 + max_streams) * 8); s.replay_rows = try A.take(max_streams * (W + 1) * 4); s.round = try A.take((6 * W + max_streams + 1 + W * c.conv_taps + W + 4 * max_streams + 3 * items * max_streams + W + 32) * 4); // + each part's 16-byte padding s.tables = try A.take(g.layers * max_streams * 8); - s.conv_cat = try A.take(g.layers * max_streams * (c.conv_taps - 1) * g.convDim() * bf); + s.conv_cat = try A.take(g.layers * s.streams() * (c.conv_taps - 1) * g.convDim() * bf); for (0..g.layers) |i| { if (c.linear(i)) { s.records[i] = .{ diff --git a/zig/src/families/qwen27/cuda/weights.zig b/zig/src/families/qwen27/cuda/weights.zig index 5c796bd43a..efe861391a 100644 --- a/zig/src/families/qwen27/cuda/weights.zig +++ b/zig/src/families/qwen27/cuda/weights.zig @@ -19,6 +19,7 @@ pub const Attention = struct { q: QLinear, k: QLinear, v: QLinear, o: QLinear, q pub const Mixer = union(enum) { delta: Delta, attention: Attention }; pub const Layer = struct { input_norm: u64, post_norm: u64, mixer: Mixer, gate: QLinear, up: QLinear, down: QLinear }; + pub const Weights = struct { gpa: std.mem.Allocator, g: c.Geometry, @@ -30,10 +31,13 @@ pub const Weights = struct { eos: [config_eos]u32 = undefined, // config.json's end tokens, then generation_config.json's eos_count: usize = 0, buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + host_buffers: std.ArrayList(cuda.HostBuffer) = .empty, // TF_QWEN27_HOST_EMBED=1: the token table, mapped pub fn deinit(w: *Weights) void { for (w.buffers.items) |*b| b.free(); w.buffers.deinit(w.gpa); + for (w.host_buffers.items) |*b| b.free(); + w.host_buffers.deinit(w.gpa); w.gpa.free(w.layers); w.* = undefined; } @@ -84,6 +88,15 @@ const Loader = struct { return L.host.items; } + /// A tensor in page-locked host memory the kernels read over PCIe (a row lookup reads a few KB a token). + fn rawHost(L: *Loader, t: Tensor) !u64 { + var b = try cuda.HostBuffer.allocMapped(L.ops.k.d, t.bytes.len); + errdefer b.free(); + try L.src.read(b.bytes, t.bytes); + try L.w.host_buffers.append(L.gpa, b); + return b.device(); + } + fn raw(L: *Loader, t: Tensor) !u64 { const ptr = try L.alloc(t.bytes.len); try L.src.upload(ptr, t.bytes); @@ -226,10 +239,13 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, layer.down = try L.tile(&ck, try join(&nb, p, "mlp.down_proj"), g.hidden, g.intermediate); layer.post_norm = try L.vector(&ck, try join(&nb, p, "post_attention_layernorm.weight"), g.hidden); } + // TF_QWEN27_HOST_EMBED=1: the token table in mapped host memory, 0.67 GiB back to the caches (a prompt row reads it over PCIe) + const host_embed = if (std.c.getenv("TF_QWEN27_HOST_EMBED")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; + const table = if (host_embed) &Loader.rawHost else &Loader.raw; w.embed = .{ - .w = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.weight", .u32, &.{ c.vocab, g.hidden / 8 })), - .s = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.scales", .bf16, &.{ c.vocab, g.hidden / 64 })), - .b = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.biases", .bf16, &.{ c.vocab, g.hidden / 64 })), + .w = try table(&L, try ck.expect(prefix ++ "model.embed_tokens.weight", .u32, &.{ c.vocab, g.hidden / 8 })), + .s = try table(&L, try ck.expect(prefix ++ "model.embed_tokens.scales", .bf16, &.{ c.vocab, g.hidden / 64 })), + .b = try table(&L, try ck.expect(prefix ++ "model.embed_tokens.biases", .bf16, &.{ c.vocab, g.hidden / 64 })), }; w.norm = try L.vector(&ck, "model.norm.weight", g.hidden); w.head = try L.tile(&ck, "lm_head", c.vocab, g.hidden); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 032f64c603..1a9a70bfb4 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -414,6 +414,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O var family_options: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments }; if (@hasField(F.Options, "slide")) family_options.slide = o.slide; if (@hasField(F.Options, "drafter")) family_options.drafter = o.drafter; + if (@hasField(F.Options, "streams")) family_options.streams = o.lanes; const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, family_options) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} with {s} glue kernels ({s})", .{ o.dir, kernels orelse "its own", @errorName(e) }); return null; From c6f544eb6685bcb549888abb5e294ad94aeea1bb Mon Sep 17 00:00:00 2001 From: "Riel St. Amand" <54451475+CanadaApollo6@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:26:48 -0400 Subject: [PATCH 4/5] qwen27 CUDA: the attention plan's parts from 16 bytes too _tail's STREAM pointer followed W i32 rows directly, so it landed off 16 bytes when W % 4 != 0 (a lone stream's small windows) and matched no captured variant. Rows, streams, items and parents each start on 16 bytes. --- zig/src/families/qwen27/cuda/round.zig | 19 +++++++++++++++---- zig/src/families/qwen27/cuda/state.zig | 4 ++-- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/zig/src/families/qwen27/cuda/round.zig b/zig/src/families/qwen27/cuda/round.zig index e3f6f30621..7c247f86f3 100644 --- a/zig/src/families/qwen27/cuda/round.zig +++ b/zig/src/families/qwen27/cuda/round.zig @@ -57,6 +57,7 @@ pub fn parentsOf(p: Part, buf: []i32) []const i32 { /// attention.plan_host: rows, (first, rows, keys, slots) per stream, items, then window parents (global rows). fn attnPlan(gpa: std.mem.Allocator, out: *std.ArrayList(i32), parts: []const Part, group: usize) !usize { for (parts, 0..) |p, s| for (p.tokens) |_| try out.append(gpa, @intCast(s)); + while (out.items.len % 4 != 0) try out.append(gpa, 0); // each part from 16 bytes, as planAt places them var start: usize = 0; for (parts) |p| { const w = p.tokens.len; @@ -77,6 +78,7 @@ fn attnPlan(gpa: std.mem.Allocator, out: *std.ArrayList(i32), parts: []const Par } } } + while (out.items.len % 4 != 0) try out.append(gpa, 0); start = 0; var buf: [tree.max_nodes]i32 = undefined; for (parts) |p| { @@ -86,6 +88,13 @@ fn attnPlan(gpa: std.mem.Allocator, out: *std.ArrayList(i32), parts: []const Par return n_items; } +/// attnPlan's parts in i32s from its start (the start on 16 bytes): rows, then streams, items and parents each from 16 bytes. +const PlanAt = struct { streams: usize, items: usize, parents: usize }; +fn planAt(W: usize, S: usize, n_items: usize) PlanAt { + const rows = std.mem.alignForward(usize, W, 4); + return .{ .streams = rows, .items = rows + 4 * S, .parents = rows + 4 * S + std.mem.alignForward(usize, 3 * n_items, 4) }; +} + /// Each stream's key and value caches as bf16 element offsets from the scratch's origin. fn offsetOf(s: *const st.Scratch, ptr: u64) i64 { return @divExact(@as(i64, @bitCast(ptr -% s.origin)), bf); @@ -175,8 +184,9 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, r.starts = s.plan + 3 * W * 4; r.tables = 0; r.offs = s.offs; - r.attn = .{ .rows = s.attn_plan, .streams = s.attn_plan + W * 4, .items = s.attn_plan + (W + 4) * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; - try t.paths(s.attn_plan + (W + 4 + 3 * n_items) * 4, s.paths, s.depths, W); + const at = planAt(W, 1, n_items); + r.attn = .{ .rows = s.attn_plan, .streams = s.attn_plan + at.streams * 4, .items = s.attn_plan + at.items * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; + try t.paths(s.attn_plan + at.parents * 4, s.paths, s.depths, W); return r; } // multi_tree_forward: positions, stream ids, plan entries and starts, ids, windows, attention's plan in one copy, each @@ -209,7 +219,8 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, r.ids = base + at_ids * 4; r.windows = base + at_windows * 4; const ab = base + at_attn * 4; - r.attn = .{ .rows = ab, .streams = ab + W * 4, .items = ab + (W + 4 * S) * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; + const at = planAt(W, S, n_items); + r.attn = .{ .rows = ab, .streams = ab + at.streams * 4, .items = ab + at.items * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; // the state pointer table and cache offsets, one copy each; every stream's conv rows stacked by layer var ptrs: std.ArrayList(u64) = .empty; defer ptrs.deinit(gpa); @@ -236,6 +247,6 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, try ops.s.synchronize(); r.tables = s.tables; r.offs = s.offs; - try t.paths(ab + (W + 4 * S + 3 * n_items) * 4, s.paths, s.depths, W); + try t.paths(ab + at.parents * 4, s.paths, s.depths, W); return r; } diff --git a/zig/src/families/qwen27/cuda/state.zig b/zig/src/families/qwen27/cuda/state.zig index 695029a483..e79467205d 100644 --- a/zig/src/families/qwen27/cuda/state.zig +++ b/zig/src/families/qwen27/cuda/state.zig @@ -271,7 +271,7 @@ pub const Scratch = struct { s.prompt_values = try A.take(R * g.kvInner() * bf); s.windows = try A.take(R * c.conv_taps * 4); s.plan = try A.take(R * 3 * 4); - s.attn_plan = try A.take((W + 4 + 3 * items + W) * 4); + s.attn_plan = try A.take((W + 4 + 3 * items + W + 6) * 4); // + its parts' 16-byte padding s.paths = try A.take(W * tri.max_nodes * 4); s.depths = try A.take(W * 4); s.offs = try A.take(g.layers * max_streams * 2 * 8); @@ -284,7 +284,7 @@ pub const Scratch = struct { s.taps = try A.take(@max(R, 2 * W) * 5 * g.hidden * bf); // a round's rows, then its kept rows gathered s.replay_table = try A.take(g.layers * (4 + max_streams) * 8); s.replay_rows = try A.take(max_streams * (W + 1) * 4); - s.round = try A.take((6 * W + max_streams + 1 + W * c.conv_taps + W + 4 * max_streams + 3 * items * max_streams + W + 32) * 4); // + each part's 16-byte padding + s.round = try A.take((6 * W + max_streams + 1 + W * c.conv_taps + W + 4 * max_streams + 3 * items * max_streams + W + 40) * 4); // + each part's 16-byte padding s.tables = try A.take(g.layers * max_streams * 8); s.conv_cat = try A.take(g.layers * s.streams() * (c.conv_taps - 1) * g.convDim() * bf); for (0..g.layers) |i| { From 99ce15afad6394832db5e78a9fd5d8f329aa5d6c Mon Sep 17 00:00:00 2001 From: "Riel St. Amand" <54451475+CanadaApollo6@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:26:48 -0400 Subject: [PATCH 5/5] qwen27 CUDA: comments within the lean checker's 120 columns --- zig/src/cuda/aot.zig | 2 +- zig/src/families/qwen27/cuda/engine.zig | 2 +- zig/src/families/qwen27/cuda/round.zig | 5 ++--- zig/src/families/qwen27/cuda/state.zig | 2 +- zig/src/families/qwen27/cuda/weights.zig | 2 +- 5 files changed, 6 insertions(+), 7 deletions(-) diff --git a/zig/src/cuda/aot.zig b/zig/src/cuda/aot.zig index fce864b3d1..1a736ed7a8 100644 --- a/zig/src/cuda/aot.zig +++ b/zig/src/cuda/aot.zig @@ -124,7 +124,7 @@ pub const Set = struct { if (!std.mem.eql(u8, v.spec.@"fn", function)) continue; if (matches(v.spec, args, consts, false)) return v; } - // else one compiled without assuming 16-byte alignment where these are aligned: the same code, valid for any value + // else one compiled without the 16-byte alignment hint: the same code, valid for any pointer for (self.variants) |*v| { if (!std.mem.eql(u8, v.spec.@"fn", function)) continue; if (matches(v.spec, args, consts, true)) return v; diff --git a/zig/src/families/qwen27/cuda/engine.zig b/zig/src/families/qwen27/cuda/engine.zig index b25a872f6e..2f24825a77 100644 --- a/zig/src/families/qwen27/cuda/engine.zig +++ b/zig/src/families/qwen27/cuda/engine.zig @@ -19,7 +19,7 @@ pub fn chunkBounds(start: usize, end: usize, size: usize, j: usize) [2]usize { return .{ start + (end - start) * j / n, start + (end - start) * (j + 1) / n }; } -/// `streams`: a shared round's windows at most (the scratch's window rows follow); `prompt_rows`: a prompt chunk's rows. +/// `streams`: a shared round's windows at most (the scratch follows); `prompt_rows`: a prompt chunk's rows. pub const Options = struct { context: usize, streams: usize = st.max_streams, prompt_rows: usize = prompt_rows }; pub const Engine = struct { diff --git a/zig/src/families/qwen27/cuda/round.zig b/zig/src/families/qwen27/cuda/round.zig index 7c247f86f3..9cc798dab6 100644 --- a/zig/src/families/qwen27/cuda/round.zig +++ b/zig/src/families/qwen27/cuda/round.zig @@ -88,7 +88,7 @@ fn attnPlan(gpa: std.mem.Allocator, out: *std.ArrayList(i32), parts: []const Par return n_items; } -/// attnPlan's parts in i32s from its start (the start on 16 bytes): rows, then streams, items and parents each from 16 bytes. +/// attnPlan's parts in i32s from its 16-byte start: rows, then streams, items and parents, each on 16 bytes. const PlanAt = struct { streams: usize, items: usize, parents: usize }; fn planAt(W: usize, S: usize, n_items: usize) PlanAt { const rows = std.mem.alignForward(usize, W, 4); @@ -189,8 +189,7 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, try t.paths(s.attn_plan + at.parents * 4, s.paths, s.depths, W); return r; } - // multi_tree_forward: positions, stream ids, plan entries and starts, ids, windows, attention's plan in one copy, each - // part from a 16-byte boundary (the captured Triton variants are specialized on their pointers' alignment) + // multi_tree_forward's tables in one copy, each part on 16 bytes: captured variants assume their pointer alignment const pad = struct { fn to16(gpa_: std.mem.Allocator, list: *std.ArrayList(i32)) !usize { while (list.items.len % 4 != 0) try list.append(gpa_, 0); diff --git a/zig/src/families/qwen27/cuda/state.zig b/zig/src/families/qwen27/cuda/state.zig index e79467205d..28366907eb 100644 --- a/zig/src/families/qwen27/cuda/state.zig +++ b/zig/src/families/qwen27/cuda/state.zig @@ -193,7 +193,7 @@ pub const Scratch = struct { attn_records: []AttnRecord, // by layer: the round's rotated keys and values gpa: std.mem.Allocator, - /// Scratch for prompt chunks of `prompt_rows` and rounds of `streams` windows of `max_rows`, over caches of `capacity` positions. + /// Scratch for `prompt_rows` prompt chunks and rounds of `streams` windows, over caches of `capacity` positions. pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: c.Geometry, prompt_rows: usize, capacity: usize, stream_count: usize) !Scratch { if (stream_count < 1 or stream_count > max_streams) return error.BadStreams; var s: Scratch = undefined; diff --git a/zig/src/families/qwen27/cuda/weights.zig b/zig/src/families/qwen27/cuda/weights.zig index efe861391a..e7c2b9b6b7 100644 --- a/zig/src/families/qwen27/cuda/weights.zig +++ b/zig/src/families/qwen27/cuda/weights.zig @@ -239,7 +239,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, layer.down = try L.tile(&ck, try join(&nb, p, "mlp.down_proj"), g.hidden, g.intermediate); layer.post_norm = try L.vector(&ck, try join(&nb, p, "post_attention_layernorm.weight"), g.hidden); } - // TF_QWEN27_HOST_EMBED=1: the token table in mapped host memory, 0.67 GiB back to the caches (a prompt row reads it over PCIe) + // TF_QWEN27_HOST_EMBED=1: the token table in mapped host memory (prompt rows read it over PCIe), 0.67 GiB freed const host_embed = if (std.c.getenv("TF_QWEN27_HOST_EMBED")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; const table = if (host_embed) &Loader.rawHost else &Loader.raw; w.embed = .{