From 8cc510c25d6e565e613aa30c1de99ef364e525e9 Mon Sep 17 00:00:00 2001 From: Diwaker Gupta <15990+diwakergupta@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:39:56 -0400 Subject: [PATCH 1/3] qwen3_5 cuda: the DFlash2 drafter reads a dense EXL3 head over its two vocabulary spans On #1465 the drafter projects all 248,320 columns of an EXL3 target's head. It now projects only the spans it proposes from, as column views of the stored matrix: exl3.Weight.columns keeps the parent's strides and moves to whole 128-column blocks, with a host test. The x3 arm shares the affine head's packing path. MoE EXL3 targets still refuse a drafter. The family's formats list exl3, and exl3.md and cuda.md no longer say the native engine serves no EXL3 pack. --- docs/recipes/cuda.md | 4 +- docs/recipes/exl3.md | 6 ++- zig/src/cuda/exl3.zig | 30 ++++++++++++++ zig/src/families/qwen3_5/cuda_draft_frame.zig | 41 ++++++++++++------- zig/src/families/qwen3_5/cuda_native.zig | 2 +- 5 files changed, 63 insertions(+), 20 deletions(-) diff --git a/docs/recipes/cuda.md b/docs/recipes/cuda.md index a9049eb7e1..473a843f4f 100644 --- a/docs/recipes/cuda.md +++ b/docs/recipes/cuda.md @@ -53,8 +53,8 @@ whose CUDA engine reads it declares `EXL3_VARIANT = "any"`, and TensorFold check downloading. The dense linear layer, its plan and its measured throughput are in [EXL3 weights](exl3.md); `python -m tensorfold.cuda.exl3.inspect MODEL_DIR` prints what a checkpoint holds. The module is at the [`v0.6.6` tag](https://github.com/ashhart/TensorFold/tree/v0.6.6/src/tensorfold/cuda/exl3) -([#42](https://github.com/ashhart/TensorFold/pull/42)). The native engine serves no EXL3 pack yet; its reference -decoder is `zig/src/core/exl3_format.zig`. +([#42](https://github.com/ashhart/TensorFold/pull/42)). The native engine's EXL3 path is in +[EXL3 weights](exl3.md#native-cuda-zig); its reference decoder is `zig/src/core/exl3_format.zig`. ## Checkpoints diff --git a/docs/recipes/exl3.md b/docs/recipes/exl3.md index e1771fbaa5..39ef8127c0 100644 --- a/docs/recipes/exl3.md +++ b/docs/recipes/exl3.md @@ -209,5 +209,7 @@ under real routing) included. The program that records them runs the Python engi NCCL (output-column cuts, input-row cuts with fp32 partials summed rank 0 first, and the experts' intermediate cut) and checks every rank's bytes and the sums the same way. -No family serves EXL3 natively yet: the Metal projection (#681) and these CUDA pieces are what a family's engine -takes, through `qlinear` and `exl3_experts`. +The Metal projection (#681) and these CUDA pieces are what a family's engine takes, through `qlinear` and +`exl3_experts`. On CUDA, `qwen3_5` serves EXL3 packs of the Qwen3.5 MoE models and of the dense +[Qwen3.8-27B](qwen3.8-27b.md#cuda), whose DFlash2 drafter reads the head through column views +(`exl3.Weight.columns`: whole 128-column blocks with the parent's strides). diff --git a/zig/src/cuda/exl3.zig b/zig/src/cuda/exl3.zig index 54176e9f64..8626eafce5 100644 --- a/zig/src/cuda/exl3.zig +++ b/zig/src/cuda/exl3.zig @@ -95,6 +95,17 @@ pub const Weight = struct { codebook: Codebook, stride_k: u64, stride_nb: u64, + + /// Columns [from, from + n) as a view of the stored matrix: whole 128-column blocks, the parent's strides. + pub fn columns(w: Weight, from: u32, n: u32) !Weight { + if (from % 128 != 0 or n % 128 != 0 or n == 0 or @as(u64, from) + n > w.n) return error.Exl3Shape; + var v = w; + v.n = n; + v.trellis = w.trellis + @as(u64, from / 128) * w.stride_nb * 4; + v.svh = w.svh + @as(u64, from) * 2; + if (w.bias != 0) v.bias = w.bias + @as(u64, from) * 2; + return v; + } }; /// The module and its kernels. @@ -331,6 +342,25 @@ test "a failed scratch allocation keeps the old scratch owned" { try std.testing.expectEqual(@as(usize, 16), slot.?.len); } +test "a column view keeps the parent's strides and moves to its first 128-column block" { + for ([_]bool{ false, true }) |strips| { + const st = strides(256, 1024, 6, strips); + const w: Weight = .{ .trellis = 0x1000, .suh = 0x2000, .svh = 0x3000, .bias = 0x4000, .k = 256, .n = 1024, .halves = 6, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; + const v = try w.columns(384, 512); + try std.testing.expectEqual(@as(u32, 512), v.n); + try std.testing.expectEqual(w.trellis + 3 * st.nb * 4, v.trellis); + try std.testing.expectEqual(w.svh + 384 * 2, v.svh); + try std.testing.expectEqual(w.bias + 384 * 2, v.bias); + try std.testing.expectEqual(w.suh, v.suh); + try std.testing.expectEqual(w.stride_k, v.stride_k); + try std.testing.expectEqual(w.stride_nb, v.stride_nb); + try std.testing.expectEqual(w, try w.columns(0, 1024)); + try std.testing.expectError(error.Exl3Shape, w.columns(64, 128)); + try std.testing.expectError(error.Exl3Shape, w.columns(0, 192)); + try std.testing.expectError(error.Exl3Shape, w.columns(896, 256)); + } +} + test "candidate selection 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 }; diff --git a/zig/src/families/qwen3_5/cuda_draft_frame.zig b/zig/src/families/qwen3_5/cuda_draft_frame.zig index bd98176ca4..909ae19faa 100644 --- a/zig/src/families/qwen3_5/cuda_draft_frame.zig +++ b/zig/src/families/qwen3_5/cuda_draft_frame.zig @@ -149,23 +149,34 @@ 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), + // the head's first span and the whole tiles holding its tail span, as views of the stored matrix + const first: layers.Proj, const tail: layers.Proj, const tail_at: usize = switch (e.head) { + .a4 => |q| blk: { + if (q.n != spans[1][1] or q.k % 64 != 0) return full(d, rows); + break :blk .{ + .{ .a4 = .{ .w = q.w, .s = q.s, .b = q.b, .n = spans[0][1], .k = q.k, .npad = q.npad } }, + .{ .a4 = .{ .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 } }, + tail_tile, + }; + }, + .x3 => |q| blk: { + // EXL3 columns come in 128-column blocks: the tail span from row 247,936, its first 96 columns dropped + const at = spans[1][0] / 128 * 128; + if (q.n != spans[1][1] or q.k % 128 != 0) return full(d, rows); + break :blk .{ .{ .x3 = try q.columns(0, spans[0][1]) }, .{ .x3 = try q.columns(at, spans[1][1] - at) }, at }; + }, + .an, .ql, .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 tail_n = spans[1][1] - tail_at; + const a = d.logits; // rows x 98,304, then rows x the tail's tiles, then the packed rows x 98,592 + const b = a + rows * spans[0][1] * 2; + const out = b + rows * tail_n * 2; + try draft.precision.project(e, first, d.x, d.xs, a, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff); + try draft.precision.project(e, tail, d.x, d.xs, b, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff); + const skip = spans[1][0] - tail_at; 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 }, + .{ .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; diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index e6917403f1..fb18ad6241 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -25,7 +25,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; From f4d84c92c6109150fe6fbdf00cbc9b9443e48b1e Mon Sep 17 00:00:00 2001 From: Diwaker Gupta <15990+diwakergupta@users.noreply.github.com> Date: Tue, 27 Oct 2026 20:46:45 -0400 Subject: [PATCH 2/3] qwen3_5 cuda: an EXL3 head's tail view takes the last 64 blocks, so it plans as the whole head does The 384-column tail planned eight K splits where the full head plans one, so its logits and some proposals changed. The view now starts at column 240,128 (exl3.wide_blocks blocks), keeping #1465's bits. A host test checks that both span views plan as the whole head does. --- zig/qwen35_draft_frame_host.zig | 14 ++++++++++++++ zig/src/cuda/exl3.zig | 7 +++++-- zig/src/families/qwen3_5/cuda_draft_frame.zig | 15 +++++++++++---- 3 files changed, 30 insertions(+), 6 deletions(-) diff --git a/zig/qwen35_draft_frame_host.zig b/zig/qwen35_draft_frame_host.zig index 2731436fd4..0a931f5f2a 100644 --- a/zig/qwen35_draft_frame_host.zig +++ b/zig/qwen35_draft_frame_host.zig @@ -359,3 +359,17 @@ test "Draft frames: enlarged frame byte counts reject overflow" { d.c.rank = std.math.maxInt(usize); try std.testing.expectError(error.Overflow, d.layout(false)); } + +test "Draft frames: an EXL3 head's span views plan their K split as the whole head does" { + const x3 = cuda.exl3; + const F = q.draft.frame; + for ([_]bool{ false, true }) |strips| { + const st = x3.strides(5120, 248320, 12, strips); + const w: x3.Weight = .{ .trellis = 0x1000, .suh = 0x2000, .svh = 0x3000, .k = 5120, .n = 248320, .halves = 12, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; + const v = try F.x3Views(w); + const whole = x3.plan(w.k, w.n, 192, 8); + try std.testing.expectEqual(whole, x3.plan(v.first.k, v.first.n, 192, 8)); + try std.testing.expectEqual(whole, x3.plan(v.tail.k, v.tail.n, 192, 8)); + try std.testing.expect(F.x3_tail_at <= 248032 and F.x3_tail_at + v.tail.n == w.n); + } +} diff --git a/zig/src/cuda/exl3.zig b/zig/src/cuda/exl3.zig index 8626eafce5..75bbe10798 100644 --- a/zig/src/cuda/exl3.zig +++ b/zig/src/cuda/exl3.zig @@ -59,13 +59,16 @@ pub fn variantIndex(halves: u32, codebook: Codebook, warps: u32) ?usize { return null; } +/// From this many 128-column blocks a shape plans one K split, so a view that wide keeps the whole matrix's bits. +pub const wide_blocks = 64; + /// What one shape launches, from the shape alone so a row's bits never depend on the row count (linear.py plan). pub fn plan(k: usize, n: usize, blocks: usize, min_tiles: usize) struct { sk: u32, wk: u32 } { const kt = k / 16; const nb = n / 128; var sk: usize = 1; - var wk: usize = if (nb >= 64) 8 else 4; - if (nb < 64) { + var wk: usize = if (nb >= wide_blocks) 8 else 4; + if (nb < wide_blocks) { while (nb * sk < blocks and sk < 64) sk *= 2; } while ((wk > 4 or sk > 1) and (kt % (sk * wk) != 0 or kt / (sk * wk) < min_tiles)) { diff --git a/zig/src/families/qwen3_5/cuda_draft_frame.zig b/zig/src/families/qwen3_5/cuda_draft_frame.zig index 909ae19faa..503e50d8f8 100644 --- a/zig/src/families/qwen3_5/cuda_draft_frame.zig +++ b/zig/src/families/qwen3_5/cuda_draft_frame.zig @@ -146,6 +146,14 @@ fn narrow(contexts: []const *draft.Ctx) bool { return true; } +/// An EXL3 head's tail view: the last wide_blocks blocks (from row 240,128), so it plans as the whole head does. +pub const x3_tail_at = spans[1][1] - cuda.exl3.wide_blocks * 128; + +/// An EXL3 head's first span and the blocks holding its tail span, as views that keep the whole head's bits. +pub fn x3Views(q: cuda.exl3.Weight) !struct { first: cuda.exl3.Weight, tail: cuda.exl3.Weight } { + return .{ .first = try q.columns(0, spans[0][1]), .tail = try q.columns(x3_tail_at, spans[1][1] - x3_tail_at) }; +} + /// 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; @@ -160,15 +168,14 @@ fn head(d: anytype, rows: usize) !?u64 { }; }, .x3 => |q| blk: { - // EXL3 columns come in 128-column blocks: the tail span from row 247,936, its first 96 columns dropped - const at = spans[1][0] / 128 * 128; if (q.n != spans[1][1] or q.k % 128 != 0) return full(d, rows); - break :blk .{ .{ .x3 = try q.columns(0, spans[0][1]) }, .{ .x3 = try q.columns(at, spans[1][1] - at) }, at }; + const v = try x3Views(q); + break :blk .{ .{ .x3 = v.first }, .{ .x3 = v.tail }, x3_tail_at }; }, .an, .ql, .bf => return full(d, rows), }; const tail_n = spans[1][1] - tail_at; - const a = d.logits; // rows x 98,304, then rows x the tail's tiles, then the packed rows x 98,592 + const a = d.logits; // rows x 98,304, then rows x the tail's columns, then the packed rows x 98,592 const b = a + rows * spans[0][1] * 2; const out = b + rows * tail_n * 2; try draft.precision.project(e, first, d.x, d.xs, a, rows, .proposal, draft.precision.dflash_projection_prompt_cutoff); From debf9f8f989708fa612d467b0df4081bd38a3f70 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:37:23 +0000 Subject: [PATCH 3/3] changelog: DFlash2 reads a dense EXL3 head over its two spans (#1507) --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c4e361db8..63194c6385 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Qwen3.8-27B EXL3 on CUDA: the DFlash2 drafter projects only the two vocabulary spans it proposes from, as column views of the stored head, instead of all 248,320 columns, skipping about 545 MB of reads a round at a 6-bit head. Both views plan the head's matrix multiply as the whole head does, so replies are the same tokens as before (144 of 144 equal on an RTX 3090 Ti with the 3.50 bpw pack), and single-stream decode there was 0.5 to 1.8% faster in the cells whose tokens per round matched. `qwen3_5` on CUDA now lists `exl3` among the formats it reads (#1507, after #1465). Thanks to @diwakergupta. + - Dory on CUDA with `TENSORFOLD_DORY_KV_PAGES=1` refuses to attach a context that is not a whole number of 256-position pages, takes its slot count from the round geometry instead of a fixed 8, and loads the paged attention kernels only when paging is on, so the default path no longer loads that module (#1387, for #1284). Thanks to @jschmied. - Gemma 4 on CUDA serves structured output: `response_format` (`json_schema`, `json_object`) and `tool_choice: "required"` or a named function now constrain its replies, with or without a drafter (#1386, for #1384). Sampled rows are masked on the GPU before the draw; greedy rows pick the first largest allowed token on the host, the same rule as the GPU pick. Requests without a grammar take the old path unchanged. Thanks to @ThomasHalwax.