From 663a64245b7fe8d8c871cb26c22a15d6183550d7 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:32:21 +0000 Subject: [PATCH 1/4] nemotron: ModelOpt checkpoints below sm_90, the lane matmul's slices reduced through a per-stream buffer The lane matmul already reduces its K slices through fp32 partials where a card has no clusters (qmmf.plan's reduce, qmmf.cu's reduce_kernel), and its FP4 and FP8 decode to bf16 for mma.sync, but the family passed it no partials buffer and refused sm_89 and below at load. Each stream's scratch now holds the most partials any dense shape writes at rows short of fusing, Ops carries it to decode and prompt rows, and the refusal goes. The reduce adds the slices in the cluster path's order, so the sums do not change. --- zig/src/families/nemotron/cuda_engine.zig | 2 +- zig/src/families/nemotron/cuda_kernels.zig | 6 +++--- zig/src/families/nemotron/cuda_state.zig | 24 +++++++++++++++++----- 3 files changed, 23 insertions(+), 9 deletions(-) diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index f287c4ced4..7c3250f686 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -334,7 +334,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_kernels.zig b/zig/src/families/nemotron/cuda_kernels.zig index 7932f1c437..b25d8f5bc0 100644 --- a/zig/src/families/nemotron/cuda_kernels.zig +++ b/zig/src/families/nemotron/cuda_kernels.zig @@ -136,7 +136,6 @@ pub const Kernels = struct { pub fn loadModelopt(k: *Kernels, ctx: *const cuda.Context) !void { if (k.modelopt != null) return; const major: u32 = @intCast(try ctx.attribute(.compute_capability_major)); - if (major < 9) return error.ModeloptNeedsSm90; // the lane matmul sums K slices in a cluster: no reduce scratch const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); var lane = try cuda.qmmf.Lane.load(k.d, major); errdefer lane.unload(); @@ -166,6 +165,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) fn go(o: Ops, f: cuda.Function, grid: [3]usize, block: u32, shared: u32, args: *cuda.Args) !void { try cuda.launch.launch(f, .{ .grid = .{ .x = u(grid[0]), .y = u(grid[1]), .z = u(grid[2]) }, .block = .{ .x = block }, .shared = shared }, o.s, args); @@ -214,7 +214,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 { @@ -231,7 +231,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. diff --git a/zig/src/families/nemotron/cuda_state.zig b/zig/src/families/nemotron/cuda_state.zig index 92bac057b9..aa05e234b4 100644 --- a/zig/src/families/nemotron/cuda_state.zig +++ b/zig/src/families/nemotron/cuda_state.zig @@ -81,7 +81,19 @@ 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 = 36; +const scratch_count = 37; + +/// 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 } }; + 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; +} /// 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 { @@ -107,6 +119,7 @@ fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { pairs * 4, items * 12, 8, pairs * 4, (pairs + 1023) / 1024 * (c.experts + 2) * 4, pairs * c.expert_width * 2, pairs * D * 4, shared_words * 4, R * kern.pocket_width * 4, + lanePartBytes(c), }; } @@ -218,6 +231,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) /// Sizes every buffer for `max_len` cache rows and `nch` attention chunk partials a row; the own sequence's first. pub fn init(d: *const cuda.Driver, c: Config, max_len: usize, nch: usize, carve: ?*cuda.Carveout) !Buffers { @@ -259,10 +273,10 @@ 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.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.lane_part, }; } From 2095d023a3516c82b963e0570b62bbd885536068 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 10:15:18 +0000 Subject: [PATCH 2/4] nemotron cuda: each prompt segment reduces through its own partials, and the loader's Ops needs no scratch A staggered prompt's segments built their Ops without a partials buffer, so below sm_90 a segment's rows short of fusing wrote the lane matmul's slices through a null pointer; each now passes its own scratch set's lane_part. The loader took Engine.ops(), which reads the buffers made after it; it now gets an Ops without one, since it uploads and packs and runs no lane matmul. --- zig/src/families/nemotron/cuda_engine.zig | 4 +++- zig/src/families/nemotron/cuda_segments.zig | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index 7c3250f686..18ffd6b0fa 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -156,7 +156,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(); e.b = try state.Buffers.init(ctx.d, e.c, e.max_len, e.nch, e.carve); errdefer e.b.deinit(); diff --git a/zig/src/families/nemotron/cuda_segments.zig b/zig/src/families/nemotron/cuda_segments.zig index 158155d965..02087c0bd3 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; seg[k] = .{ .f = f, .walk = .{ .rows = rows, .pos = e.pos + k * R } }; From fcb45bdea9c73d9c7bb1113057473c6e002c4bc2 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 10:30:23 +0000 Subject: [PATCH 3/4] tf-cuda-test modelopt: the lane rows get slice partials, which a card without clusters reduces through The fp8 and bf16 lane checks ([320, 2048], split-K 4) passed no partials buffer, so at 16 and 128 rows below sm_90 the lane's reduce wrote through a null pointer; they now pass one sized for the widest unfused call, as the fp8 and nvfp4 tests do. --- zig/tests/cuda/modelopt_tests.zig | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/zig/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig index faa25fdfdd..8812d1b7ef 100644 --- a/zig/tests/cuda/modelopt_tests.zig +++ b/zig/tests/cuda/modelopt_tests.zig @@ -187,22 +187,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| { From 0b60b397ca890da874834831af1085d91c49f665 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:38:01 +0000 Subject: [PATCH 4/4] nemotron cuda: the forward's Ops keeps the lane partials when it takes its buffers' ck8 1.0.12's Forward.init rebuilds its Ops from the bound buffers to carry their e4m3 rows (ck8), which dropped the partials buffer the engine and each segment pass in: on a card without clusters the first prompt's lane matmul reduced through a null pointer (CUDA_ERROR_LAUNCH_FAILED on an RTX 4090). It now takes the same buffers' lane_part. tf-cuda-test modelopt runs its lane rows from sm_80, the floor this branch sets, instead of sm_90. --- zig/src/families/nemotron/cuda_forward.zig | 2 +- zig/tests/cuda/modelopt_tests.zig | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) 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/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig index 9dce145122..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);