nemotron cuda: ModelOpt checkpoints below sm_90, the lane matmul's slices reduced through a per-stream buffer

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 {

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,

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);
 }

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);

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;

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,
         };
     }
 

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| {