diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index 99e6b17816..6c105dff1f 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -180,7 +180,9 @@ pub const Engine = struct { std.log.err("no captured attention kernel covers {d} chunks of 512 keys; capture one for this context", .{chunks}); return error.ContextUnsupported; }) else chunks; - e.w = try weights.load(gpa, io, e.ops(), model_dir, e.c, opts.mtp); + // ops() reads e.b, made below; the loader uploads and packs and runs no lane matmul + const load_ops: kern.Ops = .{ .k = &e.k, .s = e.stream }; + e.w = try weights.load(gpa, io, load_ops, model_dir, e.c, opts.mtp); errdefer e.w.deinit(); if (opts.vision) |lim| { var problem: []const u8 = ""; @@ -373,7 +375,7 @@ pub const Engine = struct { } pub fn ops(e: *const Engine) kern.Ops { - return .{ .k = &e.k, .s = e.stream }; + return .{ .k = &e.k, .s = e.stream, .part = e.b.lane_part }; } pub fn forward(e: *const Engine, dump: ?*Dump) Forward { diff --git a/zig/src/families/nemotron/cuda_forward.zig b/zig/src/families/nemotron/cuda_forward.zig index 2a7c6d75a1..1736897d42 100644 --- a/zig/src/families/nemotron/cuda_forward.zig +++ b/zig/src/families/nemotron/cuda_forward.zig @@ -69,7 +69,7 @@ pub const Forward = struct { .c = c, .w = w, .b = b, - .ops = .{ .k = ops.k, .s = ops.s, .ck8 = b.ck8 }, // FP8 layers' e4m3 rows in these buffers + .ops = .{ .k = ops.k, .s = ops.s, .ck8 = b.ck8, .part = b.lane_part }, // FP8 rows and lane partials in these buffers .glue = .{ .set = if (ops.k.triton) |*t| t else null, .f = &ops.k.glue, diff --git a/zig/src/families/nemotron/cuda_kernels.zig b/zig/src/families/nemotron/cuda_kernels.zig index 38cc0c48ab..c28bbb1f53 100644 --- a/zig/src/families/nemotron/cuda_kernels.zig +++ b/zig/src/families/nemotron/cuda_kernels.zig @@ -171,14 +171,14 @@ pub const Kernels = struct { /// The checkpoint's own math needs this major capability: the block-scaled FP4 mma (sm_120a, sm_121a). pub const ck_major: u32 = 12; -/// FP8 and NVFP4 paths stop below this full capability: decode's split-K sums meet in a thread-block cluster. -pub const modelopt_min: u32 = 90; +/// FP8 and NVFP4 paths stop below this full capability: they decode to bf16 for mma.sync. +pub const modelopt_min: u32 = 80; /// Why a ModelOpt checkpoint cannot load at `capability` (10 * major + minor), or null when it may. pub fn modeloptRefusal(capability: u32) ?[]const u8 { if (capability >= modelopt_min) return null; - return "FP8 and NVFP4 weights meet their split-K sums in a thread-block cluster: " ++ - "a ModelOpt checkpoint needs compute capability 9.0 or newer"; + return "FP8 and NVFP4 weights decode to bf16 for mma.sync: " ++ + "a ModelOpt checkpoint needs compute capability 8.0 or newer"; } /// qmm.split_k (cuda/qlinear.zig). @@ -199,6 +199,7 @@ fn u(x: usize) u32 { pub const Ops = struct { k: *const Kernels, s: cuda.Stream, + part: u64 = 0, // the lane matmul's slice partials where it reduces them (no clusters below sm_90) ck8: u64 = 0, // the e4m3 rows an FP8 projection reads in the checkpoint's math (its buffers' ck8) fn go(o: Ops, f: cuda.Function, grid: [3]usize, block: u32, shared: u32, args: *cuda.Args) !void { @@ -248,7 +249,7 @@ pub const Ops = struct { /// Decode rows of a projection on this stream (the affine-4 paths' bits unchanged); xs: x's 64-group sums. pub fn dense(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize) !void { - return o.linear().decode(o.s, q, .{ .x = x, .sums = xs }, out, rows); + return o.linear().decode(o.s, q, .{ .x = x, .sums = xs, .part = o.part }, out, rows); } pub fn gemvSplit(o: Ops, x: u64, xs: u64, q: Affine4, out: u64, rows: usize, sk: usize, sp: Split) !void { @@ -265,7 +266,7 @@ pub const Ops = struct { /// Prompt rows of a projection on this stream, by its format's prompt GEMM. pub fn prefillDense(o: Ops, x: u64, q: QLinear, out: u64, rows: usize) !void { - return o.linear().prompt(o.s, q, .{ .x = x }, out, rows); + return o.linear().prompt(o.s, q, .{ .x = x, .part = o.part }, out, rows); } /// nemotron_ops' rest: y (bf16, or fp32 when `y_f32`) += r x at rows `i * row_mul + row_add` for i < rows. @@ -413,7 +414,7 @@ test "items follow the Python shapes" { try std.testing.expectEqual(@as(usize, 8), maxItems(8, 130, 16)); } -test "ModelOpt's floor refuses every card below sm_90, the Ampere ones included, and admits sm_90 up" { - for ([_]u32{ 75, 80, 86, 89 }) |c| try std.testing.expect(modeloptRefusal(c) != null); - for ([_]u32{ 90, 100, 120, 121 }) |c| try std.testing.expect(modeloptRefusal(c) == null); +test "ModelOpt's floor refuses every card below sm_80 and admits Ampere up, the cards without clusters included" { + for ([_]u32{ 70, 75 }) |c| try std.testing.expect(modeloptRefusal(c) != null); + for ([_]u32{ 80, 86, 89, 90, 100, 120, 121 }) |c| try std.testing.expect(modeloptRefusal(c) == null); } diff --git a/zig/src/families/nemotron/cuda_latent.zig b/zig/src/families/nemotron/cuda_latent.zig index 634b1555ef..13db61e9de 100644 --- a/zig/src/families/nemotron/cuda_latent.zig +++ b/zig/src/families/nemotron/cuda_latent.zig @@ -357,7 +357,7 @@ pub fn run(f: *const fwd.Forward, m: W.MoE, rows: usize) !void { const Lw: u64 = c.latent; try route(f, m, b.y, rows); try f.mark(.route); - try o.linear().decode(o.s, lt.into, .{ .x = b.y }, b.lat, rows); + try o.linear().decode(o.s, lt.into, .{ .x = b.y, .part = o.part }, b.lat, rows); if (f.tp != null) try worker.send(f, lt.id, rows); // the other ranks' shares start while ours run try partials(o, lt.t, .{ .plan = b.plan, .act = b.act, .sact = b.sact, .sids = b.sids, .ck4 = b.ck4 }, b.pick, b.lat, b.y, rows, own, c, b.ymoe); try f.mark(.experts_down); @@ -385,7 +385,7 @@ pub fn run(f: *const fwd.Forward, m: W.MoE, rows: usize) !void { a.add(b.comb); a.add(int(c.latent)); try go(fns.combine, o.s, .{ rows, (c.latent + 255) / 256, 1 }, 256, 0, &a); - try o.linear().decode(o.s, lt.out, .{ .x = b.comb }, b.delta, rows); + try o.linear().decode(o.s, lt.out, .{ .x = b.comb, .part = o.part }, b.delta, rows); var z: cuda.Args = .{}; z.add(b.delta); z.add(b.ymoe + @as(u64, pairs) * Lw * 4); diff --git a/zig/src/families/nemotron/cuda_segments.zig b/zig/src/families/nemotron/cuda_segments.zig index bf32087b50..ec264ea3e9 100644 --- a/zig/src/families/nemotron/cuda_segments.zig +++ b/zig/src/families/nemotron/cuda_segments.zig @@ -85,7 +85,7 @@ pub fn prefill(e: *Engine, s: *Segments, prompt: []const u32, head: ?*Head, part for (0..c.parts) |k| { // every segment's ids go up on the caller's stream, ahead of the call's fork const rows = segs.chunkRows(c.rows, R, k); try e.ops().upload(bufs[k].p_ids, std.mem.sliceAsBytes(host[k * R ..][0..rows])); - const ops: kern.Ops = .{ .k = &e.k, .s = s.runner.stream(e.stream, k) }; + const ops: kern.Ops = .{ .k = &e.k, .s = s.runner.stream(e.stream, k), .part = bufs[k].lane_part }; var f = fwd.Forward.init(e.c, &e.w, bufs[k], ops, e.max_len, e.nch, e.sampling != null); f.untaught = e.bound.untaught; f.vision = e.vision; diff --git a/zig/src/families/nemotron/cuda_state.zig b/zig/src/families/nemotron/cuda_state.zig index 992f3ca739..289888612e 100644 --- a/zig/src/families/nemotron/cuda_state.zig +++ b/zig/src/families/nemotron/cuda_state.zig @@ -83,7 +83,18 @@ fn seqSizes(c: Config, max_len: usize) [seq_fields.len]usize { return .{ kv, kv, m[0], m[1], nm * 2 * W * cd * 2, nm * 2 * W * cd * 2, nm * 2 * W * c.mamba_heads * 4, 16, W * 4, W * c.hidden * 2, W * 4, @sizeOf(sampler.Rule) }; } -const scratch_count = 43; +/// The most fp32 slice partials the lane matmul writes when it reduces them: every dense shape, rows short of fusing. +fn lanePartBytes(c: Config) usize { + const shapes = [_][2]usize{ .{ c.projDim(), c.hidden }, .{ c.hidden, c.inner() }, .{ c.qkvDim(), c.hidden }, .{ c.hidden, c.heads * c.head_dim }, .{ c.vocab, c.hidden }, .{ c.hidden, 2 * c.hidden }, .{ c.latent, c.hidden }, .{ c.hidden, c.latent } }; + const rows = cuda.qmmf.fused_rows - 1; + var most: usize = 16; + for (shapes) |s| { + const p = cuda.qmmf.plan(rows, s[0], s[1], false); + if (p.reduce) most = @max(most, @as(usize, p.sk) * rows * s[0] * 4); + } + return most; +} +const scratch_count = 44; /// Scratch windows and chunks share on one stream: window logits, chunk io, activations, the expert plan. fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { @@ -121,7 +132,7 @@ fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { ymoe, shared_words * 4, R * kern.pocket_width * 4, R * c.latent * 2, R * c.latent * 2, R * own * sw * 2, R * own * 4, job, ck8, - ck4, + ck4, lanePartBytes(c), }; } @@ -233,6 +244,7 @@ pub const Buffers = struct { rule: u64, // the sequence's sampled rule, as sample.cu reads it shared: u64, // the shared expert's fixed plans: 32 members, then (items, counts) for 1..16 rows pocket: u64, // the pocket's activations, fp32 [rows, pocket_width] + lane_part: u64, // the lane matmul's fp32 slice partials on a card without clusters (rows short of fusing) lat: u64, // a latent MoE's rows in the latent width (bf16) comb: u64, // its routed pairs' latent sum a row (bf16) sact: u64, // its shared expert's activations (bf16) @@ -281,11 +293,11 @@ pub const Buffers = struct { /// Every scratch field, in scratchSizes' order. fn scratchPtrs(b: *Buffers) [scratch_count]*u64 { return .{ - &b.logits, &b.p_meta, &b.p_ids, &b.p_hidden, &b.p_logits, &b.p_sampled, &b.emb, &b.h[0], &b.h[1], &b.y, - &b.xs, &b.delta, &b.proj, &b.p_xc, &b.sy, &b.g, &b.gxs, &b.qkv, &b.q, &b.att, - &b.axs, &b.po, &b.pm, &b.pl, &b.part, &b.pick, &b.wts, &b.plan.members, &b.plan.items, &b.plan.counts, - &b.plan.rank, &b.plan.hist, &b.act, &b.ymoe, &b.shared, &b.pocket, &b.lat, &b.comb, &b.sact, &b.sids, - &b.job, &b.ck8, &b.ck4, + &b.logits, &b.p_meta, &b.p_ids, &b.p_hidden, &b.p_logits, &b.p_sampled, &b.emb, &b.h[0], &b.h[1], &b.y, + &b.xs, &b.delta, &b.proj, &b.p_xc, &b.sy, &b.g, &b.gxs, &b.qkv, &b.q, &b.att, + &b.axs, &b.po, &b.pm, &b.pl, &b.part, &b.pick, &b.wts, &b.plan.members, &b.plan.items, &b.plan.counts, + &b.plan.rank, &b.plan.hist, &b.act, &b.ymoe, &b.shared, &b.pocket, &b.lat, &b.comb, &b.sact, &b.sids, + &b.job, &b.ck8, &b.ck4, &b.lane_part, }; } diff --git a/zig/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig index 1e9f5162f2..e1cce35fa0 100644 --- a/zig/tests/cuda/modelopt_tests.zig +++ b/zig/tests/cuda/modelopt_tests.zig @@ -57,9 +57,9 @@ pub fn run(gpu: Gpu) !void { defer stream.deinit(); try packers(gpu, rng, &pack, stream); const capability = try gpu.ctx.capability(); - if (capability / 10 < 9) { - std.debug.print("SKIP modelopt rows: the decode lane matmul sums its split-K slices in a cluster, " ++ - "and a ModelOpt checkpoint refuses below sm_90 anyway\n", .{}); + if (capability / 10 < 8) { + std.debug.print("SKIP modelopt rows: the lane matmul's bf16 mma needs sm_80, " ++ + "and a ModelOpt checkpoint refuses below it anyway\n", .{}); return; } var lane = try qmmf.Lane.load(gpu.d, capability / 10); @@ -193,22 +193,26 @@ fn lanes(gpu: Gpu, rng: std.Random, p: *const modelopt.Packer, lin: cuda.qlinear defer y.free(); var y2 = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 2); defer y2.free(); - try lin.decode(s, q, .{ .x = dx.ptr }, y.ptr, 16); + // the lane's slice partials where it reduces them (no clusters below sm_90): the widest unfused call is 128 rows + var part = try cuda.DeviceBuffer.alloc(gpu.d, @as(usize, qmmf.splitK(n, k)) * 128 * n * 4); + defer part.free(); + const in: cuda.qlinear.Rows = .{ .x = dx.ptr, .part = part.ptr }; + try lin.decode(s, q, in, y.ptr, 16); try s.synchronize(); const wide = try check.download(gpu, y); defer a.free(wide); for (1..16) |rows| { try y2.fill8(0xff, s.handle); - try lin.decode(s, q, .{ .x = dx.ptr }, y2.ptr, rows); + try lin.decode(s, q, in, y2.ptr, rows); try s.synchronize(); try same(gpu, "decode rows at a narrower width", y2, wide[0 .. rows * n * 2]); } var worst = [2]f64{ 0, 0 }; - try lin.prompt(s, q, .{ .x = dx.ptr }, y.ptr, m); + try lin.prompt(s, q, in, y.ptr, m); try s.synchronize(); const prompt = try check.download(gpu, y); defer a.free(prompt); - try lin.prompt(s, q, .{ .x = dx.ptr }, y2.ptr, 128); + try lin.prompt(s, q, in, y2.ptr, 128); try s.synchronize(); try same(gpu, "prompt rows in a chunk of 128", y2, prompt[0 .. 128 * n * 2]); for (0..m) |r| for (0..n) |c| {