Qwen3.6-35B-A3B on CUDA: draft NVIDIA's NVFP4 checkpoint with a DFlash2 model

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.

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.
 

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");

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

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

zig/src/families/qwen3_5/cuda_native_moe.zig

@@ -7,15 +7,15 @@ 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;
 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;

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