From 7c2cbde4fa0f559d5dca65d90d3f2cc20014508c Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sun, 1 Nov 2026 01:16:08 +0000 Subject: [PATCH] qwen3_5: the --drafter openers move to cuda_drafter_open.zig, so cuda_native.zig has room again cuda_native.zig sat at exactly 600 lines after the shared-round pricing landed, so any change to it failed the 600-line rule (#1750). The DFlash2 and MTP openers, their reserve/close callbacks and the DFlash2 calibration move unchanged; Opened and openDrafter are re-exported, so nothing else changes. --- .../families/qwen3_5/cuda_drafter_open.zig | 91 +++++++++++++++++++ zig/src/families/qwen3_5/cuda_native.zig | 91 ++----------------- 2 files changed, 97 insertions(+), 85 deletions(-) create mode 100644 zig/src/families/qwen3_5/cuda_drafter_open.zig diff --git a/zig/src/families/qwen3_5/cuda_drafter_open.zig b/zig/src/families/qwen3_5/cuda_drafter_open.zig new file mode 100644 index 0000000000..f96620d30b --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_drafter_open.zig @@ -0,0 +1,91 @@ +//! --drafter over the Qwen CUDA engine: a DFlash2 draft model, or the MoE target's own MTP layer. +const std = @import("std"); +const lanes = @import("lanes"); +const native = @import("cuda_native.zig"); +const Draft = @import("cuda_draft.zig").Draft; +const Mtp = @import("cuda_mtp.zig").Mtp; +const draft_config = @import("cuda_draft_config.zig"); +const cost_adapter = native.cost_adapter; +const Lanes = native.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. +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 + 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(); + try l.e.setTaps(d.taps()); + d.target_taps = l.e.tapBytes(); + cost_adapter.identify(l, dir, .dflash); + l.calibration = d; + l.draft = d; + return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter, .prepare = .{ .ctx = l, .run = calibrateDrafter, .bytes = if (l.capacity < l.e.window_rows) try native.cacheBytes(l.e, l.e.window_rows, l.page_test) else 0 } }; +} + +/// The MTP layer beside the checkpoint in `dir` as the drafter over engine `p` (qwen3_5_moe), its ring the window. +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 + 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(); + try l.e.setTaps(d.taps()); + cost_adapter.identify(l, dir, .mtp); + try l.timeWindows(d, null); + l.mtp = d; + return .{ .drafter = d.drafter(), .ctx = d, .close = closeMtp, .reserve = reserveMtp }; +} + +fn reserveMtp(p: *anyopaque, streams: u32) anyerror!void { + const d: *Mtp = @ptrCast(@alignCast(p)); + try d.reserve(streams); +} + +fn closeMtp(p: *anyopaque) void { + const d: *Mtp = @ptrCast(@alignCast(p)); + d.deinit(); +} + +fn calibrateDrafter(p: *anyopaque) !lanes.drafter.Facts { + const l: *Lanes = @ptrCast(@alignCast(p)); + const d = l.calibration orelse return error.MissingCalibrationDrafter; + if (d.n_costs == 0) { + if (l.pool.items.len == 0) return error.MissingReservedQwenCache; + try l.timeWindows(d, l.pool.items[0]); + } + if (d.n_costs != l.e.window_rows) return error.IncompleteWindowMeasurements; + l.calibration = null; + var phases = l.startup; + phases.add(d.startup); + phases.print(); + return d.drafter().facts(); +} + +fn reserveDrafter(p: *anyopaque, streams: u32) anyerror!void { + const d: *Draft = @ptrCast(@alignCast(p)); + try d.reserve(streams); +} + +fn closeDrafter(p: *anyopaque) void { + const d: *Draft = @ptrCast(@alignCast(p)); + d.deinit(); +} diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 1665dea3e2..f26bbf6e41 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -11,7 +11,6 @@ pub const cost_adapter = @import("cuda_cost_adapter.zig"); const draws = @import("draw.zig"); const Draft = @import("cuda_draft.zig").Draft; const cuda_rope = @import("cuda_rope.zig"); -const draft_config = @import("cuda_draft_config.zig"); const prefill = @import("cuda_prefill.zig"); const Fill = prefill.Fill; const reuse = @import("cuda_reuse.zig"); @@ -83,7 +82,7 @@ pub const Loaded = struct { all_on_pages: bool = false, }; -fn cacheBytes(e: *const E.Engine, cap: usize, test_tokens: usize) !usize { +pub fn cacheBytes(e: *const E.Engine, cap: usize, test_tokens: usize) !usize { return std.math.add(usize, e.streamBytes(cap), try cache_mod.pageExtra(e.g, cap, test_tokens, e.kv8)); } @@ -191,87 +190,9 @@ 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. -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 - 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(); - try l.e.setTaps(d.taps()); - d.target_taps = l.e.tapBytes(); - cost_adapter.identify(l, dir, .dflash); - l.calibration = d; - l.draft = d; - return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter, .prepare = .{ .ctx = l, .run = calibrateDrafter, .bytes = if (l.capacity < l.e.window_rows) try cacheBytes(l.e, l.e.window_rows, l.page_test) else 0 } }; -} - -/// The MTP layer beside the checkpoint in `dir` as the drafter over engine `p` (qwen3_5_moe), its ring the window. -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 - 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(); - try l.e.setTaps(d.taps()); - cost_adapter.identify(l, dir, .mtp); - try l.timeWindows(d, null); - l.mtp = d; - return .{ .drafter = d.drafter(), .ctx = d, .close = closeMtp, .reserve = reserveMtp }; -} - -fn reserveMtp(p: *anyopaque, streams: u32) anyerror!void { - const d: *Mtp = @ptrCast(@alignCast(p)); - try d.reserve(streams); -} - -fn closeMtp(p: *anyopaque) void { - const d: *Mtp = @ptrCast(@alignCast(p)); - d.deinit(); -} - -fn calibrateDrafter(p: *anyopaque) !lanes.drafter.Facts { - const l: *Lanes = @ptrCast(@alignCast(p)); - const d = l.calibration orelse return error.MissingCalibrationDrafter; - if (d.n_costs == 0) { - if (l.pool.items.len == 0) return error.MissingReservedQwenCache; - try l.timeWindows(d, l.pool.items[0]); - } - if (d.n_costs != l.e.window_rows) return error.IncompleteWindowMeasurements; - l.calibration = null; - var phases = l.startup; - phases.add(d.startup); - phases.print(); - return d.drafter().facts(); -} - -fn reserveDrafter(p: *anyopaque, streams: u32) anyerror!void { - const d: *Draft = @ptrCast(@alignCast(p)); - try d.reserve(streams); -} - -fn closeDrafter(p: *anyopaque) void { - const d: *Draft = @ptrCast(@alignCast(p)); - d.deinit(); -} +const drafter_open = @import("cuda_drafter_open.zig"); +pub const Opened = drafter_open.Opened; +pub const openDrafter = drafter_open.openDrafter; /// A request this engine refuses, in words, or with no engine (null) its open; null: none of its own. pub fn explain(ctx: ?*anyopaque, err: anyerror) ?[]const u8 { @@ -307,7 +228,7 @@ fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { return .{ .rule = cuda.sample.Rule.from(lanes.gpu_rule.rule(s)), .position = position }; } -const Lanes = struct { +pub const Lanes = struct { timing_on: bool = false, trace: ?*round_timing.Trace = null, gpa: std.mem.Allocator, @@ -398,7 +319,7 @@ const Lanes = struct { } /// Each width timed on a reset cache after its captured launch shape is warm. - fn timeWindows(self: *Lanes, d: anytype, reserved: ?*E.Cache) !void { + pub fn timeWindows(self: *Lanes, d: anytype, reserved: ?*E.Cache) !void { self.startup.add(try cost_adapter.measure(self, d, reserved)); }