Qwen3.8-27B on CUDA: RTX 3090 (sm_86) support, the scratch sized by --parallel

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

@@ -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];

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,13 +169,15 @@ 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);
     }
 
     /// 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;

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

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

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(),