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/6] 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 4b486922f..f7aa6a4c1 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/6] 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 b8e58d2e8..fce864b3d 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 2b2d4e269..33ac29cf2 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 4cbe06c07..47e4d4672 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/6] 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 a92416ce7..aab471404 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 02e365952..08ad3fc9a 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 52576d7cd..b25a872f6 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 4e4964be4..867ce1ee6 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 8a5f86fad..e73796e3d 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 0397ea706..5d75100bf 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 6e3f362fd..aa576026f 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 33ac29cf2..e3f6f3062 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 47e4d4672..695029a48 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 5c796bd43..efe861391 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 032f64c60..1a9a70bfb 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/6] 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 e3f6f3062..7c247f86f 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 695029a48..e79467205 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/6] 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 fce864b3d..1a736ed7a 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 b25a872f6..2f24825a7 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 7c247f86f..9cc798dab 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 e79467205..28366907e 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 efe861391..e7c2b9b6b 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 = .{ From 083532cf86f085c91927a26f13e71dc19169b1ea Mon Sep 17 00:00:00 2001 From: "Riel St. Amand" <54451475+CanadaApollo6@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:46:28 -0400 Subject: [PATCH 6/6] qwen27 CUDA: a pre-sm_90 decode matmul (qmm_lane), one launch for weights sharing an input Before sm_90, decode matmuls take qmm_lane.cu: the weights as the mma's A operand, so 1-16 rows fill a tile, and each tile's K slices added in slice order, the bits of 0.6.6's qmm.cu on the same card. Weights over one input (GDN qkv/z/b/a, attention q/k/v, gate/up) share one segmented launch, each at its own K split, so every output equals its lone launch. Wider rounds take row tiles of 16-64, 128 rows a turn. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/qmm_lane.cu | 509 +++++++++++++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen27/cuda/forward.zig | 12 +- zig/src/families/qwen27/cuda/kernels.zig | 132 +++++- 5 files changed, 637 insertions(+), 18 deletions(-) create mode 100644 zig/kernels/cuda/qmm_lane.cu diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 676bda8c8..f670922c1 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -25,6 +25,7 @@ const kernels = [_]Kernel{ .{ .name = "nemotron_ops", .flags = &.{"-O3"} }, // ours: Nemotron's serial feed, routed plan and rests .{ .name = "affine4_pack", .flags = &.{"-O3"} }, // ours: the MLX affine-4 repack into qlinear's tiles .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA + .{ .name = "qmm_lane", .flags = &.{"-O3"} }, // pre-sm_90 decode matmuls, K slices in slice order .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "qwen_ops", .flags = &.{"-O3"} }, // ours: the 27B's batched state copies, taps, drafter norms and rotary .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue diff --git a/zig/kernels/cuda/qmm_lane.cu b/zig/kernels/cuda/qmm_lane.cu new file mode 100644 index 000000000..0b5300a2d --- /dev/null +++ b/zig/kernels/cuda/qmm_lane.cu @@ -0,0 +1,509 @@ +// Pre-sm_90 4-bit decode matmuls (no clusters): a tile's K slices added in slice order, 0.6.6's qmm.cu bits. + +#include +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace { + +using namespace qmm_frag; + +// Tile shapes: BM rows by BN columns a block, WM x WN warps, each warp (BM / WM) x (BN / WN). +template +struct Tile { + static constexpr int THREADS = WM * WN * 32; + static constexpr int MT = BM / WM / 16; // m16 tiles a warp + static constexpr int NT = BN / WN / 8; // n8 tiles a warp + static constexpr int ROW = GS * 2; // bytes of one input row a group + static constexpr int CHUNKS = ROW / 16; + static constexpr int X = BM * ROW; // stage bytes: inputs, + static constexpr int W = BN * GS / 2; // weights, + static constexpr int S = BN * 2; // scales, biases (bf16), + static constexpr int XS = BM * 4; // and input sums (fp32) + static constexpr int STAGE = X + W + 2 * S + XS; + static constexpr int PARTIALS = MT * NT * 4 * THREADS * 4; // a K slice's partial, parked for the cluster sum + static constexpr int SMEM = STAGES * STAGE > PARTIALS ? STAGES * STAGE : PARTIALS; +}; + +// A tile's K slice: acc = fma(xs, b, fma(p, s, acc)) a group, in order from zero; every launch form runs these steps. +template +__device__ __forceinline__ void tile_slice(float (&acc)[Tile::MT] + [Tile::NT][4], + unsigned char* buf, const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, + const uint32_t* __restrict__ w, const __nv_bfloat16* __restrict__ scales, + const __nv_bfloat16* __restrict__ biases, int M, int KG, int per, int g0, int m0, int n0, int npad, int ldx) { + using T = Tile; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + // one group's inputs, weights, scales, biases and input sums into stage s + auto load = [&](int s, int g, int m0) { + unsigned char* p = stage(s); + for (int c = tid; c < BM * T::CHUNKS; c += T::THREADS) { + const int r = c / T::CHUNKS, ch = c % T::CHUNKS; + const int row = min(m0 + r, M - 1); + cp16z(p + r * T::ROW + swz(r, ch) * 16, x + static_cast(row) * ldx + g * GS + ch * 8, + m0 + r < M); + } + unsigned char* pw = p + T::X; + constexpr int TILE_BYTES = 64 * GS / 2; // one stored 64-column tile's group block + for (int c = tid; c < T::W / 16; c += T::THREADS) { + const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + const size_t tile = static_cast(n0 / 64 + t) * KG + g; + cp16(pw + c * 16, reinterpret_cast(w) + tile * TILE_BYTES + off * 16); + } + unsigned char* ps = pw + T::W; + for (int c = tid; c < 2 * (T::S / 16); c += T::THREADS) { + const int which = c / (T::S / 16), off = c % (T::S / 16); + const __nv_bfloat16* src = (which ? biases : scales) + static_cast(g) * npad + n0 + off * 8; + cp16(ps + which * T::S + off * 16, src); + } + float* px = reinterpret_cast(ps + 2 * T::S); + for (int r = tid; r < BM; r += T::THREADS) + cp4(px + r, xs + static_cast(min(m0 + r, M - 1)) * KG + g); + }; + + { +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f; + // acc = fma(xs, b, fma(p, s, acc)) per output, the order every row gets + auto epilogue = [&](float (&acc_)[T::MT][T::NT][4], const float (&d_)[T::MT][T::NT][4], + const float (&s_)[T::NT][2], const float (&b_)[T::NT][2], const float (&x_)[T::MT][2]) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc_[i][j][e] = __fmaf_rn(x_[i][e >> 1], b_[j][e & 1], __fmaf_rn(d_[i][j][e], s_[j][e & 1], + acc_[i][j][e])); + }; + float dp[T::MT][T::NT][4], sp[T::NT][2], bp[T::NT][2], xp[T::MT][2]; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < per) load(s, g0 + s, m0); + commit(); + } + for (int it = 0; it < per; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) load(next % STAGES, g0 + next, m0); + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W); + const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S); + uint32_t words[T::NT][GS / 32]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int v = 0; v < GS / 32; ++v) + words[j][v] = pw[((wn * T::NT + j) * 32 + lane) * (GS / 32) + v]; + float d[T::MT][T::NT][4]; +#pragma unroll + for (int kt = 0; kt < GS / 16; ++kt) { + uint32_t a[T::MT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = kt * 2 + (lane >> 4); + ldmatrix4(a[i], p + r * T::ROW + swz(r, ch) * 16); + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const uint32_t b0 = pair(words[j][kt / 2], (kt & 1) * 8); + const uint32_t b1 = pair(words[j][kt / 2], (kt & 1) * 8 + 4); +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[i][j], a[i], b0, b1); + else mma(d[i][j], a[i], b0, b1); + } + } + if (PIPE && kt == 0 && it > 0) epilogue(acc, dp, sp, bp, xp); // the previous group's, between MMAs + } + float sv[T::NT][2], bv[T::NT][2], xv[T::MT][2]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; + const __nv_bfloat162 s2 = *reinterpret_cast(ps + col); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + BN + col); + sv[j][0] = __low2float(s2); sv[j][1] = __high2float(s2); + bv[j][0] = __low2float(b2); bv[j][1] = __high2float(b2); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int row = wm * (BM / WM) + i * 16 + (lane >> 2); + xv[i][0] = px[row]; + xv[i][1] = px[row + 8]; + } + if constexpr (PIPE) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) dp[i][j][e] = d[i][j][e]; + xp[i][0] = xv[i][0]; xp[i][1] = xv[i][1]; + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + sp[j][0] = sv[j][0]; sp[j][1] = sv[j][1]; bp[j][0] = bv[j][0]; bp[j][1] = bv[j][1]; + } + } else { + epilogue(acc, d, sv, bv, xv); + } + } + if (PIPE && per > 0) epilogue(acc, dp, sp, bp, xp); + wait<0>(); + __syncthreads(); + } +} + +template +__device__ __forceinline__ void qmm_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, + void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, int group) { + using T = Tile; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int KG = K / GS, per = KG / SK; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y, slice = blockIdx.z, g0 = slice * per; + { + float acc[T::MT][T::NT][4]; + tile_slice(acc, buf, x, xs, w, scales, biases, M, KG, per, g0, m0, n0, npad, + ldx); + if constexpr (CLUSTER) { +#if __CUDA_ARCH__ >= 900 + // K slices of a tile form a cluster: slice 0 adds peers' partials in slice order, as reduce_kernel does + constexpr int E = T::MT * T::NT * 4; + auto cluster = cooperative_groups::this_cluster(); + float* mine = reinterpret_cast(buf); + if (slice != 0) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) mine[((i * T::NT + j) * 4 + e) * T::THREADS + tid] = acc[i][j][e]; + } + cluster.sync(); + if (slice == 0) { + for (int peer = 1; peer < SK; ++peer) { + const float* theirs = cluster.map_shared_rank(mine, peer); +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc[i][j][e] = acc[i][j][e] + theirs[((i * T::NT + j) * 4 + e) * T::THREADS + tid]; + } + } + cluster.sync(); // peers keep their memory until slice 0 has read it + if (slice != 0) return; +#else + __trap(); // clusters are 9.0+; qmm_cuda sends slices to reduce_kernel +#endif + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = m0 + wm * (BM / WM) + i * 16 + (lane >> 2) + h * 8; + if (row >= M) continue; + const float v0 = acc[i][j][2 * h], v1 = acc[i][j][2 * h + 1]; + if (SK > 1 && !CLUSTER) { + float* dst = part + (static_cast(slice) * M + row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else if (F32) { + float* dst = reinterpret_cast(out) + static_cast(row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else { + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * N + col; + if (col + 1 < N && (N & 1) == 0) + *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1); + else { + if (col < N) dst[0] = __float2bfloat16_rn(v0); + if (col + 1 < N) dst[1] = __float2bfloat16_rn(v1); + } + } + } + } + } +} + +// The K slices' partial sums added in slice order, then rounded once. +template +__device__ __forceinline__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ out, long long total, int SK) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total) return; + float acc = part[i]; + for (int s = 1; s < SK; ++s) acc = acc + part[s * total + i]; + if (F32) reinterpret_cast(out)[i] = acc; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); +} + +// tile_slice with the weights as the mma's A and 8-token input tiles as its B: the same steps, so tile_slice's bits. +template +__device__ __forceinline__ void tile_slice_swap(float (&acc)[TT][4], unsigned char* buf, + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, int M, int KG, int per, + int g0, int n0, int npad, int ldx) { + constexpr int BM = 16; + using T = Tile; + static_assert(T::NT == 2 && T::MT == 1, "a warp: one m16 tile of 16 columns"); + const int tid = threadIdx.x % T::THREADS, lane = tid & 31, wn = tid >> 5, gq = lane >> 2, t = lane & 3; + // G > 1: the block holds G slices of a tile, a group of T::THREADS each with its own stages, synced alone + auto sync = [] { + if constexpr (G == 1) __syncthreads(); + else asm volatile("bar.sync %0, %1;" :: "r"(1 + static_cast(threadIdx.x) / T::THREADS), "n"(T::THREADS)); + }; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + // a group's stage: each thread's 16-byte copies, its sources found once and stepped a group at a time + static_assert(T::THREADS == BM * T::CHUNKS && T::W == T::THREADS * 16 && 2 * (T::S / 16) <= T::THREADS, + "one copy of each kind a thread"); + const int xr = tid / T::CHUNKS, xc = tid % T::CHUNKS; + const bool xread = xr < M; + const int xoff = xr * T::ROW + swz(xr, xc) * 16; + const __nv_bfloat16* xsrc = x + static_cast(min(xr, M - 1)) * ldx + static_cast(g0) * GS + xc * 8; + const unsigned char* wsrc = reinterpret_cast(w) + + (static_cast(n0 / 64) * KG + g0) * (64 * GS / 2) + tid * 16; + const int sw = tid / (T::S / 16), so = tid % (T::S / 16); + const bool sread = tid < 2 * (T::S / 16); + const __nv_bfloat16* ssrc = (sw ? biases : scales) + static_cast(g0) * npad + n0 + so * 8; + const int soff = T::X + T::W + sw * T::S + so * 16; + const bool xsread = tid < BM; + const float* xssrc = xs + static_cast(min(tid, M - 1)) * KG + g0; + const int xsoff = T::X + T::W + 2 * T::S + tid * 4; + auto load = [&](int s, int i) { // group g0 + i into stage s + unsigned char* p = stage(s); + cp16z(p + xoff, xsrc + i * GS, xread); + cp16(p + T::X + tid * 16, wsrc + static_cast(i) * (64 * GS / 2)); + if (sread) cp16(p + soff, ssrc + static_cast(i) * npad); + if (xsread) cp4(p + xsoff, xssrc + i); + }; + +#pragma unroll + for (int u = 0; u < TT; ++u) acc[u][0] = acc[u][1] = acc[u][2] = acc[u][3] = 0.0f; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < per) load(s, s); + commit(); + } + const int c0 = wn * 16 + gq, c1 = c0 + 8; // the lane's columns in the block + for (int it = 0; it < per; ++it) { + wait(); + sync(); + const int next = it + STAGES - 1; + if (next < per) load(next % STAGES, next); + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W); + const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S); + uint32_t words[2][GS / 32]; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int v = 0; v < GS / 32; ++v) words[j][v] = pw[((wn * 2 + j) * 32 + lane) * (GS / 32) + v]; + float d[TT][4]; +#pragma unroll + for (int kt = 0; kt < GS / 16; ++kt) { + uint32_t a[4]; + const int r = (lane & 7) + ((lane >> 3) & 1) * 8; + ldmatrix4(a, p + r * T::ROW + swz(r, kt * 2 + (lane >> 4)) * 16); + const int sh = (kt & 1) * 8; + const uint32_t wa[4] = {pair(words[0][kt / 2], sh), pair(words[1][kt / 2], sh), + pair(words[0][kt / 2], sh + 4), pair(words[1][kt / 2], sh + 4)}; +#pragma unroll + for (int u = 0; u < TT; ++u) { + if (kt == 0) mma0(d[u], wa, a[u], a[u + 2]); + else mma(d[u], wa, a[u], a[u + 2]); + } + } + const float s0 = __bfloat162float(ps[c0]), s1 = __bfloat162float(ps[c1]); + const float b0 = __bfloat162float(ps[BN + c0]), b1 = __bfloat162float(ps[BN + c1]); +#pragma unroll + for (int u = 0; u < TT; ++u) { + const float x0 = px[u * 8 + 2 * t], x1 = px[u * 8 + 2 * t + 1]; + acc[u][0] = __fmaf_rn(x0, b0, __fmaf_rn(d[u][0], s0, acc[u][0])); + acc[u][1] = __fmaf_rn(x1, b0, __fmaf_rn(d[u][1], s0, acc[u][1])); + acc[u][2] = __fmaf_rn(x0, b1, __fmaf_rn(d[u][2], s1, acc[u][2])); + acc[u][3] = __fmaf_rn(x1, b1, __fmaf_rn(d[u][3], s1, acc[u][3])); + } + } + wait<0>(); + sync(); +} + +// Several weights of one input in one launch, each at its own K slices; a tile's last block adds them in slice order. +constexpr int SEGS = 4; +struct Seg { + const uint32_t* w; + const __nv_bfloat16* scales; + const __nv_bfloat16* biases; + void* out; + float* part; + int N, npad, SK, blocks, sync; // blocks: tiles x SK; sync: its tiles' first counter +}; +struct Segs { + Seg s[SEGS]; + int count; +}; + +template +__device__ __forceinline__ void put(void* out, size_t i, float v) { + if (F32) reinterpret_cast(out)[i] = v; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(v); +} + +// TT 0: tile_slice, 1-2: tile_slice_swap's 8-token tiles; G > 1: a block's G slices, summed in shared memory in order. +template +__device__ __forceinline__ void qmm_seg_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const Segs segs, int* __restrict__ sync, + int M, int K, int ldx) { + using T = Tile; + extern __shared__ __align__(128) unsigned char buf[]; + __shared__ int last; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + int b = blockIdx.x; + Seg sg = segs.s[0]; +#pragma unroll + for (int i = 1; i < SEGS; ++i) + if (i < segs.count && b >= sg.blocks / G) { // G > 1: a block a tile + b -= sg.blocks / G; + sg = segs.s[i]; + } + const int KG = K / GS, per = KG / sg.SK, tiles = sg.blocks / sg.SK; + const int slice = G > 1 ? static_cast(threadIdx.x) / T::THREADS : b / tiles; + const int tile = G > 1 ? b : b - slice * tiles, n0 = tile * BN; + constexpr int E = TT ? TT * 4 : T::MT * T::NT * 4; // the lane's outputs + float v[E]; + int rows[E], cols[E]; + if constexpr (TT == 0) { + float acc[T::MT][T::NT][4]; + tile_slice(acc, buf, x, xs, sg.w, sg.scales, sg.biases, M, KG, per, + slice * per, 0, n0, sg.npad, ldx); +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int q = (i * T::NT + j) * 4 + e; + v[q] = acc[i][j][e]; + rows[q] = wm * (BM / WM) + i * 16 + (lane >> 2) + (e >> 1) * 8; + cols[q] = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2 + (e & 1); + } + } else { + static_assert(WM == 1 && BM == 16, "the swapped form: one 16-row tile, a warp's 16 columns"); + float acc[TT][4]; + tile_slice_swap(acc, buf + (G > 1 ? slice * STAGES * T::STAGE : 0), x, xs, sg.w, + sg.scales, sg.biases, M, KG, per, slice * per, n0, sg.npad, ldx); +#pragma unroll + for (int u = 0; u < TT; ++u) +#pragma unroll + for (int e = 0; e < 4; ++e) { + v[u * 4 + e] = acc[u][e]; + rows[u * 4 + e] = u * 8 + 2 * (lane & 3) + (e & 1); + cols[u * 4 + e] = n0 + wn * 16 + (lane >> 2) + (e >> 1) * 8; + } + } + const int N = sg.N; + const size_t MN = static_cast(M) * N; + if constexpr (G > 1) { + float* red = reinterpret_cast(buf); // slice s's lane values at [s - 1][q][thread] + const int lt = threadIdx.x % T::THREADS; + __syncthreads(); // every slice's stages done + if (slice > 0) +#pragma unroll + for (int q = 0; q < E; ++q) red[((slice - 1) * E + q) * T::THREADS + lt] = v[q]; + __syncthreads(); + if (slice > 0) return; +#pragma unroll + for (int q = 0; q < E; ++q) { + float a = v[q]; +#pragma unroll + for (int s2 = 1; s2 < G; ++s2) a = a + red[((s2 - 1) * E + q) * T::THREADS + lt]; + if (rows[q] < M && cols[q] < N) put(sg.out, static_cast(rows[q]) * N + cols[q], a); + } + return; + } +#pragma unroll + for (int q = 0; q < E; ++q) { + if (rows[q] >= M || cols[q] >= N) continue; + const size_t at = static_cast(rows[q]) * N + cols[q]; + if (sg.SK == 1) put(sg.out, at, v[q]); + else sg.part[slice * MN + at] = v[q]; + } + if (sg.SK == 1) return; + __threadfence(); // this slice's partials, then its count + __syncthreads(); + if (tid == 0) last = atomicAdd(sync + sg.sync + tile, 1) == sg.SK - 1; + __syncthreads(); + if (!last) return; + __threadfence(); +#pragma unroll + for (int q = 0; q < E; ++q) { + if (rows[q] >= M || cols[q] >= N) continue; + const size_t at = static_cast(rows[q]) * N + cols[q]; + float p[8]; // every slice's load in flight, then added in order +#pragma unroll + for (int s = 0; s < 8; ++s) p[s] = s < sg.SK ? __ldcg(sg.part + s * MN + at) : 0.f; + float a = p[0]; +#pragma unroll + for (int s = 1; s < 8; ++s) + if (s < sg.SK) a = a + p[s]; + put(sg.out, at, a); + } + if (tid == 0) sync[sg.sync + tile] = 0; +} +} // namespace + +// Segmented decode launch (rows <= 16): TT 8-token tiles, G K slices a block, STAGES 3 or 4. +#define TF_SEG(ST, TT, G) \ + extern "C" __global__ void __launch_bounds__(128 * G) tf_qmm_seg_s##ST##_t##TT##_g##G( \ + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const Segs segs, int* __restrict__ sync, \ + int M, int K, int ldx) { \ + qmm_seg_kernel<64, 16, 64, 1, 4, ST, false, TT, G>(x, xs, segs, sync, M, K, ldx); \ + } +TF_SEG(3, 1, 1) TF_SEG(3, 2, 1) TF_SEG(3, 1, 2) TF_SEG(3, 2, 2) TF_SEG(3, 1, 4) TF_SEG(3, 2, 4) +TF_SEG(4, 1, 1) TF_SEG(4, 2, 1) TF_SEG(4, 1, 2) TF_SEG(4, 2, 2) TF_SEG(4, 1, 4) TF_SEG(4, 2, 4) +#undef TF_SEG + +// Wider rounds: BM-row tiles, K slices into part (SK, M, N) for tf_qmm_reduce, or straight to bf16 with one slice. +#define TF_LANE(BM) \ + extern "C" __global__ void __launch_bounds__(128) tf_qmm_lane_m##BM( \ + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, \ + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, \ + void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, \ + int group) { \ + qmm_kernel<64, BM, 64, 1, 4, 4, false, false>(x, xs, w, scales, biases, out, part, M, N, K, SK, npad, ldx, group); \ + } +TF_LANE(16) TF_LANE(32) TF_LANE(64) +#undef TF_LANE + +extern "C" __global__ void tf_qmm_reduce(const float* __restrict__ part, void* __restrict__ out, long long total, int SK) { + reduce_kernel(part, out, total, SK); +} + diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index ac6c02628..5a56f6501 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -24,6 +24,7 @@ pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_a pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes else &.{}; pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; pub const lane_gemv: []const u8 = if (available) &Blob("fatbin_lane_gemv").bytes else &.{}; +pub const qmm_lane: []const u8 = if (available) &Blob("fatbin_qmm_lane").bytes else &.{}; pub const affine4_pack: []const u8 = if (available) &Blob("fatbin_affine4_pack").bytes else &.{}; pub const qwen_ops: []const u8 = if (available) &Blob("fatbin_qwen_ops").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; diff --git a/zig/src/families/qwen27/cuda/forward.zig b/zig/src/families/qwen27/cuda/forward.zig index 867ce1ee6..32914175a 100644 --- a/zig/src/families/qwen27/cuda/forward.zig +++ b/zig/src/families/qwen27/cuda/forward.zig @@ -97,10 +97,7 @@ pub const Forward = struct { switch (layer.mixer) { .delta => |d| { const rec = s.records[i]; - try f.ops.dense(s.y, s.xs, d.qkv, rec.qkv, W); - try f.ops.dense(s.y, s.xs, d.z, s.z, W); - try f.ops.dense(s.y, s.xs, d.b, s.b, W); - try f.ops.dense(s.y, s.xs, d.a, s.a, W); + try f.ops.denseMany(s.y, s.xs, &.{ d.qkv, d.z, d.b, d.a }, &.{ rec.qkv, s.z, s.b, s.a }, W); const sid: ?u64 = if (r.multi) r.sids else null; try f.t.gdnPre(rec.qkv, r.conv(parts, s, g_, i, lin), d.conv, r.windows, sid, s.a, s.b, d.a_log, d.dt_bias, rec.q, rec.k, rec.v, rec.g, rec.beta, W, f.gdnShape()); const state = if (r.multi) 0 else parts[0].seq.rec[i]; @@ -112,9 +109,7 @@ pub const Forward = struct { }, .attention => |a| { const rec = s.attn_records[i]; - try f.ops.dense(s.y, s.xs, a.q, s.qg, W); - try f.ops.dense(s.y, s.xs, a.k, s.key, W); - try f.ops.dense(s.y, s.xs, a.v, rec.v, W); + try f.ops.denseMany(s.y, s.xs, &.{ a.q, a.k, a.v }, &.{ s.qg, s.key, rec.v }, W); try f.t.attnPrep(s.qg, s.key, a.q_norm, a.k_norm, r.pos, w.inv_freq, s.q_rot, rec.k, W, f.attnShape(), c.eps); try f.t.attention(s.q_rot, rec.k, rec.v, s.origin, r.offsets(i, att), r.attn, .{ .o = s.part_o, .m = s.part_m, .l = s.part_l }, s.attn, f.attnShape()); try f.t.gateMul(s.attn, s.qg, s.gated, s.gated_xs, W, g_.query_heads, c.head_dim); @@ -123,8 +118,7 @@ pub const Forward = struct { }, } x = try f.norm(x, s.r, layer.post_norm, W); - try f.ops.dense(s.y, s.xs, layer.gate, s.gate, W); - try f.ops.dense(s.y, s.xs, layer.up, s.up, W); + try f.ops.denseMany(s.y, s.xs, &.{ layer.gate, layer.up }, &.{ s.gate, s.up }, W); try f.t.swiglu(s.gate, s.up, s.act, s.act_xs, W, g_.intermediate); try f.ops.dense(s.act, s.act_xs, layer.down, s.pending, W); pending = s.pending; diff --git a/zig/src/families/qwen27/cuda/kernels.zig b/zig/src/families/qwen27/cuda/kernels.zig index f7aa6a4c1..04fe87dec 100644 --- a/zig/src/families/qwen27/cuda/kernels.zig +++ b/zig/src/families/qwen27/cuda/kernels.zig @@ -30,7 +30,12 @@ pub const GemvPart = extern struct { w: u64, scales: u64, biases: u64, out: u64, /// gdn.cu's Pending (the tree kernel's rows still to fold): all null here, commits replay instead. pub const Pending = extern struct { k: u64 = 0, v: u64 = 0, g: u64 = 0, beta: u64 = 0, rows: u64 = 0, row_stride: c_int = 0, counts: u64 = 0, count_stride: c_int = 0 }; +/// qmm_lane.cu's Seg and Segs (the segmented launch), passed by value: four weights at most. +pub const Seg = extern struct { w: u64, scales: u64, biases: u64, out: u64, part: u64, n: c_int, npad: c_int, sk: c_int, blocks: c_int, sync: c_int }; +pub const Segs = extern struct { s: [4]Seg, count: c_int }; + comptime { + std.debug.assert(@sizeOf(Seg) == 64 and @sizeOf(Segs) == 264); std.debug.assert(@sizeOf(Part) == 56 and @sizeOf(Parts) == 232 and @sizeOf(GemvPart) == 48 and @sizeOf(Pending) == 64); } @@ -38,6 +43,14 @@ pub const group_smem: u32 = 35328; // LaneTile<64, 16, 64, 1, 4, 8>::SMEM /// qmm_group's GB10 tiles 3-5 for wider rounds: rows a tile, columns a tile, threads, LaneTile::SMEM. pub const Wide = struct { bm: usize, bn: usize, threads: u32, smem: u32 }; pub const wide = [_]Wide{ .{ .bm = 32, .bn = 64, .threads = 128, .smem = 26112 }, .{ .bm = 64, .bn = 64, .threads = 128, .smem = 43008 }, .{ .bm = 64, .bn = 128, .threads = 256, .smem = 39168 } }; +/// qmm_lane.cu's Tile<64, BM, 64, 1, 4, STAGES>: a stage's bytes at 16 rows; SMEM at 4 stages for 16, 32, 64 rows. +pub const seg_stage: u32 = 16 * 128 + 64 * 32 + 2 * 64 * 2 + 16 * 4; +pub const lane_smem = [_]u32{ 17664, 26112, 43008 }; +/// Rows a wide lane launch takes at most (wider rounds run in turns), so K-slice partials fit `part`. +pub const lane_rows: usize = 128; +/// split_k stops once tiles x slices reach 192, so a weight's slices hold fewer than 384 column tiles of partials. +pub const part_bytes: usize = lane_rows * 384 * 64 * 4; +pub const sync_counters: usize = 1 << 14; pub const prefill_mm_smem: u32 = 41984; // Tile<64, 128, 128, 2, 2, 2>::SMEM: prompt_tile 9, two stages pub const pattn_smem: u32 = 98304; // two 32-key slots of 256 dims beside eight warps' staged queries @@ -49,11 +62,17 @@ pub const QLinear = struct { w: u64, s: u64, b: u64, n: usize, k: usize, npad: u pub const Kernels = struct { d: *const cuda.Driver, - mods: [15]cuda.Module, + mods: [16]cuda.Module, triton: cuda.aot.Set, group: cuda.Function, wide: [wide.len]cuda.Function, // tiles 3-5: rows <= 32, <= 64, more gemv: cuda.Function, + seg: [2][2][3]cuda.Function, // qmm_lane's segmented launch by stages (3, 4), 8-token tiles (1, 2), slices a block (1, 2, 4) + lane: [3]cuda.Function, // qmm_lane's row tiles of 16, 32 and 64 for wider rounds + reduce: cuda.Function, + part: cuda.DeviceBuffer, // K-slice partials of the pre-sm_90 launches + sync: cuda.DeviceBuffer, // the segmented launch's finished-slice counters, zero between launches + sms: usize, prefill_mm: cuda.Function, pattn: cuda.Function, trees: [cuda.kernels.gdn_variants.len]cuda.Function, // tree_kernel by gdn_variants: [0] the chain one @@ -84,7 +103,7 @@ pub const Kernels = struct { var k: Kernels = undefined; k.d = d; const kk = cuda.kernels; - const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.prefill_attention, kk.gdn, kk.gdn_prefill, kk.affine4_pack, kk.torch_argmax, kk.torch_topk, kk.torch_pointwise, kk.torch_indexing, kk.torch_movement, kk.torch_nemotron_constants, kk.qwen_ops, kk.torch_norm_rotary }; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.prefill_attention, kk.gdn, kk.gdn_prefill, kk.affine4_pack, kk.torch_argmax, kk.torch_topk, kk.torch_pointwise, kk.torch_indexing, kk.torch_movement, kk.torch_nemotron_constants, kk.qwen_ops, kk.torch_norm_rotary, kk.qmm_lane }; var loaded: usize = 0; errdefer for (k.mods[0..loaded]) |*m| m.unload(); for (images, 0..) |img, i| { @@ -108,12 +127,31 @@ pub const Kernels = struct { k.inv_freq = try k.mods[13].function("tf_inv_freq"); k.rotary = try k.mods[13].function("tf_rotary"); k.torch_rms = try k.mods[14].function("tf_dflash_rms5120_bf16_fma_kernel"); + inline for (0..2) |si| inline for (0..2) |ti| inline for (0..3) |gi| { + const st: u32 = 3 + si; + const g: u32 = 1 << gi; + const f = &k.seg[si][ti][gi]; + f.* = try k.mods[15].function(std.fmt.comptimePrint("tf_qmm_seg_s{d}_t{d}_g{d}", .{ st, ti + 1, g })); + try f.allowDynamicShared(if (g > 1) g * st * seg_stage else @max(st * seg_stage, 4096)); + }; + for (&k.lane, [_][:0]const u8{ "tf_qmm_lane_m16", "tf_qmm_lane_m32", "tf_qmm_lane_m64" }, lane_smem) |*f, name, smem| { + f.* = try k.mods[15].function(name); + try f.allowDynamicShared(smem); + } + k.reduce = try k.mods[15].function("tf_qmm_reduce"); const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); // gdn_prefill_cuda: a block takes 128 value rows, or 64 when there are fewer heads than SMs k.prompt_rows = if (heads >= sms) 128 else 64; k.prompt_chain = try k.mods[5].function(if (k.prompt_rows == 128) sym.chain128 else sym.chain64); k.triton = try cuda.aot.Set.load(gpa, io, d, ctx.device, triton_dir); errdefer k.triton.deinit(); + k.sms = sms; + k.clusters = try ctx.attribute(.compute_capability_major) >= 9; + k.part = try cuda.DeviceBuffer.alloc(d, if (k.clusters) 0 else part_bytes); + errdefer k.part.free(); + k.sync = try cuda.DeviceBuffer.alloc(d, if (k.clusters) 0 else sync_counters * 4); + errdefer k.sync.free(); + if (!k.clusters) try k.sync.fill32(0, null); try k.group.allowDynamicShared(group_smem); for (k.wide, wide) |f, t| try f.allowDynamicShared(t.smem); try k.gemv.allowDynamicShared(group_smem); @@ -123,7 +161,6 @@ 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; @@ -133,6 +170,8 @@ pub const Kernels = struct { pub fn deinit(k: *Kernels) void { k.triton.deinit(); + k.part.free(); + k.sync.free(); for (&k.mods) |*m| m.unload(); } }; @@ -169,15 +208,44 @@ 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 (!o.k.clusters and !(sk > 1 and rows <= o.k.gemv_rows)) return o.lane3090(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); } + /// Weights over one input: before sm_90 and up to 16 rows, a segmented launch a K split, each weight at its own. + pub fn denseMany(o: Ops, x: u64, xs: u64, qs: []const QLinear, outs: []const u64, rows: usize) !void { + if (o.k.clusters or rows > 16 or (o.k.gemv_rows > 0 and rows <= o.k.gemv_rows)) { + for (qs, outs) |q, out| try o.dense(x, xs, q, out, rows); + return; + } + var done: u8 = 0; + for (qs, 0..) |q0, i| { + if (done & (@as(u8, 1) << @intCast(i)) != 0) continue; + const sk = splitK(q0.n, q0.k); + var segs: Segs = std.mem.zeroes(Segs); + var blocks: usize = 0; + var sync: usize = 0; + var part: u64 = o.k.part.ptr; + for (qs[i..], outs[i..], i..) |q, out, j| { + if (splitK(q.n, q.k) != sk or q.k != q0.k or segs.count == 4) continue; + done |= @as(u8, 1) << @intCast(j); + const tiles = (q.n + 63) / 64; + segs.s[@intCast(segs.count)] = .{ .w = q.w, .scales = q.s, .biases = q.b, .out = out, .part = if (sk > 1) part else 0, .n = int(q.n), .npad = int(q.npad), .sk = int(sk), .blocks = int(tiles * sk), .sync = int(sync) }; + segs.count += 1; + blocks += tiles * sk; + sync += tiles; + part += sk * rows * q.n * 4; + } + if (sk > 1 and sync > sync_counters) return error.TooManyTiles; + try o.segLaunch(x, xs, segs, blocks, sk, rows, q0.k); + } + } + /// 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 (!o.k.clusters and !(sk > 1 and rows <= o.k.gemv_rows)) return o.lane3090(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); } @@ -202,14 +270,60 @@ 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 { + /// Before sm_90 (no clusters): qmm_lane's segmented launch to 16 rows, else row tiles 128 rows a turn. + fn lane3090(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + if (rows <= 16) return o.seg(x, xs, q, out, rows, sk); 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); + while (r < rows) : (r += lane_rows) { + try o.lane(x + r * q.k * 2, xs + r * (q.k / 64) * 4, q, out + r * q.n * 2, @min(lane_rows, rows - r), sk); } } + fn seg(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + const tiles = (q.n + 63) / 64; + const blocks = tiles * sk; + if (sk > 1 and tiles > sync_counters) return error.TooManyTiles; + var segs: Segs = std.mem.zeroes(Segs); + segs.count = 1; + segs.s[0] = .{ .w = q.w, .scales = q.s, .biases = q.b, .out = out, .part = if (sk > 1) o.k.part.ptr else 0, .n = int(q.n), .npad = int(q.npad), .sk = int(sk), .blocks = int(blocks), .sync = 0 }; + return o.segLaunch(x, xs, segs, blocks, sk, rows, q.k); + } + + fn segLaunch(o: Ops, x: u64, xs: u64, segs: Segs, blocks: usize, sk: usize, rows: usize, k: usize) !void { + // qmm_seg_cuda: past 4 blocks an SM, 3 stages; 2 or 4 slices side by side in a block + const st: u32 = if (blocks > 4 * o.k.sms) 3 else 4; + const g: u32 = if (sk == 2 or sk == 4) @intCast(sk) else 1; + const f = o.k.seg[st - 3][if (rows <= 8) 0 else 1][if (g == 1) 0 else if (g == 2) 1 else 2]; + var a: cuda.Args = .{}; + a.add(x); + a.add(xs); + a.add(segs); + a.add(o.k.sync.ptr); + for ([_]usize{ rows, k, k }) |v| a.add(int(v)); + const smem = if (g > 1) g * st * seg_stage else @max(st * seg_stage, 4096); + try cuda.launch.launch(f, .{ .grid = .{ .x = u(blocks / g) }, .block = .{ .x = 128 * g }, .shared = smem }, o.s, &a); + } + + fn lane(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + const i: usize = if (rows <= 16) 0 else if (rows <= 32) 1 else 2; + const bm: usize = @as(usize, 16) << @intCast(i); + const rows_t = (rows + bm - 1) / bm; + // a group's inputs stay near 12 MB of L2 while its blocks sweep the column tiles + const group = @max(1, @min(rows_t, (12 << 20) / (bm * q.k * 2))); + var a: cuda.Args = .{}; + for ([_]u64{ x, xs, q.w, q.s, q.b, out, if (sk > 1) o.k.part.ptr else 0 }) |v| a.add(v); + for ([_]usize{ rows, q.n, q.k, sk, q.npad, q.k, group }) |v| a.add(int(v)); + try cuda.launch.launch(o.k.lane[i], .{ .grid = .{ .x = u(rows_t * ((q.n + 63) / 64)), .z = u(sk) }, .block = .{ .x = 128 }, .shared = lane_smem[i] }, o.s, &a); + if (sk == 1) return; + const total = rows * q.n; + var r: cuda.Args = .{}; + r.add(o.k.part.ptr); + r.add(out); + r.add(@as(i64, @intCast(total))); + r.add(int(sk)); + try cuda.launch.launch(o.k.reduce, .{ .grid = .{ .x = u((total + 255) / 256) }, .block = .{ .x = 256 } }, o.s, &r); + } + /// 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;