From 07cfa2f5806ca3aceefb8421af8b5d6069d33569 Mon Sep 17 00:00:00 2001 From: Diwaker Gupta <15990+diwakergupta@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:06:20 -0400 Subject: [PATCH 1/4] feat(qwen3_5): enable dense EXL3 checkpoints and DFlash2 drafter on CUDA Admit non-MoE qwen3_5 EXL3 checkpoints in cuda_geometry.zig, load dense SwiGLU MLP projections (gate_proj, up_proj, down_proj) in cuda_exl3_weights.zig while guarding grouped-expert workspace allocation on g.moe(), register exl3 in cuda_native.zig formats, and allow attaching the 4-bit DFlash2 drafter to EXL3 targets. --- zig/qwen35_formats_host.zig | 11 ++++++--- .../families/qwen3_5/cuda_exl3_weights.zig | 23 +++++++++++++------ zig/src/families/qwen3_5/cuda_geometry.zig | 2 +- zig/src/families/qwen3_5/cuda_native.zig | 4 +--- 4 files changed, 26 insertions(+), 14 deletions(-) diff --git a/zig/qwen35_formats_host.zig b/zig/qwen35_formats_host.zig index 5545e681eb..139147b72f 100644 --- a/zig/qwen35_formats_host.zig +++ b/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,4 @@ 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); } + diff --git a/zig/src/families/qwen3_5/cuda_exl3_weights.zig b/zig/src/families/qwen3_5/cuda_exl3_weights.zig index 8c333de8b7..b58580ae72 100644 --- a/zig/src/families/qwen3_5/cuda_exl3_weights.zig +++ b/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; } @@ -276,7 +278,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), + } }; } } diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig index 382bdc3f55..4e3ea04aa1 100644 --- a/zig/src/families/qwen3_5/cuda_geometry.zig +++ b/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); diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 5151b9e95c..2fd9236355 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -24,7 +24,7 @@ pub const slide = @import("cuda_slide.zig"); pub const linked = slide.linked; pub const model_type = "qwen3_5"; -pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const formats: []const []const u8 = &.{ "mlx-q4g64", "exl3" }; pub const default_context: i64 = 32768; pub const max_segments: u32 = 1; pub const prompt_rows: u32 = E.max_rows; @@ -161,7 +161,6 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void, prepare: struct { ctx: *anyopaque, run: *const fn (*anyopaque) anyerror!lanes.drafter.Facts, bytes: usize } } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); - if (l.e.g.exl3) return error.Exl3DrafterUnsupported; if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; const d = try Draft.init(gpa, io, l.e, dir); errdefer d.deinit(); @@ -177,7 +176,6 @@ pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []con pub fn openMtp(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); - if (l.e.g.exl3) return error.Exl3DrafterUnsupported; if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; const d = try Mtp.init(gpa, io, l.e, dir, l.capacity); errdefer d.deinit(); From 9bdb08d3da2a053a1ae8c19ecaae20819325ae29 Mon Sep 17 00:00:00 2001 From: Diwaker Gupta <15990+diwakergupta@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:06:20 -0400 Subject: [PATCH 2/4] perf(exl3): use 64-row tensor-core prompt tiles and slice DFlash2 lm_head spans Preload exl3_candidate during Dense.reserve and route Dense.prompt for rows > 16 to the 64-row tensor-core prompt kernel (tf_exl3_prompt_2) when compatible. Zero-fill split-K atomic counters once on allocation/growth since the kernel resets counters before exiting. Slice .x3 lm_head projections in cuda_draft_frame.zig across the two active vocabulary spans (0..98304 and 248064..248320). --- zig/src/cuda/exl3.zig | 40 ++++++++++--- zig/src/cuda/exl3_candidate.zig | 4 +- zig/src/families/qwen3_5/cuda_draft_frame.zig | 59 ++++++++++++------- 3 files changed, 74 insertions(+), 29 deletions(-) diff --git a/zig/src/cuda/exl3.zig b/zig/src/cuda/exl3.zig index 54176e9f64..aad1d57290 100644 --- a/zig/src/cuda/exl3.zig +++ b/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; diff --git a/zig/src/cuda/exl3_candidate.zig b/zig/src/cuda/exl3_candidate.zig index c894f33096..b4fdad48b2 100644 --- a/zig/src/cuda/exl3_candidate.zig +++ b/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)) }; } }; diff --git a/zig/src/families/qwen3_5/cuda_draft_frame.zig b/zig/src/families/qwen3_5/cuda_draft_frame.zig index bd98176ca4..5f98c93873 100644 --- a/zig/src/families/qwen3_5/cuda_draft_frame.zig +++ b/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. From 723b2247e64fca77950b69d746a8ae745648d132 Mon Sep 17 00:00:00 2001 From: Diwaker Gupta <15990+diwakergupta@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:06:21 -0400 Subject: [PATCH 3/4] perf(qwen3_5): map EXL3 embed_tokens in host memory and union disjoint scratch Pin BF16 embed_tokens.weight in mapped host memory via HostBuffer.allocMapped in cuda_exl3_weights.zig so embedRows gathers token rows over PCIe without occupying device arena memory, and overlap the sequentially executed DeltaNet mixer, full-attention mixer, and dense MLP scratch regions in Engine.scratch. --- zig/qwen35_formats_host.zig | 13 ++++++++++++ zig/src/families/qwen3_5/cuda_engine.zig | 21 +++++++++++++++++-- .../families/qwen3_5/cuda_exl3_weights.zig | 9 +++++--- 3 files changed, 38 insertions(+), 5 deletions(-) diff --git a/zig/qwen35_formats_host.zig b/zig/qwen35_formats_host.zig index 139147b72f..ae9df7a001 100644 --- a/zig/qwen35_formats_host.zig +++ b/zig/qwen35_formats_host.zig @@ -260,3 +260,16 @@ test "format coexistence: EXL3 BF16 prompt: existing rows and partial final slic 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); +} diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 9c7898e017..24e4fbdbc3 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/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); diff --git a/zig/src/families/qwen3_5/cuda_exl3_weights.zig b/zig/src/families/qwen3_5/cuda_exl3_weights.zig index b58580ae72..a196119591 100644 --- a/zig/src/families/qwen3_5/cuda_exl3_weights.zig +++ b/zig/src/families/qwen3_5/cuda_exl3_weights.zig @@ -233,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; From 00a81f04b4342ff89fd8671b88b1783da4c0a6fd Mon Sep 17 00:00:00 2001 From: Diwaker Gupta <15990+diwakergupta@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:06:21 -0400 Subject: [PATCH 4/4] perf(qwen3_5): require long_match (20 tokens) for 27B root copies Set facts.copy_match to lanes.windows.long_match (20) alongside facts.root_copies on 27B dense (hidden == 5120, layers == 64, experts == 0) so short 5-token suffix matches do not preempt DFlash2 drafting. --- zig/src/families/qwen3_5/cuda_native.zig | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 2fd9236355..99ad2f870a 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -129,7 +129,10 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .backend = l.backend(), .facts = blk: { var facts = laneFacts(@intCast(e.window_rows), @intCast(e.verify_rows)); - facts.root_copies = e.g.hidden == 5120 and e.g.layers == 64 and e.g.experts == 0; + if (e.g.hidden == 5120 and e.g.layers == 64 and e.g.experts == 0) { + facts.root_copies = true; + facts.copy_match = lanes.windows.long_match; + } break :blk facts; }, .rows = @intCast(e.window_rows), @@ -176,6 +179,7 @@ pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []con pub fn openMtp(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); + if (l.e.g.exl3) return error.Exl3DrafterUnsupported; if (l.caches.count() != 0 or l.pool.items.len != 0) return error.TapsAfterStreams; const d = try Mtp.init(gpa, io, l.e, dir, l.capacity); errdefer d.deinit();