qwen3_5 cuda: serve dense EXL3 checkpoints with DFlash2, tensor-core prompt tiles, and host-mapped embeddings

zig/qwen35_formats_host.zig

@@ -43,19 +43,23 @@ fn config(quant: []const u8, moe: bool, pocket: bool) ![]u8 {
 const exl3 = "\"quantization_config\":{\"quant_method\":\"exl3\"}";
 const ct = "\"quantization_config\":{\"quant_method\":\"compressed-tensors\"}";
 
-test "format coexistence: EXL3 stays MoE only and ModelOpt keeps its admission and refusals" {
+test "format coexistence: EXL3 admits MoE and dense while ModelOpt keeps its admission and refusals" {
     const x = try config(exl3, true, false);
     defer t.allocator.free(x);
     const gx = try geo.parse(t.allocator, x);
-    try t.expect(gx.exl3);
+    try t.expect(gx.exl3 and gx.moe());
     try t.expect(!gx.kinds.modelopt);
+    const d = try config(exl3, false, false);
+    defer t.allocator.free(d);
+    const gd = try geo.parse(t.allocator, d);
+    try t.expect(gd.exl3 and !gd.moe() and !gd.kinds.modelopt);
     const m = try config(ct, false, false);
     defer t.allocator.free(m);
     const gm = try geo.parse(t.allocator, m);
     try t.expect(gm.kinds.modelopt and gm.kinds.ct and !gm.exl3);
     try t.expectEqual(@as(u8, 16), gm.bits.embed);
     for ([_]struct { q: []const u8, moe: bool, pocket: bool, err: anyerror }{
-        .{ .q = exl3, .moe = false, .pocket = false, .err = error.UnsupportedQwenQuantization },
+        .{ .q = exl3, .moe = false, .pocket = true, .err = error.UnsupportedQwenQuantization },
         .{ .q = exl3, .moe = true, .pocket = true, .err = error.UnsupportedQwenQuantization },
         .{ .q = ct, .moe = true, .pocket = false, .err = error.UnsupportedQwenQuantization },
         .{ .q = ct, .moe = false, .pocket = true, .err = error.UnsupportedLearningFormat },
@@ -255,3 +259,17 @@ test "format coexistence: EXL3 BF16 prompt: 4096 rows reach both GDN scalar proj
 test "format coexistence: EXL3 BF16 prompt: existing rows and partial final slices keep every argument and row" {
     for ([_]usize{ 0, 1, 16, 17, 127, 128, 129, 2049, 4095 }) |rows| try bf16Prompt(rows);
 }
+
+test "format coexistence: disjoint mixer and dense MLP scratch regions share one arena union" {
+    const c = try config(exl3, false, false);
+    defer t.allocator.free(c);
+    var e: Engine = undefined;
+    e.g = try geo.parse(t.allocator, c);
+    e.arena = .{ .d = undefined, .ptr = 0x100000, .len = 1 << 30 };
+    e.arena_at = 0;
+    const measured = try e.scratch(.measure);
+    e.arena_at = 0;
+    try t.expectEqual(measured, try e.scratch(.allocate));
+    try t.expectEqual(e.y, e.gate);
+    try t.expectEqual(e.z, e.qraw);
+}

zig/src/cuda/exl3.zig

@@ -219,8 +219,28 @@ pub const Dense = struct {
     /// Grow every scratch buffer to what 128 rows of `w` need, so a call inside a graph capture allocates nothing.
     pub fn reserve(de: *Dense, w: Weight) !void {
         if (plan(w.k, w.n, 192, 8).sk > 1) _ = try de.scratch(&de.z, Linear.zBytes(w, 128));
-        _ = try de.scratch(&de.counters, Linear.counterBytes(w));
+        _ = try de.zeroedCounters(Linear.counterBytes(w));
         _ = try de.scratch(&de.xh, 128 * w.k * 2);
+        if (candidateWeight(w)) |_| {
+            if (de.candidate_k == null or de.candidate_k.?.k_max < w.k) {
+                if (kernels.available) {
+                    var mod = try Module.load(de.d, kernels.exl3_candidate);
+                    errdefer mod.unload();
+                    const replacement = try candidate.Kernels.resolve(de.d, mod, w.k);
+                    if (de.candidate_k) |*k| k.deinit();
+                    if (de.candidate_mod) |*old| old.unload();
+                    de.candidate_mod = mod;
+                    de.candidate_k = replacement;
+                }
+            }
+        } else |_| {}
+    }
+
+    fn zeroedCounters(de: *Dense, bytes: usize) !u64 {
+        if (de.counters != null and de.counters.?.len >= bytes) return de.counters.?.ptr;
+        const ptr = try de.scratch(&de.counters, bytes);
+        try de.counters.?.fill8(0, null);
+        return ptr;
     }
 
     /// Decode rows: one call (the kernel takes 1..128), any count above in 128-row blocks.
@@ -229,16 +249,18 @@ pub const Dense = struct {
         return de.blocks(s, w, x, x_dtype, rows, y, y_dtype);
     }
 
-    /// Prompt rows: 128-row blocks of `matmul`, so a row's bits never depend on where its chunk starts.
+    /// Prompt rows: 64-row tensor-core prompt tiles when compatible, else 128-row blocks of `matmul`.
     pub fn prompt(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void {
-        if (de.implementation == .tiles) return de.tiles(s, w, x, x_dtype, rows, y, y_dtype, true);
+        if (de.implementation == .tiles or (rows > 16 and x_dtype == .bf16 and y_dtype != .f16 and (candidateWeight(w) catch null) != null))
+            return de.tiles(s, w, x, x_dtype, rows, y, y_dtype, true);
         return de.blocks(s, w, x, x_dtype, rows, y, y_dtype);
     }
     pub fn candidateWeight(w: Weight) !candidate.Matrix {
         const st = strides(w.k, w.n, w.halves, false);
-        if (w.halves % 2 != 0 or w.bias != 0 or w.stride_k != st.k or w.stride_nb != st.nb)
+        const word_row = @as(u64, w.halves) * 4;
+        if (w.halves % 2 != 0 or w.bias != 0 or word_row == 0 or w.stride_k % word_row != 0 or w.stride_k < st.k or w.stride_nb != st.nb)
             return error.CandidateFormat;
-        const m: candidate.Matrix = .{ .trellis = w.trellis, .suh = w.suh, .svh = w.svh, .k = w.k, .n = w.n, .bits = w.halves / 2, .codebook = @fromBackingInt(@intCast(@backingInt(w.codebook))) };
+        const m: candidate.Matrix = .{ .trellis = w.trellis, .suh = w.suh, .svh = w.svh, .k = w.k, .n = w.n, .bits = w.halves / 2, .codebook = @fromBackingInt(@intCast(@backingInt(w.codebook))), .tiles_n = @intCast(w.stride_k / word_row) };
         try m.check();
         return m;
     }
@@ -264,14 +286,13 @@ pub const Dense = struct {
     fn blocks(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void {
         const p = plan(w.k, w.n, 192, 8);
         const z: u64 = if (p.sk > 1) try de.scratch(&de.z, Linear.zBytes(w, 128)) else 0;
-        const counters = try de.scratch(&de.counters, Linear.counterBytes(w));
+        const counters: u64 = if (p.sk > 1) try de.zeroedCounters(Linear.counterBytes(w)) else 0;
         const xe: u64 = if (x_dtype == .f32) 4 else 2;
         const ye: u64 = if (y_dtype == .f32) 4 else 2;
         var r: usize = 0;
         while (r < rows) : (r += 128) {
             const m = @min(rows - r, 128);
             const xh = try de.scratch(&de.xh, m * w.k * 2);
-            try de.counters.?.fill8(0, s.handle);
             try de.lin.matmul(s, w, x + @as(u64, r) * w.k * xe, x_dtype, m, y + @as(u64, r) * w.n * ye, y_dtype, xh, z, counters);
         }
     }
@@ -331,10 +352,13 @@ test "a failed scratch allocation keeps the old scratch owned" {
     try std.testing.expectEqual(@as(usize, 16), slot.?.len);
 }
 
-test "candidate selection refuses layouts it cannot represent" {
+test "candidate selection admits sliced column spans and refuses layouts it cannot represent" {
     const st = strides(128, 256, 8, false);
     var w: Weight = .{ .trellis = 1, .suh = 2, .svh = 3, .k = 128, .n = 256, .halves = 8, .codebook = .mcg, .stride_k = st.k, .stride_nb = st.nb };
     try std.testing.expectEqual(@as(u32, 4), (try Dense.candidateWeight(w)).bits);
+    w.n = 128; // column-sliced span sharing the parent matrix's stride_k
+    try std.testing.expectEqual(@as(usize, 16), (try Dense.candidateWeight(w)).tiles_n);
+    w.n = 256;
     w.halves = 7;
     try std.testing.expectError(error.CandidateFormat, Dense.candidateWeight(w));
     w.halves = 8;

zig/src/cuda/exl3_candidate.zig

@@ -22,6 +22,7 @@ pub const Matrix = struct {
     n: usize,
     bits: u32,
     codebook: Codebook,
+    tiles_n: usize = 0,
 
     /// Whole widths 1-8 over whole 128-blocks; exl3.cu takes no half widths and no partial blocks.
     pub fn check(m: Matrix) !void {
@@ -33,7 +34,8 @@ pub const Matrix = struct {
     }
 
     fn args(m: Matrix) MatrixArgs {
-        return .{ .k = @intCast(m.k), .n = @intCast(m.n), .bits = m.bits, .tiles_n = @intCast(m.n / 16), .split = @intCast(split(m.k, m.n)) };
+        const tn = if (m.tiles_n != 0) m.tiles_n else m.n / 16;
+        return .{ .k = @intCast(m.k), .n = @intCast(m.n), .bits = m.bits, .tiles_n = @intCast(tn), .split = @intCast(split(m.k, m.n)) };
     }
 };
 

zig/src/families/qwen3_5/cuda_draft_frame.zig

@@ -149,26 +149,45 @@ fn narrow(contexts: []const *draft.Ctx) bool {
 /// The rows' draft logits over the spans' columns, packed in d.logits after the two projections (null: not this head).
 fn head(d: anytype, rows: usize) !?u64 {
     const e = d.e;
-    const q = switch (e.head) {
-        .a4 => |q| q,
-        .an, .ql => return full(d, rows),
-        .x3, .bf => return full(d, rows),
-    };
-    if (q.n != spans[1][1] or q.k % 64 != 0) return full(d, rows);
-    const first: cuda.qlinear.Affine4 = .{ .w = q.w, .s = q.s, .b = q.b, .n = spans[0][1], .k = q.k, .npad = q.npad };
-    const tail: cuda.qlinear.Affine4 = .{ .w = q.w + tail_tile * q.k / 2, .s = q.s + tail_tile * 2, .b = q.b + tail_tile * 2, .n = q.n - tail_tile, .k = q.k, .npad = q.npad };
-    const a = d.logits; // rows x 98,304, then rows x 320, then the packed rows x 98,592
-    const b = a + rows * first.n * 2;
-    const out = b + rows * tail.n * 2;
-    try draft.precision.project(e, layers.Proj{ .a4 = first }, d.x, d.xs, a, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff);
-    try draft.precision.project(e, layers.Proj{ .a4 = tail }, d.x, d.xs, b, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff);
-    const skip = spans[1][0] - tail_tile;
-    const copies = [2]cuda.abi.Memcpy2D{
-        .{ .src = a, .src_pitch = first.n * 2, .dst = out, .dst_pitch = span_width * 2, .width = first.n * 2, .height = rows },
-        .{ .src = b + skip * 2, .src_pitch = tail.n * 2, .dst = out + first.n * 2, .dst_pitch = span_width * 2, .width = (tail.n - skip) * 2, .height = rows },
-    };
-    for (&copies) |*m| try e.d.check(e.d.api.cuMemcpy2DAsync_v2(m, e.stream.handle), "cuMemcpy2DAsync");
-    return out;
+    switch (e.head) {
+        .a4 => |q| {
+            if (q.n != spans[1][1] or q.k % 64 != 0) return full(d, rows);
+            const first: cuda.qlinear.Affine4 = .{ .w = q.w, .s = q.s, .b = q.b, .n = spans[0][1], .k = q.k, .npad = q.npad };
+            const tail: cuda.qlinear.Affine4 = .{ .w = q.w + tail_tile * q.k / 2, .s = q.s + tail_tile * 2, .b = q.b + tail_tile * 2, .n = q.n - tail_tile, .k = q.k, .npad = q.npad };
+            const a = d.logits; // rows x 98,304, then rows x 320, then the packed rows x 98,592
+            const b = a + rows * first.n * 2;
+            const out = b + rows * tail.n * 2;
+            try draft.precision.project(e, layers.Proj{ .a4 = first }, d.x, d.xs, a, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff);
+            try draft.precision.project(e, layers.Proj{ .a4 = tail }, d.x, d.xs, b, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff);
+            const skip = spans[1][0] - tail_tile;
+            const copies = [2]cuda.abi.Memcpy2D{
+                .{ .src = a, .src_pitch = first.n * 2, .dst = out, .dst_pitch = span_width * 2, .width = first.n * 2, .height = rows },
+                .{ .src = b + skip * 2, .src_pitch = tail.n * 2, .dst = out + first.n * 2, .dst_pitch = span_width * 2, .width = (tail.n - skip) * 2, .height = rows },
+            };
+            for (&copies) |*m| try e.d.check(e.d.api.cuMemcpy2DAsync_v2(m, e.stream.handle), "cuMemcpy2DAsync");
+            return out;
+        },
+        .x3 => |q| {
+            const tail_x3 = spans[1][0] / 128 * 128;
+            if (q.n != spans[1][1] or q.k % 128 != 0 or spans[0][1] % 128 != 0 or (q.n - tail_x3) % 128 != 0) return full(d, rows);
+            const first: cuda.exl3.Weight = .{ .trellis = q.trellis, .suh = q.suh, .svh = q.svh, .bias = q.bias, .k = q.k, .n = @intCast(spans[0][1]), .halves = q.halves, .codebook = q.codebook, .stride_k = q.stride_k, .stride_nb = q.stride_nb };
+            const tail_n: usize = q.n - tail_x3;
+            const tail: cuda.exl3.Weight = .{ .trellis = q.trellis + (tail_x3 / 128) * q.stride_nb * 4, .suh = q.suh, .svh = q.svh + tail_x3 * 2, .bias = if (q.bias != 0) q.bias + tail_x3 * 2 else 0, .k = q.k, .n = @intCast(tail_n), .halves = q.halves, .codebook = q.codebook, .stride_k = q.stride_k, .stride_nb = q.stride_nb };
+            const a = d.logits;
+            const b = a + rows * spans[0][1] * 2;
+            const out = b + rows * tail_n * 2;
+            try draft.precision.project(e, layers.Proj{ .x3 = first }, d.x, d.xs, a, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff);
+            try draft.precision.project(e, layers.Proj{ .x3 = tail }, d.x, d.xs, b, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff);
+            const skip = spans[1][0] - tail_x3;
+            const copies = [2]cuda.abi.Memcpy2D{
+                .{ .src = a, .src_pitch = spans[0][1] * 2, .dst = out, .dst_pitch = span_width * 2, .width = spans[0][1] * 2, .height = rows },
+                .{ .src = b + skip * 2, .src_pitch = tail_n * 2, .dst = out + spans[0][1] * 2, .dst_pitch = span_width * 2, .width = (tail_n - skip) * 2, .height = rows },
+            };
+            for (&copies) |*m| try e.d.check(e.d.api.cuMemcpy2DAsync_v2(m, e.stream.handle), "cuMemcpy2DAsync");
+            return out;
+        },
+        .an, .ql, .bf => return full(d, rows),
+    }
 }
 
 /// The rows' draft logits over the whole vocabulary, in d.logits.

zig/src/families/qwen3_5/cuda_engine.zig

@@ -105,6 +105,7 @@ pub const Engine = struct {
     embed_s: u64,
     embed_b: u64,
     embed_bits: u8, // the embedding's rows' width: 4 or 8 (16: bf16 rows, an EXL3 checkpoint's)
+    host_embed: ?cuda.HostBuffer = null, // pinned device-mapped host storage for uncompressed bf16 embeddings
     kit: ?*x3w.Kit = null, // an EXL3 checkpoint's kernels (g.exl3)
     head: Proj,
     final_norm: u64,
@@ -183,6 +184,7 @@ pub const Engine = struct {
         e.learning = null;
         e.mlp = null;
         e.taught = 0;
+        e.host_embed = null;
         e.kit = null;
         e.d = d;
         e.g = g;
@@ -278,6 +280,7 @@ pub const Engine = struct {
     pub fn deinit(e: *Engine) void {
         recurrent.shutdown(e);
         if (e.kit) |k| k.deinit(e.gpa);
+        if (e.host_embed) |*he| he.free();
         e.pockets.deinit();
         var it = e.graphs.valueIterator();
         while (it.next()) |x| x.deinit();
@@ -340,16 +343,30 @@ pub const Engine = struct {
         const R = max_rows;
         const wide = @max(g.valueWidth(), g.attnWidth(), g.hidden); // y and mix: either mixer's rows
         inline for (.{ "h", "x", "r" }) |name| @field(e, name) = try t.alloc(R * g.hidden * 2);
+        e.xs = try t.alloc(R * (@max(g.intermediate, wide) / 64) * 4);
+        // DeltaNet mixer, full-attention mixer, and dense MLP run sequentially on e.stream and never overlap:
+        const union_total = total;
+        const union_at = e.arena_at;
         inline for (.{ "y", "mix", "z", "v" }) |name| @field(e, name) = try t.alloc(R * wide * 2);
         inline for (.{ "q", "k" }) |name| @field(e, name) = try t.alloc(R * g.k_heads * geometry.dk * 2);
         e.qkv = try t.alloc(R * g.convWidth() * 2);
-        inline for (.{ "gate", "up", "act" }) |name| @field(e, name) = try t.alloc(R * g.intermediate * 2);
-        e.xs = try t.alloc(R * (@max(g.intermediate, wide) / 64) * 4);
         inline for (.{ "a", "b", "beta" }) |name| @field(e, name) = try t.alloc(R * g.v_heads * 2);
         e.g_ = try t.alloc(R * g.v_heads * 4);
+        const delta_total = total;
+        const delta_at = e.arena_at;
+        total = union_total;
+        e.arena_at = union_at;
+        inline for (.{ 0, 1 }) |_| _ = try t.alloc(R * wide * 2); // y and mix
         e.qraw = try t.alloc(R * 2 * g.attnWidth() * 2);
         e.qn = try t.alloc(R * g.attnWidth() * 2);
         inline for (.{ "kraw", "vraw", "kn" }) |name| @field(e, name) = try t.alloc(R * g.kv_heads * head_dim * 2);
+        const attn_total = total;
+        const attn_at = e.arena_at;
+        total = union_total;
+        e.arena_at = union_at;
+        inline for (.{ "gate", "up", "act" }) |name| @field(e, name) = try t.alloc(R * g.intermediate * 2);
+        total = @max(total, @max(delta_total, attn_total));
+        e.arena_at = @max(e.arena_at, @max(delta_at, attn_at));
         // the decode partials: verify rows only (a prompt chunk's attention runs cuda_prompt_attention's tiles)
         e.pm = try t.alloc(most_verify * g.q_heads * max_chunks * 4);
         e.pl = try t.alloc(most_verify * g.q_heads * max_chunks * 4);

zig/src/families/qwen3_5/cuda_exl3_weights.zig

@@ -31,12 +31,14 @@ pub const Kit = struct {
         k.b16 = try b16.Kernels.load(d);
         errdefer k.b16.unload();
         k.grouped = .{ .d = d, .kernels = &k.experts };
-        // the grouped experts' scratch grows on first use: allocate it before any graph capture, for the widest call
-        const pairs = engine.max_rows * g.slots();
-        const items = cuda.grouped.maxItems(pairs, g.experts + 1, 16);
-        const up = try cuda.exl3_grouped.Grouped.bytes(g.hidden, g.intermediate, items, try ex.plan(g.hidden, g.intermediate, true), 2);
-        const down = try cuda.exl3_grouped.Grouped.bytes(g.intermediate, g.hidden, items, try ex.plan(g.intermediate, g.hidden, false), 1);
-        k.grouped.storage = try cuda.DeviceBuffer.alloc(d, @max(up, down));
+        if (g.moe()) {
+            // grouped-expert scratch grows on first use: allocate it before graph capture, for the widest call
+            const pairs = engine.max_rows * g.slots();
+            const items = cuda.grouped.maxItems(pairs, g.experts + 1, 16);
+            const up = try cuda.exl3_grouped.Grouped.bytes(g.hidden, g.intermediate, items, try ex.plan(g.hidden, g.intermediate, true), 2);
+            const down = try cuda.exl3_grouped.Grouped.bytes(g.intermediate, g.hidden, items, try ex.plan(g.intermediate, g.hidden, false), 1);
+            k.grouped.storage = try cuda.DeviceBuffer.alloc(d, @max(up, down));
+        }
         return k;
     }
 
@@ -231,12 +233,15 @@ fn all(s: *Sink, ck: *core.Checkpoint) !void {
     const g = e.g;
     const kit = e.kit;
     const m = "model.language_model.";
-    // embeddings stay bf16 rows (embedRows); the head is an EXL3 projection
-    const embed = try plain(s, ck, "{s}embed_tokens.weight", .{m}, .bf16, &.{ g.vocab, g.hidden });
+    // embeddings stay bf16 rows (embedRows) in mapped host memory; the head is an EXL3 projection
+    var ebuf: [200]u8 = undefined;
+    const et = try ck.expect(try name(&ebuf, "{s}embed_tokens.weight", .{m}), .bf16, &.{ g.vocab, g.hidden });
     const head = try dense(s, kit, ck, "", "lm_head", g.vocab, g.hidden);
     const final = try norm1(s, ck, "{s}norm.weight", .{m}, g.hidden);
     if (!s.measure) {
-        e.embed_w = embed;
+        e.host_embed = try cuda.HostBuffer.allocMapped(e.d, et.bytes.len);
+        @memcpy(e.host_embed.?.bytes, et.bytes);
+        e.embed_w = try e.host_embed.?.device();
         e.embed_s = 0;
         e.embed_b = 0;
         e.embed_bits = 16;
@@ -276,7 +281,14 @@ fn all(s: *Sink, ck: *core.Checkpoint) !void {
                 .k_norm = try norm1(s, ck, "{s}self_attn.k_norm.weight", .{p}, geometry.head_dim),
             } };
         }
-        l.mlp = .{ .moe = .{ .router = try router(s, ck, p), .experts = try experts(s, ck, p) } };
+        l.mlp = if (g.moe())
+            .{ .moe = .{ .router = try router(s, ck, p), .experts = try experts(s, ck, p) } }
+        else
+            .{ .dense = .{
+                .gate = try dense(s, kit, ck, p, "mlp.gate_proj", g.intermediate, g.hidden),
+                .up = try dense(s, kit, ck, p, "mlp.up_proj", g.intermediate, g.hidden),
+                .down = try dense(s, kit, ck, p, "mlp.down_proj", g.hidden, g.intermediate),
+            } };
     }
 }
 

zig/src/families/qwen3_5/cuda_geometry.zig

@@ -325,7 +325,7 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry {
         return g;
     }
     if (exl3) {
-        if (!moe or g.pocket != 0) return error.UnsupportedQwenQuantization;
+        if (g.pocket != 0) return error.UnsupportedQwenQuantization;
         return g;
     }
     g.bits = try widths(gpa, root, try object(qv.?), g);