qwen27 CUDA: a pre-sm_90 decode matmul, ~52% faster serving on an RTX 3090 (exact; stacked on #3)

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

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;

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,19 +195,22 @@ 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;
                 if (!std.mem.eql(u8, p.type, "fp32")) return false;
             },
             .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;

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 &.{};

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

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

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;
@@ -295,7 +289,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 +323,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 +334,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 +372,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];