From 50a41a6d4179b7d1d0a87edf79bf4d26f4fb243e Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:54:30 +0000 Subject: [PATCH] feat(qwen3_5): let the MoE target draft with a DFlash2 model On CUDA the Qwen3.6 MoE target wired --drafter to its MTP layer alone, though the engine's DFlash2 drafter already reads the SGLang recipe's incoai/Qwen3.6-35B-A3B-DFlash2. One openDrafter now serves both families: a directory whose config.json has dflash_config opens DFlash2, as every directory does under a dense target, and any other opens the MoE target's MTP layer. Both openers return one Opened struct whose prepare is optional, which cuda_drafted.zig now accepts. On a GB10 one build drafts with either; every drafted reply equals its plain run, and DFlash2 decodes 130.3 tok/s at one stream on a 4.2K-token prompt against MTP's 98.6. --- CHANGELOG.md | 2 ++ docs/recipes/qwen3.6-moe.md | 14 ++++++++++ zig/src/families/qwen3_5/cuda.zig | 1 + .../families/qwen3_5/cuda_draft_config.zig | 27 +++++++++++++++++++ zig/src/families/qwen3_5/cuda_native.zig | 23 +++++++++++++--- zig/src/families/qwen3_5/cuda_native_moe.zig | 4 +-- zig/src/native/cuda_drafted.zig | 7 ++++- 7 files changed, 72 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b7b990c84f..c90c3796a3 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.6-35B-A3B on CUDA: `--drafter` also takes a DFlash2 draft model, chosen by the `dflash_config` in its `config.json`; the checkpoint's own directory still drafts with its MTP layer. With `incoai/Qwen3.6-35B-A3B-DFlash2` over `nvidia/Qwen3.6-35B-A3B-NVFP4`, one GB10 decodes 130 tok/s at one stream on a 4.2K-token prompt (MTP: 99), and every drafted reply equals its plain run. + - A fresh-start session disk now clears the `.tfs` and `.tmp` files earlier serving processes left behind, under an exclusive lock on its directory, so restarts no longer fill the disk. Ported by @jayleaton from @mikolaj92's downstream fix, corroborated by @Jenpo. - GLM-5.3-Flash on CUDA: `TF_GLM_KV_SHARE=n` with `TF_KV_PAGES=on` keeps prompts as the KV pool's pages in the shared prompt store (the one `TF_QWEN_KV_SHARE` uses), with at most n states past them, on one GPU and at tensor parallel. Where main's flat store has no room for a long state (GB10s at TP=3 with the 1M window), a 64k message sent again resumes in about 1.7 s instead of prefilling for 36 s; token-id prompts resume at every width (#1510). Unset, nothing changes. diff --git a/docs/recipes/qwen3.6-moe.md b/docs/recipes/qwen3.6-moe.md index 1afebc9d7a..8ce01d93e0 100644 --- a/docs/recipes/qwen3.6-moe.md +++ b/docs/recipes/qwen3.6-moe.md @@ -153,6 +153,20 @@ The checkpoint keeps its MTP layer in bf16 in its own shards. Without `mtp-4bit. quantizes that layer to MLX affine-4 at load, as Nemotron's ModelOpt drafter does; this changes only which rows a round proposes, never a reply, and the startup log says so. With the file present, the file is read instead. +`--drafter` can also name a DFlash2 draft model: a directory whose `config.json` has `dflash_config` drafts with it, +any other with the target's MTP layer. `incoai/Qwen3.6-35B-A3B-DFlash2` (revision `51ef7b69`, the SGLang DFlash +recipe's draft) reads eight target layers and is quantized to affine-4 at load: + +```bash +hf download incoai/Qwen3.6-35B-A3B-DFlash2 --revision 51ef7b6923ad6c14cb1bb41c37a9041446496ab6 +tensorfold serve nvidia/Qwen3.6-35B-A3B-NVFP4 --drafter incoai/Qwen3.6-35B-A3B-DFlash2 +``` + +On the fleet benchmark's 4.2K-token prompt (`bench.py`, 512 tokens), one build with either drafter, it decodes 130.3 +tok/s at one stream against MTP's 98.6, and gives 131.3 / 161.4 / 198.0 tok/s aggregate at 2 / 4 / 8 users against +124.7 / 170.6 / 208.4: it proposes more rows a round and keeps fewer of them, so it leads at one and two streams and +MTP leads under load. Every drafted reply equals its `"draft": false` run with either drafter. + The rules come from @gilby's Python port ([#345](https://github.com/ashhart/TensorFold/pull/345)): bf16 activations for the W4A16 layers, the shared expert last in the table, and the MTP layer drafting through the target's head. diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 7ee736f47a..137195286b 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -30,6 +30,7 @@ test { _ = geometry; _ = @import("draw.zig"); _ = @import("cuda_draft_weights.zig"); + _ = @import("cuda_draft_config.zig"); _ = draft; _ = vision; _ = @import("cuda_vision_prompt.zig"); diff --git a/zig/src/families/qwen3_5/cuda_draft_config.zig b/zig/src/families/qwen3_5/cuda_draft_config.zig index 24553d2db4..71d7425148 100644 --- a/zig/src/families/qwen3_5/cuda_draft_config.zig +++ b/zig/src/families/qwen3_5/cuda_draft_config.zig @@ -91,6 +91,17 @@ pub const Config = struct { } }; +/// Whether `dir`'s config.json declares a DFlash2 draft (a top-level dflash_config); no readable object: not one. +pub fn declared(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) bool { + const path = std.fs.path.join(gpa, &.{ dir, "config.json" }) catch return false; + defer gpa.free(path); + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)) catch return false; + defer gpa.free(bytes); + const doc = std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}) catch return false; + defer doc.deinit(); + return doc.value == .object and doc.value.object.get("dflash_config") != null; +} + /// A size, count or divisor: positive, so a zero can't reach a remainder or an empty allocation. fn int(o: std.json.ObjectMap, key: []const u8) !usize { const v = try whole(o, key); @@ -112,3 +123,19 @@ fn num(o: std.json.ObjectMap, key: []const u8) !f64 { else => error.BadDraftConfig, }; } + +test "a drafter directory is DFlash2's by its config's dflash_config; a checkpoint's own config or none is not" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(root); + try std.testing.expect(!declared(a, io, root)); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = "{\"model_type\":\"qwen3_5_moe\",\"text_config\":{}}" }); + try std.testing.expect(!declared(a, io, root)); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = "{\"model_type\":\"qwen3\",\"dflash_config\":{\"block_size\":8}}" }); + try std.testing.expect(declared(a, io, root)); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = "[1]" }); + try std.testing.expect(!declared(a, io, root)); +} diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index fadd01b1a8..64dcc32d7e 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -10,6 +10,7 @@ pub const window_costs = @import("cuda_window_costs.zig"); pub const cost_adapter = @import("cuda_cost_adapter.zig"); const draws = @import("draw.zig"); const Draft = @import("cuda_draft.zig").Draft; +const draft_config = @import("cuda_draft_config.zig"); const prefill = @import("cuda_prefill.zig"); const Fill = prefill.Fill; const reuse = @import("cuda_reuse.zig"); @@ -29,7 +30,7 @@ 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; -/// --drafter loads the DFlash2 draft model (openDrafter); the checkpoints have no head of their own. +/// --drafter loads a DFlash2 draft model (openDrafter), or on the MoE target its MTP layer; DFlash2 has no own head. pub const takes_drafter = true; /// A request's pages hold one more forward past its prompt and reply (TF_QWEN_KV_PAGES). pub const window_rows: u32 = pages_mod.window_rows; @@ -177,8 +178,24 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { const l: *Lanes = @ptrCast(@alignCast(p)); try reservation.target(E.Cache, E.Part, l, streams); // with a pool: the first window, the shapes on its lanes } +/// A drafter opened over the engine; DFlash2's `prepare` times the target's windows once the pool is reserved. +pub const Opened = 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 } = null, +}; + +/// --drafter's directory over engine `p`: a DFlash2 draft (its config's dflash_config), else the MoE's MTP layer. +pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !Opened { + const l: *Lanes = @ptrCast(@alignCast(p)); + if (!l.e.g.moe() or draft_config.declared(gpa, io, dir)) return openDFlash(gpa, io, p, dir, bits); + return openMtp(gpa, io, p, dir, bits); +} + /// The DFlash2 drafter in `dir` over engine `p`, 4-bit only, and the target's windows timed for the round loop. -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 } } { +fn openDFlash(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !Opened { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); if (l.e.g.exl3 and l.e.g.moe()) return error.Exl3DrafterUnsupported; // a dense checkpoint drafts as an MLX one @@ -194,7 +211,7 @@ pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []con } /// The MTP layer beside the checkpoint in `dir` as the drafter over engine `p` (qwen3_5_moe), its ring the window. -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 } { +fn openMtp(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !Opened { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); if (l.e.g.exl3 and l.e.g.moe()) return error.Exl3DrafterUnsupported; // a dense checkpoint drafts as an MLX one diff --git a/zig/src/families/qwen3_5/cuda_native_moe.zig b/zig/src/families/qwen3_5/cuda_native_moe.zig index 3461c51bc0..87c9cfbfbf 100644 --- a/zig/src/families/qwen3_5/cuda_native_moe.zig +++ b/zig/src/families/qwen3_5/cuda_native_moe.zig @@ -7,7 +7,7 @@ pub const formats: []const []const u8 = &.{ "mlx-q4g64", "modelopt-mixed", "mode pub const default_context = native.default_context; pub const max_segments = native.max_segments; pub const prompt_rows = native.prompt_rows; -/// --drafter names the checkpoint's own directory: its MTP layer (mtp-4bit.safetensors) drafts chains. +/// --drafter names the checkpoint's own directory (its MTP layer drafts chains) or a DFlash2 draft model's. pub const takes_drafter = true; pub const Options = native.Options; pub const open = native.open; @@ -15,7 +15,7 @@ pub const attach = native.attach; pub const pagedLayout = native.pagedLayout; // TF_QWEN_KV_PAGES: the same pool as the dense models' pub const fits = native.fits; pub const window_rows = native.window_rows; -pub const openDrafter = native.openMtp; +pub const openDrafter = native.openDrafter; pub const explain = native.explain; /// The dense family's kept prompt states, its open refusals' words and its image tower: the MLP is the only change. pub const snap_bytes = native.snap_bytes; diff --git a/zig/src/native/cuda_drafted.zig b/zig/src/native/cuda_drafted.zig index 2d54af6a0c..a0307fb5e5 100644 --- a/zig/src/native/cuda_drafted.zig +++ b/zig/src/native/cuda_drafted.zig @@ -70,7 +70,12 @@ pub fn wrap(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, loaded: s.facts = h.drafted.facts(loaded.facts); s.stream_bytes = h.drafted.streamBytes(loaded.stream_bytes); s.held = h; - if (@hasField(@TypeOf(opened), "prepare")) s.prepare = .{ .ctx = opened.prepare.ctx, .run = opened.prepare.run, .bytes = if (@hasField(@TypeOf(opened.prepare), "bytes")) opened.prepare.bytes else 0 }; + if (@hasField(@TypeOf(opened), "prepare")) { + // a family's prepare, or an optional one when only some drafters time windows after the pool (not Qwen's MTP) + const P = @TypeOf(opened.prepare); + const maybe: ?(if (@typeInfo(P) == .optional) @typeInfo(P).optional.child else P) = opened.prepare; + if (maybe) |p| s.prepare = .{ .ctx = p.ctx, .run = p.run, .bytes = if (@hasField(@TypeOf(p), "bytes")) p.bytes else 0 }; + } return s; }