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/cuda/aot.zig b/zig/src/cuda/aot.zig index b8e58d2e8e..1a736ed7a8 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 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; } 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/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..2f24825a77 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 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 { 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/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; 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 2b2d4e2697..9cc798dab6 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. @@ -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 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); + 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); @@ -117,7 +126,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 +137,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; @@ -175,30 +184,42 @@ 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, starts, ids, windows, attention's plan in one copy + // 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); + 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 }; + 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); @@ -211,7 +232,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 }); @@ -225,6 +246,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 4cbe06c07e..28366907eb 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_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; + 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 @@ -263,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); @@ -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) * 4); + 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 * 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..e7c2b9b6b7 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 (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 = .{ - .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;