From 310e5e04d48105ed2340d0bc21cd85220dd18b25 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] feat(qwen3_5): serve dense EXL3 checkpoints with the DFlash2 drafter on CUDA Admit non-MoE qwen3_5 EXL3 checkpoints (pocketed ones stay refused), load the dense SwiGLU MLP as EXL3 linears, allocate the grouped-expert scratch only for MoE geometry, and let openDrafter attach DFlash2 to EXL3 targets. The drafter reads the target's EXL3 head over its two vocabulary spans through exl3.Weight.columns, a view of whole 128-column blocks with the parent's strides, sharing the packing path with the affine head. Docs: the 27B recipe's native EXL3 section, 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 ++- docs/recipes/qwen3.8-27b.md | 25 ++++++++--- zig/qwen35_formats_host.zig | 10 +++-- zig/src/cuda/exl3.zig | 30 ++++++++++++++ zig/src/families/qwen3_5/cuda_draft_frame.zig | 41 ++++++++++++------- .../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 +- 9 files changed, 107 insertions(+), 38 deletions(-) diff --git a/docs/recipes/cuda.md b/docs/recipes/cuda.md index 945c2c3113..b1556d6461 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..064cf77562 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#native-exl3-checkpoints), whose DFlash2 drafter reads the head through column views +(`exl3.Weight.columns`: whole 128-column blocks with the parent's strides). diff --git a/docs/recipes/qwen3.8-27b.md b/docs/recipes/qwen3.8-27b.md index fcc4cdb43a..e90829b4ec 100644 --- a/docs/recipes/qwen3.8-27b.md +++ b/docs/recipes/qwen3.8-27b.md @@ -3,8 +3,8 @@ The native engine serves `TensorFold/Qwen3.8-27B-MLX-4bit` on Metal, qualified on M5 Max and M3 Ultra, from `zig/src/families/qwen27/`. It drafts with DFlash2 (`zig/src/families/qwen27/dflash/`) only when `--drafter` names the drafter; without it the model decodes plain. On CUDA the native engine serves the MLX affine checkpoints (2, 3, -4, 5, 6 and 8 bits in groups of 64) from `zig/src/families/qwen3_5/cuda_*.zig`, with DFlash2 drafts and -`--kv-cache fp8`; it reads no NVFP4 or EXL3 checkpoint yet. +4, 5, 6 and 8 bits in groups of 64) and turboderp's EXL3 packs ([native EXL3](#native-exl3-checkpoints)) from +`zig/src/families/qwen3_5/cuda_*.zig`, with DFlash2 drafts and `--kv-cache fp8`; it reads no NVFP4 checkpoint yet. ```bash tensorfold pull TensorFold/Qwen3.8-27B-MLX-4bit @@ -90,9 +90,24 @@ quality comparisons are TBD [release-0.3.5]. ## CUDA -This section is the Python engine's (0.6.6). The native CUDA engine reads only MLX affine checkpoints so far: the -NVFP4 and EXL3 readers below have not been ported to it (1.1.0 carries the shared EXL3 module of #832, which the -27B does not read through yet). +### Native EXL3 checkpoints + +The native CUDA engine reads turboderp's EXL3 packs of the model (`turboderp/Qwen3.8-27B-exl3`, a branch per size, +Apache-2.0 like the base model) through the shared EXL3 module ([EXL3 weights](exl3.md)), on one GPU, and drafts +with `z-lab/Qwen3.8-27B-DFlash2` (Apache-2.0) at `--drafter-bits 4`: + +```bash +python -c "from huggingface_hub import snapshot_download as d; d('turboderp/Qwen3.8-27B-exl3', revision='3.50bpw', local_dir='qwen27b-exl3-3.50bpw')" +tensorfold pull z-lab/Qwen3.8-27B-DFlash2 +tensorfold serve qwen27b-exl3-3.50bpw --drafter z-lab/Qwen3.8-27B-DFlash2 --kv-cache fp8 +``` + +Every projection, the dense MLP's included, is an EXL3 linear; the embeddings stay bf16 rows and the head is the +pack's EXL3 matrix. Prompt and verify rows take the row-invariant EXL3 linear, so drafted replies equal +`"draft": false` ones. The drafter reads the head over its draft vocabulary as two column views of the stored +matrix, whole 128-column blocks: rows 0-98,303 and 247,936-248,319, of which it keeps the last 288. + +The sections below are the Python engine's (0.6.6): the native engine has not taken its NVFP4 reader. On CUDA the 27B serves NVFP4 and EXL3 checkpoints, and the MLX 4-bit checkpoint as the portable option: the same files a Mac serves, and the only format two ranks read. `tensorfold serve` loads the checkpoint you name; it picks diff --git a/zig/qwen35_formats_host.zig b/zig/qwen35_formats_host.zig index a70119e4a4..c957d6caaa 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 }, 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_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 82987bfa1f..6a4fea26d9 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();