qwen3_5 cuda: the DFlash2 drafter reads a dense EXL3 head over its two vocabulary spans

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.

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
 

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

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

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)) {
@@ -95,6 +98,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 +345,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 };

zig/src/families/qwen3_5/cuda_draft_frame.zig

@@ -146,26 +146,44 @@ 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;
-    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: {
+            if (q.n != spans[1][1] or q.k % 128 != 0) return full(d, rows);
+            const v = try x3Views(q);
+            break :blk .{ .{ .x3 = v.first }, .{ .x3 = v.tail }, x3_tail_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 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);
+    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;

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;