From 1be6dc00b243b1c692d5ccc06e52d74143e4773d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:03:35 +0000 Subject: [PATCH 01/35] Windows: a platform layer and host builds that pass - zig/src/platform: `sys`, imported by every module (zig/build/sys.zig wires it). On POSIX it is std.c, std.posix and std.DynLib unchanged; on Windows win_c.zig gives libc's surface over the UCRT and Win32 (integer descriptors, directory descriptors, readdir, pread/pwrite, high-resolution sleep, monotonic clock, /tmp and /proc/self/exe mapped), win_posix.zig the server's sockets over Winsock, console events as SIGINT/SIGTERM, and LoadLibraryW for DynLib. - std.c and std.DynLib calls go through `platform` (mechanical). - build.zig routes Windows targets to the CUDA build. - Server: socket closes and writes are explicit, the key file's mode check is skipped on NTFS (ACLs), descriptor limits are a no-op. - shard_edit: the shard's map is released before the rename, since Windows refuses to replace a mapped file. - Tests that need symlink privilege or mode bits skip on Windows. zig build test-cuda-host passes on x86_64-windows. --- build.zig | 6 +- zig/build/sys.zig | 41 ++ zig/src/core/checkpoint_metal.zig | 23 +- zig/src/core/direct_io.zig | 25 +- zig/src/core/heat.zig | 13 +- zig/src/core/lanes/cost_cache.zig | 5 +- zig/src/core/lanes/grammar.zig | 3 +- zig/src/core/lanes/learned_dirs.zig | 105 ++-- zig/src/core/lanes/learned_disk.zig | 6 +- zig/src/core/learned_fault_test.zig | 41 +- zig/src/core/learned_index_orphan_test.zig | 19 +- zig/src/core/learned_part_test.zig | 45 +- zig/src/core/paced_read.zig | 25 +- zig/src/core/pair.zig | 15 +- zig/src/core/prompt_cache_fake.zig | 11 +- zig/src/core/prompt_cache_learn_test.zig | 29 +- zig/src/core/prompt_imprint.zig | 133 ++-- zig/src/core/prompt_imprint_fault_test.zig | 19 +- zig/src/core/prompt_imprint_test.zig | 31 +- zig/src/core/shard_edit.zig | 57 +- zig/src/core/snapshot_file.zig | 21 +- zig/src/core/template/methods.zig | 15 +- zig/src/cuda/cublaslt.zig | 5 +- zig/src/cuda/driver.zig | 7 +- zig/src/cuda/nccl.zig | 5 +- zig/src/fabric/mailbox.zig | 33 +- zig/src/fabric/mcdma.zig | 3 +- zig/src/fabric/mcdma_abi.zig | 5 +- zig/src/fabric/registration.zig | 45 +- zig/src/fabric/verbs.zig | 7 +- zig/src/fabric/verbs_abi.zig | 11 +- zig/src/fabric/words.zig | 19 +- zig/src/families/flashnext/call_log.zig | 3 +- zig/src/families/flashnext/engine.zig | 19 +- zig/src/families/flashnext/follow.zig | 9 +- zig/src/families/flashnext/host_test.zig | 9 +- zig/src/families/flashnext/learned.zig | 19 +- zig/src/families/flashnext/replay.zig | 9 +- zig/src/families/flashnext/rows_load.zig | 3 +- zig/src/families/flashnext/tp.zig | 41 +- zig/src/families/glm/checks.zig | 15 +- zig/src/families/glm/engine.zig | 51 +- zig/src/families/glm/ep.zig | 23 +- zig/src/families/glm/ep_control.zig | 11 +- zig/src/families/glm/learned_pair.zig | 11 +- zig/src/families/glm/load_plan.zig | 5 +- zig/src/families/glm/mirror.zig | 7 +- zig/src/families/glm/names.zig | 7 +- zig/src/families/glm/page_cache.zig | 35 +- zig/src/families/glm/plain_rounds.zig | 5 +- zig/src/families/glm/pool.zig | 3 +- zig/src/families/glm/prompt.zig | 3 +- zig/src/families/glm/slots.zig | 7 +- zig/src/families/glm/snapshot_file.zig | 29 +- zig/src/families/glm/weights.zig | 1 + zig/src/families/glm/weights_fix.zig | 19 +- zig/src/families/glm/weights_test.zig | 7 +- zig/src/families/nemotron/backend.zig | 3 +- zig/src/families/nemotron/cuda_engine.zig | 3 +- zig/src/families/nemotron/cuda_native.zig | 5 +- zig/src/families/nemotron/cuda_seq_graphs.zig | 3 +- zig/src/families/nemotron/forward.zig | 3 +- zig/src/families/nemotron/gpu_round.zig | 5 +- zig/src/families/nemotron/learned_prompts.zig | 21 +- zig/src/families/nemotron/learner.zig | 3 +- zig/src/families/nemotron/living.zig | 5 +- zig/src/families/nemotron/model.zig | 3 +- zig/src/families/nemotron/pocket.zig | 7 +- zig/src/families/nemotron/pocket_dump.zig | 7 +- zig/src/families/nemotron/whole.zig | 5 +- zig/src/families/qwen27/model.zig | 3 +- zig/src/hip/driver.zig | 5 +- zig/src/hip/runtime.zig | 5 +- zig/src/hip/symbols.zig | 3 +- zig/src/learned_warning_test.zig | 3 +- zig/src/metal/clock.zig | 7 +- zig/src/metal/device.zig | 7 +- zig/src/native/cuda.zig | 3 +- zig/src/native/flashnext_host.zig | 3 +- zig/src/native/metal.zig | 15 +- zig/src/native/qwen27_lanes.zig | 3 +- zig/src/native/qwen27_served.zig | 3 +- zig/src/platform/sys.zig | 47 ++ zig/src/platform/win_c.zig | 577 ++++++++++++++++++ zig/src/platform/win_posix.zig | 279 +++++++++ zig/src/server/auth.zig | 50 +- zig/src/server/http_conn.zig | 3 +- zig/src/server/late_system_test.zig | 21 +- zig/src/server/listener.zig | 3 +- zig/src/server/live.zig | 3 +- zig/src/server/log.zig | 5 +- zig/src/server/request_log.zig | 8 +- zig/src/server/serve.zig | 6 +- zig/src/server/server.zig | 5 +- zig/tests/cuda/pack_compare.zig | 3 +- zig/tests/cuda/train_pack.zig | 3 +- zig/tests/fabric_library.zig | 3 +- zig/tests/flashnext_bench.zig | 3 +- zig/tests/flashnext_reuse.zig | 5 +- zig/tests/flashnext_run.zig | 137 ++--- zig/tests/flashnext_think.zig | 5 +- zig/tests/glm_attn_probe.zig | 7 +- zig/tests/glm_dsa_prompt_check.zig | 5 +- zig/tests/glm_dsa_run.zig | 9 +- zig/tests/glm_rows_speed.zig | 7 +- zig/tests/glm_run.zig | 69 +-- zig/tests/linktest.zig | 9 +- zig/tests/linktest_probe.zig | 3 +- zig/tests/linktest_suite.zig | 3 +- zig/tests/mm_bench.zig | 5 +- zig/tests/moe_bench.zig | 9 +- zig/tests/nemotron_prefill_check.zig | 3 +- zig/tests/pocket_fit_check.zig | 3 +- zig/tests/qwen27/dflash_graph.zig | 13 +- zig/tests/qwen27/dflash_run.zig | 3 +- zig/tests/qwen27/dflash_served.zig | 3 +- zig/tests/qwen27/reuse.zig | 3 +- zig/tests/server/golden_check.zig | 5 +- zig/tests/train_check.zig | 5 +- zig/tests/weight_read_check.zig | 3 +- 120 files changed, 1883 insertions(+), 762 deletions(-) create mode 100644 zig/build/sys.zig create mode 100644 zig/src/platform/sys.zig create mode 100644 zig/src/platform/win_c.zig create mode 100644 zig/src/platform/win_posix.zig diff --git a/build.zig b/build.zig index 08765bd1df..01a2a30bb4 100644 --- a/build.zig +++ b/build.zig @@ -12,6 +12,7 @@ const counters_build = @import("zig/build/metal_counters.zig"); const exl3_build = @import("zig/build/exl3.zig"); const exlr8_build = @import("zig/build/exlr8.zig"); const golden_build = @import("zig/build/golden.zig"); +const sys_build = @import("zig/build/sys.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -153,13 +154,13 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&cost_run.step); switch (target.result.os.tag) { .macos => metalTargets(b, target, optimize, draft_ids, build_options, test_step), - .linux => cuda_build.targets(b, target, optimize, draft_ids, build_options), + .linux, .windows => cuda_build.targets(b, target, optimize, draft_ids, build_options), else => {}, } const xgrammar = grammar_build.option(b); dist_build.targets(b, draft_ids, build_options, release_version, xgrammar); cuda_build.hostTests(b, draft_ids, build_options, test_step, target.result.os.tag != .macos); - if (target.result.os.tag == .linux) { + if (target.result.os.tag == .linux or target.result.os.tag == .windows) { // the engine seam's tests (lane_host.zig and its siblings), which the macOS branch runs from its own api module const lanes_mod = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const api_mod = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes_mod }} }); @@ -189,6 +190,7 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&warning_run.step); b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); + sys_build.wire(b); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/sys.zig b/zig/build/sys.zig new file mode 100644 index 0000000000..b34be717ed --- /dev/null +++ b/zig/build/sys.zig @@ -0,0 +1,41 @@ +//! `sys` (zig/src/platform/sys.zig) for every module the build compiles: one walk of the step graph after it is +//! declared, so no module's declaration needs to name it. + +const std = @import("std"); + +pub fn wire(b: *std.Build) void { + var seen_steps: std.AutoHashMapUnmanaged(*std.Build.Step, void) = .empty; + var seen_mods: std.AutoHashMapUnmanaged(*std.Build.Module, void) = .empty; + var per_target: std.StringHashMapUnmanaged(*std.Build.Module) = .empty; + var w: Walk = .{ .b = b, .seen_steps = &seen_steps, .seen_mods = &seen_mods, .per_target = &per_target }; + for (b.top_level_steps.values()) |top| w.step(&top.step); +} + +const Walk = struct { + b: *std.Build, + seen_steps: *std.AutoHashMapUnmanaged(*std.Build.Step, void), + seen_mods: *std.AutoHashMapUnmanaged(*std.Build.Module, void), + per_target: *std.StringHashMapUnmanaged(*std.Build.Module), + + fn step(w: *Walk, s: *std.Build.Step) void { + if ((w.seen_steps.getOrPut(w.b.allocator, s) catch @panic("OOM")).found_existing) return; + if (s.cast(std.Build.Step.Compile)) |compile| w.module(compile.root_module, compile.root_module.resolved_target); + for (s.dependencies.items) |dep| w.step(dep); + } + + fn module(w: *Walk, m: *std.Build.Module, inherited: ?std.Build.ResolvedTarget) void { + if ((w.seen_mods.getOrPut(w.b.allocator, m) catch @panic("OOM")).found_existing) return; + const target = m.resolved_target orelse inherited orelse w.b.graph.host; + // only TensorFold's own sources see it; the sys module itself and generated modules do not + if (m.root_source_file) |root| if (root == .src_path and !m.import_table.contains("sys") and !std.mem.endsWith(u8, root.src_path.sub_path, "platform/sys.zig")) + m.addImport("sys", w.sys(target)); + for (m.import_table.values()) |dep| w.module(dep, target); + } + + fn sys(w: *Walk, target: std.Build.ResolvedTarget) *std.Build.Module { + const key = target.result.zigTriple(w.b.allocator) catch @panic("OOM"); + const slot = w.per_target.getOrPut(w.b.allocator, key) catch @panic("OOM"); + if (!slot.found_existing) slot.value_ptr.* = w.b.createModule(.{ .root_source_file = w.b.path("zig/src/platform/sys.zig"), .target = target, .link_libc = true }); + return slot.value_ptr.*; + } +}; diff --git a/zig/src/core/checkpoint_metal.zig b/zig/src/core/checkpoint_metal.zig index aabd70f340..45939c6bb5 100644 --- a/zig/src/core/checkpoint_metal.zig +++ b/zig/src/core/checkpoint_metal.zig @@ -1,5 +1,6 @@ //! The shared safetensors index read into Metal buffers: one shared buffer a shard, each tensor at its data offset. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("safetensors.zig"); @@ -33,7 +34,7 @@ pub const Deferred = struct { file: usize, at: usize, bytes: usize, dtype: DType pub const Later = *const fn (name: []const u8, rank: usize) bool; /// A shard file that deferred tensors stay in, and its address-only views' reservation once one is asked for. -const File = struct { fd: std.c.fd_t, end: usize = 0, views: ?[]align(std.heap.page_size_min) u8 = null }; +const File = struct { fd: platform.c.fd_t, end: usize = 0, views: ?[]align(std.heap.page_size_min) u8 = null }; /// Kept tensors pack at this alignment; reads this large or larger run on 8 threads. pub const alignment = 256; @@ -63,14 +64,14 @@ pub const Checkpoint = struct { /// As addFileSelected; with `later`, the tensors it names stay in the file for `read` and the rest pack into one buffer. pub fn addFileSplit(self: *Checkpoint, device: mtl.Device, path: [:0]const u8, prefix: []const u8, selected: ?[]const u8, later: ?Later) !void { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) { std.log.err("cannot open {s}", .{path}); return error.OpenFailed; } var held = false; defer if (!held) { - _ = std.c.close(fd); + _ = platform.c.close(fd); }; var head: [8]u8 = undefined; try readAll(fd, &head, 0); @@ -78,7 +79,7 @@ pub const Checkpoint = struct { const header = try self.allocator.alloc(u8, header_len); defer self.allocator.free(header); try readAll(fd, header, 8); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); const data_start = 8 + header_len; const data_len: usize = @as(usize, @intCast(end)) - data_start; @@ -114,7 +115,7 @@ pub const Checkpoint = struct { } /// The kept entries packed into one buffer, the `later` ones recorded for `read`; `held` once the file stays open. - fn split(self: *Checkpoint, device: mtl.Device, fd: std.c.fd_t, held: *bool, data_start: usize, entries: st.Header, prefix: []const u8, later: Later) !void { + fn split(self: *Checkpoint, device: mtl.Device, fd: platform.c.fd_t, held: *bool, data_start: usize, entries: st.Header, prefix: []const u8, later: Later) !void { var arena = std.heap.ArenaAllocator.init(self.allocator); defer arena.deinit(); const names = try arena.allocator().alloc([]const u8, entries.count()); @@ -184,7 +185,7 @@ pub const Checkpoint = struct { /// Close the files the deferred tensors live in, once load has read them. pub fn closeFiles(self: *Checkpoint) void { for (self.files.items) |*file| if (file.fd >= 0) { - _ = std.c.close(file.fd); + _ = platform.c.close(file.fd); file.fd = -1; }; } @@ -300,16 +301,16 @@ pub const Stage = struct { } }; -fn readAll(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readAll(fd: platform.c.fd_t, dest: []u8, at: usize) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } } -const Part = struct { fd: std.c.fd_t, dest: []u8, at: usize, failed: bool = false }; +const Part = struct { fd: platform.c.fd_t, dest: []u8, at: usize, failed: bool = false }; fn readPart(part: *Part) void { readAll(part.fd, part.dest, part.at) catch { @@ -318,7 +319,7 @@ fn readPart(part: *Part) void { } /// pread `dest.len` bytes from `at` on 8 threads (page cache or disk, the copy is the load's main cost). -fn readParallel(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readParallel(fd: platform.c.fd_t, dest: []u8, at: usize) !void { const n_threads = 8; var parts: [n_threads]Part = undefined; var threads: [n_threads]?std.Thread = @splat(null); @@ -335,7 +336,7 @@ fn readParallel(fd: std.c.fd_t, dest: []u8, at: usize) !void { } /// One tensor's bytes: small ones on this thread, large ones on 8. -fn readInto(fd: std.c.fd_t, dest: []u8, at: usize) !void { +fn readInto(fd: platform.c.fd_t, dest: []u8, at: usize) !void { return if (dest.len >= parallel_bytes) readParallel(fd, dest, at) else readAll(fd, dest, at); } diff --git a/zig/src/core/direct_io.zig b/zig/src/core/direct_io.zig index c06f213388..2f4cc546f6 100644 --- a/zig/src/core/direct_io.zig +++ b/zig/src/core/direct_io.zig @@ -1,6 +1,7 @@ //! Positional reads that skip the page cache (O_DIRECT on Linux): 4 KiB-aligned spans into aligned buffers. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); /// Offset, length and buffer alignment a direct read needs (the logical block size of every NVMe we run on). @@ -20,22 +21,22 @@ pub fn fits(bytes: usize) usize { } pub const File = struct { - fd: std.c.fd_t, + fd: platform.c.fd_t, direct: bool, /// `path` read-only, direct when its file system allows it, else through the page cache. pub fn open(path: [:0]const u8) !File { if (builtin.os.tag == .linux) { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }); if (fd >= 0) return .{ .fd = fd, .direct = true }; } - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }); if (fd < 0) return error.FileNotFound; return .{ .fd = fd, .direct = false }; } pub fn close(f: *File) void { - _ = std.c.close(f.fd); + _ = platform.c.close(f.fd); f.* = undefined; } @@ -46,11 +47,11 @@ pub const File = struct { } }; -fn sysPread(_: *anyopaque, fd: std.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { - return std.c.pread(fd, buf, len, @intCast(offset)); +fn sysPread(_: *anyopaque, fd: platform.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { + return platform.c.pread(fd, buf, len, @intCast(offset)); } -fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: *const fn (*anyopaque, std.c.fd_t, [*]u8, usize, u64) isize) ![]u8 { +fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: *const fn (*anyopaque, platform.c.fd_t, [*]u8, usize, u64) isize) ![]u8 { const lo = std.mem.alignBackward(u64, offset, alignment); const want: usize = @intCast(offset + len - lo); const whole = span(offset, len); @@ -62,8 +63,8 @@ fn readAt(f: File, buf: Buffer, offset: u64, len: usize, ctx: *anyopaque, pread: const ask = whole - at; const n = pread(ctx, f.fd, buf.ptr + at, ask, lo + @as(u64, @intCast(at))); if (n < 0) { - if (std.c.errno(n) == .INTR) continue; - std.log.err("read of {d} bytes at {d} failed: {t}", .{ ask, lo + at, std.c.errno(n) }); + if (platform.c.errno(n) == .INTR) continue; + std.log.err("read of {d} bytes at {d} failed: {t}", .{ ask, lo + at, platform.c.errno(n) }); return error.ReadFailed; } if (n == 0) return error.EndOfFile; @@ -89,7 +90,7 @@ test "direct reads return the same bytes as buffered reads" { var f = try File.open("/proc/self/exe"); defer f.close(); // The reference uses the page cache, so its unaligned reads do not inherit O_DIRECT. - var plain: File = .{ .fd = std.c.open("/proc/self/exe", .{ .ACCMODE = .RDONLY, .CLOEXEC = true }), .direct = false }; + var plain: File = .{ .fd = platform.c.open("/proc/self/exe", .{ .ACCMODE = .RDONLY, .CLOEXEC = true }), .direct = false }; try std.testing.expect(plain.fd >= 0); defer plain.close(); const buf = try gpa.alignedAlloc(u8, .fromByteUnits(alignment), 1 << 20); @@ -97,7 +98,7 @@ test "direct reads return the same bytes as buffered reads" { const want = try gpa.alloc(u8, 1 << 19); defer gpa.free(want); for ([_][2]usize{ .{ 0, 1 }, .{ 1, 4095 }, .{ 4095, 2 }, .{ 12345, 100000 }, .{ 8192, 1 << 19 } }) |c| { - const n = std.c.pread(plain.fd, want.ptr, c[1], @intCast(c[0])); + const n = platform.c.pread(plain.fd, want.ptr, c[1], @intCast(c[0])); try std.testing.expectEqual(@as(isize, @intCast(c[1])), n); try std.testing.expectEqualSlices(u8, want[0..c[1]], try f.read(buf, c[0], c[1])); } @@ -111,7 +112,7 @@ const Script = struct { calls: usize = 0, bad: bool = false, - fn pread(ctx: *anyopaque, _: std.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { + fn pread(ctx: *anyopaque, _: platform.c.fd_t, buf: [*]u8, len: usize, offset: u64) isize { const s: *Script = @ptrCast(@alignCast(ctx)); if (s.calls < s.offs.len) s.offs[s.calls] = offset; s.calls += 1; diff --git a/zig/src/core/heat.zig b/zig/src/core/heat.zig index dc024d653d..eec72a6348 100644 --- a/zig/src/core/heat.zig +++ b/zig/src/core/heat.zig @@ -1,5 +1,6 @@ //! Opt-in wait before a prompt chunk when the hottest thermal zone is above a band. const std = @import("std"); +const platform = @import("sys"); pub const default_root = "/sys/class/thermal"; @@ -129,7 +130,7 @@ fn noteNoZone(root: []const u8) void { } fn envSpan(name: [*:0]const u8) ?[]const u8 { - const p = std.c.getenv(name) orelse return null; + const p = platform.c.getenv(name) orelse return null; const s = std.mem.span(p); if (s.len == 0) return null; return s; @@ -194,13 +195,15 @@ test "the chunk runs under the high band, and a hot reading waits until the low try std.testing.expectEqual(@as(u32, 1), stuck.sleeps); // a cancel ends the wait after the pause it came in } -extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: i32) i32; -extern "c" fn unsetenv(name: [*:0]const u8) i32; +const env = if (@import("builtin").os.tag == .windows) @import("sys").c else struct { + extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: i32) i32; + extern "c" fn unsetenv(name: [*:0]const u8) i32; +}; fn putEnv(name: [*:0]const u8, value: ?[*:0]const u8) void { if (value) |v| { - _ = setenv(name, v, 1); - } else _ = unsetenv(name); + _ = env.setenv(name, v, 1); + } else _ = env.unsetenv(name); } test "TF_HEAT_HIGH and TF_HEAT_LOW set the bands together, and TF_HEAT_ROOT the zones' folder" { diff --git a/zig/src/core/lanes/cost_cache.zig b/zig/src/core/lanes/cost_cache.zig index 1b04fe8fde..279aca3931 100644 --- a/zig/src/core/lanes/cost_cache.zig +++ b/zig/src/core/lanes/cost_cache.zig @@ -1,5 +1,6 @@ //! Timed costs are measured once per engine build, chip and model shape, then read back by later loads. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; const magic = "TFWC2\n"; @@ -53,7 +54,7 @@ fn hashPart(h: *std.crypto.hash.sha2.Sha256, part: []const u8) void { } fn path(a: Allocator, k: [32]u8, suffix: []const u8) !?[]u8 { - const home = std.c.getenv("HOME") orelse return null; + const home = platform.c.getenv("HOME") orelse return null; return try std.fmt.allocPrint(a, "{s}/.cache/tensorfold/window-costs/{s}{s}", .{ std.mem.span(home), std.fmt.bytesToHex(k, .lower), suffix }); } @@ -78,7 +79,7 @@ pub fn save(comptime T: type, a: Allocator, io: std.Io, k: [32]u8, value: T) voi const p = (path(a, k, ".bin") catch return) orelse return; defer a.free(p); std.Io.Dir.cwd().createDirPath(io, std.fs.path.dirname(p) orelse return) catch return; - const temp = std.fmt.allocPrint(a, "{s}.{d}.tmp", .{ p, std.c.getpid() }) catch return; + const temp = std.fmt.allocPrint(a, "{s}.{d}.tmp", .{ p, platform.c.getpid() }) catch return; defer a.free(temp); const bytes = std.mem.concat(a, u8, &.{ magic, std.mem.asBytes(&value) }) catch return; defer a.free(bytes); diff --git a/zig/src/core/lanes/grammar.zig b/zig/src/core/lanes/grammar.zig index 371c2ac4d8..2fe9ae3d87 100644 --- a/zig/src/core/lanes/grammar.zig +++ b/zig/src/core/lanes/grammar.zig @@ -1,6 +1,7 @@ //! Structured output: libtfgrammar (xgrammar, loaded at the first structured request) and a stream's grammar walk. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Allocator = std.mem.Allocator; @@ -37,7 +38,7 @@ fn load(io: std.Io, a: Allocator) ?Fns { library, }; var dl = for (names) |name| { - break std.DynLib.open(name) catch continue; + break platform.DynLib.open(name) catch continue; } else return null; var f: Fns = undefined; const info = @typeInfo(Fns).@"struct"; diff --git a/zig/src/core/lanes/learned_dirs.zig b/zig/src/core/lanes/learned_dirs.zig index 780a997a5c..cdf9648438 100644 --- a/zig/src/core/lanes/learned_dirs.zig +++ b/zig/src/core/lanes/learned_dirs.zig @@ -1,5 +1,6 @@ //! Bounded-name learned directories: both ranks plan the same identity deletions before either writes. const std = @import("std"); +const platform = @import("sys"); pub fn rootOf(dir: []const u8) []const u8 { return std.fs.path.dirname(dir) orelse "."; } @@ -13,23 +14,23 @@ fn isCurrent(keep: []const u8, id: u64) bool { } pub fn next(root: []const u8, keep: []const u8, after: ?u64) !?u64 { var buf: [1100]u8 = undefined; - const d = std.c.opendir(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0)) orelse return error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0)) orelse return error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); var out: ?u64 = null; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (name.len != 16) continue; const id = std.fmt.parseInt(u64, name, 16) catch continue; if (id == std.math.maxInt(u64)) return error.LearnedDirectoryType; - if (ent.type != std.c.DT.DIR) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.DIR) return error.LearnedDirectoryType; if (isCurrent(keep, id) or (after != null and id <= after.?)) continue; if (out == null or id < out.?) out = id; } return out; } pub fn unlink(file: [:0]const u8) !void { - const rc = std.c.unlink(file); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.unlink(file); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } fn stateKey(name: []const u8) ?u64 { @@ -52,38 +53,38 @@ fn identityName(name: []const u8) bool { } fn unlinkAt(fd: c_int, name: [*:0]const u8) !void { - const rc = std.c.unlinkat(fd, name, 0); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.unlinkat(fd, name, 0); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } /// Only the configured root may be an alias; cleanup never follows an identity or file link. pub fn cleanupParts(root: []const u8) !void { var buf: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0), .{ .ACCMODE = .RDONLY, .DIRECTORY = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0), .{ .ACCMODE = .RDONLY, .DIRECTORY = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const d = std.c.fdopendir(fd) orelse { - _ = std.c.close(fd); + const d = platform.c.fdopendir(fd) orelse { + _ = platform.c.close(fd); return error.LearnedDirectoryRead; }; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); - if (ent.type == std.c.DT.REG and std.mem.eql(u8, name, "clock.part")) { + if (ent.type == platform.c.DT.REG and std.mem.eql(u8, name, "clock.part")) { try unlinkAt(fd, @ptrCast(&ent.name)); continue; } if (!identityName(name)) continue; - if (ent.type != std.c.DT.DIR) return error.LearnedDirectoryType; - const sub_fd = std.c.openat(fd, @ptrCast(&ent.name), .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + if (ent.type != platform.c.DT.DIR) return error.LearnedDirectoryType; + const sub_fd = platform.c.openat(fd, @ptrCast(&ent.name), .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (sub_fd < 0) return error.LearnedDirectoryRead; - const sub = std.c.fdopendir(sub_fd) orelse { - _ = std.c.close(sub_fd); + const sub = platform.c.fdopendir(sub_fd) orelse { + _ = platform.c.close(sub_fd); return error.LearnedDirectoryRead; }; { - defer _ = std.c.closedir(sub); - while (std.c.readdir(sub)) |file| { - if (file.type == std.c.DT.REG and partName(std.mem.sliceTo(&file.name, 0))) try unlinkAt(sub_fd, @ptrCast(&file.name)); + defer _ = platform.c.closedir(sub); + while (platform.c.readdir(sub)) |file| { + if (file.type == platform.c.DT.REG and partName(std.mem.sliceTo(&file.name, 0))) try unlinkAt(sub_fd, @ptrCast(&file.name)); } } } @@ -91,15 +92,15 @@ pub fn cleanupParts(root: []const u8) !void { /// Payloads have no token metadata to recover a lost index entry; only exact owned, unreferenced names go. pub fn removeUnindexed(dir: [:0]const u8, metas: anytype) !void { - const fd = std.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const d = std.c.fdopendir(fd) orelse { - _ = std.c.close(fd); + const d = platform.c.fdopendir(fd) orelse { + _ = platform.c.close(fd); return error.LearnedDirectoryRead; }; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { - if (ent.type != std.c.DT.REG) continue; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { + if (ent.type != platform.c.DT.REG) continue; const name = std.mem.sliceTo(&ent.name, 0); const key = stateKey(name) orelse continue; const indexed = for (metas) |x| { @@ -110,29 +111,29 @@ pub fn removeUnindexed(dir: [:0]const u8, metas: anytype) !void { } } pub fn bytes(dir: [:0]const u8) !u64 { - const d = std.c.opendir(dir) orelse return if (std.c.errno(-1) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(dir) orelse return if (platform.c.errno(-1) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); var n: u64 = 0; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) continue; - if (ent.type != std.c.DT.REG) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.REG) return error.LearnedDirectoryType; var buf: [1400]u8 = undefined; const f = try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0); - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.LearnedDirectoryRead; - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - _ = std.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); + _ = platform.c.close(fd); if (end < 0) return error.LearnedDirectoryRead; n = try std.math.add(u64, n, @intCast(end)); } return n; } pub fn fileBytes(file: [:0]const u8) !u64 { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.LearnedDirectoryRead; return @intCast(end); } @@ -156,16 +157,16 @@ pub fn total(root: []const u8) !u64 { pub fn used(root: []const u8, id: u64) !u64 { var buf: [1100]u8 = undefined; const f = try std.fmt.bufPrintSentinel(&buf, "{s}/{x:0>16}/used", .{ root, id }, 0); - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; - defer _ = std.c.close(fd); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.LearnedDirectoryRead; + defer _ = platform.c.close(fd); var n: u64 = 0; const data = std.mem.asBytes(&n); var done: usize = 0; while (done < data.len) { - const count = std.c.read(fd, data.ptr + done, data.len - done); + const count = platform.c.read(fd, data.ptr + done, data.len - done); if (count < 0) { - if (std.c.errno(count) == .INTR) continue; + if (platform.c.errno(count) == .INTR) continue; return error.LearnedDirectoryRead; } if (count == 0) return 0; // A torn recency hint is oldest; it cannot make a learned payload valid. @@ -178,23 +179,23 @@ pub fn removeOther(root: []const u8, keep: []const u8, id: u64) !void { var buf: [1100]u8 = undefined; const dir = try path(&buf, root, id); _ = try bytes(dir); - const d = std.c.opendir(dir) orelse return if (std.c.errno(-1) == .NOENT) {} else error.LearnedDirectoryRead; - defer _ = std.c.closedir(d); - while (std.c.readdir(d)) |ent| { + const d = platform.c.opendir(dir) orelse return if (platform.c.errno(-1) == .NOENT) {} else error.LearnedDirectoryRead; + defer _ = platform.c.closedir(d); + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) continue; - if (ent.type != std.c.DT.REG) return error.LearnedDirectoryType; + if (ent.type != platform.c.DT.REG) return error.LearnedDirectoryType; var f: [1400]u8 = undefined; try unlink(try std.fmt.bufPrintSentinel(&f, "{s}/{s}", .{ dir, name }, 0)); } - const rc = std.c.rmdir(dir); - if (rc != 0 and std.c.errno(rc) != .NOENT) return error.LearnedUnlink; + const rc = platform.c.rmdir(dir); + if (rc != 0 and platform.c.errno(rc) != .NOENT) return error.LearnedUnlink; } test "acknowledged deletion refuses a directory and accepts an absent file" { var buf: [128]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-unlink-{d}", .{std.c.getpid()}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-unlink-{d}", .{platform.c.getpid()}, 0); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); try std.testing.expectError(error.LearnedUnlink, unlink(dir)); var missing: [160]u8 = undefined; try unlink(try std.fmt.bufPrintSentinel(&missing, "{s}/absent", .{dir}, 0)); diff --git a/zig/src/core/lanes/learned_disk.zig b/zig/src/core/lanes/learned_disk.zig index 9ca5fbd4c0..d6e275fae1 100644 --- a/zig/src/core/lanes/learned_disk.zig +++ b/zig/src/core/lanes/learned_disk.zig @@ -1,5 +1,6 @@ //! Reserve learned-state writes above a free-disk floor, with bounded backoff after refusals or write failures. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const Count = if (builtin.os.tag == .linux) u64 else u32; const Stat = extern struct { @@ -19,14 +20,15 @@ const Stat = extern struct { extern "c" fn statvfs(path: [*:0]const u8, result: *Stat) c_int; pub fn available(dir: [:0]const u8) ?u64 { + if (builtin.os.tag == .windows) return platform.c.diskAvailable(dir); var stat: Stat = undefined; if (statvfs(dir, &stat) != 0) return null; return std.math.mul(u64, stat.frsize, stat.bavail) catch null; } fn now() ?u64 { - var t: std.c.timespec = undefined; - if (std.c.clock_gettime(std.c.CLOCK.MONOTONIC, &t) != 0 or t.sec < 0) return null; + var t: platform.c.timespec = undefined; + if (platform.c.clock_gettime(platform.c.CLOCK.MONOTONIC, &t) != 0 or t.sec < 0) return null; return @intCast(t.sec); } diff --git a/zig/src/core/learned_fault_test.zig b/zig/src/core/learned_fault_test.zig index fefd560d99..98d9c9209b 100644 --- a/zig/src/core/learned_fault_test.zig +++ b/zig/src/core/learned_fault_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const Imprint = @import("prompt_imprint.zig").Imprint; const Fake = @import("prompt_cache_fake.zig").Fake; @@ -14,14 +15,14 @@ const Disk = struct { } }; fn root(buf: []u8, suffix: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-learn-fault-{s}-{d}", .{ suffix, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-learn-fault-{s}-{d}", .{ suffix, platform.c.getpid() }); } fn put(dir: []const u8, name: []const u8, count: usize) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.Create; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(count)) != 0) return error.Create; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(count)) != 0) return error.Create; } pub fn learnOne(im: *Imprint, prompt: []const u32) !void { @@ -61,9 +62,9 @@ pub fn recallOne(im: *Imprint, prompt: []const u32) !void { fn appendIndex(dir: []const u8, bytes: []const u8) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/index", .{dir}, 0), .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/index", .{dir}, 0), .{ .ACCMODE = .WRONLY, .APPEND = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try @import("prompt_imprint.zig").writeAll(fd, bytes); } @@ -125,9 +126,9 @@ test "invalid complete records salvage the valid prefix but unknown versions rem } else |err| try std.testing.expectEqual(error.ImprintRead, err); try std.testing.expectEqual(@as(u64, 104), try im.indexScratchBytes()); var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var actual: [44]u8 = undefined; try std.testing.expect(@import("prompt_imprint.zig").readAt(fd, &actual, 60)); try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(&record), &actual); @@ -206,9 +207,9 @@ test "short future-version prefixes preserve index bytes and payloads on open an defer im.deinit(); if (prefix) try learnOne(&im, &first); var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try @import("prompt_imprint.zig").writeAll(fd, std.mem.sliceAsBytes(&future)[0..cut]); var expected: [96]u8 = undefined; const len = (if (prefix) @as(usize, 60) else 0) + cut; @@ -225,13 +226,13 @@ test "short future-version prefixes preserve index bytes and payloads on open an var actual: [96]u8 = undefined; try std.testing.expect(@import("prompt_imprint.zig").readAt(fd, actual[0..len], 0)); try std.testing.expectEqualSlices(u8, expected[0..len], actual[0..len]); - const payload = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/00000000000000af.bin", .{im.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const payload = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/00000000000000af.bin", .{im.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (payload < 0) return error.FuturePayloadRemoved; - _ = std.c.close(payload); + _ = platform.c.close(payload); if (prefix) { - const kept = std.c.open(try Fake.file(&path, im.dir, Imprint.keyOf(first[0..6])), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const kept = platform.c.open(try Fake.file(&path, im.dir, Imprint.keyOf(first[0..6])), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (kept < 0) return error.ValidPrefixPayloadRemoved; - _ = std.c.close(kept); + _ = platform.c.close(kept); } } }; @@ -272,14 +273,14 @@ test "an index replacement failure retains its metadata and removes its partial var idx_buf: [1200]u8 = undefined; const idx = try std.fmt.bufPrintSentinel(&idx_buf, "{s}/index", .{im.dir}, 0); try dirs.unlink(idx); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(idx, 0o700)); - defer _ = std.c.rmdir(idx); // the tree removal goes two levels down + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(idx, 0o700)); + defer _ = platform.c.rmdir(idx); // the tree removal goes two levels down try std.testing.expectError(error.ImprintWrite, im.remove(Imprint.keyOf(&.{ 1, 2 }))); try std.testing.expect(im.has(Imprint.keyOf(&.{ 1, 2 }))); const part = try std.fmt.bufPrintSentinel(&path, "{s}/index.part", .{im.dir}, 0); - const fd = std.c.open(part, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(part, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd >= 0) { - _ = std.c.close(fd); + _ = platform.c.close(fd); return error.PartRemains; } } @@ -315,8 +316,8 @@ test "unlink failure is observable and missing files are idempotently removed" { const r = try root(&buf, "unlink"); var path: [150]u8 = undefined; const dir = try std.fmt.bufPrintSentinel(&path, "{s}", .{r}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); try std.testing.expectError(error.LearnedUnlink, dirs.unlink(dir)); var absent: [180]u8 = undefined; try dirs.unlink(try std.fmt.bufPrintSentinel(&absent, "{s}/absent", .{dir}, 0)); diff --git a/zig/src/core/learned_index_orphan_test.zig b/zig/src/core/learned_index_orphan_test.zig index 7b361f0907..6b210bb469 100644 --- a/zig/src/core/learned_index_orphan_test.zig +++ b/zig/src/core/learned_index_orphan_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const imprint = @import("prompt_imprint.zig"); const Imprint = imprint.Imprint; const Fake = @import("prompt_cache_fake.zig").Fake; @@ -26,15 +27,15 @@ fn learn(im: *Imprint, prompt: []const u32) !void { try std.testing.expect(im.has(Imprint.keyOf(prompt[0..6]))); } fn exists(file: [:0]const u8) bool { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return false; - _ = std.c.close(fd); + _ = platform.c.close(fd); return true; } fn put(file: [:0]const u8, bytes: []const u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(file, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imprint.writeAll(fd, bytes); } @@ -42,7 +43,7 @@ test "reopen sweeps only unindexed owned payloads left by old post-tear learning const a = std.testing.allocator; for (0..3) |kind| { var buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-lost-{d}-{d}", .{ std.c.getpid(), kind }); + const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-lost-{d}-{d}", .{ platform.c.getpid(), kind }); defer @import("prompt_imprint_test.zig").rmTree(root); const first = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8 }; var later = [_]u32{ 11, 12, 13, 14, 15, 16, 17, 18 }; @@ -57,14 +58,14 @@ test "reopen sweeps only unindexed owned payloads left by old post-tear learning later[0] += 10; } var path: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{im.dir}, 0), .{ .ACCMODE = .RDWR }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var original: [240]u8 = undefined; try std.testing.expect(imprint.readAt(fd, &original, 0)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 60)); - if (std.c.lseek(fd, 60, std.c.SEEK.SET) != 60) return error.TestFile; + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 60)); + if (platform.c.lseek(fd, 60, platform.c.SEEK.SET) != 60) return error.TestFile; const key = Imprint.keyOf(&.{ 31, 32 }); const torn = [_]u32{ 0x494d5052, 2, @truncate(key), @truncate(key >> 32), 1, if (kind == 0) 512 else 2, 0, 100, 0 }; if (kind != 1) try imprint.writeAll(fd, std.mem.sliceAsBytes(&torn)); diff --git a/zig/src/core/learned_part_test.zig b/zig/src/core/learned_part_test.zig index c013e5de1b..2f7ab53027 100644 --- a/zig/src/core/learned_part_test.zig +++ b/zig/src/core/learned_part_test.zig @@ -1,22 +1,23 @@ const std = @import("std"); +const platform = @import("sys"); const Imprint = @import("prompt_imprint.zig").Imprint; const life = @import("learned_fault_test.zig"); fn scratch(buf: []u8, tag: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-learn-part-{s}-{d}", .{ tag, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-learn-part-{s}-{d}", .{ tag, platform.c.getpid() }); } fn put(dir: []const u8, name: []const u8, count: usize) !void { var buf: [1200]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(count)) != 0) return error.TestFile; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(count)) != 0) return error.TestFile; } fn exists(dir: []const u8, name: []const u8) bool { var buf: [1200]u8 = undefined; - const fd = std.c.open(std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0) catch return false, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0) catch return false, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); if (fd < 0) return false; - _ = std.c.close(fd); + _ = platform.c.close(fd); return true; } @@ -95,30 +96,30 @@ test "owned-name links and directories stay untouched and foreign directories ar const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}/target", .{root}, 0); var link_buf: [1200]u8 = undefined; const link = try std.fmt.bufPrintSentinel(&link_buf, "{s}/00000000000000af.bin.part", .{im.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, link)); - defer _ = std.c.unlink(link); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, link)); + defer _ = platform.c.unlink(link); var dir_buf: [1200]u8 = undefined; const dir = try std.fmt.bufPrintSentinel(&dir_buf, "{s}/index.part", .{im.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(dir, 0o700)); - defer _ = std.c.rmdir(dir); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(dir, 0o700)); + defer _ = platform.c.rmdir(dir); var foreign_buf: [1200]u8 = undefined; const foreign = try std.fmt.bufPrintSentinel(&foreign_buf, "{s}/foreign", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(foreign, 0o700)); - defer _ = std.c.rmdir(foreign); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(foreign, 0o700)); + defer _ = platform.c.rmdir(foreign); try put(foreign, "index.part", 3); var reopened = try Imprint.open(a, root, 1, 1 << 20); defer reopened.deinit(); var actual: [1200]u8 = undefined; - const n = std.c.readlink(link, &actual, actual.len); + const n = platform.c.readlink(link, &actual, actual.len); try std.testing.expectEqual(@as(isize, @intCast(target.len)), n); try std.testing.expectEqualSlices(u8, target, actual[0..@intCast(n)]); - const directory = std.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); + const directory = platform.c.open(dir, .{ .ACCMODE = .RDONLY, .DIRECTORY = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); try std.testing.expect(directory >= 0); - _ = std.c.close(directory); + _ = platform.c.close(directory); try std.testing.expect(exists(root, "target")); try std.testing.expect(exists(foreign, "index.part")); var marker_buf: [1400]u8 = undefined; - _ = std.c.unlink(try std.fmt.bufPrintSentinel(&marker_buf, "{s}/index.part", .{foreign}, 0)); + _ = platform.c.unlink(try std.fmt.bufPrintSentinel(&marker_buf, "{s}/index.part", .{foreign}, 0)); } test "an identity-directory symlink refuses open and preserves foreign payloads" { @@ -132,12 +133,12 @@ test "an identity-directory symlink refuses open and preserves foreign payloads" } var target_buf: [1200]u8 = undefined; const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}/foreign", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(target, 0o700)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(target, 0o700)); try put(target, "00000000000000af.bin", 7); var link_buf: [1200]u8 = undefined; const link = try std.fmt.bufPrintSentinel(&link_buf, "{s}/0000000000000001", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, link)); - defer _ = std.c.unlink(link); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, link)); + defer _ = platform.c.unlink(link); try std.testing.expectError(error.ImprintRead, Imprint.open(a, root, 1, 1 << 20)); try std.testing.expect(exists(target, "00000000000000af.bin")); try std.testing.expect(!exists(target, "used")); @@ -154,11 +155,11 @@ test "an explicitly configured root alias still cleans only its real identity ch try put(seeded.dir, "index.part", 4096); } var alias_buf: [128]u8 = undefined; - const alias = try std.fmt.bufPrintSentinel(&alias_buf, "/tmp/tf-learn-root-alias-{d}", .{std.c.getpid()}, 0); + const alias = try std.fmt.bufPrintSentinel(&alias_buf, "/tmp/tf-learn-root-alias-{d}", .{platform.c.getpid()}, 0); var target_buf: [128]u8 = undefined; const target = try std.fmt.bufPrintSentinel(&target_buf, "{s}", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.symlink(target, alias)); - defer _ = std.c.unlink(alias); + try std.testing.expectEqual(@as(c_int, 0), platform.c.symlink(target, alias)); + defer _ = platform.c.unlink(alias); var reopened = try Imprint.open(a, alias, 1, 1 << 20); defer reopened.deinit(); try std.testing.expect(!exists(reopened.dir, "index.part")); diff --git a/zig/src/core/paced_read.zig b/zig/src/core/paced_read.zig index b23dd15c4a..0b72289b1f 100644 --- a/zig/src/core/paced_read.zig +++ b/zig/src/core/paced_read.zig @@ -1,5 +1,6 @@ //! Checkpoint reads that keep out of the page cache and wait for free memory: a load never outruns the kernel's reclaim. const std = @import("std"); +const platform = @import("sys"); /// Offsets, lengths and buffers of uncached reads align to the VM page (F_NOCACHE only skips the cache for aligned reads). pub const alignment = 16384; @@ -16,17 +17,17 @@ pub fn stage(gpa: std.mem.Allocator, bytes: usize) !Stage { pub fn floor() u64 { var mem: u64 = 0; var len: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; + if (platform.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; return @max(8 << 30, mem / 100 * 3); } pub fn freeBytes() u64 { var pages: u32 = 0; var len: usize = @sizeOf(u32); - if (std.c.sysctlbyname("vm.page_free_count", &pages, &len, null, 0) != 0) return std.math.maxInt(u64); + if (platform.c.sysctlbyname("vm.page_free_count", &pages, &len, null, 0) != 0) return std.math.maxInt(u64); var size: u64 = 0; var len2: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.pagesize", &size, &len2, null, 0) != 0) size = alignment; + if (platform.c.sysctlbyname("hw.pagesize", &size, &len2, null, 0) != 0) size = alignment; return @as(u64, pages) * size; } @@ -35,13 +36,13 @@ pub fn waitForRoom(limit_ms: u32) void { const want = floor(); var waited: u32 = 0; while (waited < limit_ms and freeBytes() < want) : (waited += 2) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 2 * std.time.ns_per_ms }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 2 * std.time.ns_per_ms }; + _ = platform.c.nanosleep(&ts, null); } } /// `dest.len` bytes at `at` of `fd` (opened with F_NOCACHE) through page-aligned reads into `buf`, pausing for free memory. -pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { +pub fn read(fd: platform.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { var done: usize = 0; while (done < dest.len) { waitForRoom(10_000); @@ -52,7 +53,7 @@ pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { const span = std.mem.alignForward(usize, skip + want, alignment); var got: usize = 0; while (got < skip + want) { // a read past the file's end comes back short once the bytes asked for are in - const n = std.c.pread(fd, buf.ptr + got, span - got, @intCast(lo + got)); + const n = platform.c.pread(fd, buf.ptr + got, span - got, @intCast(lo + got)); if (n <= 0) return error.ShortRead; got += @intCast(n); } @@ -64,17 +65,17 @@ pub fn read(fd: std.c.fd_t, dest: []u8, at: u64, buf: Stage) !void { test "aligned spans copy the bytes asked for, at any offset and length" { const gpa = std.testing.allocator; var name_buf: [64]u8 = undefined; - const name = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-paced-read-{d}", .{std.c.getpid()}, 0); - const fd = std.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const name = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-paced-read-{d}", .{platform.c.getpid()}, 0); + const fd = platform.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); try std.testing.expect(fd >= 0); defer { - _ = std.c.close(fd); - _ = std.c.unlink(name); + _ = platform.c.close(fd); + _ = platform.c.unlink(name); } const data = try gpa.alloc(u8, 3 * alignment + 777); defer gpa.free(data); for (data, 0..) |*b, i| b.* = @truncate(i *% 131 +% 7); - try std.testing.expectEqual(@as(isize, @intCast(data.len)), std.c.write(fd, data.ptr, data.len)); + try std.testing.expectEqual(@as(isize, @intCast(data.len)), platform.c.write(fd, data.ptr, data.len)); const buf = try stage(gpa, 2 * alignment); defer gpa.free(buf); for ([_][2]usize{ .{ 0, 1 }, .{ 5, 100 }, .{ alignment - 3, 10 }, .{ 1, 2 * alignment + 50 }, .{ 3 * alignment, 777 }, .{ 9, data.len - 9 } }) |c| { diff --git a/zig/src/core/pair.zig b/zig/src/core/pair.zig index b7d6daed68..d6fc64cd68 100644 --- a/zig/src/core/pair.zig +++ b/zig/src/core/pair.zig @@ -1,5 +1,6 @@ //! Two Macs' tensor-parallel exchanges in a command buffer over MCDMA, and rank 0's commands to rank 1. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const sources = @import("kernel_sources"); @@ -54,7 +55,7 @@ threadlocal var raised = false; /// This thread at the user-interactive QoS once: busy CPUs must not delay the pair's sends and commands. pub fn raise() void { if (raised) return; - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); raised = true; } @@ -230,16 +231,16 @@ pub const Pair = struct { /// Wait until this Mac's word at `off` reaches `value`: spins, then short sleeps (`bounded`: 60 s at most). fn reach(p: *Pair, off: usize, value: u64, bounded: bool) !void { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); var spins: usize = 0; while (@atomicLoad(u64, p.word64(off), .acquire) < value) { if (p.stop.load(.acquire)) return error.PairStopped; if (p.failed.load(.acquire)) return error.PairLinkFailed; - if (bounded and std.c.mach_absolute_time() - t0 > 60 * 24_000_000) return error.PairPeerSilent; + if (bounded and platform.c.mach_absolute_time() - t0 > 60 * 24_000_000) return error.PairPeerSilent; spins += 1; if (spins > 200_000) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } } @@ -254,8 +255,8 @@ pub const Pair = struct { if (!reached(posted, next_x) or !reached(p.queued.load(.acquire), next_x)) { idle += 1; if (idle > 1_000_000) { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); continue; } diff --git a/zig/src/core/prompt_cache_fake.zig b/zig/src/core/prompt_cache_fake.zig index 6f6754d0b6..734ecafae9 100644 --- a/zig/src/core/prompt_cache_fake.zig +++ b/zig/src/core/prompt_cache_fake.zig @@ -1,5 +1,6 @@ //! A host-memory snapshot family shared by the prompt cache's tests and the learned-disk fault tests. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const imprint = @import("prompt_imprint.zig"); const Snapshots = pc.Snapshots; @@ -96,7 +97,7 @@ pub const Fake = struct { } var path: [512]u8 = undefined; - _ = std.c.unlink(file(&path, dir, key) catch return); + _ = platform.c.unlink(file(&path, dir, key) catch return); } pub fn file(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { return std.fmt.bufPrintSentinel(buf, "{s}/{x:0>16}.bin", .{ dir, key }, 0); @@ -106,17 +107,17 @@ pub const Fake = struct { f.writes += 1; if (f.fail_write) return error.WriteFailed; var path: [512]u8 = undefined; - const fd = std.c.open(try file(&path, dir, key), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try file(&path, dir, key), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.WriteFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imprint.writeAll(fd, std.mem.asBytes(@as(*State, @ptrCast(@alignCast(saved))))); } fn readFn(ptr: *anyopaque, dir: [:0]const u8, key: u64, at: u32) anyerror!Saved { const f = of(ptr); var path: [512]u8 = undefined; - const fd = std.c.open(try file(&path, dir, key), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(try file(&path, dir, key), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.ReadFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const st = try f.gpa.create(State); errdefer f.gpa.destroy(st); if (!imprint.readAt(fd, std.mem.asBytes(st), 0) or st.at != at) return error.ReadFailed; diff --git a/zig/src/core/prompt_cache_learn_test.zig b/zig/src/core/prompt_cache_learn_test.zig index 73eb857f2a..26f9f3f28e 100644 --- a/zig/src/core/prompt_cache_learn_test.zig +++ b/zig/src/core/prompt_cache_learn_test.zig @@ -1,5 +1,6 @@ //! The prompt cache's learned-state tests: harness states kept on disk, the cap, the disk floor and peers' halves. const std = @import("std"); +const platform = @import("sys"); const pc = @import("prompt_cache.zig"); const imprint = @import("prompt_imprint.zig"); const rmTree = @import("prompt_imprint_test.zig").rmTree; @@ -14,7 +15,7 @@ test "a learned harness state outlives its store: a fresh session on a new one r defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-test-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-test-{d}", .{platform.c.getpid()}); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const harness = [_]u32{ 7, 7, 7, 7, 7, 7 }; // shared by every session: a state at 5 reads it all (lookahead 1) const key = imprint.Imprint.keyOf(&harness); @@ -50,7 +51,7 @@ test "past the learned-state cap the least recently used state is forgotten, and defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-cap-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-cap-{d}", .{platform.c.getpid()}); defer rmTree(root); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const one = [_]u32{ 7, 7, 7, 7, 7, 7 }; @@ -67,8 +68,8 @@ test "past the learned-state cap the least recently used state is forgotten, and try std.testing.expectEqual(fresh(&p2), f.pass(&s, &p2, try s.begin(a, &p2, 7, &.{5}, &.{}, null, &.{}))); try std.testing.expect(im.has(imprint.Imprint.keyOf(&two)) and !im.has(imprint.Imprint.keyOf(&one))); var path: [512]u8 = undefined; - try std.testing.expect(std.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&one))) != 0); // its file went too - _ = std.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&two))); // a file lost behind the store's back + try std.testing.expect(platform.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&one))) != 0); // its file went too + _ = platform.c.unlink(try Fake.file(&path, im.dir, imprint.Imprint.keyOf(&two))); // a file lost behind the store's back var g: Fake = .{ .gpa = gpa }; var t = Store.init(gpa, g.learned(), rules, 1 << 20); defer t.deinit(); @@ -87,7 +88,7 @@ test "learned states are prompt arithmetic: a request that decodes rows below on defer arena.deinit(); const a = arena.allocator(); var root_buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-spans-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&root_buf, "/tmp/tf-learn-spans-{d}", .{platform.c.getpid()}); defer rmTree(root); const rules: pc.Rules = .{ .lookahead = 1, .min_prompt = 0, .min_gap = 1 }; const harness = [_]u32{ 7, 7, 7, 7, 7, 7 }; @@ -138,7 +139,7 @@ test "Store learning preserves the disk floor, backs off writes across keys, and var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-backoff-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-backoff-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 11, 1 << 20); defer im.deinit(); @@ -180,7 +181,7 @@ test "learning reserves the existing index rewrite above the disk floor and resu var arena = std.heap.ArenaAllocator.init(a); defer arena.deinit(); var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-index-floor-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-learn-index-floor-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 12, 1 << 20); defer im.deinit(); @@ -213,7 +214,7 @@ test "actual Store learn writes neither half on peer refusal, backs off, and rem }; const a = std.testing.allocator; var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-refuse-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-refuse-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 51, 1 << 20); defer im.deinit(); @@ -242,9 +243,9 @@ test "actual Store learn writes neither half on peer refusal, backs off, and rem try std.testing.expectEqual(@as(u64, 0), im.admission.reserved); var file_path: [512]u8 = undefined; const file = try Fake.file(&file_path, im.dir, imprint.Imprint.keyOf(prompt[0..6])); - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd >= 0) { - _ = std.c.close(fd); + _ = platform.c.close(fd); return error.OrphanedHalf; } } @@ -263,16 +264,16 @@ test "under-cap disk pressure selects the same LRU victim before either half is }; const a = std.testing.allocator; var path: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-victim-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&path, "/tmp/tf-pair-victim-{d}", .{platform.c.getpid()}); defer rmTree(root); var im = try imprint.Imprint.open(a, root, 52, 1 << 20); defer im.deinit(); try im.add(imprint.Imprint.keyOf(&.{ 1, 2 }), 1, &.{ 1, 2 }, &.{}, 16); var victim_path: [512]u8 = undefined; - const victim_fd = std.c.open(try Fake.file(&victim_path, im.dir, imprint.Imprint.keyOf(&.{ 1, 2 })), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + const victim_fd = platform.c.open(try Fake.file(&victim_path, im.dir, imprint.Imprint.keyOf(&.{ 1, 2 })), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (victim_fd < 0) return error.Create; - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(victim_fd, 16)); - _ = std.c.close(victim_fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(victim_fd, 16)); + _ = platform.c.close(victim_fd); im.admission = .{ .floor = 100, .free = Disk.free, .clock = Disk.clock }; const prompt = [_]u32{ 7, 7, 7, 7, 7, 7, 9, 10 }; Disk.space = 100 + 105 + 256 + (6 * 4) + (try im.indexScratchBytes()) - 8; diff --git a/zig/src/core/prompt_imprint.zig b/zig/src/core/prompt_imprint.zig index 8fee9da5ad..024a8a9b6a 100644 --- a/zig/src/core/prompt_imprint.zig +++ b/zig/src/core/prompt_imprint.zig @@ -1,5 +1,6 @@ //! Prompt-cache states on disk (Imprint): a learned harness state outlives the server and serves later fresh sessions. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; /// One state on disk: the tokens it read (position plus lookahead), the chunk starts below it, and its bytes. @@ -74,10 +75,10 @@ pub const Imprint = struct { /// Reserve the current index's size too: LRU eviction rewrites it through a temporary file. pub fn indexScratchBytes(m: *const Imprint) !u64 { var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) 0 else error.ImprintRead; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) 0 else error.ImprintRead; + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.ImprintRead; return @intCast(end); } @@ -119,10 +120,10 @@ pub const Imprint = struct { pub fn pendingWrite(m: *const Imprint) !?u64 { var buf: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/pending", .{m.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) null else error.ImprintRead; - defer _ = std.c.close(fd); - if (std.c.lseek(fd, 0, std.c.SEEK.END) != 48) return error.ImprintRead; + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/pending", .{m.dir}, 0), .{ .ACCMODE = .RDONLY, .NOFOLLOW = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) null else error.ImprintRead; + defer _ = platform.c.close(fd); + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != 48) return error.ImprintRead; var data: [48]u8 = undefined; if (!readAt(fd, &data, 0)) return error.ImprintRead; var digest: [32]u8 = undefined; @@ -135,19 +136,19 @@ pub const Imprint = struct { var buf: [1100]u8 = undefined; var dest: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&buf, "{s}/pending.part", .{m.dir}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var data: [48]u8 = undefined; std.mem.writeInt(u64, data[0..8], 0x494d5052494e5431, .little); std.mem.writeInt(u64, data[8..16], key, .little); std.crypto.hash.sha2.Sha256.hash(data[0..16], data[16..48], .{}); try writeAll(fd, &data); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; } - if (std.c.rename(tmp, try std.fmt.bufPrintSentinel(&dest, "{s}/pending", .{m.dir}, 0)) != 0) return error.ImprintWrite; + if (platform.c.rename(tmp, try std.fmt.bufPrintSentinel(&dest, "{s}/pending", .{m.dir}, 0)) != 0) return error.ImprintWrite; } pub fn clearPending(m: *const Imprint) !void { var buf: [1100]u8 = undefined; @@ -159,16 +160,16 @@ pub const Imprint = struct { if (m.has(key)) return; const below = cut(starts, at); var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true, .NOFOLLOW = true, .NONBLOCK = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .APPEND = true, .CREAT = true, .NOFOLLOW = true, .NONBLOCK = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); _ = indexSize(fd) catch return error.ImprintWrite; - var end = std.c.lseek(fd, 0, std.c.SEEK.END); + var end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0) return error.ImprintWrite; if (end != m.index_end) { try m.reload(fd); if (m.has(key)) return; - end = std.c.lseek(fd, 0, std.c.SEEK.END); + end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end < 0 or end != m.index_end) return error.ImprintWrite; } const t = try m.gpa.dupe(u32, tokens); @@ -176,9 +177,9 @@ pub const Imprint = struct { const s = try m.gpa.dupe(u32, below); errdefer m.gpa.free(s); try m.metas.ensureUnusedCapacity(m.gpa, 1); - errdefer _ = std.c.ftruncate(fd, end); + errdefer _ = platform.c.ftruncate(fd, end); try record(fd, key, at, tokens, below, bytes); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; m.index_end = @as(u64, @intCast(end)) + HEAD * 4 + 4 * (tokens.len + below.len); m.clock += 1; m.metas.appendAssumeCapacity(.{ .key = key, .at = at, .tokens = t, .starts = s, .bytes = bytes, .used = m.clock }); @@ -190,15 +191,15 @@ pub const Imprint = struct { var path: [1100]u8 = undefined; var part: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&part, "{s}/index.part", .{m.dir}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.ImprintWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); for (m.metas.items, 0..) |x, j| if (j != i) try record(fd, x.key, x.at, x.tokens, x.starts, x.bytes); - if (std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.fsync(fd) != 0) return error.ImprintWrite; } - if (std.c.rename(tmp, try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0)) != 0) return error.ImprintWrite; + if (platform.c.rename(tmp, try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0)) != 0) return error.ImprintWrite; free(m.gpa, m.metas.orderedRemove(i)); m.index_end = 0; for (m.metas.items) |x| m.index_end += HEAD * 4 + 4 * (x.tokens.len + x.starts.len); @@ -207,9 +208,9 @@ pub const Imprint = struct { /// Keep a valid prefix across torn or invalid suffixes; unknown versions and I/O errors refuse learning. fn load(m: *Imprint) !void { var path: [1100]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .NONBLOCK = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return if (std.c.errno(fd) == .NOENT) {} else error.ImprintRead; - defer _ = std.c.close(fd); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&path, "{s}/index", .{m.dir}, 0), .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .NONBLOCK = true }, @as(platform.c.mode_t, 0)); + if (fd < 0) return if (platform.c.errno(fd) == .NOENT) {} else error.ImprintRead; + defer _ = platform.c.close(fd); try m.readIndex(fd); } @@ -300,20 +301,20 @@ fn indexSize(fd: c_int) !u64 { if (@import("builtin").os.tag == .linux) { const linux = std.os.linux; var sx: linux.Statx = undefined; - if (std.c.statx(fd, "", linux.AT.EMPTY_PATH, .{ .TYPE = true, .SIZE = true }, &sx) != 0) return error.ImprintRead; - if (!sx.mask.TYPE or !sx.mask.SIZE or @as(u32, sx.mode) & std.c.S.IFMT != std.c.S.IFREG) return error.ImprintRead; + if (platform.c.statx(fd, "", linux.AT.EMPTY_PATH, .{ .TYPE = true, .SIZE = true }, &sx) != 0) return error.ImprintRead; + if (!sx.mask.TYPE or !sx.mask.SIZE or @as(u32, sx.mode) & platform.c.S.IFMT != platform.c.S.IFREG) return error.ImprintRead; if (sx.size > @as(u64, std.math.maxInt(i64))) return error.ImprintRead; return sx.size; } else { - var stat: std.c.Stat = undefined; - if (std.c.fstat(fd, &stat) != 0 or stat.mode & std.c.S.IFMT != std.c.S.IFREG or stat.size < 0) return error.ImprintRead; + var stat: platform.c.Stat = undefined; + if (platform.c.fstat(fd, &stat) != 0 or stat.mode & platform.c.S.IFMT != platform.c.S.IFREG or stat.size < 0) return error.ImprintRead; return @intCast(stat.size); } } fn repairTail(fd: c_int, whole: u64, observed: u64) !void { - if (std.c.lseek(fd, 0, std.c.SEEK.END) != observed) return error.ImprintRead; - if (std.c.ftruncate(fd, @intCast(whole)) != 0 or std.c.fsync(fd) != 0) return error.ImprintWrite; + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != observed) return error.ImprintRead; + if (platform.c.ftruncate(fd, @intCast(whole)) != 0 or platform.c.fsync(fd) != 0) return error.ImprintWrite; } fn ordered(starts: []const u32) bool { @@ -358,9 +359,9 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { var dirs: std.ArrayList(Dir) = .empty; defer dirs.deinit(gpa); var buf: [1100]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return 0) orelse return 0; + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return 0) orelse return 0; var held: u64 = 0; - while (std.c.readdir(d)) |ent| { + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (!identityName(name) or std.mem.endsWith(u8, keep, name)) continue; var sub: [1100]u8 = undefined; @@ -370,7 +371,7 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { held += x.bytes; dirs.append(gpa, x) catch break; } - _ = std.c.closedir(d); + _ = platform.c.closedir(d); std.mem.sort(Dir, dirs.items, {}, struct { fn older(_: void, a: Dir, b: Dir) bool { return a.used < b.used; @@ -381,7 +382,7 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { var sub: [1100]u8 = undefined; const path = std.fmt.bufPrintSentinel(&sub, "{s}/{s}", .{ root, &x.name }, 0) catch continue; _ = files(path, true); - if (std.c.rmdir(path) == 0) held -|= x.bytes; + if (platform.c.rmdir(path) == 0) held -|= x.bytes; } return held; } @@ -389,27 +390,27 @@ fn sweep(gpa: Allocator, root: []const u8, keep: []const u8, room: u64) u64 { /// The bytes of the regular files directly in `dir`, each unlinked when `remove`. fn files(dir: []const u8, remove: bool) u64 { var buf: [1100]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{dir}, 0) catch return 0) orelse return 0; - defer _ = std.c.closedir(d); + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{dir}, 0) catch return 0) orelse return 0; + defer _ = platform.c.closedir(d); var n: u64 = 0; - while (std.c.readdir(d)) |ent| { - if (ent.type != std.c.DT.REG) continue; + while (platform.c.readdir(d)) |ent| { + if (ent.type != platform.c.DT.REG) continue; var p: [1400]u8 = undefined; const f = std.fmt.bufPrintSentinel(&p, "{s}/{s}", .{ dir, std.mem.sliceTo(&ent.name, 0) }, 0) catch continue; - const fd = std.c.open(f, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(f, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) continue; - const end = std.c.lseek(fd, 0, std.c.SEEK.END); - _ = std.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); + _ = platform.c.close(fd); if (end > 0) n += @intCast(end); - if (remove) _ = std.c.unlink(f); + if (remove) _ = platform.c.unlink(f); } return n; } fn readWord(path: [:0]const u8) u64 { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return 0; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var w: u64 = 0; return if (readAt(fd, std.mem.asBytes(&w), 0)) w else 0; } @@ -417,19 +418,19 @@ fn readWord(path: [:0]const u8) u64 { fn writeWord(path: [:0]const u8, w: u64) void { var buf: [1200]u8 = undefined; const part = std.fmt.bufPrintSentinel(&buf, "{s}.part", .{path}, 0) catch return; - const fd = std.c.open(part, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(part, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true, .NOFOLLOW = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return; - defer _ = std.c.close(fd); - defer _ = std.c.unlink(part); + defer _ = platform.c.close(fd); + defer _ = platform.c.unlink(part); writeAll(fd, std.mem.asBytes(&w)) catch return; - if (std.c.fsync(fd) != 0) return; - _ = std.c.rename(part, path); + if (platform.c.fsync(fd) != 0) return; + _ = platform.c.rename(part, path); } pub fn writeAll(fd: c_int, bytes: []const u8) !void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes.ptr + done, bytes.len - done); + const n = platform.c.write(fd, bytes.ptr + done, bytes.len - done); if (n <= 0) return error.ImprintWrite; done += @intCast(n); } @@ -438,7 +439,7 @@ pub fn writeAll(fd: c_int, bytes: []const u8) !void { pub fn readAt(fd: c_int, dest: []u8, at: u64) bool { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return false; done += @intCast(n); } @@ -448,13 +449,13 @@ pub fn readAt(fd: c_int, dest: []u8, at: u64) bool { /// This Mac's OS build (sysctl kern.osversion): its Metal compiler can change a kernel's bits. pub fn osBuild(buf: []u8) []const u8 { var n: usize = buf.len; - if (std.c.sysctlbyname("kern.osversion", buf.ptr, &n, null, 0) != 0) return ""; + if (platform.c.sysctlbyname("kern.osversion", buf.ptr, &n, null, 0) != 0) return ""; return std.mem.sliceTo(buf[0..n], 0); } /// Where learned states live unless --learn-dir says: $HOME/.cache/tensorfold/learned. pub fn defaultRoot(a: Allocator) ![]const u8 { - const home = std.c.getenv("HOME") orelse return error.ImprintDir; + const home = platform.c.getenv("HOME") orelse return error.ImprintDir; return std.fmt.allocPrint(a, "{s}/.cache/tensorfold/learned", .{std.mem.span(home)}); } @@ -464,11 +465,11 @@ fn makePath(gpa: Allocator, dir: [:0]const u8) !void { defer gpa.free(tmp); for (tmp[1..], 1..) |ch, i| if (ch == '/') { tmp[i] = 0; - _ = std.c.mkdir(tmp.ptr, 0o700); + _ = platform.c.mkdir(tmp.ptr, 0o700); tmp[i] = '/'; }; - const rc = std.c.mkdir(tmp.ptr, 0o700); - if (rc != 0 and std.c.errno(rc) != .EXIST) return error.ImprintDir; + const rc = platform.c.mkdir(tmp.ptr, 0o700); + if (rc != 0 and platform.c.errno(rc) != .EXIST) return error.ImprintDir; } test { @@ -479,19 +480,19 @@ test { test "a failed tail truncation keeps the bytes and refuses repair" { var buf: [128]u8 = undefined; - const path = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-index-repair-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(path); - const writer = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const path = try std.fmt.bufPrintSentinel(&buf, "/tmp/tf-index-repair-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(path); + const writer = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (writer < 0) return error.TestFile; { - defer _ = std.c.close(writer); + defer _ = platform.c.close(writer); try writeAll(writer, "record"); } - const reader = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const reader = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (reader < 0) return error.TestFile; - defer _ = std.c.close(reader); + defer _ = platform.c.close(reader); try std.testing.expectError(error.ImprintWrite, repairTail(reader, 0, 6)); - try std.testing.expectEqual(@as(i64, 6), std.c.lseek(reader, 0, std.c.SEEK.END)); + try std.testing.expectEqual(@as(i64, 6), platform.c.lseek(reader, 0, platform.c.SEEK.END)); var actual: [6]u8 = undefined; try std.testing.expect(readAt(reader, &actual, 0)); try std.testing.expectEqualStrings("record", &actual); diff --git a/zig/src/core/prompt_imprint_fault_test.zig b/zig/src/core/prompt_imprint_fault_test.zig index 418f69e0b5..a1d954bfd2 100644 --- a/zig/src/core/prompt_imprint_fault_test.zig +++ b/zig/src/core/prompt_imprint_fault_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const Imprint = @import("prompt_imprint.zig").Imprint; const imp = @import("prompt_imprint.zig"); const dirs = @import("lanes").learned_dirs; @@ -41,12 +42,12 @@ const Tracked = struct { } }; fn scratch(buf: []u8, suffix: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-imprint-fault-{s}-{d}", .{ suffix, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-imprint-fault-{s}-{d}", .{ suffix, platform.c.getpid() }); } fn word(path: [:0]const u8, n: u64) !void { - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); try imp.writeAll(fd, std.mem.asBytes(&n)); } test "Imprint failed open releases each allocation once after another identity scan" { @@ -58,8 +59,8 @@ test "Imprint failed open releases each allocation once after another identity s defer old.deinit(); var path: [256]u8 = undefined; const unreadable = try std.fmt.bufPrintSentinel(&path, "{s}/unreadable", .{old.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(unreadable, 0o700)); - defer _ = std.c.rmdir(unreadable); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(unreadable, 0o700)); + defer _ = platform.c.rmdir(unreadable); var tracked: Tracked = .{}; try std.testing.expectError(error.ImprintRead, Imprint.open(tracked.allocator(), root, 2, 1 << 20)); try std.testing.expectEqual(@as(usize, 0), tracked.duplicate_frees); @@ -92,8 +93,8 @@ test "failed atomic stamp replacement preserves the previous complete used word" var path: [256]u8 = undefined; try word(try std.fmt.bufPrintSentinel(&path, "{s}/used", .{m.dir}, 0), 42); const part = try std.fmt.bufPrintSentinel(&path, "{s}/used.part", .{m.dir}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(part, 0o700)); - defer _ = std.c.rmdir(part); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(part, 0o700)); + defer _ = platform.c.rmdir(part); var reopened = try Imprint.open(a, root, 1, 1 << 20); defer reopened.deinit(); try std.testing.expectEqual(@as(u64, 42), try dirs.used(root, 1)); @@ -108,7 +109,7 @@ test "used stamp reads still refuse an actual I/O error" { var path: [256]u8 = undefined; const used = try std.fmt.bufPrintSentinel(&path, "{s}/used", .{m.dir}, 0); try dirs.unlink(used); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(used, 0o700)); - defer _ = std.c.rmdir(used); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(used, 0o700)); + defer _ = platform.c.rmdir(used); try std.testing.expectError(error.LearnedDirectoryRead, dirs.used(root, 1)); } diff --git a/zig/src/core/prompt_imprint_test.zig b/zig/src/core/prompt_imprint_test.zig index 66bdabd15e..f3e856e21b 100644 --- a/zig/src/core/prompt_imprint_test.zig +++ b/zig/src/core/prompt_imprint_test.zig @@ -1,5 +1,6 @@ //! Learned prompt states on disk: reopen, eviction order, removal and the sweep of other identities. const std = @import("std"); +const platform = @import("sys"); const imp = @import("prompt_imprint.zig"); const Imprint = imp.Imprint; @@ -7,38 +8,38 @@ const Imprint = imp.Imprint; pub fn rmTree(root: []const u8) void { var buf: [512]u8 = undefined; const r = std.fmt.bufPrintSentinel(&buf, "{s}", .{root}, 0) catch return; - const d = std.c.opendir(r) orelse return; - while (std.c.readdir(d)) |ent| { + const d = platform.c.opendir(r) orelse return; + while (platform.c.readdir(d)) |ent| { const name = std.mem.sliceTo(&ent.name, 0); if (name[0] == '.') continue; var sub: [512]u8 = undefined; const p = std.fmt.bufPrintSentinel(&sub, "{s}/{s}", .{ root, name }, 0) catch continue; - if (std.c.unlink(p) == 0) continue; - if (std.c.opendir(p)) |s| { - while (std.c.readdir(s)) |f| { + if (platform.c.unlink(p) == 0) continue; + if (platform.c.opendir(p)) |s| { + while (platform.c.readdir(s)) |f| { const fname = std.mem.sliceTo(&f.name, 0); if (fname[0] == '.') continue; var fb: [600]u8 = undefined; - _ = std.c.unlink(std.fmt.bufPrintSentinel(&fb, "{s}/{s}", .{ p, fname }, 0) catch continue); + _ = platform.c.unlink(std.fmt.bufPrintSentinel(&fb, "{s}/{s}", .{ p, fname }, 0) catch continue); } - _ = std.c.closedir(s); + _ = platform.c.closedir(s); } - _ = std.c.rmdir(p); + _ = platform.c.rmdir(p); } - _ = std.c.closedir(d); - _ = std.c.rmdir(r); + _ = platform.c.closedir(d); + _ = platform.c.rmdir(r); } fn scratch(buf: []u8, tag: []const u8) ![]const u8 { - return std.fmt.bufPrint(buf, "/tmp/tf-imprint-{s}-{d}", .{ tag, std.c.getpid() }); + return std.fmt.bufPrint(buf, "/tmp/tf-imprint-{s}-{d}", .{ tag, platform.c.getpid() }); } /// A file of `n` bytes at `dir`/`name`. fn fill(dir: []const u8, name: []const u8, n: usize) !void { var buf: [512]u8 = undefined; - const fd = std.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&buf, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var zeros: [256]u8 = @splat(0); var left = n; while (left > 0) : (left -= @min(left, zeros.len)) try imp.writeAll(fd, zeros[0..@min(left, zeros.len)]); @@ -46,8 +47,8 @@ fn fill(dir: []const u8, name: []const u8, n: usize) !void { fn exists(path: []const u8) bool { var buf: [512]u8 = undefined; - const d = std.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{path}, 0) catch return false) orelse return false; - _ = std.c.closedir(d); + const d = platform.c.opendir(std.fmt.bufPrintSentinel(&buf, "{s}", .{path}, 0) catch return false) orelse return false; + _ = platform.c.closedir(d); return true; } diff --git a/zig/src/core/shard_edit.zig b/zig/src/core/shard_edit.zig index b340a40523..dff4f9dd5a 100644 --- a/zig/src/core/shard_edit.zig +++ b/zig/src/core/shard_edit.zig @@ -1,5 +1,6 @@ //! Replaced tensors written into a checkpoint's own shards, index and config: in place when sizes hold, else by rename. const std = @import("std"); +const platform = @import("sys"); const st = @import("safetensors.zig"); const Allocator = std.mem.Allocator; const Io = std.Io; @@ -200,18 +201,18 @@ fn inPlace(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacement) } var arena_state = std.heap.ArenaAllocator.init(gpa); defer arena_state.deinit(); - const fd = std.c.open(try arena_state.allocator().dupeSentinel(u8, path, 0), .{ .ACCMODE = .WRONLY }); + const fd = platform.c.open(try arena_state.allocator().dupeSentinel(u8, path, 0), .{ .ACCMODE = .WRONLY }); if (fd < 0) return error.ShardWrite; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); for (edits, places[0..edits.len]) |r, at| { var done: usize = 0; while (done < r.bytes.len) { - const n = std.c.pwrite(fd, r.bytes.ptr + done, @min(r.bytes.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pwrite(fd, r.bytes.ptr + done, @min(r.bytes.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.ShardWrite; done += @intCast(n); } } - if (std.c.fsync(fd) != 0) return error.ShardWrite; + if (platform.c.fsync(fd) != 0) return error.ShardWrite; return true; } @@ -221,7 +222,8 @@ pub fn rewrite(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacem defer arena_state.deinit(); const arena = arena_state.allocator(); var file = try st.File.open(gpa, io, path); - defer file.close(io); + var mapped = true; + defer if (mapped) file.close(io); for (edits) |r| if (!file.names.contains(r.name)) return error.MissingTensor; const raw = file.map.memory; const original = try std.json.parseFromSliceLeaky(Value, arena, raw[8..file.data], .{}); @@ -257,7 +259,11 @@ pub fn rewrite(gpa: Allocator, io: Io, path: []const u8, edits: []const Replacem const spaces = try arena.alloc(u8, padded - text.len); @memset(spaces, ' '); try pieces.insertSlice(arena, 0, &.{ &head, text, spaces }); - try writeFile(arena, path, pieces.items); + const tmp = try writeTemp(arena, path, pieces.items); + // the pieces are written; the map goes before the rename, since Windows refuses to replace a mapped file + file.close(io); + mapped = false; + try commit(arena, tmp, path); return @as(i64, @intCast(offset)) - @as(i64, @intCast(before)); } @@ -325,23 +331,32 @@ fn writeJson(arena: Allocator, path: []const u8, v: Value) !void { /// `parts` to `path` through `path`.tmp, synced, then renamed over it; a reader sees the old file or the new one. fn writeFile(arena: Allocator, path: []const u8, parts: []const []const u8) !void { + try commit(arena, try writeTemp(arena, path, parts), path); +} + +/// `parts` written and synced to `path`.tmp, which is returned. +fn writeTemp(arena: Allocator, path: []const u8, parts: []const []const u8) ![:0]const u8 { const tmp = try std.fmt.allocPrintSentinel(arena, "{s}.tmp", .{path}, 0); - const dest = try arena.dupeSentinel(u8, path, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o644)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o644)); if (fd < 0) return error.ShardWrite; - errdefer _ = std.c.unlink(tmp); - { - defer _ = std.c.close(fd); - for (parts) |p| try put(fd, p); - if (std.c.fsync(fd) != 0) return error.ShardWrite; - } - if (std.c.rename(tmp, dest) != 0) return error.ShardWrite; + errdefer _ = platform.c.unlink(tmp); + defer _ = platform.c.close(fd); + for (parts) |p| try put(fd, p); + if (platform.c.fsync(fd) != 0) return error.ShardWrite; + return tmp; +} + +/// `tmp` renamed over `path`; a failed rename removes `tmp`. +fn commit(arena: Allocator, tmp: [:0]const u8, path: []const u8) !void { + errdefer _ = platform.c.unlink(tmp); + const dest = try arena.dupeSentinel(u8, path, 0); + if (platform.c.rename(tmp, dest) != 0) return error.ShardWrite; } fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.ShardWrite; done += @intCast(n); } @@ -518,13 +533,19 @@ test "a folder whose shard is a symlink or a hard link is refused for learning; defer tmp.cleanup(); try tmp.dir.createDirPath(io, "blobs"); for ([_][]const u8{ "config.json", index_name, "model-00001-of-00002.safetensors", "blobs/b2", "notes.txt" }) |name| try tmp.dir.writeFile(io, .{ .sub_path = name, .data = "x" }); - try tmp.dir.symLink(io, "model-00001-of-00002.safetensors", "readme-link", .{}); + // Windows grants symlinks only with Developer Mode or elevation + tmp.dir.symLink(io, "model-00001-of-00002.safetensors", "readme-link", .{}) catch |e| return if (e == error.PermissionDenied and @import("builtin").os.tag == .windows) error.SkipZigTest else e; const dir = try std.fs.path.join(arena, &.{ ".zig-cache", "tmp", &tmp.sub_path }); try std.testing.expect((try linked(arena, io, dir)) == null); try tmp.dir.symLink(io, "blobs/b2", "model-00002-of-00002.safetensors", .{}); try std.testing.expectEqualStrings("model-00002-of-00002.safetensors", (try linked(arena, io, dir)).?); try tmp.dir.deleteFile(io, "model-00002-of-00002.safetensors"); - try tmp.dir.hardLink("blobs/b2", tmp.dir, "model-00002-of-00002.safetensors", io, .{}); + if (@import("builtin").os.tag == .windows) { + // std's Io has no hard links on Windows; the platform's link (CreateHardLinkW) is what pull uses there + const old = try std.fs.path.joinZ(arena, &.{ dir, "blobs", "b2" }); + const new = try std.fs.path.joinZ(arena, &.{ dir, "model-00002-of-00002.safetensors" }); + try std.testing.expectEqual(@as(c_int, 0), platform.c.link(old, new)); + } else try tmp.dir.hardLink("blobs/b2", tmp.dir, "model-00002-of-00002.safetensors", io, .{}); try std.testing.expectEqualStrings("model-00002-of-00002.safetensors", (try linked(arena, io, dir)).?); } diff --git a/zig/src/core/snapshot_file.zig b/zig/src/core/snapshot_file.zig index 609916ed99..f56997b862 100644 --- a/zig/src/core/snapshot_file.zig +++ b/zig/src/core/snapshot_file.zig @@ -1,5 +1,6 @@ //! Learned snapshots carry a SHA-256 of their entire header and payload; exact length is required on admission. const std = @import("std"); +const platform = @import("sys"); const Hash = std.crypto.hash.sha2.Sha256; @@ -15,7 +16,7 @@ pub fn write(fd: c_int, header: []const u8, payload: []const u8) !void { pub fn read(fd: c_int, header: []const u8, payload: []u8) !void { const length = std.math.add(usize, header.len, payload.len) catch return error.SnapshotRead; const total = std.math.add(usize, length, 32) catch return error.SnapshotRead; - if (std.c.lseek(fd, 0, std.c.SEEK.END) != total) return error.SnapshotRead; + if (platform.c.lseek(fd, 0, platform.c.SEEK.END) != total) return error.SnapshotRead; if (!get(fd, payload, header.len)) return error.SnapshotRead; var stored: [32]u8 = undefined; if (!get(fd, &stored, length)) return error.SnapshotRead; @@ -27,21 +28,21 @@ pub fn read(fd: c_int, header: []const u8, payload: []u8) !void { test "payload bit flips, header changes, truncation, legacy files and extra bytes are refused" { var path: [128]u8 = undefined; - const name = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-snapshot-test-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(name); - const fd = std.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const name = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-snapshot-test-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(name); + const fd = platform.c.open(name, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var output: [4]u8 = undefined; try write(fd, "head", "body"); try read(fd, "head", &output); try std.testing.expectEqualStrings("body", &output); try std.testing.expectError(error.SnapshotRead, read(fd, "dead", &output)); - try std.testing.expectEqual(@as(isize, 1), std.c.pwrite(fd, "X", 1, 5)); + try std.testing.expectEqual(@as(isize, 1), platform.c.pwrite(fd, "X", 1, 5)); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 8)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 8)); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); - _ = std.c.lseek(fd, 0, std.c.SEEK.SET); + _ = platform.c.lseek(fd, 0, platform.c.SEEK.SET); try write(fd, "head", "body"); try put(fd, "x"); try std.testing.expectError(error.SnapshotRead, read(fd, "head", &output)); @@ -50,7 +51,7 @@ test "payload bit flips, header changes, truncation, legacy files and extra byte fn put(fd: c_int, bytes: []const u8) !void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30)); + const n = platform.c.write(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -58,7 +59,7 @@ fn put(fd: c_int, bytes: []const u8) !void { fn get(fd: c_int, bytes: []u8, at: usize) bool { var done: usize = 0; while (done < bytes.len) { - const n = std.c.pread(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, bytes.ptr + done, @min(bytes.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return false; done += @intCast(n); } diff --git a/zig/src/core/template/methods.zig b/zig/src/core/template/methods.zig index 2ec14e09e1..ecccb1497e 100644 --- a/zig/src/core/template/methods.zig +++ b/zig/src/core/template/methods.zig @@ -1,5 +1,6 @@ //! Python methods templates call on strings, dicts and lists, the loop helpers, and the template globals. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const v = @import("value.zig"); const uni = @import("unicode.zig"); @@ -432,14 +433,22 @@ pub fn callFunc(rt: *Rt, f: v.Func, pos: []const Value, kw: []const Kw) Error!Va } const Tm = extern struct { sec: c_int, min: c_int, hour: c_int, mday: c_int, mon: c_int, year: c_int, wday: c_int, yday: c_int, isdst: c_int, gmtoff: c_long, zone: ?[*:0]const u8 }; -extern "c" fn localtime_r(t: *const std.c.time_t, out: *Tm) ?*Tm; +const localtime_r = if (builtin.os.tag == .windows) localtimeWindows else struct { + extern "c" fn localtime_r(t: *const platform.c.time_t, out: *Tm) ?*Tm; +}.localtime_r; +/// The UCRT's localtime_s fills the nine ISO fields; gmtoff and zone are POSIX's and stay unset. +extern "c" fn _localtime64_s(out: *Tm, t: *const i64) c_int; +fn localtimeWindows(t: *const platform.c.time_t, out: *Tm) ?*Tm { + out.* = std.mem.zeroes(Tm); + return if (_localtime64_s(out, t) == 0) out else null; +} extern "c" fn strftime(buf: [*]u8, max: usize, fmt: [*:0]const u8, tm: *const Tm) usize; /// datetime.now().strftime(): %f, %z and %Z are Python's own (naive time), the rest is the C library's. fn strftimeNow(rt: *Rt, fmt: []const u8) Error![]const u8 { if (!builtin.link_libc) return rt.fail("strftime_now needs libc", .{}); - var tv: std.c.timeval = undefined; - _ = std.c.gettimeofday(&tv, null); + var tv: platform.c.timeval = undefined; + _ = platform.c.gettimeofday(&tv, null); var tm: Tm = undefined; if (localtime_r(&tv.sec, &tm) == null) return rt.fail("localtime failed", .{}); var pre: std.ArrayList(u8) = .empty; diff --git a/zig/src/cuda/cublaslt.zig b/zig/src/cuda/cublaslt.zig index cc733b8537..9e1f154172 100644 --- a/zig/src/cuda/cublaslt.zig +++ b/zig/src/cuda/cublaslt.zig @@ -1,6 +1,7 @@ //! cuBLASLt opened at run time: descriptors, layouts, heuristics and matmul, with fp32 compute only (no TF32 or fast modes). const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Status = c_int; @@ -54,7 +55,7 @@ pub const Api = struct { }; pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, /// The CUDA 13 soname first, then the unversioned link name. @@ -69,7 +70,7 @@ pub const Library = struct { } pub fn openPath(path: []const u8) Error!Library { - var lib = std.DynLib.open(path) catch return error.LibraryUnavailable; + var lib = platform.DynLib.open(path) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 0a860eac92..1d644f3742 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -1,21 +1,22 @@ //! The CUDA driver opened at run time: libcuda.so.1's versioned entry points in one table, and checked calls. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Error = error{ DriverUnavailable, MissingSymbol, CudaFailed, OutOfDeviceMemory, NotReady, NotFound, Invalid }; pub const Driver = struct { - lib: std.DynLib, + lib: platform.DynLib, api: abi.Api, pub fn open() Error!Driver { - return openPath("libcuda.so.1"); + return openPath(platform.cuda_driver); } /// Resolves every field of `abi.Api` by its exact name; a missing symbol refuses the whole driver. pub fn openPath(path: []const u8) Error!Driver { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); var api: abi.Api = undefined; const info = @typeInfo(abi.Api).@"struct"; diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..5bc478c6b5 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -1,6 +1,7 @@ //! NCCL opened at run time for tensor-parallel ranks: communicators and the collectives the lane engines use. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Result = c_int; @@ -30,7 +31,7 @@ pub const Api = struct { }; pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, pub fn open() Error!Library { @@ -38,7 +39,7 @@ pub const Library = struct { } pub fn openPath(path: []const u8) Error!Library { - var lib = std.DynLib.open(path) catch return error.LibraryUnavailable; + var lib = platform.DynLib.open(path) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; diff --git a/zig/src/fabric/mailbox.zig b/zig/src/fabric/mailbox.zig index 3010d50af6..49b08f4b4b 100644 --- a/zig/src/fabric/mailbox.zig +++ b/zig/src/fabric/mailbox.zig @@ -1,5 +1,6 @@ //! One mcdma-rpcd mailbox (protocol 1): the mapped halves, a client's ordered call and a service's request and reply. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const layout = @import("layout.zig"); const words = @import("words.zig"); @@ -48,32 +49,32 @@ pub const Mailbox = struct { fn openShm(name: []const u8) Error!Mailbox { var path: [64]u8 = undefined; const z = std.mem.printSentinel(&path, "/mcdma-rpc.{s}", .{name}, 0) catch return error.InvalidName; - const flags: c_int = @bitCast(std.c.O{ .ACCMODE = .RDWR }); - const fd = if (builtin.os.tag == .macos) std.c.shm_open(z.ptr, flags, @as(c_uint, 0)) else std.c.shm_open(z.ptr, flags, 0); + const flags: c_int = @bitCast(platform.c.O{ .ACCMODE = .RDWR }); + const fd = if (builtin.os.tag == .macos) platform.c.shm_open(z.ptr, flags, @as(c_uint, 0)) else platform.c.shm_open(z.ptr, flags, 0); if (fd < 0) return error.NoMailbox; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); return mapFd(fd); } /// The listen end's mailbox: POSIX shared memory on a Mac, the file /dev/shm/mcdma-rpc.NAME on Linux; MCDMA_RPC_BOX_DIR names a test directory. pub fn openListen(name: []const u8) Error!Mailbox { if (!layout.validName(name)) return error.InvalidName; - const env = std.c.getenv("MCDMA_RPC_BOX_DIR"); + const env = platform.c.getenv("MCDMA_RPC_BOX_DIR"); if (env == null and builtin.os.tag == .macos) return openShm(name); const dir = if (env) |d| std.mem.span(d) else "/dev/shm"; var path: [256]u8 = undefined; const z = std.mem.printSentinel(&path, "{s}/mcdma-rpc.{s}", .{ dir, name }, 0) catch return error.InvalidName; - const fd = std.c.open(z.ptr, .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .CLOEXEC = true }); + const fd = platform.c.open(z.ptr, .{ .ACCMODE = .RDWR, .NOFOLLOW = true, .CLOEXEC = true }); if (fd < 0) return error.NoMailbox; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); return mapFd(fd); } /// Map a mailbox the daemon created; it must be this user's and owner-only, as the daemons make them. - fn mapFd(fd: std.c.fd_t) Error!Mailbox { - var st: std.c.Stat = undefined; - if (std.c.fstat(fd, &st) != 0) return error.NoMailbox; - if (st.uid != std.c.getuid() or st.mode & 0o077 != 0) return error.UnsafeMailbox; + fn mapFd(fd: platform.c.fd_t) Error!Mailbox { + var st: platform.c.Stat = undefined; + if (platform.c.fstat(fd, &st) != 0) return error.NoMailbox; + if (st.uid != platform.c.getuid() or st.mode & 0o077 != 0) return error.UnsafeMailbox; const size: usize = @intCast(st.size); if (size < 2 * layout.segment) return error.BadSizes; const mem = std.posix.mmap(null, size, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0) catch return error.NoMailbox; @@ -210,14 +211,14 @@ pub fn allocFake(req: usize, rep: usize) ![]align(page) u8 { } test "a listen mailbox on this Mac is POSIX shared memory, as MCDMA 2.0.0's Mac listen end makes it" { - if (builtin.os.tag != .macos or std.c.getenv("MCDMA_RPC_BOX_DIR") != null) return error.SkipZigTest; + if (builtin.os.tag != .macos or platform.c.getenv("MCDMA_RPC_BOX_DIR") != null) return error.SkipZigTest; const name = "/mcdma-rpc.tf-zig-shm-test"; - const flags: c_int = @bitCast(std.c.O{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true }); - const fd = std.c.shm_open(name, flags, @as(c_uint, 0o600)); + const flags: c_int = @bitCast(platform.c.O{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true }); + const fd = platform.c.shm_open(name, flags, @as(c_uint, 0o600)); try std.testing.expect(fd >= 0); - defer _ = std.c.shm_unlink(name); - defer _ = std.c.close(fd); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 2 * layout.segment)); + defer _ = platform.c.shm_unlink(name); + defer _ = platform.c.close(fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 2 * layout.segment)); const mem = try std.posix.mmap(null, 2 * layout.segment, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); std.mem.writeInt(u64, mem[layout.sizes..][0..8], layout.segment, .little); std.mem.writeInt(u64, mem[layout.sizes + 8 ..][0..8], layout.segment, .little); diff --git a/zig/src/fabric/mcdma.zig b/zig/src/fabric/mcdma.zig index 290bcc1c61..d3292477a9 100644 --- a/zig/src/fabric/mcdma.zig +++ b/zig/src/fabric/mcdma.zig @@ -1,5 +1,6 @@ //! Registered windows over first-party MCDMA: one progress context per port, with staged sources held through flush. const std = @import("std"); +const platform = @import("sys"); const abi = @import("mcdma_abi.zig"); const rma = @import("rdma.zig"); const words = @import("words.zig"); @@ -367,7 +368,7 @@ test "an endpoint gives the library's placing thread user-interactive QoS unless for ([_]?[*:0]const u8{ null, "0", "1" }) |set| { _ = if (set) |v| setenv("MCDMA_FABRIC_QOS", v, 1) else unsetenv("MCDMA_FABRIC_QOS"); try std.testing.expectError(error.LibraryUnavailable, Endpoint.create(std.testing.allocator, "/nonexistent/libmcdma-fabric.dylib", config)); - try std.testing.expectEqualStrings(if (set) |v| std.mem.span(v) else "1", if (std.c.getenv("MCDMA_FABRIC_QOS")) |v| std.mem.span(v) else "(unset)"); + try std.testing.expectEqualStrings(if (set) |v| std.mem.span(v) else "1", if (platform.c.getenv("MCDMA_FABRIC_QOS")) |v| std.mem.span(v) else "(unset)"); } } diff --git a/zig/src/fabric/mcdma_abi.zig b/zig/src/fabric/mcdma_abi.zig index 0b4bcf062c..feeacbb75c 100644 --- a/zig/src/fabric/mcdma_abi.zig +++ b/zig/src/fabric/mcdma_abi.zig @@ -1,5 +1,6 @@ //! First-party libmcdma-fabric ABI 1, loaded only from an explicit absolute library path. const std = @import("std"); +const platform = @import("sys"); pub const Context = opaque {}; pub const Peer = opaque {}; @@ -27,7 +28,7 @@ pub const MaxLinks = *const fn () callconv(.c) c_uint; pub const ws_room: usize = 64; pub const Library = struct { - handle: std.DynLib, + handle: platform.DynLib, api: Api, write_signal: ?WriteSignal = null, connect_links: ?ConnectLinks = null, @@ -36,7 +37,7 @@ pub const Library = struct { /// No install-path or loader-search fallback; the caller supplies the inspected user-directory library. pub fn open(path: [:0]const u8) Error!Library { if (!std.fs.path.isAbsolute(path)) return error.InvalidPath; - var handle = std.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; + var handle = platform.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; errdefer handle.close(); var api: Api = undefined; inline for (@typeInfo(Api).@"struct".field_names) |name| { diff --git a/zig/src/fabric/registration.zig b/zig/src/fabric/registration.zig index 2bbaf084b4..11a18e92cc 100644 --- a/zig/src/fabric/registration.zig +++ b/zig/src/fabric/registration.zig @@ -1,13 +1,14 @@ //! A service's registration with its listen daemon ("MODE poll"), and the daemons' STATUS lines, as protocol 1 defines. const std = @import("std"); +const platform = @import("sys"); const layout = @import("layout.zig"); pub const Error = error{ InvalidName, NoDaemon, Busy, Refused, Io }; /// The registration channel: the daemon sends nothing after OK until the link drops (BYE or a closed socket). pub const Registration = struct { - in_fd: std.c.fd_t, - out_fd: std.c.fd_t, + in_fd: platform.c.fd_t, + out_fd: platform.c.fd_t, /// Register with the listen daemon at MCDMA_RPCD_SOCKET, else /tmp/mcdma-rpcd.NAME.sock. pub fn attach(name: []const u8) Error!Registration { @@ -35,21 +36,21 @@ pub const Registration = struct { /// Whether the daemon still holds the registration: anything readable means BYE or a closed socket. pub fn alive(r: *const Registration) bool { - var fds = [_]std.c.pollfd{.{ .fd = r.in_fd, .events = std.c.POLL.IN, .revents = 0 }}; - const n = std.c.poll(&fds, 1, 0); + var fds = [_]platform.c.pollfd{.{ .fd = r.in_fd, .events = platform.c.POLL.IN, .revents = 0 }}; + const n = platform.c.poll(&fds, 1, 0); return n == 0; } pub fn close(r: *Registration) void { - _ = std.c.close(r.in_fd); - if (r.out_fd != r.in_fd) _ = std.c.close(r.out_fd); + _ = platform.c.close(r.in_fd); + if (r.out_fd != r.in_fd) _ = platform.c.close(r.out_fd); r.* = undefined; } }; /// The listen daemon's socket for link `name`, honouring MCDMA_RPCD_SOCKET as mcdma-rpcd does. pub fn socketPath(buf: []u8, name: []const u8) ![:0]const u8 { - if (std.c.getenv("MCDMA_RPCD_SOCKET")) |env| { + if (platform.c.getenv("MCDMA_RPCD_SOCKET")) |env| { const s = std.mem.span(env); if (s.len > 0) return std.mem.printSentinel(buf, "{s}", .{s}, 0); } @@ -57,35 +58,35 @@ pub fn socketPath(buf: []u8, name: []const u8) ![:0]const u8 { } /// Connect to the daemon's local Unix socket; the only socket the fabric opens, and never TCP. -fn connectUnix(path: [:0]const u8) ?std.c.fd_t { - var addr: std.c.sockaddr.un = .{ .path = undefined }; +fn connectUnix(path: [:0]const u8) ?platform.c.fd_t { + var addr: platform.c.sockaddr.un = .{ .path = undefined }; if (path.len >= addr.path.len) return null; @memset(&addr.path, 0); @memcpy(addr.path[0..path.len], path); - const fd = std.c.socket(std.c.AF.UNIX, std.c.SOCK.STREAM, 0); + const fd = platform.c.socket(platform.c.AF.UNIX, platform.c.SOCK.STREAM, 0); if (fd < 0) return null; - if (std.c.connect(fd, @ptrCast(&addr), @sizeOf(std.c.sockaddr.un)) != 0) { - _ = std.c.close(fd); + if (platform.c.connect(fd, @ptrCast(&addr), @sizeOf(platform.c.sockaddr.un)) != 0) { + _ = platform.c.close(fd); return null; } return fd; } -fn writeAll(fd: std.c.fd_t, bytes: []const u8) Error!void { +fn writeAll(fd: platform.c.fd_t, bytes: []const u8) Error!void { var done: usize = 0; while (done < bytes.len) { - const n = std.c.write(fd, bytes[done..].ptr, bytes.len - done); + const n = platform.c.write(fd, bytes[done..].ptr, bytes.len - done); if (n <= 0) return error.Io; done += @intCast(n); } } /// Read one line a byte at a time, as the daemons write them; the newline is dropped. -fn readLine(fd: std.c.fd_t, buf: []u8) Error![]const u8 { +fn readLine(fd: platform.c.fd_t, buf: []u8) Error![]const u8 { var n: usize = 0; while (n < buf.len) { var c: [1]u8 = undefined; - if (std.c.read(fd, &c, 1) != 1) return error.Io; + if (platform.c.read(fd, &c, 1) != 1) return error.Io; if (c[0] == '\n') return buf[0..n]; buf[n] = c[0]; n += 1; @@ -139,9 +140,9 @@ pub fn parsePeer(line: []const u8) ?Peer { return p; } -fn pipePair() ![2]std.c.fd_t { - var fds: [2]std.c.fd_t = undefined; - if (std.c.pipe(&fds) != 0) return error.Io; +fn pipePair() ![2]platform.c.fd_t { + var fds: [2]platform.c.fd_t = undefined; + if (platform.c.pipe(&fds) != 0) return error.Io; return fds; } @@ -149,14 +150,14 @@ test "registration answers OK, ERR busy or anything else, and BYE ends it" { for ([_][]const u8{ "OK\n", "ERR busy\n", "ERR unknown command\n" }, 0..) |answer, i| { const to_daemon = try pipePair(); const from_daemon = try pipePair(); - defer for ([_]std.c.fd_t{ to_daemon[0], to_daemon[1], from_daemon[0], from_daemon[1] }) |fd| { - _ = std.c.close(fd); + defer for ([_]platform.c.fd_t{ to_daemon[0], to_daemon[1], from_daemon[0], from_daemon[1] }) |fd| { + _ = platform.c.close(fd); }; try writeAll(from_daemon[1], answer); const r: Registration = .{ .in_fd = from_daemon[0], .out_fd = to_daemon[1] }; const got = r.handshake(); var sent: [16]u8 = undefined; - try std.testing.expectEqual(@as(isize, 10), std.c.read(to_daemon[0], &sent, sent.len)); + try std.testing.expectEqual(@as(isize, 10), platform.c.read(to_daemon[0], &sent, sent.len)); try std.testing.expectEqualStrings("MODE poll\n", sent[0..10]); switch (i) { 0 => { diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index 1da50d386d..9b9353c554 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -1,5 +1,6 @@ //! The system's verbs library by dlopen (Apple's librdma, Linux libibverbs): every entry point resolved by name. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const abi = @import("verbs_abi.zig"); @@ -52,13 +53,13 @@ const GidEx = *const fn (ctx: Ctx, port: u32, index: u32, entry: *abi.GidEntry, const InOrder = *const fn (qp: *abi.Qp, op: c_int, flags: u32) callconv(.c) c_int; pub const Verbs = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, gid_ex: ?GidEx = null, in_order: ?InOrder = null, pub fn open(path: ?[:0]const u8) Error!Verbs { - var lib = std.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; + var lib = platform.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; errdefer lib.close(); var api: Api = undefined; const info = @typeInfo(Api).@"struct"; @@ -195,7 +196,7 @@ test "a data-device GID scan skips zero entries and prefers link-local over an I } test "verbs smoke: open the device list and close it (TF_FABRIC_VERBS=1)" { - const flag = std.c.getenv("TF_FABRIC_VERBS") orelse return error.SkipZigTest; + const flag = platform.c.getenv("TF_FABRIC_VERBS") orelse return error.SkipZigTest; if (flag[0] != '1') return error.SkipZigTest; var v = Verbs.open(null) catch |err| { std.debug.print("verbs: {s} unavailable ({s})\n", .{ default_path, @errorName(err) }); diff --git a/zig/src/fabric/verbs_abi.zig b/zig/src/fabric/verbs_abi.zig index 1941fc3157..801fa9d80a 100644 --- a/zig/src/fabric/verbs_abi.zig +++ b/zig/src/fabric/verbs_abi.zig @@ -1,5 +1,6 @@ //! The verbs ABI (rdma-core's infiniband/verbs.h, which Apple's SDK ships for librdma): the structs and constants we use. const std = @import("std"); +const platform = @import("sys"); pub const Device = extern struct { _ops: [2]?*anyopaque, @@ -22,8 +23,8 @@ pub const Qp = extern struct { qp_num: u32, state: c_int, qp_type: c_int, - mutex: std.c.pthread_mutex_t, - cond: std.c.pthread_cond_t, + mutex: platform.c.pthread_mutex_t, + cond: platform.c.pthread_cond_t, events_completed: u32, }; @@ -33,8 +34,8 @@ pub const Cq = extern struct { cq_context: ?*anyopaque, handle: u32, cqe: c_int, - mutex: std.c.pthread_mutex_t, - cond: std.c.pthread_cond_t, + mutex: platform.c.pthread_mutex_t, + cond: platform.c.pthread_cond_t, comp_events_completed: u32, async_events_completed: u32, }; @@ -59,7 +60,7 @@ pub const Context = extern struct { cmd_fd: c_int, async_fd: c_int, num_comp_vectors: c_int, - mutex: std.c.pthread_mutex_t, + mutex: platform.c.pthread_mutex_t, abi_compat: ?*anyopaque, }; diff --git a/zig/src/fabric/words.zig b/zig/src/fabric/words.zig index 53d41a0f58..f9c749fc72 100644 --- a/zig/src/fabric/words.zig +++ b/zig/src/fabric/words.zig @@ -1,5 +1,6 @@ //! Ordered mailbox words: MCDMA's libmcdma-rpc (ABI 1, rpc/mcdma_rpc.h) by dlopen, or its same rules in Zig. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const layout = @import("layout.zig"); @@ -30,20 +31,20 @@ pub fn load(w: *const u64) u64 { /// CLOCK_MONOTONIC in nanoseconds, the clock libmcdma-rpc waits on. pub fn nowNs() u64 { - var ts: std.c.timespec = undefined; - _ = std.c.clock_gettime(.MONOTONIC, &ts); + var ts: platform.c.timespec = undefined; + _ = platform.c.clock_gettime(.MONOTONIC, &ts); return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); } /// Sleep about `ns` nanoseconds (the library's short step between polls). pub fn pause(ns: u64) void { - const ts: std.c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + _ = platform.c.nanosleep(&ts, null); } pub fn checkedNowNs() error{ClockUnavailable}!u64 { - var ts: std.c.timespec = undefined; - if (std.c.clock_gettime(.MONOTONIC, &ts) != 0 or ts.sec < 0 or ts.nsec < 0 or ts.nsec >= std.time.ns_per_s) return error.ClockUnavailable; + var ts: platform.c.timespec = undefined; + if (platform.c.clock_gettime(.MONOTONIC, &ts) != 0 or ts.sec < 0 or ts.nsec < 0 or ts.nsec >= std.time.ns_per_s) return error.ClockUnavailable; return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); } @@ -81,7 +82,7 @@ else /// libmcdma-rpc opened at run time; the three Metal functions are optional, as the header says. pub const Library = struct { - lib: std.DynLib, + lib: platform.DynLib, abi_fn: AbiFn, wait_word: WaitFn, store_word: StoreFn, @@ -92,7 +93,7 @@ pub const Library = struct { /// Open `path`, else MCDMA_RPC_LIBRARY, else the install prefix; refuse any ABI but 1. pub fn open(path: ?[:0]const u8) Error!Library { if (path) |p| return openPath(p); - if (std.c.getenv("MCDMA_RPC_LIBRARY")) |env| { + if (platform.c.getenv("MCDMA_RPC_LIBRARY")) |env| { const p = std.mem.span(env); if (p.len > 0) return openPath(p); } @@ -104,7 +105,7 @@ pub const Library = struct { } pub fn openPath(path: [:0]const u8) Error!Library { - var lib = std.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; + var lib = platform.DynLib.openZ(path.ptr) catch return error.LibraryUnavailable; errdefer lib.close(); var l: Library = .{ .lib = lib, diff --git a/zig/src/families/flashnext/call_log.zig b/zig/src/families/flashnext/call_log.zig index 387ec89220..5afdca5da9 100644 --- a/zig/src/families/flashnext/call_log.zig +++ b/zig/src/families/flashnext/call_log.zig @@ -1,5 +1,6 @@ //! FZ_CALL_LOG: a prompt pass's calls (rows, parts, wall and GPU milliseconds), logged in one line after its first token. const std = @import("std"); +const platform = @import("sys"); const N = 16; @@ -17,7 +18,7 @@ pub const CallLog = struct { if (c.n == N) return; c.rows[c.n] = @intCast(rows); c.parts[c.n] = @intCast(parts); - c.wall[c.n] = @floatCast(@as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6); // mach ticks: 125/3 ns + c.wall[c.n] = @floatCast(@as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6); // mach ticks: 125/3 ns c.gpu[c.n] = @floatCast((g1 - g0) * 1e3); c.n += 1; } diff --git a/zig/src/families/flashnext/engine.zig b/zig/src/families/flashnext/engine.zig index c9f8481249..9ff39b2845 100644 --- a/zig/src/families/flashnext/engine.zig +++ b/zig/src/families/flashnext/engine.zig @@ -1,6 +1,7 @@ //! Flash Next served from the replay engine: one reply at a time, its prompt through prompt chunks (the MTP head's //! keys for every prompt row), then GPU-side rounds whose drafts the depth rule sets; tokens leave as rounds are read. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fz = @import("replay.zig"); const config = @import("config.zig"); @@ -151,8 +152,8 @@ pub const Engine = struct { e.copy_long = 6; e.segments = true; e.mark_taps = true; - e.pair_min = tp_settings.pairMin(if (std.c.getenv("FZ_PAIR_MIN")) |v| std.mem.span(v) else null, PAIR_MIN); // rank 1 takes rank 0's - e.call_log = std.c.getenv("FZ_CALL_LOG") != null; + e.pair_min = tp_settings.pairMin(if (platform.c.getenv("FZ_PAIR_MIN")) |v| std.mem.span(v) else null, PAIR_MIN); // rank 1 takes rank 0's + e.call_log = platform.c.getenv("FZ_CALL_LOG") != null; e.passed = null; e.snap_pool = .{}; e.handoff_ns = 0; @@ -177,8 +178,8 @@ pub const Engine = struct { r.event = try device.sharedEvent(); if (dump_dir) |d| try r.compile(d) else try r.compileChecked(); r.sel = try Select.init(r, MAXR); - r.lane_new = std.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (the recorded bits) - if (std.c.getenv("FZ_GDN")) |v| { + r.lane_new = platform.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (the recorded bits) + if (platform.c.getenv("FZ_GDN")) |v| { r.gdn_pipe = try fz.gdn_step.compile(r, false); if (v[0] == '2') r.gdn_kept = try fz.gdn_step.compile(r, true); // GPU-side rounds keep one state a layer } @@ -211,7 +212,7 @@ pub const Engine = struct { }; e.identity_hash = std.hash.Wyhash.hash(0x6964, identity); const owned_pack_dir = if (dump_dir == null) blk: { - const home = if (std.c.getenv("HOME")) |v| std.mem.span(v) else ""; + const home = if (platform.c.getenv("HOME")) |v| std.mem.span(v) else ""; const root = try pack_cache.defaultRoot(arena, home); break :blk try pack_cache.ensure(gpa, io, model_dir, root, identity, PackBuilder{}); } else null; @@ -434,7 +435,7 @@ pub const Engine = struct { e.tsel = try GSelect.init(r, 0); e.msel = try GSelect.init(r, 0); // speed-up mode: TP=2 with the peer the settings name (tp.zig); both ranks then run the same replies in lockstep - const tp_path: ?[]const u8 = speed_up orelse if (std.c.getenv("TF_FLASHNEXT_TP")) |p| std.mem.span(p) else null; + const tp_path: ?[]const u8 = speed_up orelse if (platform.c.getenv("TF_FLASHNEXT_TP")) |p| std.mem.span(p) else null; if (tp_path != null and !uniform) return error.SpeedUpNeeds6Bit32; // TP splits the 6/32 kernels only if (tp_path) |path| r.tp = try fz.Tp2.init(arena, r.device, path); if (r.tp == null) try fz.rows_load.attach( @@ -504,8 +505,8 @@ pub const Engine = struct { const pool = mtl.objc.Pool.push(); defer pool.pop(); if (r.tp) |tp| if (tp.rank == 0) { // served speed-up mode: rank 1 runs the same request (follow) - const t0 = std.c.mach_absolute_time(); - defer e.handoff_ns = @intCast((std.c.mach_absolute_time() - t0) * 125 / 3); // mach ticks are 125/3 ns on Apple silicon + const t0 = platform.c.mach_absolute_time(); + defer e.handoff_ns = @intCast((platform.c.mach_absolute_time() - t0) * 125 / 3); // mach ticks are 125/3 ns on Apple silicon try follow_mod.send(e, tp, prompt, from, marks, max_tokens, eos, depth, budget); if (!try tp.waitAck()) return error.PeerNotResumed; }; @@ -522,7 +523,7 @@ pub const Engine = struct { var calls: CallLog = .{}; while (at < prompt.len) { if (try e.agree(out.cancelled(out.ctx))) return .{ .reason = .cancelled }; - const t_call = std.c.mach_absolute_time(); + const t_call = platform.c.mach_absolute_time(); const g_call = m.gpu_seconds; if (r.tp) |tp| { // speed-up mode: the call's rows split across the two Macs, running on past up to MARKS marks const pair_taps = e.mark_taps and m.marks != null; diff --git a/zig/src/families/flashnext/follow.zig b/zig/src/families/flashnext/follow.zig index 4255faeb6e..482870c098 100644 --- a/zig/src/families/flashnext/follow.zig +++ b/zig/src/families/flashnext/follow.zig @@ -1,5 +1,6 @@ //! Speed-up mode's rank 1: rank 0's requests run here in step, and this Mac keeps its own states under rank 0's names. const std = @import("std"); +const platform = @import("sys"); const engine = @import("engine.zig"); const snapshot = @import("snapshot.zig"); const fz = @import("replay.zig"); @@ -9,7 +10,7 @@ const Engine = engine.Engine; /// Rank 0's requests as they come, until its empty request; between them a free buffer is readied inside the budget. pub fn run(e: *Engine) !void { while (try one(e, null)) |_| { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const got = e.snap_pool.ready(e.r.device, engine.NEXT_TURN, .{ .ctx = e.r.tp.?, .check = requestWaiting }, e.peer_budget -| e.peer_held -| e.snap_pool.spare()); if (got.cap > 0) std.log.info("speed-up rank 1: readied {d} of {d} MiB for the next save in {d:.1} ms", .{ got.touched >> 20, got.cap >> 20, msSince(t0) }); } @@ -21,7 +22,7 @@ fn requestWaiting(ctx: *anyopaque) bool { } fn msSince(t0: u64) f64 { - return @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6; // mach ticks are 125/3 ns on Apple silicon + return @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e6; // mach ticks are 125/3 ns on Apple silicon } /// Rank 0's request for `one`: a 16-word head, then the prompt, marks, dropped states' keys and the budget. @@ -65,7 +66,7 @@ pub fn one(e: *Engine, out: ?engine.Out) !?engine.Result { const from = head[12]; if (head[15] == 0) return error.TpProtocol; // rank 0 sends its positive pair minimum with every request e.pair_min = head[15]; - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const ok = from == 0 or blk: { const st = e.peer_kept.get(engine.keyOf(prompt[0 .. from + 1])) orelse break :blk false; snapshot.restore(e, st) catch break :blk false; @@ -123,7 +124,7 @@ const Reply = struct { } fn marked(ctx: *anyopaque, at: usize) void { const r: *Reply = @ptrCast(@alignCast(ctx)); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); defer { r.kept += 1; r.keep_ns += msSince(t0) * 1e6; diff --git a/zig/src/families/flashnext/host_test.zig b/zig/src/families/flashnext/host_test.zig index bc198ce4ad..642812ce54 100644 --- a/zig/src/families/flashnext/host_test.zig +++ b/zig/src/families/flashnext/host_test.zig @@ -1,5 +1,6 @@ //! Host-only contracts exercise strict text configuration, affine six-bit packing and borrowed-view shapes. const std = @import("std"); +const platform = @import("sys"); const posix = std.posix; const family = @import("flashnext.zig"); const st = @import("../../core/safetensors.zig"); @@ -298,12 +299,12 @@ test "the prompt's checked-in text resolves and patches from another working dir const sources = @import("kernel_sources").flashnext_gen; var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); - const old_fd = std.c.open(".", .{}); // O_RDONLY via the flags struct's default - defer _ = std.c.close(old_fd); + const old_fd = platform.c.open(".", .{}); // O_RDONLY via the flags struct's default + defer _ = platform.c.close(old_fd); const tmpz = try std.fmt.allocPrintSentinel(std.testing.allocator, ".zig-cache/tmp/{s}", .{tmp.sub_path}, 0); defer std.testing.allocator.free(tmpz); - try std.testing.expect(std.c.chdir(tmpz) == 0); - defer _ = std.c.fchdir(old_fd); + try std.testing.expect(platform.c.chdir(tmpz) == 0); + defer _ = platform.c.fchdir(old_fd); for (table.entries) |e| { var text: []const u8 = ""; for (sources.sources) |s| { diff --git a/zig/src/families/flashnext/learned.zig b/zig/src/families/flashnext/learned.zig index 5d1485abdf..262f44f447 100644 --- a/zig/src/families/flashnext/learned.zig +++ b/zig/src/families/flashnext/learned.zig @@ -1,5 +1,6 @@ //! Flash Next's learned prompt states (--learn): a kept state's file on disk, and the probe keying this build's bits. const std = @import("std"); +const platform = @import("sys"); const ks = @import("kernel_sources"); const snapshot = @import("snapshot.zig"); const roles_gen = @import("roles_gen.zig"); @@ -21,24 +22,24 @@ pub fn path(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { pub fn writeFile(st: *const State, file: [:0]const u8) !void { var tmp_buf: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head: Head = .{ .magic = FILE_MAGIC, .layout = @backingInt(st.layout), .at = st.at, .bytes = st.bytes, .hist = st.hist }; try put(fd, std.mem.asBytes(&head)); try put(fd, st.buf.contents()[0..st.bytes]); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } /// `file`'s state after `at` tokens, in a buffer from the engine's pool; another position, layout or size is refused. pub fn readFile(e: *Engine, gpa: std.mem.Allocator, at: usize, file: [:0]const u8) !*State { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: Head = undefined; try get(fd, std.mem.asBytes(&head), 0); const layout: u32 = @backingInt(snapshot.layoutOf(e)); @@ -53,7 +54,7 @@ pub fn readFile(e: *Engine, gpa: std.mem.Allocator, at: usize, file: [:0]const u fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -62,7 +63,7 @@ fn put(fd: c_int, b: []const u8) !void { fn get(fd: c_int, b: []u8, at: u64) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.SnapshotRead; done += @intCast(n); } diff --git a/zig/src/families/flashnext/replay.zig b/zig/src/families/flashnext/replay.zig index 834c0a732a..7021f8c6b1 100644 --- a/zig/src/families/flashnext/replay.zig +++ b/zig/src/families/flashnext/replay.zig @@ -1,6 +1,7 @@ //! Flash Next on our Metal runtime: the Python engine's recorded kernels (a dump from tools/zig/flashnext_dump.py) //! replayed from Zig, with our own kernels for wide experts, prompt chunks, block selection and GPU-side rounds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); const roles_gen = @import("roles_gen.zig"); @@ -43,7 +44,7 @@ pub const Buf = struct { return .{ .b = x.b, .off = x.off + bytes }; } }; -pub const Entry = struct { fd: std.c.fd_t, at: usize, len: usize }; +pub const Entry = struct { fd: platform.c.fd_t, at: usize, len: usize }; pub const Variant = struct { inputs: [][]const u8, outputs: [][]const u8, meta: [][]const u8, pipe: mtl.Pipeline, file: []const u8 = "", name: []const u8 = "", text: []const u8 = "" }; pub const Site = struct { v: *Variant, grid: mtl.Size, tg: mtl.Size }; @@ -814,10 +815,10 @@ const hc_up_source = \\ ; -pub fn readAll(fd: std.c.fd_t, dest: []u8, at: usize) !void { +pub fn readAll(fd: platform.c.fd_t, dest: []u8, at: usize) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } @@ -915,7 +916,7 @@ pub const Run = struct { /// Every tensor of a safetensors file in the index, at its absolute offset. pub fn indexFile(r: *Run, path: [:0]const u8) !void { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; var head: [8]u8 = undefined; try readAll(fd, &head, 0); diff --git a/zig/src/families/flashnext/rows_load.zig b/zig/src/families/flashnext/rows_load.zig index cf8e89cfbb..1932830cd1 100644 --- a/zig/src/families/flashnext/rows_load.zig +++ b/zig/src/families/flashnext/rows_load.zig @@ -1,5 +1,6 @@ //! Rows mode without tensor units (0.6.6's pre-M5 decode): every dense lane also as MLX-layout members for rows_w. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const replay = @import("replay.zig"); pub const rows_w = @import("rows_w.zig"); @@ -14,7 +15,7 @@ const KO = 6144; // the attention and DeltaNet output projections' inputs /// Whether decode reads its dense projections on the row kernels: no tensor units and FZ_M3_LANES unset. pub fn on(device: mtl.Device) bool { - return !device.tensorUnits() and std.c.getenv("FZ_M3_LANES") == null; + return !device.tensorUnits() and platform.c.getenv("FZ_M3_LANES") == null; } /// `d` on the run's encoder at `rows` rows, unless the role is knocked out. diff --git a/zig/src/families/flashnext/tp.zig b/zig/src/families/flashnext/tp.zig index e5718126d2..3fe8c3a514 100644 --- a/zig/src/families/flashnext/tp.zig +++ b/zig/src/families/flashnext/tp.zig @@ -1,5 +1,6 @@ //! Speed-up mode for Flash Next: two Macs with the whole model each, half the work each, over MCDMA. Decode: each rank computes every routed and the shared expert for its half of the window's rows (`ownRows`) and those rows' MoE branch with one Mac's sums; after each target layer's experts the GPU posts a sequence; the host sends the branch rows (bf16) to the peer, and the next layer's combine waits for the flag the peer's message sets, then updates every row's streams (one Mac's bits). DeltaNet layers split their heads: each rank's out-projection partial (fp32), exchanged and added in rank order. Prefill: a chunk's rows split across the two Macs (replay.zig `chunkPair`); the host sends each layer's handoff into the peer's slot for that layer and signals the layer's flag, which the peer's GPU waits for. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const host_wait = @import("tp_host_wait.zig"); @@ -341,9 +342,9 @@ pub const Tp2 = struct { .jobs = try gpa.alloc(Job, JOBS), }; t.one.slice(i32, 1)[0] = 1; - t.trace = std.c.getenv("TF_TP_TRACE") != null; - t.local = std.c.getenv("TF_TP_LOCAL") != null; - t.stats = std.c.getenv("TF_TP_STATS") != null; + t.trace = platform.c.getenv("TF_TP_TRACE") != null; + t.local = platform.c.getenv("TF_TP_LOCAL") != null; + t.stats = platform.c.getenv("TF_TP_STATS") != null; // both ranks up before the first round: a word each way try rd.signal(t.peer, HELLO, 1); while (@atomicLoad(u64, t.word64(HELLO), .acquire) < 1) std.atomic.spinLoopHint(); @@ -529,11 +530,11 @@ pub const Tp2 = struct { /// Served, rank 0: rank 1's answer for the request just sent: true when it restored the same resume state. pub fn waitAck(t: *Tp2) !bool { - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (true) { const v = @atomicLoad(u64, t.word64(ACK_FLAG), .acquire); if (v >> 1 >= t.req) return v >> 1 == t.req and v & 1 != 0; - if (std.c.mach_absolute_time() - t0 > 7_200_000_000) return error.TpPeerSilent; // 300 s: rank 1 may still be loading + if (platform.c.mach_absolute_time() - t0 > 7_200_000_000) return error.TpPeerSilent; // 300 s: rank 1 may still be loading std.atomic.spinLoopHint(); } } @@ -551,8 +552,8 @@ pub const Tp2 = struct { if (t.quitting.load(.acquire)) return null; spins += 1; if (spins > 100_000) { // idle: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 50_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } return .{ .head = @ptrCast(t.win.ptr + REQ), .tokens = @ptrCast(@alignCast(t.win.ptr + REQ + 64)) }; @@ -565,11 +566,11 @@ pub const Tp2 = struct { try t.rd.signal(t.peer, CTRL_FLAG, (t.ctrl << 6) | (@as(u64, depth & 31) << 1) | @intFromBool(quit)); return .{ .quit = quit, .depth = depth }; } - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (true) { const v = @atomicLoad(u64, t.word64(CTRL_FLAG), .acquire); if (v >> 6 >= t.ctrl) return .{ .quit = quit or (v >> 6 == t.ctrl and v & 1 != 0), .depth = @intCast((v >> 1) & 31) }; - if (std.c.mach_absolute_time() - t0 > 240_000_000) return error.TpPeerSilent; // 10 s + if (platform.c.mach_absolute_time() - t0 > 240_000_000) return error.TpPeerSilent; // 10 s std.atomic.spinLoopHint(); } } @@ -604,7 +605,7 @@ pub const Tp2 = struct { /// for the peer's, by the flag its message sets) or a prefill send (its writes, then its flag). While idle it /// watches the last exchange: 10 s without the peer's answer fails the link. fn service(t: *Tp2, first: u32) void { - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send const posted = t.word32(SYNC + GPU * 4); const flag = t.word64(FLAG); var seq = first; @@ -614,7 +615,7 @@ pub const Tp2 = struct { var w = @atomicLoad(u32, posted, .acquire); while (!postReached(w, seq) or !reached(t.queued.load(.acquire), seq)) { if (t.stop.load(.acquire)) return; - if (want_at > 0 and !reached(@truncate(@atomicLoad(u64, flag, .acquire)), @truncate(want)) and !t.failed.load(.acquire) and std.c.mach_absolute_time() - want_at > 240_000_000) { // 10 s of 24 MHz ticks + if (want_at > 0 and !reached(@truncate(@atomicLoad(u64, flag, .acquire)), @truncate(want)) and !t.failed.load(.acquire) and platform.c.mach_absolute_time() - want_at > 240_000_000) { // 10 s of 24 MHz ticks t.fail(seq, error.PeerSilent); t.land(want); } @@ -622,7 +623,7 @@ pub const Tp2 = struct { w = @atomicLoad(u32, posted, .acquire); } const job = t.jobs[seq % JOBS]; - const seen = if (t.stats) std.c.mach_absolute_time() else 0; + const seen = if (t.stats) platform.c.mach_absolute_time() else 0; if ((t.local and job.kind != .send) or t.failed.load(.acquire)) { // a failed link drains: the GPU never hangs if (job.kind != .send) t.land(job.value); seq +%= 1; @@ -638,11 +639,11 @@ pub const Tp2 = struct { } if (job.kind != .send) { want = x; - want_at = std.c.mach_absolute_time(); + want_at = platform.c.mach_absolute_time(); if (t.failed.load(.acquire)) t.land(x); } if (t.stats) { - const held = std.c.mach_absolute_time() - seen; // ticks: 41.67 ns each on Apple silicon + const held = platform.c.mach_absolute_time() - seen; // ticks: 41.67 ns each on Apple silicon t.held_ns += held; t.held_max = @max(t.held_max, held); t.held_kind[@backingInt(job.kind)][0] += held; @@ -705,13 +706,13 @@ test "the host service follows the GPU's posts across the 27-bit post wrap and t for (0..8) |_| { // local serving lands each exchange's flag once its post is seen x +%= 1; const seq = t.queue(.{ .kind = .exchange, .value = x }); - const ts: std.c.timespec = .{ .sec = 0, .nsec = 2_000_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 2_000_000 }; + _ = platform.c.nanosleep(&ts, null); try std.testing.expect(@as(u32, @truncate(@atomicLoad(u64, t.word64(FLAG), .acquire))) != x); // not before its post @atomicStore(u32, t.word32(SYNC + GPU * 4), (seq << 5) | 4, .release); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); while (@as(u32, @truncate(@atomicLoad(u64, t.word64(FLAG), .acquire))) != x) { - if (std.c.mach_absolute_time() - t0 > 24_000_000) return error.ServiceStuck; // 1 s + if (platform.c.mach_absolute_time() - t0 > 24_000_000) return error.ServiceStuck; // 1 s std.atomic.spinLoopHint(); } } @@ -729,8 +730,8 @@ test "partNext and hostWait across the 32-bit wrap: the next exchange's slot, an } const Peer = struct { // the peer's message landing a call's flag, a little later fn land(w: *u64, v: u64) void { - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000_000 }; + _ = platform.c.nanosleep(&ts, null); @atomicStore(u64, w, v, .release); } }; diff --git a/zig/src/families/glm/checks.zig b/zig/src/families/glm/checks.zig index 98404ca22a..f1a92ef8c3 100644 --- a/zig/src/families/glm/checks.zig +++ b/zig/src/families/glm/checks.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's dev-time checks on a loaded engine: captures, traces, teacher forcing and knock-out profiles. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("state.zig"); const fwd = @import("forward.zig"); @@ -27,9 +28,9 @@ pub fn capture(e: *Engine, prompt: []const u32, path: []const u8) !void { fwd.head(&x, b.enc, e.sc.hidden, dump.at(x.dump_at), e.sc.picks, n); try e.finish(b.cb, b.enc); fwd.flipKda(&x); - const file = std.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; std.debug.print("captured {d} rows ({d} bytes) in {s}\n", .{ n, bytes, path }); } @@ -47,8 +48,8 @@ pub fn trace(e: *Engine, prompt: []const u32, reply: []const u32, steps: usize, @memcpy(Engine.u32s(e.prompt_ids.at(@as(usize, P) * 4), total - P), reply[0 .. total - P]); const most = @as(usize, st.max_rows) * c.hidden * 2 * (2 + 4 * @as(usize, c.run)) + @as(usize, st.max_rows) * c.vocab * 2; const dump = try e.arena.buffer(most); - const file = std.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(e.gpa, "{s}", .{path}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); var x = e.ctx(); var at: u32 = 0; var calls: usize = 0; @@ -62,7 +63,7 @@ pub fn trace(e: *Engine, prompt: []const u32, reply: []const u32, steps: usize, try e.finish(b.cb, b.enc); fwd.flipKda(&x); const bytes = x.dump_at + @as(usize, n) * c.vocab * 2; - if (std.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; + if (platform.c.fwrite(dump.addr(), 1, bytes, file) != bytes) return error.WriteFailed; at += n; } std.debug.print("traced {d} calls ({d} prompt rows, {d} steps) in {s}\n", .{ calls, P, total - P, path }); @@ -106,7 +107,7 @@ pub fn profile(e: *Engine, depth: u32, reps: usize, only: bool, parts: bool) !vo const c = &e.c; const D = c.hidden; const d: u32 = if (e.w.mtp == null) 0 else @min(depth, st.max_rows - 1); - const wide = if (std.c.getenv("GLM_PROFILE_ROWS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 1 else 1; // a copy window's rows, without the MTP head + const wide = if (platform.c.getenv("GLM_PROFILE_ROWS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 1 else 1; // a copy window's rows, without the MTP head const R = if (d == 0) std.math.clamp(wide, 1, st.max_rows) else d + 1; const pool = mtl.objc.Pool.push(); defer pool.pop(); diff --git a/zig/src/families/glm/engine.zig b/zig/src/families/glm/engine.zig index e1ffe5ac41..e030c7200a 100644 --- a/zig/src/families/glm/engine.zig +++ b/zig/src/families/glm/engine.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash on Metal: one greedy reply at a time, the prompt in 16-row windows or tensor-unit chunks, then drafted rounds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -95,7 +96,7 @@ pub const Engine = struct { /// The checkpoint in `dir`, caches for `cap` tokens; GLM_LAYERS=N: the first N layers only; GLM_EP=settings: half the experts. pub fn load(gpa: std.mem.Allocator, dir: []const u8, cap: u32) !*Engine { - return loadWith(gpa, dir, cap, if (std.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false); + return loadWith(gpa, dir, cap, if (platform.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false); } /// `load` with expert parallel over the link in `ep_path` (this Mac's settings), or on one Mac when null. @@ -104,23 +105,23 @@ pub const Engine = struct { errdefer gpa.destroy(e); const pool = mtl.objc.Pool.push(); defer pool.pop(); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); e.gpa = gpa; e.ev = 0; e.residency = null; e.gpu = .{ 0, 0 }; - e.fused_route = if (std.c.getenv("GLM_ROUTE")) |v| v[0] != '0' else true; - e.draft_vocab = if (std.c.getenv("GLM_DRAFT_VOCAB")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; + e.fused_route = if (platform.c.getenv("GLM_ROUTE")) |v| v[0] != '0' else true; + e.draft_vocab = if (platform.c.getenv("GLM_DRAFT_VOCAB")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; e.committed = 0; e.ep = null; e.pr = null; e.trace_last = null; e.margins = null; e.chunk_rows = prompt_mod.max_rows; - e.cut = if (std.c.getenv("GLM_CUTS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; - e.copy_min = if (std.c.getenv("GLM_COPY")) |v| std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0, 0, 8) else 0; - e.rank_log = if (std.c.getenv("GLM_RANKS")) |v| v[0] == '1' else false; - if (std.c.getenv("GLM_CHUNK")) |v| e.chunk_rows = std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch prompt_mod.max_rows, st.max_rows + 1, prompt_mod.max_rows); + e.cut = if (platform.c.getenv("GLM_CUTS")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; + e.copy_min = if (platform.c.getenv("GLM_COPY")) |v| std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0, 0, 8) else 0; + e.rank_log = if (platform.c.getenv("GLM_RANKS")) |v| v[0] == '1' else false; + if (platform.c.getenv("GLM_CHUNK")) |v| e.chunk_rows = std.math.clamp(std.fmt.parseInt(u32, std.mem.span(v), 10) catch prompt_mod.max_rows, st.max_rows + 1, prompt_mod.max_rows); e.ep_arena = .init(gpa); errdefer e.ep_arena.deinit(); e.device = try mtl.Device.init(); @@ -138,17 +139,17 @@ pub const Engine = struct { e.draft_vocab -= e.draft_vocab % 4; // the head's kernel takes four rows a simdgroup const model = try load_plan.modelHash(gpa, dir, f.bytes[0..f.size]); e.model_hash = model; - if (std.c.getenv("GLM_LAYERS")) |v| try cfg.subset(&e.c, std.fmt.parseInt(u32, std.mem.span(v), 10) catch return error.BadLayerCount); - if (std.c.getenv("GLM_MTP")) |v| if (v[0] == '0') { // no MTP head: no drafts, its ~4 GB left on disk + if (platform.c.getenv("GLM_LAYERS")) |v| try cfg.subset(&e.c, std.fmt.parseInt(u32, std.mem.span(v), 10) catch return error.BadLayerCount); + if (platform.c.getenv("GLM_MTP")) |v| if (v[0] == '0') { // no MTP head: no drafts, its ~4 GB left on disk e.c.mtp = 0; }; const link: ?ep_mod.Settings = if (ep_path) |sp| blk: { const sf = try mtl.MappedFile.open(try e.ep_arena.allocator().dupeSentinel(u8, sp, 0)); defer sf.deinit(); const s = try ep_mod.readSettings(e.ep_arena.allocator(), sf.bytes[0..sf.size]); - const by_rows = if (std.c.getenv("GLM_EP_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "experts") else true; + const by_rows = if (platform.c.getenv("GLM_EP_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "experts") else true; if (by_rows) try cfg.splitRows(&e.c, s.rank, 2) else try cfg.split(&e.c, s.rank, 2); - const tp = if (std.c.getenv("GLM_TP")) |v| v[0] != '0' else true; // TP2: each Mac half the KDA heads (GLM_TP=0: off) + const tp = if (platform.c.getenv("GLM_TP")) |v| v[0] != '0' else true; // TP2: each Mac half the KDA heads (GLM_TP=0: off) if (by_rows and tp) try cfg.splitHeads(&e.c, s.rank, 2); if (e.c.tp > 1) { // TP2's head halves: the whole vocabulary (its kernel's pitch); no host logits to rank e.draft_vocab = e.c.vocab; @@ -162,7 +163,7 @@ pub const Engine = struct { gpa.destroy(e.k); } try frags.check(e.device, e.queue, gpa); // the fragment layout nax.h gives GLM's prompt and attention kernels - const chunked = if (std.c.getenv("GLM_PROMPT")) |v| v[0] != '0' else true; + const chunked = if (platform.c.getenv("GLM_PROMPT")) |v| v[0] != '0' else true; e.limit = try load_plan.loadLimit(e.device); const limit = e.limit.bytes; var plan_bytes = try load_plan.planBytes(gpa, e.device, dir, &e.c); @@ -176,7 +177,7 @@ pub const Engine = struct { plan_bytes = lean_bytes; } } - if (std.c.getenv("GLM_DRY") != null) return error.DryRun; + if (platform.c.getenv("GLM_DRY") != null) return error.DryRun; e.arena = .{ .device = e.device, .gpa = gpa }; errdefer e.arena.deinit(); const both = try st.init(&e.arena, &e.c, cap); @@ -195,9 +196,9 @@ pub const Engine = struct { return error.OverMemoryLimit; } { // free memory short of the load: a cached copy of the checkpoint holds it, the paced reads would wait (#583) - const dropped_at = std.c.mach_absolute_time(); + const dropped_at = platform.c.mach_absolute_time(); if (page_cache.dropForLoad(gpa, dir, plan_bytes + e.arena.bytes)) |d| { - if (d.dropped > 0) std.log.info("glm: dropped {d:.1} GiB of the checkpoint from the file cache, the {d:.1} GiB free memory lacked for the load ({d:.2} s)", .{ @as(f64, @floatFromInt(d.dropped)) / (1 << 30), @as(f64, @floatFromInt(d.deficit)) / (1 << 30), @as(f64, @floatFromInt(std.c.mach_absolute_time() - dropped_at)) / 24e6 }); + if (d.dropped > 0) std.log.info("glm: dropped {d:.1} GiB of the checkpoint from the file cache, the {d:.1} GiB free memory lacked for the load ({d:.2} s)", .{ @as(f64, @floatFromInt(d.dropped)) / (1 << 30), @as(f64, @floatFromInt(d.deficit)) / (1 << 30), @as(f64, @floatFromInt(platform.c.mach_absolute_time() - dropped_at)) / 24e6 }); } else |err| std.log.warn("glm: could not drop the checkpoint's cached pages ({s}); the load reads beside them", .{@errorName(err)}); } e.w = try wts.load(gpa, e.device, dir, &e.c, 16, false); @@ -216,7 +217,7 @@ pub const Engine = struct { e.ep = try ep_mod.Ep.init(gpa, e.device, s, me); } // opt-in: wiring 181 GB leaves macOS nothing to reclaim if another model shares the Mac (Flash Next runs without) - if (std.c.getenv("GLM_RESIDENCY") == null) {} else if (e.device.residencySet(e.w.buffers.items.len + e.arena.buffers.items.len)) |set| { + if (platform.c.getenv("GLM_RESIDENCY") == null) {} else if (e.device.residencySet(e.w.buffers.items.len + e.arena.buffers.items.len)) |set| { for (e.w.buffers.items) |b| set.add(b); for (e.arena.buffers.items) |b| set.add(b); set.commit(); @@ -228,13 +229,13 @@ pub const Engine = struct { if (e.residency) |set| e.keepalive_sets[0] = set; e.keepalive_target = .{ .queue = e.queue, .sets = if (e.residency != null) &e.keepalive_sets else &.{} }; // a process's first prompt pays a one-time cost: every rank of a pair pays it here at the same step (GLM_NO_WARMUP: skip) - if (e.pr != null and cap >= 128 and std.c.getenv("GLM_NO_WARMUP") == null) { + if (e.pr != null and cap >= 128 and platform.c.getenv("GLM_NO_WARMUP") == null) { var ids: [64]u32 = undefined; for (&ids, 0..) |*t, i| t.* = @intCast(1000 + i); var quiet: Quiet = .{}; _ = try e.generate(&ids, 2, &.{}, 0, .{ .ctx = &quiet, .prefilled = Quiet.prefilled, .tokens = Quiet.tokens, .cancelled = Quiet.cancelled }); } - e.load_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) / 24e6; + e.load_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) / 24e6; return e; } @@ -324,7 +325,7 @@ pub const Engine = struct { e.ev += 1; cb.signal(e.event, e.ev); cb.commit(); - e.committed = std.c.mach_absolute_time(); + e.committed = platform.c.mach_absolute_time(); if (spin) { while (@backingInt(cb.status()) < @backingInt(mtl.types.CommandBufferStatus.completed)) std.atomic.spinLoopHint(); } else cb.wait(); @@ -408,7 +409,7 @@ pub const Engine = struct { e.s.reset(); @memcpy(u32s(e.prompt_ids, prompt.len), prompt); var x = e.ctx(); - const t_start = std.c.mach_absolute_time(); + const t_start = platform.c.mach_absolute_time(); // the prompt: windows of up to 16 rows, the MTP head taking each row with the token after it var at: u32 = 0; var last_n: u32 = 1; @@ -456,7 +457,7 @@ pub const Engine = struct { at += n; } e.s.pos = P; - const t_prompt = std.c.mach_absolute_time(); + const t_prompt = platform.c.mach_absolute_time(); out.prefilled(out.ctx); var res: Result = .{ .reason = .length, .prompt_seconds = @as(f64, @floatFromInt(t_prompt - t_start)) / 24e6 }; var tok = u32s(e.sc.picks, 1)[0]; @@ -469,7 +470,7 @@ pub const Engine = struct { if (max_tokens <= 1) return res; if (d == 0 and e.ep == null and e.copy_min == 0 and e.margins == null) { try plain_rounds.run(e, &x, tok, max_tokens, eos, out, &res); - res.decode_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t_prompt)) / 24e6; + res.decode_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t_prompt)) / 24e6; return res; } var emitted: usize = 1; @@ -490,7 +491,7 @@ pub const Engine = struct { u32s(e.sc.next, 1)[0] = tok; var last_end: f64 = 0; while (true) { - const t_enc = std.c.mach_absolute_time(); + const t_enc = platform.c.mach_absolute_time(); const b = e.begin(); if (window > 0) { // the last window's kept rows (the prompt's windows kept all theirs) fwd.keepKda(&x, b.enc, window, keep); @@ -584,7 +585,7 @@ pub const Engine = struct { break; } } - res.decode_seconds = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t_prompt)) / 24e6; + res.decode_seconds = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t_prompt)) / 24e6; return res; } }; diff --git a/zig/src/families/glm/ep.zig b/zig/src/families/glm/ep.zig index 74ea506c26..1fd8a94b9f 100644 --- a/zig/src/families/glm/ep.zig +++ b/zig/src/families/glm/ep.zig @@ -1,5 +1,6 @@ //! Expert parallel over two Macs (MCDMA): half the routed experts each, a MoE call's (row, slot) outputs swapped, all combined on both. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fabric = @import("fabric"); const settings = @import("../flashnext/tp_settings.zig"); @@ -120,8 +121,8 @@ pub const Ep = struct { t.ctl = try control.Control.init(gpa, rd, CONTROL, &t.failed); errdefer t.ctl.deinit(gpa); _ = try t.ctl.hello(me); // both Macs up, running the same thing, before the first exchange - t.trace = std.c.getenv("GLM_EP_TRACE") != null; - if (std.c.getenv("GLM_EP_DELAY_US")) |v| t.delay_ns = 1000 * (std.fmt.parseInt(u64, std.mem.span(v), 10) catch 0); + t.trace = platform.c.getenv("GLM_EP_TRACE") != null; + if (platform.c.getenv("GLM_EP_DELAY_US")) |v| t.delay_ns = 1000 * (std.fmt.parseInt(u64, std.mem.span(v), 10) catch 0); t.thread = try std.Thread.spawn(.{}, service, .{t}); std.log.info("expert parallel: rank {d} of 2 holds experts {d}-{d}, connected", .{ t.rank, me.own_lo, me.own_hi - 1 }); return t; @@ -307,7 +308,7 @@ pub const Ep = struct { /// Each posted exchange's entries to the peer, then its flag; 10 s without an answer fails the link and lands every flag. fn service(t: *Ep) void { - _ = std.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send + _ = platform.c.pthread_set_qos_class_self_np(.USER_INTERACTIVE, 0); // busy CPUs must not hold back a send const posted = t.word32(POSTED); const flag = t.word64(FLAG); var x: u64 = 1; @@ -320,26 +321,26 @@ pub const Ep = struct { t.st.watch(@atomicLoad(u64, flag, .acquire), x - 1); if (t.stop.load(.acquire)) return; if (waiting and @atomicLoad(u64, flag, .acquire) >= want) waiting = false; - if (waiting and !t.failed.load(.acquire) and std.c.mach_absolute_time() - want_at > 240_000_000) { + if (waiting and !t.failed.load(.acquire) and platform.c.mach_absolute_time() - want_at > 240_000_000) { t.fail(want, error.PeerSilent); t.land(want); waiting = false; } spins += 1; if (spins > 1_000_000) { // idle: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } - const seen = std.c.mach_absolute_time(); + const seen = platform.c.mach_absolute_time(); const n: usize = @atomicLoad(u32, t.word32(COUNT + 4 * @as(usize, @intCast(x % 2))), .acquire); const theirs: u32 = @bitCast(@as([*]const i32, @ptrCast(@alignCast(t.lists.contents() + COUNTS)))[1]); t.st.posted(x, seen, @intCast(n), theirs); t.st.watch(@atomicLoad(u64, flag, .acquire), x); if (t.trace) std.debug.print("EP rank{d} exchange {d}: {d} entries\n", .{ t.rank, x, n }); if (t.delay_ns > 0 and n * ENTRY > 1 << 20) { - const ts: std.c.timespec = .{ .sec = @intCast(t.delay_ns / std.time.ns_per_s), .nsec = @intCast(t.delay_ns % std.time.ns_per_s) }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = @intCast(t.delay_ns / std.time.ns_per_s), .nsec = @intCast(t.delay_ns % std.time.ns_per_s) }; + _ = platform.c.nanosleep(&ts, null); } if (t.failed.load(.acquire) or n * ENTRY > SLOT) { if (n * ENTRY > SLOT) t.fail(x, error.BadCount); @@ -349,7 +350,7 @@ pub const Ep = struct { t.land(x); }; want = x; - want_at = std.c.mach_absolute_time(); + want_at = platform.c.mach_absolute_time(); waiting = true; t.sent += 1; t.held_ticks += want_at - seen; @@ -394,7 +395,7 @@ pub const Ep = struct { /// The flag's new value seen now; every exchange up to `upto` whose post and landing are both known is counted. fn watch(s: *Stats, flag: u64, upto: u64) void { if (flag > s.flag_seen) { - const now = std.c.mach_absolute_time(); + const now = platform.c.mach_absolute_time(); var v = s.flag_seen + 1; while (v <= flag) : (v += 1) s.flag_t[v % 64] = now; s.flag_seen = flag; diff --git a/zig/src/families/glm/ep_control.zig b/zig/src/families/glm/ep_control.zig index 420d7411d4..4e69e3e39b 100644 --- a/zig/src/families/glm/ep_control.zig +++ b/zig/src/families/glm/ep_control.zig @@ -1,5 +1,6 @@ //! The expert-parallel pair's control protocol over a link's window: identities, requests and stop decisions, each acked. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const Rdma = fabric.rdma.Rdma; @@ -54,7 +55,7 @@ pub const Request = struct { max_tokens: usize, depth: usize, eos: []const u32, pub const Command = union(enum) { request: Request, lanes: struct { kind: u32, words: []const u32 } }; fn ticks() u64 { - return std.c.mach_absolute_time(); // 24 MHz on Apple silicon + return platform.c.mach_absolute_time(); // 24 MHz on Apple silicon } pub const Control = struct { @@ -96,8 +97,8 @@ pub const Control = struct { if (bounded and ticks() - t0 > c.limit) return error.EpPeerSilent; spins += 1; if (spins > 100_000) { // a long wait: back off - const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000 }; - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 50_000 }; + _ = platform.c.nanosleep(&ts, null); } else std.atomic.spinLoopHint(); } } @@ -295,8 +296,8 @@ test "early cancels: rank 0 waits for rank 1's acks, nothing is overwritten or s }; var e0: ?anyerror = undefined; const th = try std.Thread.spawn(.{}, Rank0.run, .{ &p.c[0], &e0 }); - const ts: std.c.timespec = .{ .sec = 0, .nsec = 200_000_000 }; // rank 1 held back while rank 0 runs ahead - _ = std.c.nanosleep(&ts, null); + const ts: platform.c.timespec = .{ .sec = 0, .nsec = 200_000_000 }; // rank 1 held back while rank 0 runs ahead + _ = platform.c.nanosleep(&ts, null); for (prompts, 0..) |pr, i| { const r = (try p.c[1].waitRequest()).?; try std.testing.expectEqualSlices(u32, &pr, r.prompt); diff --git a/zig/src/families/glm/learned_pair.zig b/zig/src/families/glm/learned_pair.zig index 4d15eb992f..b1c8c57321 100644 --- a/zig/src/families/glm/learned_pair.zig +++ b/zig/src/families/glm/learned_pair.zig @@ -1,5 +1,6 @@ //! A peer's disk reservation covers its snapshot half until rank 0 has committed or removed the shared index entry. const std = @import("std"); +const platform = @import("sys"); const disk = @import("lanes").learned_disk; pub const State = struct { admission: disk.Admission = .{}, @@ -98,12 +99,12 @@ test "each rank includes actual other-identity files in its independent cap befo var filename: [1400]u8 = undefined; const current = try std.fmt.bufPrintSentinel(&one, "{s}/0000000000000001", .{root}, 0); const other = try std.fmt.bufPrintSentinel(&two, "{s}/0000000000000002", .{root}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(current, 0o700)); - try std.testing.expectEqual(@as(c_int, 0), std.c.mkdir(other, 0o700)); - const fd = std.c.open(try std.fmt.bufPrintSentinel(&filename, "{s}/state.bin", .{other}, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(std.c.mode_t, 0o600)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(current, 0o700)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.mkdir(other, 0o700)); + const fd = platform.c.open(try std.fmt.bufPrintSentinel(&filename, "{s}/state.bin", .{other}, 0), .{ .ACCMODE = .WRONLY, .CREAT = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.TestCreate; - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 96)); - _ = std.c.close(fd); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 96)); + _ = platform.c.close(fd); var rank0: State = .{ .admission = .{ .floor = 0, .free = Fake.free, .clock = Fake.clock }, .cap = 200 }; var rank1: State = .{ .admission = .{ .floor = 0, .free = Fake.free, .clock = Fake.clock }, .cap = 100 }; try std.testing.expectEqual(@as(u64, 0), rank0.need(current, 8, 2)); diff --git a/zig/src/families/glm/load_plan.zig b/zig/src/families/glm/load_plan.zig index 53a0dddbcd..94d2b98b68 100644 --- a/zig/src/families/glm/load_plan.zig +++ b/zig/src/families/glm/load_plan.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's load plan: the checkpoint's identity, the bytes its weights and caches take, the chunk height, the limit. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -137,8 +138,8 @@ pub fn limitFor(ram: u64, working_set: u64, text: ?[]const u8) error{BadLoadLimi pub fn loadLimit(device: mtl.Device) error{BadLoadLimit}!Limit { var mem: u64 = 0; var len: usize = @sizeOf(u64); - if (std.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; // RAM unknown: a zero limit refuses every load - const text: ?[]const u8 = if (std.c.getenv("GLM_LOAD_LIMIT_GB")) |v| std.mem.span(v) else null; + if (platform.c.sysctlbyname("hw.memsize", &mem, &len, null, 0) != 0) mem = 0; // RAM unknown: a zero limit refuses every load + const text: ?[]const u8 = if (platform.c.getenv("GLM_LOAD_LIMIT_GB")) |v| std.mem.span(v) else null; const l = limitFor(mem, device.maxWorkingSet(), text) catch |err| { std.log.err("glm: GLM_LOAD_LIMIT_GB={s} is not a load limit: give a number of GB above 0, or unset it for 70% of RAM", .{text.?}); return err; diff --git a/zig/src/families/glm/mirror.zig b/zig/src/families/glm/mirror.zig index 6a3a7dbb70..4c9d862f9f 100644 --- a/zig/src/families/glm/mirror.zig +++ b/zig/src/families/glm/mirror.zig @@ -1,5 +1,6 @@ //! Speed-up mode's slot commands: rank 0 sends each slot operation before it runs it; rank 1 runs the same, in order. const std = @import("std"); +const platform = @import("sys"); const eng = @import("engine.zig"); const slots_mod = @import("slots.zig"); const st = @import("state.zig"); @@ -192,10 +193,10 @@ fn learned(sl: *Slots, k: Kind, w: []const u32) !void { } pub fn fileBytes(file: [:0]const u8) u64 { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return 0; - defer _ = std.c.close(fd); - const n = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const n = platform.c.lseek(fd, 0, platform.c.SEEK.END); return if (n > 0) @intCast(n) else 0; } diff --git a/zig/src/families/glm/names.zig b/zig/src/families/glm/names.zig index abea94eabb..3058b447e8 100644 --- a/zig/src/families/glm/names.zig +++ b/zig/src/families/glm/names.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's checkpoint index in either layout (the original or mlx-lm's), each tensor by its original name. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("../../core/safetensors.zig"); const load_fix = @import("weights_fix.zig"); @@ -45,15 +46,15 @@ pub fn read(gpa: std.mem.Allocator, arena: std.mem.Allocator, dir: []const u8, m for (files.keys(), 0..) |name, si| { const full = try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ dir, name }, 0); try shards.append(gpa, full); - const fd = std.c.open(full, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(full, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: [8]u8 = undefined; try load_fix.readAll(fd, &head, 0); const hlen = std.mem.readInt(u64, &head, .little); const header = try arena.alloc(u8, hlen); try load_fix.readAll(fd, header, 8); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); const data = 8 + hlen; const doc2 = try std.json.parseFromSliceLeaky(std.json.Value, arena, header, .{}); var eit = doc2.object.iterator(); diff --git a/zig/src/families/glm/page_cache.zig b/zig/src/families/glm/page_cache.zig index ab64f39d75..cdd9fb2ab9 100644 --- a/zig/src/families/glm/page_cache.zig +++ b/zig/src/families/glm/page_cache.zig @@ -1,5 +1,6 @@ //! A checkpoint's pages in the file cache, dropped before a load: F_NOCACHE reads keep out of it but evict nothing. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const paced_read = @import("../../core/paced_read.zig"); @@ -33,15 +34,15 @@ pub fn dropCached(gpa: std.mem.Allocator, dir: []const u8, want: u64, d: *Droppe const page: usize = std.heap.pageSize(); for (files.keys()) |name| { if (d.dropped >= want) break; - const fd = std.c.open(try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0), .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end <= 0) continue; const len: usize = @intCast(end); - const map = std.c.mmap(null, len, .{ .READ = true }, .{ .TYPE = .SHARED }, fd, 0); - if (map == std.c.MAP_FAILED) return error.MapFailed; - defer _ = std.c.munmap(@alignCast(map), len); + const map = platform.c.mmap(null, len, .{ .READ = true }, .{ .TYPE = .SHARED }, fd, 0); + if (map == platform.c.MAP_FAILED) return error.MapFailed; + defer _ = platform.c.munmap(@alignCast(map), len); const vec = try a.alloc(u8, (len + page - 1) / page); var cached: u64 = 0; if (mincore(map, len, vec.ptr) == 0) for (vec) |v| { @@ -57,31 +58,31 @@ pub fn dropCached(gpa: std.mem.Allocator, dir: []const u8, want: u64, d: *Droppe test "a checkpoint's cached pages are dropped only as far as wanted, and a shard with none is left alone" { const gpa = std.testing.allocator; var name_buf: [64]u8 = undefined; - const dir = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-drop-cached-{d}", .{std.c.getpid()}, 0); - _ = std.c.mkdir(dir, 0o700); - defer _ = std.c.rmdir(dir); + const dir = try std.fmt.bufPrintSentinel(&name_buf, "/tmp/tf-drop-cached-{d}", .{platform.c.getpid()}, 0); + _ = platform.c.mkdir(dir, 0o700); + defer _ = platform.c.rmdir(dir); var path_buf: [128]u8 = undefined; const index = try std.fmt.bufPrintSentinel(&path_buf, "{s}/model.safetensors.index.json", .{dir}, 0); const text = "{\"weight_map\": {\"a\": \"s.safetensors\", \"b\": \"s.safetensors\"}}"; var shard_buf: [128]u8 = undefined; const shard = try std.fmt.bufPrintSentinel(&shard_buf, "{s}/s.safetensors", .{dir}, 0); for ([_][:0]const u8{ index, shard }, [_][]const u8{ text, "" }) |file, body| { - const fd = std.c.open(file, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); try std.testing.expect(fd >= 0); - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); if (body.len > 0) { - _ = std.c.write(fd, body.ptr, body.len); + _ = platform.c.write(fd, body.ptr, body.len); continue; } const data = try gpa.alloc(u8, 4 << 20); defer gpa.free(data); @memset(data, 7); - try std.testing.expectEqual(@as(isize, @intCast(data.len)), std.c.write(fd, data.ptr, data.len)); - _ = std.c.fsync(fd); - _ = std.c.pread(fd, data.ptr, data.len, 0); // read through the file cache + try std.testing.expectEqual(@as(isize, @intCast(data.len)), platform.c.write(fd, data.ptr, data.len)); + _ = platform.c.fsync(fd); + _ = platform.c.pread(fd, data.ptr, data.len, 0); // read through the file cache } - defer _ = std.c.unlink(index); - defer _ = std.c.unlink(shard); + defer _ = platform.c.unlink(index); + defer _ = platform.c.unlink(shard); var none: Dropped = .{ .deficit = 0 }; try dropCached(gpa, dir, 0, &none); // nothing wanted: no shard touched, the pages stay cached try std.testing.expectEqual(@as(u64, 0), none.dropped); diff --git a/zig/src/families/glm/plain_rounds.zig b/zig/src/families/glm/plain_rounds.zig index c5f2114916..3289fdc513 100644 --- a/zig/src/families/glm/plain_rounds.zig +++ b/zig/src/families/glm/plain_rounds.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's plain greedy rounds in generate: the next round committed before this one's pick is read. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const fwd = @import("forward.zig"); const Fence = @import("../../core/fence.zig").Fence; @@ -12,7 +13,7 @@ const Pending = struct { cb: mtl.CommandBuffer, t_enc: u64, committed: u64 }; /// One plain round at `pos`, its token from `from` (null: the ids as set), its pick into `to`; after a round it waits on the fence, not the event. fn plainRound(e: *Engine, x: *fwd.Ctx, from: ?Ref, to: Ref, pos: u32) Pending { - const t_enc = std.c.mach_absolute_time(); + const t_enc = platform.c.mach_absolute_time(); const fence = e.round_fence.?; const cb = e.queue.commandBuffer(); if (from == null and e.ev > 0) cb.waitFor(e.event, e.ev); @@ -31,7 +32,7 @@ fn plainRound(e: *Engine, x: *fwd.Ctx, from: ?Ref, to: Ref, pos: u32) Pending { e.ev += 1; b.cb.signal(e.event, e.ev); b.cb.commit(); - e.committed = std.c.mach_absolute_time(); + e.committed = platform.c.mach_absolute_time(); return .{ .cb = b.cb, .t_enc = t_enc, .committed = e.committed }; } diff --git a/zig/src/families/glm/pool.zig b/zig/src/families/glm/pool.zig index fb14809842..56f36077f0 100644 --- a/zig/src/families/glm/pool.zig +++ b/zig/src/families/glm/pool.zig @@ -1,5 +1,6 @@ //! One Mac's streams share the engine's MLA cache, a contiguous region each; kept prefixes hold blocks until needed. const std = @import("std"); +const platform = @import("sys"); const st = @import("state.zig"); const snapshot = @import("snapshot.zig"); const Engine = @import("engine.zig").Engine; @@ -41,7 +42,7 @@ pub const Pool = struct { /// Whether the slots share one pool: on one Mac (a pair's ranks mirror slot commands), unless GLM_POOL=0. pub fn on(e: *const Engine) bool { if (e.ep != null or e.followsPeer()) return false; - const v = std.c.getenv("GLM_POOL") orelse return true; + const v = platform.c.getenv("GLM_POOL") orelse return true; return !std.mem.eql(u8, std.mem.span(v), "0"); } diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index 4f73d916b0..4be9a124cf 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -1,5 +1,6 @@ //! Prompt chunks: matmuls on the M5 tensor units at the row kernels' arithmetic (core/affine_mm.zig), the rest on the row kernels. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -83,7 +84,7 @@ pub fn init(gpa: std.mem.Allocator, arena: *st.Arena, device: mtl.Device, c: *co /// The prompt's sparse MLA on the tensor units (GLM_SPARSE_NAX=0: the row kernel, its output beside its queries). fn sparseNax() bool { - return if (std.c.getenv("GLM_SPARSE_NAX")) |v| v[0] != '0' else true; + return if (platform.c.getenv("GLM_SPARSE_NAX")) |v| v[0] != '0' else true; } /// The bytes `init` takes from the arena for `rows`-row chunks. diff --git a/zig/src/families/glm/slots.zig b/zig/src/families/glm/slots.zig index 90aa566962..fc97750612 100644 --- a/zig/src/families/glm/slots.zig +++ b/zig/src/families/glm/slots.zig @@ -1,5 +1,6 @@ //! Stream slots on one Mac: a stream's caches each; prompt chunks, shared windows, keeps and drafts run on them. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("state.zig"); const fwd = @import("forward.zig"); @@ -76,7 +77,7 @@ pub const Slots = struct { .last = try e.arena.buffer(@as(usize, e.c.hidden) * 2), }; const plane = @as(usize, st.max_rows) * e.c.hidden * 2; - var sl: Slots = .{ .gpa = gpa, .e = e, .slots = slots, .log = std.c.getenv("GLM_WINDOWS") != null, .rows = try e.arena.buffer(plane), .ids = try e.arena.buffer(st.max_rows * 4), .lasts = try e.arena.buffer(plane), .picks = try e.arena.buffer(st.max_rows * 4), .ring = try e.arena.buffer(ring_len * 4) }; + var sl: Slots = .{ .gpa = gpa, .e = e, .slots = slots, .log = platform.c.getenv("GLM_WINDOWS") != null, .rows = try e.arena.buffer(plane), .ids = try e.arena.buffer(st.max_rows * 4), .lasts = try e.arena.buffer(plane), .picks = try e.arena.buffer(st.max_rows * 4), .ring = try e.arena.buffer(ring_len * 4) }; if (pool_mod.on(e)) pool_mod.setUp(&sl); return sl; } @@ -247,7 +248,7 @@ pub const Slots = struct { const e = sl.e; var segs: [st.max_rows]fwd.Seg = undefined; if (wins.len == 0 or wins.len > segs.len) return error.WindowOutOfStep; - const t_in = std.c.mach_absolute_time(); + const t_in = platform.c.mach_absolute_time(); sl.settleAll(); try sl.launch(); // the GPU runs the keeps and drafts while the window encodes const ids = Engine.u32s(e.sc.ids, st.max_rows); @@ -274,7 +275,7 @@ pub const Slots = struct { fwd.head(&x, enc, e.sc.hidden, e.sc.logits, e.sc.picks, total); sl.ahead = .{ 0, 0 }; try sl.flushBy(true); - const ret = std.c.mach_absolute_time(); + const ret = platform.c.mach_absolute_time(); if (sl.log) { const first = if (sl.ahead[1] > 0) sl.ahead[0] else e.gpu[0]; std.log.info("glm window: streams {d} rows {d} gpu {d:.2} ms gap {d:.2} ms ahead {d:.2} join {d:.3} host {d:.2} encode {d:.2} ms", .{ wins.len, total, (e.gpu[1] - e.gpu[0]) * 1e3, (first - sl.last_end) * 1e3, (sl.ahead[1] - sl.ahead[0]) * 1e3, if (sl.ahead[1] > 0) (e.gpu[0] - sl.ahead[1]) * 1e3 else 0, @as(f64, @floatFromInt(t_in -| sl.last_ret)) / 24e3, @as(f64, @floatFromInt(e.committed - t_in)) / 24e3 }); diff --git a/zig/src/families/glm/snapshot_file.zig b/zig/src/families/glm/snapshot_file.zig index 901b5d90b6..b8b058d7d3 100644 --- a/zig/src/families/glm/snapshot_file.zig +++ b/zig/src/families/glm/snapshot_file.zig @@ -1,45 +1,46 @@ //! GLM learned halves use a versioned checksum envelope; legacy, corrupt and partial files never enter a cache. const std = @import("std"); +const platform = @import("sys"); const checked = @import("../../core/snapshot_file.zig"); const magic: u32 = 0x474c4d54; pub const Head = extern struct { magic: u32 = magic, version: u32 = 1, at: u32, reserved: u32 = 0, bytes: u64 }; pub fn write(file: [:0]const u8, at: u32, payload: []const u8) !void { var tmp_buf: [1200]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head: Head = .{ .at = at, .bytes = payload.len }; try checked.write(fd, std.mem.asBytes(&head), payload); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } pub fn read(file: [:0]const u8, at: u32, payload: []u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const expected: Head = .{ .at = at, .bytes = payload.len }; var head: Head = undefined; - if (std.c.pread(fd, std.mem.asBytes(&head).ptr, @sizeOf(Head), 0) != @sizeOf(Head) or !std.mem.eql(u8, std.mem.asBytes(&head), std.mem.asBytes(&expected))) return error.SnapshotRead; + if (platform.c.pread(fd, std.mem.asBytes(&head).ptr, @sizeOf(Head), 0) != @sizeOf(Head) or !std.mem.eql(u8, std.mem.asBytes(&head), std.mem.asBytes(&expected))) return error.SnapshotRead; try checked.read(fd, std.mem.asBytes(&head), payload); } test "both GLM half files refuse a changed body, legacy header, truncation and another position" { var path: [160]u8 = undefined; - const file = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-glm-file-{d}", .{std.c.getpid()}, 0); - defer _ = std.c.unlink(file); + const file = try std.fmt.bufPrintSentinel(&path, "/tmp/tf-glm-file-{d}", .{platform.c.getpid()}, 0); + defer _ = platform.c.unlink(file); try write(file, 128, "body"); var out: [4]u8 = undefined; try read(file, 128, &out); try std.testing.expectEqualStrings("body", &out); try std.testing.expectError(error.SnapshotRead, read(file, 129, &out)); - const fd = std.c.open(file, .{ .ACCMODE = .RDWR }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDWR }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - try std.testing.expectEqual(@as(isize, 1), std.c.pwrite(fd, "X", 1, @sizeOf(Head) + 1)); + defer _ = platform.c.close(fd); + try std.testing.expectEqual(@as(isize, 1), platform.c.pwrite(fd, "X", 1, @sizeOf(Head) + 1)); try std.testing.expectError(error.SnapshotRead, read(file, 128, &out)); - try std.testing.expectEqual(@as(c_int, 0), std.c.ftruncate(fd, 16)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.ftruncate(fd, 16)); try std.testing.expectError(error.SnapshotRead, read(file, 128, &out)); } diff --git a/zig/src/families/glm/weights.zig b/zig/src/families/glm/weights.zig index 5beb758f14..d85bebec5b 100644 --- a/zig/src/families/glm/weights.zig +++ b/zig/src/families/glm/weights.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's weights in the Python family's layout (stacked projections and experts, repacked mixes), read by parallel preads. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const st = @import("../../core/safetensors.zig"); const cfg = @import("config.zig"); diff --git a/zig/src/families/glm/weights_fix.zig b/zig/src/families/glm/weights_fix.zig index ec0861fe6e..3d12f8b75b 100644 --- a/zig/src/families/glm/weights_fix.zig +++ b/zig/src/families/glm/weights_fix.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash weights as they load: parallel reads into place, and rewrites (mixes, router, conv, norms, columns). const std = @import("std"); +const platform = @import("sys"); const st = @import("../../core/safetensors.zig"); const paced_read = @import("../../core/paced_read.zig"); @@ -12,16 +13,16 @@ pub const Transform = struct { src: Src, dst: [*]u8, fix: Fix, extra: [3]?Src = /// One tensor's bytes read straight into their place in a weights buffer. pub const Copy = struct { shard: u32, off: u64, len: u64, dst: [*]u8 }; -pub fn readAll(fd: std.c.fd_t, dest: []u8, at: u64) !void { +pub fn readAll(fd: platform.c.fd_t, dest: []u8, at: u64) !void { var done: usize = 0; while (done < dest.len) { - const n = std.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + const n = platform.c.pread(fd, dest.ptr + done, dest.len - done, @intCast(at + done)); if (n <= 0) return error.ShortRead; done += @intCast(n); } } -pub fn transform(gpa: std.mem.Allocator, fds: []std.c.fd_t, t: Transform) !void { +pub fn transform(gpa: std.mem.Allocator, fds: []platform.c.fd_t, t: Transform) !void { const raw = try gpa.alloc(u8, t.src.len); defer gpa.free(raw); try readAll(fds[t.src.shard], raw, t.src.off); @@ -110,7 +111,7 @@ test "fp32 values that are bf16 narrow in place; any other value is refused" { } const Pool = struct { - fds: []std.c.fd_t, + fds: []platform.c.fd_t, jobs: []const Copy, next: std.atomic.Value(usize) = .init(0), failed: std.atomic.Value(bool) = .init(false), @@ -132,15 +133,15 @@ const Pool = struct { /// Every copy on `threads` threads, uncached, paced by free memory (a large load otherwise fills the page cache too). pub fn runCopies(gpa: std.mem.Allocator, shards: []const [:0]const u8, copies: []const Copy, transforms: []const Transform, threads: usize) !void { - const fds = try gpa.alloc(std.c.fd_t, shards.len); + const fds = try gpa.alloc(platform.c.fd_t, shards.len); defer gpa.free(fds); for (shards, fds) |path, *fd| { - fd.* = std.c.open(path, .{ .ACCMODE = .RDONLY }); + fd.* = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd.* < 0) return error.OpenFailed; - _ = std.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE + _ = platform.c.fcntl(fd.*, 48, @as(c_int, 1)); // F_NOCACHE } defer for (fds) |fd| { - _ = std.c.close(fd); + _ = platform.c.close(fd); }; var pool: Pool = .{ .fds = fds, .jobs = copies }; const workers = try gpa.alloc(?std.Thread, threads); @@ -153,7 +154,7 @@ pub fn runCopies(gpa: std.mem.Allocator, shards: []const [:0]const u8, copies: [ const Tx = struct { gpa: std.mem.Allocator, jobs: []const Transform, - fds: []std.c.fd_t, + fds: []platform.c.fd_t, next: std.atomic.Value(usize) = .init(0), failed: std.atomic.Value(bool) = .init(false), fn run(x: *@This()) void { diff --git a/zig/src/families/glm/weights_test.zig b/zig/src/families/glm/weights_test.zig index 87e0c7fd01..92dc8ea962 100644 --- a/zig/src/families/glm/weights_test.zig +++ b/zig/src/families/glm/weights_test.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's loader on synthetic checkpoints: one sparse shard whose header names every tensor at its width. const std = @import("std"); +const platform = @import("sys"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); const quant = @import("quant.zig"); @@ -156,10 +157,10 @@ const Fixture = struct { std.mem.writeInt(u64, &len, head.items.len, .little); try tmp.dir.writeFile(io, .{ .sub_path = "model-00001.safetensors", .data = try std.mem.concat(f.a, u8, &.{ &len, head.items }) }); const path = try std.fmt.allocPrintSentinel(f.a, "{s}/model-00001.safetensors", .{dir}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.TestFile; - defer _ = std.c.close(fd); - if (std.c.ftruncate(fd, @intCast(8 + head.items.len + at)) != 0) return error.TestFile; + defer _ = platform.c.close(fd); + if (platform.c.ftruncate(fd, @intCast(8 + head.items.len + at)) != 0) return error.TestFile; } }; diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index f52e5a4a6c..ee3c374e4c 100644 --- a/zig/src/families/nemotron/backend.zig +++ b/zig/src/families/nemotron/backend.zig @@ -1,5 +1,6 @@ //! The lane core's Metal backend for Nemotron: streams' rows in one forward a round, drawn tokens in GPU slots. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const Fence = @import("../../core/fence.zig").Fence; @@ -89,7 +90,7 @@ pub const Metal = struct { .tokens = try m.device.buffer(ring * 4, opts), .prompt = try m.device.buffer(@max(o.capacity, 16) * 4, opts), .fence = try Fence.init(m.device), - .trace = std.c.getenv("TF_SLIDE_TRACE") != null, + .trace = platform.c.getenv("TF_SLIDE_TRACE") != null, }; if (o.chunk > fused_rows) b.wide = try prefill.Prefill.init(m, @min(o.chunk, o.capacity)); if (o.drafts and m.weights.mtp != null) b.head = try mtp.Head.init(m.device, m.config, &m.kernels, &m.weights, &b.pool, max_lanes, o.capacity); diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index f287c4ced4..7483f02db6 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -1,6 +1,7 @@ //! The Nemotron CUDA engine (engine.py): weights, kernels and one sequence's buffers; windows and commits. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const Config = @import("config.zig").Config; const kern = @import("cuda_kernels.zig"); @@ -130,7 +131,7 @@ pub const Engine = struct { const e = try gpa.create(Engine); errdefer gpa.destroy(e); e.* = .{ .gpa = gpa, .io = io, .ctx = ctx, .stream = undefined, .k = undefined, .w = undefined, .b = undefined, .c = undefined, .max_len = 0, .pinned = undefined, .history = undefined }; - e.trace = std.c.getenv("TF_SLIDE_TRACE") != null; + e.trace = platform.c.getenv("TF_SLIDE_TRACE") != null; e.c = try Config.read(gpa, io, model_dir); if (opts.segments < 1 or opts.segments > segs.MAX) return error.BadSegments; e.segments = opts.segments; diff --git a/zig/src/families/nemotron/cuda_native.zig b/zig/src/families/nemotron/cuda_native.zig index 64594342c4..f09ad552e6 100644 --- a/zig/src/families/nemotron/cuda_native.zig +++ b/zig/src/families/nemotron/cuda_native.zig @@ -1,5 +1,6 @@ //! Nemotron's CUDA engine, MTP head and lane backend behind the native family interface. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); @@ -134,7 +135,7 @@ fn release(p: *anyopaque) void { /// TF_CUDA_CARVEOUT=1: the display carveout KV planes go to first (cuda/carveout.zig); null when it is off. fn openCarveout(gpa: std.mem.Allocator, ctx: *const cuda.Context) !?*cuda.Carveout { const bytes = (try cuda.carveout.requested(false, env("TF_CUDA_CARVEOUT"), env("TF_CUDA_CARVEOUT_MIB"))) orelse return null; - const card: [*:0]const u8 = std.c.getenv("TF_DRM_CARD") orelse cuda.carveout.default_card; + const card: [*:0]const u8 = platform.c.getenv("TF_DRM_CARD") orelse cuda.carveout.default_card; const c = try gpa.create(cuda.Carveout); errdefer gpa.destroy(c); c.* = cuda.Carveout.open(ctx.d, card, bytes) catch |err| { @@ -151,5 +152,5 @@ fn closeCarveout(gpa: std.mem.Allocator, c: *cuda.Carveout) void { } fn env(name: [:0]const u8) ?[]const u8 { - return std.mem.span(std.c.getenv(name) orelse return null); + return std.mem.span(platform.c.getenv(name) orelse return null); } diff --git a/zig/src/families/nemotron/cuda_seq_graphs.zig b/zig/src/families/nemotron/cuda_seq_graphs.zig index 9e51e1556b..09fb6965a4 100644 --- a/zig/src/families/nemotron/cuda_seq_graphs.zig +++ b/zig/src/families/nemotron/cuda_seq_graphs.zig @@ -1,6 +1,7 @@ //! Graphs a sequence other than the engine's own captures on its own buffers, once each, and replays after. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); /// One sequence's captured graphs, by the step they run (a window's or a head level's id). @@ -27,6 +28,6 @@ pub const Stats = struct { /// Whether TF_GRAPH_LOG=1 asks for a line a request: its captures, their time, its replays and eager windows. pub fn logFromEnv() bool { - const v = std.c.getenv("TF_GRAPH_LOG") orelse return false; + const v = platform.c.getenv("TF_GRAPH_LOG") orelse return false; return std.mem.eql(u8, std.mem.span(v), "1"); } diff --git a/zig/src/families/nemotron/forward.zig b/zig/src/families/nemotron/forward.zig index 7dc93e2de2..10751fbbf7 100644 --- a/zig/src/families/nemotron/forward.zig +++ b/zig/src/families/nemotron/forward.zig @@ -1,5 +1,6 @@ //! One Nemotron-H forward over streams' row segments, encoded as FusedDecode's kernels with their geometry. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); @@ -376,7 +377,7 @@ var split_mode: enum(u8) { unread, on, off } = .unread; /// TF_COOP_SPLIT=0 keeps one threadgroup a column tile, as the Python engine dispatches. fn splitOn() bool { if (split_mode == .unread) { - const v = std.c.getenv("TF_COOP_SPLIT"); + const v = platform.c.getenv("TF_COOP_SPLIT"); split_mode = if (v != null and v.?[0] == '0') .off else .on; } return split_mode == .on; diff --git a/zig/src/families/nemotron/gpu_round.zig b/zig/src/families/nemotron/gpu_round.zig index f8db78584b..6515ec7750 100644 --- a/zig/src/families/nemotron/gpu_round.zig +++ b/zig/src/families/nemotron/gpu_round.zig @@ -1,5 +1,6 @@ //! A lone greedy stream's rounds with no host wait: the GPU accepts, keeps rows and drafts; the host reads behind it. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const st = @import("state.zig"); @@ -349,7 +350,7 @@ pub fn run(gpa: std.mem.Allocator, b: *Metal, s: *lanes.Stream, o: Options, hook var mtp_known: usize = mtp0 + 1; var queued: u64 = 0; // rounds submitted var read: u64 = 0; // rounds whose records the host took - const log = std.c.getenv("TF_ROUND_LOG") != null; + const log = platform.c.getenv("TF_ROUND_LOG") != null; var depth = Depth{ .levels = o.depth, .most = o.depth, .fixed = o.fixed }; var windows: [ring]usize = undefined; // each queued round's verify rows var levels = depth.levels; // the levels the last queued head drafts (the next verify's rows - 1) @@ -502,7 +503,7 @@ pub fn run(gpa: std.mem.Allocator, b: *Metal, s: *lanes.Stream, o: Options, hook if (o.level_bench) { // every level live (zero bars), each chain alone in its command buffer r.o.stop = false; - if (std.c.getenv("TF_HEAD_GEOMETRY")) |g| b.head.?.geometry = std.fmt.parseInt(usize, std.mem.span(g), 10) catch 0; + if (platform.c.getenv("TF_HEAD_GEOMETRY")) |g| b.head.?.geometry = std.fmt.parseInt(usize, std.mem.span(g), 10) catch 0; var med: [2]f64 = undefined; for ([_]usize{ 1, @min(8, o.depth) }, 0..) |n, k| { var ms: [9]f64 = undefined; diff --git a/zig/src/families/nemotron/learned_prompts.zig b/zig/src/families/nemotron/learned_prompts.zig index 3f0cdfb1f6..7e7edc8818 100644 --- a/zig/src/families/nemotron/learned_prompts.zig +++ b/zig/src/families/nemotron/learned_prompts.zig @@ -1,5 +1,6 @@ //! Nemotron's learned prompt states (--learn): a kept state's file on disk, and the probe that keys this build's bits. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const lanes = @import("lanes"); const snapshot = @import("snapshot.zig"); @@ -20,23 +21,23 @@ pub fn path(buf: []u8, dir: []const u8, key: u64) ![:0]const u8 { pub fn writeFile(snap: *const Snap, file: [:0]const u8) !void { var tmp_buf: [1100]u8 = undefined; const tmp = try std.fmt.bufPrintSentinel(&tmp_buf, "{s}.part", .{file}, 0); - const fd = std.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const fd = platform.c.open(tmp, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.SnapshotWrite; - errdefer _ = std.c.unlink(tmp); + errdefer _ = platform.c.unlink(tmp); { - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); const head = [HEAD]u32{ FILE_MAGIC, snap.at, @intFromBool(snap.head), @truncate(snap.bytes), @truncate(snap.bytes >> 32) }; try checked.write(fd, std.mem.sliceAsBytes(&head), snap.buf.contents()[0..snap.bytes]); - if (std.c.fsync(fd) != 0) return error.SnapshotWrite; + if (platform.c.fsync(fd) != 0) return error.SnapshotWrite; } - if (std.c.rename(tmp, file) != 0) return error.SnapshotWrite; + if (platform.c.rename(tmp, file) != 0) return error.SnapshotWrite; } /// `file`'s state into `snap`, whose buffer holds its bytes; a file of another position, layout or size is refused. pub fn readFile(snap: *Snap, file: [:0]const u8) !void { - const fd = std.c.open(file, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(file, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.SnapshotRead; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var head: [HEAD]u32 = undefined; try get(fd, std.mem.sliceAsBytes(&head), 0); const bytes = @as(u64, head[3]) | @as(u64, head[4]) << 32; @@ -47,7 +48,7 @@ pub fn readFile(snap: *Snap, file: [:0]const u8) !void { fn put(fd: c_int, b: []const u8) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); + const n = platform.c.write(fd, b.ptr + done, @min(b.len - done, 1 << 30)); if (n <= 0) return error.SnapshotWrite; done += @intCast(n); } @@ -56,7 +57,7 @@ fn put(fd: c_int, b: []const u8) !void { fn get(fd: c_int, b: []u8, at: u64) !void { var done: usize = 0; while (done < b.len) { - const n = std.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); + const n = platform.c.pread(fd, b.ptr + done, @min(b.len - done, 1 << 30), @intCast(at + done)); if (n <= 0) return error.SnapshotRead; done += @intCast(n); } @@ -127,7 +128,7 @@ pub fn snapRead(ptr: *anyopaque, dir: [:0]const u8, key: u64, at: u32) anyerror! /// --learn: learned state `key`'s file removed. pub fn snapForget(_: *anyopaque, dir: [:0]const u8, key: u64) void { var buf: [1100]u8 = undefined; - _ = std.c.unlink(path(&buf, dir, key) catch return); + _ = platform.c.unlink(path(&buf, dir, key) catch return); } test "a learned state's file name is its key in hex" { diff --git a/zig/src/families/nemotron/learner.zig b/zig/src/families/nemotron/learner.zig index 0607eb90c1..69d70cb360 100644 --- a/zig/src/families/nemotron/learner.zig +++ b/zig/src/families/nemotron/learner.zig @@ -1,5 +1,6 @@ //! Sliding Weights' lessons, one learner for every backend over a host that gives the trainer, weights and files. const std = @import("std"); +const platform = @import("sys"); const adapters = @import("learn_shape.zig"); const choice = @import("choice.zig"); const router = @import("router.zig"); @@ -96,7 +97,7 @@ pub fn Learner(comptime H: type) type { pocketed: bool = false, // the last lesson put neurons in the pocket pub fn init(gpa: std.mem.Allocator, io: std.Io, b: H.Ctx) Self { - return .{ .gpa = gpa, .io = io, .b = b, .whole = std.c.getenv("TF_SLIDE_ROUTED") == null, .edits = if (has_edits) Edits.init(gpa) else {}, .method = if (has_edits) whole_lessons.method() else .pocket, .probe = if (has_edits) Probes.init(gpa) else {}, .pockets = if (has_pocket) Pockets.init(gpa) else {} }; + return .{ .gpa = gpa, .io = io, .b = b, .whole = platform.c.getenv("TF_SLIDE_ROUTED") == null, .edits = if (has_edits) Edits.init(gpa) else {}, .method = if (has_edits) whole_lessons.method() else .pocket, .probe = if (has_edits) Probes.init(gpa) else {}, .pockets = if (has_pocket) Pockets.init(gpa) else {} }; } pub fn deinit(l: *Self) void { diff --git a/zig/src/families/nemotron/living.zig b/zig/src/families/nemotron/living.zig index 1f0e8b073b..f479c17c2d 100644 --- a/zig/src/families/nemotron/living.zig +++ b/zig/src/families/nemotron/living.zig @@ -1,5 +1,6 @@ //! Saved lessons as shards beside the model's, a file a lesson (its blocks and router); both backends pass in `core`. const std = @import("std"); +const platform = @import("sys"); const Config = @import("config.zig").Config; const shape = @import("learn_shape.zig"); @@ -27,7 +28,7 @@ pub fn Io(comptime core: type) type { while (true) : (k += 1) { const path = try std.fs.path.joinZ(gpa, &.{ dir, try fileName(&buf, k) }); defer gpa.free(path); - if (std.c.unlink(path) != 0) break; + if (platform.c.unlink(path) != 0) break; } return bytes; } @@ -127,7 +128,7 @@ pub fn Io(comptime core: type) type { pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !?Saved { const old = try std.fs.path.joinZ(gpa, &.{ dir, legacy }); defer gpa.free(old); - if (std.c.access(old, std.c.F_OK) == 0) std.log.warn("slide: {s} holds lessons saved by an earlier version; they are not read", .{legacy}); + if (platform.c.access(old, platform.c.F_OK) == 0) std.log.warn("slide: {s} holds lessons saved by an earlier version; they are not read", .{legacy}); var lessons: std.ArrayList(Lesson) = .empty; errdefer { for (lessons.items) |*l| l.file.close(io); diff --git a/zig/src/families/nemotron/model.zig b/zig/src/families/nemotron/model.zig index 8ff4dae263..740846b96d 100644 --- a/zig/src/families/nemotron/model.zig +++ b/zig/src/families/nemotron/model.zig @@ -1,5 +1,6 @@ //! A loaded Nemotron-H model: config, checkpoint buffers, kernel-layout weights and compiled pipelines. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ckpt = @import("../../core/checkpoint_metal.zig"); const shards = @import("../../core/checkpoint.zig"); @@ -84,7 +85,7 @@ pub const Model = struct { if (mtp) { const path = try std.fmt.allocPrintSentinel(allocator, "{s}/mtp-4bit.safetensors", .{dir}, 0); defer allocator.free(path); - if (std.c.access(path, 0) == 0) try m.checkpoint.addFile(m.device, path, "mtp."); + if (platform.c.access(path, 0) == 0) try m.checkpoint.addFile(m.device, path, "mtp."); } const drafting = m.checkpoint.has("mtp.layers.0.eh_proj.weight"); m.draft_ids = if (drafting) try draft_ids.load(allocator, m.config.vocab) else &.{}; diff --git a/zig/src/families/nemotron/pocket.zig b/zig/src/families/nemotron/pocket.zig index ff24ac5ee9..90232ef763 100644 --- a/zig/src/families/nemotron/pocket.zig +++ b/zig/src/families/nemotron/pocket.zig @@ -1,5 +1,6 @@ //! A pocket of new neurons beside each MoE layer's shared expert, where facts are learned; nothing the model had moves. const std = @import("std"); +const platform = @import("sys"); const examples = @import("lesson_examples.zig"); const lesson_route = @import("lesson_route.zig"); const pocket_fit = @import("pocket_fit.zig"); @@ -52,7 +53,7 @@ const Refit = struct { slot: usize, neuron: usize }; /// Whether environment variable `name` is set to `value`. fn set(name: [*:0]const u8, value: []const u8) bool { - const v = std.c.getenv(name) orelse return false; + const v = platform.c.getenv(name) orelse return false; return std.mem.eql(u8, std.mem.span(v), value); } @@ -136,7 +137,7 @@ pub fn Pocket(comptime T: type) type { read: reading.Seen = .{}, // what the last read pass read, which a lesson begun again over the same keeps pub fn init(gpa: std.mem.Allocator) Self { - const asked: Mode = if (std.c.getenv("TF_SLIDE_POCKET_KEY")) |v| std.meta.stringToEnum(Mode, std.mem.span(v)) orelse .auto else .auto; + const asked: Mode = if (platform.c.getenv("TF_SLIDE_POCKET_KEY")) |v| std.meta.stringToEnum(Mode, std.mem.span(v)) orelse .auto else .auto; return .{ .gpa = gpa, .mode = asked, .strict = !set("TF_SLIDE_POCKET_LEADS", "off"), .keeping = !set("TF_SLIDE_POCKET_KEEPS", "off"), .valuing = set("TF_SLIDE_POCKET_VALUES", "on"), .silencing = set("TF_SLIDE_POCKET_SAME", "silent"), .worst = set("TF_SLIDE_POCKET_REPLAY", "worst"), .joining = set("TF_SLIDE_POCKET_GROUP", "joint") }; } @@ -360,7 +361,7 @@ pub fn Pocket(comptime T: type) type { const d = c.hidden; const n = p.slots(); const m = p.labels.items.len; - if (std.c.getenv("TF_SLIDE_POCKET_DUMP")) |dir| pocket_dump.write(p, std.mem.span(dir), p.learned_before, d) catch |e| std.log.warn("slide pocket: the read rows were not written: {t}", .{e}); + if (platform.c.getenv("TF_SLIDE_POCKET_DUMP")) |dir| pocket_dump.write(p, std.mem.span(dir), p.learned_before, d) catch |e| std.log.warn("slide pocket: the read rows were not written: {t}", .{e}); const kinds = try p.gpa.alloc(Kind, 3 * m); defer p.gpa.free(kinds); for (p.labels.items, 0..) |l, i| for (0..3) |k| { diff --git a/zig/src/families/nemotron/pocket_dump.zig b/zig/src/families/nemotron/pocket_dump.zig index 1653589631..b1f5d152a6 100644 --- a/zig/src/families/nemotron/pocket_dump.zig +++ b/zig/src/families/nemotron/pocket_dump.zig @@ -1,11 +1,12 @@ //! A lesson's read rows written out for offline study (TF_SLIDE_POCKET_DUMP=dir): inputs, labels and examples. const std = @import("std"); +const platform = @import("sys"); /// `bytes` written to `path`. fn put(path: [:0]const u8, bytes: []const u8) !void { - const f = std.c.fopen(path, "wb") orelse return error.DumpOpen; - defer _ = std.c.fclose(f); - if (std.c.fwrite(bytes.ptr, 1, bytes.len, f) != bytes.len) return error.DumpWrite; + const f = platform.c.fopen(path, "wb") orelse return error.DumpOpen; + defer _ = platform.c.fclose(f); + if (platform.c.fwrite(bytes.ptr, 1, bytes.len, f) != bytes.len) return error.DumpWrite; } /// Pocket `p`'s read rows under `dir` for `tag`: bf16 [read, slots, D], then their labels and examples as text. diff --git a/zig/src/families/nemotron/whole.zig b/zig/src/families/nemotron/whole.zig index 4fa4a9fc01..b8f6fc8a99 100644 --- a/zig/src/families/nemotron/whole.zig +++ b/zig/src/families/nemotron/whole.zig @@ -1,5 +1,6 @@ //! Whole-mode lessons: pocket neurons (default), a weight edit (TF_SLIDE_EDIT) or a probe (TF_SLIDE_PROBE). const std = @import("std"); +const platform = @import("sys"); const examples = @import("lesson_examples.zig"); const adapters = @import("learn_shape.zig"); @@ -10,8 +11,8 @@ pub const Method = enum { pocket, edit, probe }; /// The method the environment picks for whole lessons. pub fn method() Method { - if (std.c.getenv("TF_SLIDE_PROBE") != null) return .probe; - if (std.c.getenv("TF_SLIDE_EDIT") != null) return .edit; + if (platform.c.getenv("TF_SLIDE_PROBE") != null) return .probe; + if (platform.c.getenv("TF_SLIDE_EDIT") != null) return .edit; return .pocket; } diff --git a/zig/src/families/qwen27/model.zig b/zig/src/families/qwen27/model.zig index 53771c8eda..428f764854 100644 --- a/zig/src/families/qwen27/model.zig +++ b/zig/src/families/qwen27/model.zig @@ -1,5 +1,6 @@ //! A native text model keeps immutable checkpoint bytes and prepared affine layouts under one owner. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const core = @import("core"); const cfg = @import("config.zig"); @@ -71,7 +72,7 @@ pub const Model = struct { m.frame = try Frame.init(m.device, c, rows); errdefer m.frame.deinit(); m.ane = null; - const columns = if (std.c.getenv("TF_ANE_COLUMNS")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else ane_columns; + const columns = if (platform.c.getenv("TF_ANE_COLUMNS")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else ane_columns; const share = columns > 0 and m.device.tensorUnits() and mtl.ane.available() and m.weights.uniformMlp(); const fits = if (share_fits) |f| !share or f(ane_mlp.Share.bytes(c, columns)) else true; if (!fits) std.log.info("qwen27: the Neural Engine's prompt share ({d} MiB of fp16 weights) would take this Mac past its load limit; prompts run on the GPU alone", .{ane_mlp.Share.bytes(c, columns) >> 20}); diff --git a/zig/src/hip/driver.zig b/zig/src/hip/driver.zig index 0982c96660..7fa1ef9cca 100644 --- a/zig/src/hip/driver.zig +++ b/zig/src/hip/driver.zig @@ -1,5 +1,6 @@ //! The first HIP admission boundary, with signatures from ROCm's hip_runtime_api.h; no CUDA symbol aliases or GPU kernels. const std = @import("std"); +const platform = @import("sys"); pub const Error = error{ DriverUnavailable, MissingSymbol, HipFailed, InvalidDeviceCount }; pub const Api = struct { @@ -8,11 +9,11 @@ pub const Api = struct { }; pub const Driver = struct { - lib: std.DynLib, + lib: platform.DynLib, api: Api, pub fn openPath(path: []const u8) Error!Driver { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); const api = try @import("symbols.zig").resolve(Api, &lib); try initialize(api); diff --git a/zig/src/hip/runtime.zig b/zig/src/hip/runtime.zig index 8f5527bd84..2583cf6bbd 100644 --- a/zig/src/hip/runtime.zig +++ b/zig/src/hip/runtime.zig @@ -1,10 +1,11 @@ //! Complete model-free HIP runtime admission; partial symbol tables never escape. const std = @import("std"); +const platform = @import("sys"); const abi = @import("abi.zig"); pub const Error = error{ DriverUnavailable, MissingSymbol, HipFailed, Invalid }; pub const Runtime = struct { - lib: std.DynLib, + lib: platform.DynLib, api: abi.Api, pub fn open() Error!Runtime { @@ -12,7 +13,7 @@ pub const Runtime = struct { } pub fn openPath(path: []const u8) Error!Runtime { - var lib = std.DynLib.open(path) catch return error.DriverUnavailable; + var lib = platform.DynLib.open(path) catch return error.DriverUnavailable; errdefer lib.close(); const api = try @import("symbols.zig").resolve(abi.Api, &lib); try check(api.hipInit(0)); diff --git a/zig/src/hip/symbols.zig b/zig/src/hip/symbols.zig index e334f716cb..54c2c14bc6 100644 --- a/zig/src/hip/symbols.zig +++ b/zig/src/hip/symbols.zig @@ -1,7 +1,8 @@ //! One symbol-resolution path for admission fixtures and the full runtime. const std = @import("std"); +const platform = @import("sys"); -pub fn resolve(comptime Api: type, lib: *std.DynLib) error{MissingSymbol}!Api { +pub fn resolve(comptime Api: type, lib: *platform.DynLib) error{MissingSymbol}!Api { var api: Api = undefined; const info = @typeInfo(Api).@"struct"; inline for (info.field_names, info.field_types) |name, T| { diff --git a/zig/src/learned_warning_test.zig b/zig/src/learned_warning_test.zig index cb9a9346b1..382ac2659e 100644 --- a/zig/src/learned_warning_test.zig +++ b/zig/src/learned_warning_test.zig @@ -1,4 +1,5 @@ const std = @import("std"); +const platform = @import("sys"); const pc = @import("core/prompt_cache.zig"); const Imprint = @import("core/prompt_imprint.zig").Imprint; const Fake = @import("core/prompt_cache_fake.zig").Fake; @@ -56,7 +57,7 @@ fn expectedWarnings(n: usize) !void { pub fn main() !void { const a = std.heap.page_allocator; var buf: [128]u8 = undefined; - const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-warning-{d}", .{std.c.getpid()}); + const root = try std.fmt.bufPrint(&buf, "/tmp/tf-learn-warning-{d}", .{platform.c.getpid()}); defer @import("core/prompt_imprint_test.zig").rmTree(root); { var im = try Imprint.open(a, root, 1, 1 << 20); diff --git a/zig/src/metal/clock.zig b/zig/src/metal/clock.zig index 6461c8e227..9e0572c597 100644 --- a/zig/src/metal/clock.zig +++ b/zig/src/metal/clock.zig @@ -1,13 +1,14 @@ //! Host time in seconds on the clock Metal's GPUStartTime and GPUEndTime use (mach_absolute_time). const std = @import("std"); +const platform = @import("sys"); var scale: f64 = 0; pub fn seconds() f64 { if (scale == 0) { - var info: std.c.mach_timebase_info_data = undefined; - _ = std.c.mach_timebase_info(&info); + var info: platform.c.mach_timebase_info_data = undefined; + _ = platform.c.mach_timebase_info(&info); scale = @as(f64, @floatFromInt(info.numer)) / @as(f64, @floatFromInt(info.denom)) / 1e9; } - return @as(f64, @floatFromInt(std.c.mach_absolute_time())) * scale; + return @as(f64, @floatFromInt(platform.c.mach_absolute_time())) * scale; } diff --git a/zig/src/metal/device.zig b/zig/src/metal/device.zig index 69d67e108c..353cd6d8ce 100644 --- a/zig/src/metal/device.zig +++ b/zig/src/metal/device.zig @@ -1,5 +1,6 @@ //! The GPU device: queues, buffers (zero-copy over mmap too), events and residency sets. const std = @import("std"); +const platform = @import("sys"); const objc = @import("objc.zig"); const types = @import("types.zig"); const command = @import("command.zig"); @@ -139,10 +140,10 @@ pub const MappedFile = struct { size: usize, pub fn open(path: [*:0]const u8) !MappedFile { - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }); if (fd < 0) return error.OpenFailed; - defer _ = std.c.close(fd); - const end = std.c.lseek(fd, 0, std.c.SEEK.END); + defer _ = platform.c.close(fd); + const end = platform.c.lseek(fd, 0, platform.c.SEEK.END); if (end <= 0) return error.EmptyFile; const size: usize = @intCast(end); const page = std.heap.pageSize(); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..50af1bd955 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,5 +1,6 @@ //! CUDA family registration, device setup, memory admission and the native lane host. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); @@ -44,7 +45,7 @@ fn chipClass(a: Allocator, capability: u32) ?[]const u8 { } fn getenv(name: [:0]const u8) ?[]const u8 { - return std.mem.span(std.c.getenv(name) orelse return null); + return std.mem.span(platform.c.getenv(name) orelse return null); } /// A variable's whole number: null when unset or empty, an error for anything else that is not one. diff --git a/zig/src/native/flashnext_host.zig b/zig/src/native/flashnext_host.zig index cdba9b907b..539964a682 100644 --- a/zig/src/native/flashnext_host.zig +++ b/zig/src/native/flashnext_host.zig @@ -1,6 +1,7 @@ //! Flash Next behind the native server: one reply at a time on the replay engine (prompt chunks, then GPU-side //! rounds), greedy only; the engine's thread owns the GPU, requests wait in arrival order (foreground first). const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -420,7 +421,7 @@ const Snaps = struct { /// --learn: learned state `key`'s file removed. fn forget(_: *anyopaque, dir: [:0]const u8, key: u64) void { var buf: [1100]u8 = undefined; - _ = std.c.unlink(learned.path(&buf, dir, key) catch return); + _ = platform.c.unlink(learned.path(&buf, dir, key) catch return); } }; diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 39fe531408..8370f45ed7 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -1,5 +1,6 @@ //! The Metal engines: lane core (Nemotron, Qwen3.5-2B, GLM, Qwen3.8-27B), Flash Next replay. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -260,9 +261,9 @@ fn shardStamps(h: *std.hash.Wyhash, a: Allocator, io: std.Io, dir: []const u8) ! /// A safetensors file's header (its tensors' names, shapes and offsets), in `a`. fn safetensorsHeader(a: Allocator, dir: []const u8, name: []const u8) ![]const u8 { const path = try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0); - const fd = std.c.open(path, .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY }, @as(platform.c.mode_t, 0)); if (fd < 0) return error.FileNotFound; - defer _ = std.c.close(fd); + defer _ = platform.c.close(fd); var n: u64 = 0; if (!api.prompt_imprint.readAt(fd, std.mem.asBytes(&n), 0) or n > 1 << 24) return error.BadHeader; const head = try a.alloc(u8, @intCast(n)); @@ -305,7 +306,7 @@ fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ /// Flash Next uses the dump named by TF_FLASHNEXT_DUMP, or the checked-in kernels when that variable is unset. fn openFlashNext(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { - const dump: ?[]const u8 = if (std.c.getenv("TF_FLASHNEXT_DUMP")) |d| std.mem.span(d) else null; + const dump: ?[]const u8 = if (platform.c.getenv("TF_FLASHNEXT_DUMP")) |d| std.mem.span(d) else null; const native = modelContext(a, io, o.dir); const window: i64 = o.context orelse native; if (window < 0 or (native > 0 and window > native)) { @@ -373,7 +374,7 @@ fn openGlmDsa(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[ } var ram: u64 = 0; var len: usize = 8; - if (std.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; + if (platform.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; const pool = mtl.objc.Pool.push(); defer pool.pop(); const h = glm_dsa.open(gpa, io, o.dir, @intCast(window), settings, @max(1, @min(o.lanes, 16)), ram / 100 * 66, o.prompt_cache_gib) catch |e| { @@ -469,10 +470,10 @@ test "a shard rewritten in place, same size and bytes, gives the learned states try std.testing.expectEqual(first.final(), notes.final()); std.Io.sleep(io, .fromMilliseconds(5), .awake) catch {}; const path = try std.fmt.allocPrintSentinel(a, "{s}/model-00001-of-00001.safetensors", .{dir}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY }, @as(std.c.mode_t, 0)); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY }, @as(platform.c.mode_t, 0)); try std.testing.expect(fd >= 0); - try std.testing.expectEqual(@as(isize, 7), std.c.pwrite(fd, "weights", 7, 0)); // as Sliding Weights writes a tensor of the same size - _ = std.c.close(fd); + try std.testing.expectEqual(@as(isize, 7), platform.c.pwrite(fd, "weights", 7, 0)); // as Sliding Weights writes a tensor of the same size + _ = platform.c.close(fd); var moved = std.hash.Wyhash.init(0); try shardStamps(&moved, a, io, dir); try std.testing.expect(first.final() != moved.final()); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index e09a508a8c..e7da14ef17 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -1,5 +1,6 @@ //! The 27B's lane core backend: a stream a runner slot, shared rounds across slots, DFlash2 trees for slot 0. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const api = @import("engine_api"); const tf = @import("tensorfold"); @@ -125,7 +126,7 @@ pub const Metal = struct { const drafting = b.draft != null; const shared = b.runner.slots > 1; // a lone stream drafts DFlash2 trees, as in 1.0.5: copies land less often there (TF_QWEN27_COPIES=1 opts in) - const lone_copies = if (std.c.getenv("TF_QWEN27_COPIES")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; + const lone_copies = if (platform.c.getenv("TF_QWEN27_COPIES")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; return .{ .exact_width = if (drafting or shared) max_rows else 1, .mtp = drafting, .speculate = drafting, .speculate_early = false, .drafts = b.nodes, .draft_probabilities = true, .window_costs = b.costs[0..b.timed], .hidden_rows = shared, .streams_exact = true, .batch_rows = if (shared) batch_rows else max_rows, .max_streams = b.runner.slots, .lone_copies = lone_copies }; } diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 09da59e53f..514cecdc8d 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -1,5 +1,6 @@ //! Qwen3.8-27B on the lane host: a stream a runner slot in shared rounds, DFlash2 drafts for slot 0's stream. const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const mtl = @import("metal"); const api = @import("engine_api"); @@ -289,7 +290,7 @@ pub fn attach(gpa: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q.df h.host.decoded_rows = true; // a reply's rows prefill with the bits decoding gave them, with the cache on or off // a lone greedy request takes the serial loop on tensor units, where it is faster (TF_QWEN27_LANES overrides) const units = model.device.tensorUnits(); - const serial = if (std.c.getenv("TF_QWEN27_LANES")) |v| !std.mem.eql(u8, std.mem.span(v), "1") else units; + const serial = if (platform.c.getenv("TF_QWEN27_LANES")) |v| !std.mem.eql(u8, std.mem.span(v), "1") else units; if (h.draft != null and serial) h.host.lone = .{ .ctx = h, .run = Host.lone }; h.warm = .{ .queue = model.queue }; // the ready slots are wired before the cache's budget reads the footprint, and the later ones count as taken diff --git a/zig/src/platform/sys.zig b/zig/src/platform/sys.zig new file mode 100644 index 0000000000..ee614d7507 --- /dev/null +++ b/zig/src/platform/sys.zig @@ -0,0 +1,47 @@ +//! The operating system as TensorFold calls it: `c` is libc's surface (std.c on POSIX, win_c.zig on Windows) and +//! `DynLib` opens shared libraries (std.DynLib on POSIX, LoadLibraryW on Windows). Every module imports this as `sys`. + +const std = @import("std"); +const builtin = @import("builtin"); + +pub const windows = builtin.os.tag == .windows; + +pub const c = if (windows) @import("win_c.zig") else std.c; + +/// std.posix, or its sockets, signals and limits over Winsock and console events. +pub const posix = if (windows) @import("win_posix.zig") else std.posix; + +pub const DynLib = if (windows) WinDynLib else std.DynLib; + +/// The CUDA driver's file name: nvcuda.dll comes with the display driver on Windows. +pub const cuda_driver = if (windows) "nvcuda.dll" else "libcuda.so.1"; + +extern "kernel32" fn LoadLibraryW(name: [*:0]const u16) callconv(.winapi) ?*anyopaque; +extern "kernel32" fn GetProcAddress(module: *anyopaque, name: [*:0]const u8) callconv(.winapi) ?*anyopaque; +extern "kernel32" fn FreeLibrary(module: *anyopaque) callconv(.winapi) c_int; + +const WinDynLib = struct { + module: *anyopaque, + + pub const Error = error{ FileNotFound, NameTooLong, InvalidWtf8 }; + + pub fn open(path: []const u8) Error!WinDynLib { + var buf: [1024:0]u16 = undefined; + const n = std.unicode.wtf8ToWtf16Le(buf[0 .. buf.len - 1], path) catch return error.InvalidWtf8; + buf[n] = 0; + return .{ .module = LoadLibraryW(buf[0..n :0].ptr) orelse return error.FileNotFound }; + } + + pub fn openZ(path: [*:0]const u8) Error!WinDynLib { + return open(std.mem.span(path)); + } + + pub fn close(self: *WinDynLib) void { + _ = FreeLibrary(self.module); + } + + pub fn lookup(self: *WinDynLib, comptime T: type, name: [:0]const u8) ?T { + const p = GetProcAddress(self.module, name.ptr) orelse return null; + return @ptrCast(@alignCast(p)); + } +}; diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig new file mode 100644 index 0000000000..82b4dc7b3f --- /dev/null +++ b/zig/src/platform/win_c.zig @@ -0,0 +1,577 @@ +//! The libc surface TensorFold calls through `sys.c`, on Windows: integer descriptors over the UCRT (files opened by +//! CreateFileW with every share mode, so unlink and rename behave as on POSIX), directory descriptors as remembered +//! paths, readdir over FindFirstFileW, and POSIX's names for the rest. Each call keeps libc's contract: -1 and errno. + +const std = @import("std"); + +pub const fd_t = c_int; +pub const mode_t = c_uint; +pub const off_t = i64; +pub const pid_t = c_int; +pub const time_t = i64; +pub const E = std.c.E; +pub const errno = std.c.errno; +pub const timespec = extern struct { sec: time_t, nsec: c_long }; + +pub const SEEK = struct { + pub const SET: c_int = 0; + pub const CUR: c_int = 1; + pub const END: c_int = 2; +}; +pub const F_OK: c_int = 0; +pub const R_OK: c_int = 4; +pub const W_OK: c_int = 2; +pub const X_OK: c_int = 0; // Windows has no execute bit; existence is the answer +pub const AT = struct { + pub const FDCWD: fd_t = -100; + pub const REMOVEDIR: u32 = 0x200; +}; +pub const DT = struct { + pub const UNKNOWN: u8 = 0; + pub const DIR: u8 = 4; + pub const REG: u8 = 8; + pub const LNK: u8 = 10; +}; +pub const S = struct { + pub const IFMT: mode_t = 0o170000; + pub const IFDIR: mode_t = 0o040000; + pub const IFREG: mode_t = 0o100000; + pub const IFLNK: mode_t = 0o120000; +}; +pub const RLIM = struct { + pub const INFINITY: u64 = std.math.maxInt(u64); +}; +pub const timeval = extern struct { sec: time_t, usec: c_long }; +pub const CLOCK = struct { + pub const REALTIME: c_int = 0; + pub const MONOTONIC: c_int = 1; +}; + +/// Linux's open flags by name; DIRECT asks for unbuffered reads (FILE_FLAG_NO_BUFFERING), the rest map or are moot. +pub const O = packed struct(u32) { + ACCMODE: enum(u2) { RDONLY = 0, WRONLY = 1, RDWR = 2 } = .RDONLY, + CREAT: bool = false, + EXCL: bool = false, + NOCTTY: bool = false, + TRUNC: bool = false, + APPEND: bool = false, + NONBLOCK: bool = false, + DSYNC: bool = false, + ASYNC: bool = false, + DIRECT: bool = false, + LARGEFILE: bool = false, + DIRECTORY: bool = false, + NOFOLLOW: bool = false, + NOATIME: bool = false, + CLOEXEC: bool = false, + SYNC: bool = false, + PATH: bool = false, + TMPFILE: bool = false, + _: u13 = 0, +}; + +// ---- Win32 and UCRT entry points ---- + +const HANDLE = *anyopaque; +const BOOL = c_int; +const invalid_handle: isize = -1; + +const OVERLAPPED = extern struct { internal: usize = 0, internal_high: usize = 0, offset: u32, offset_high: u32, event: ?HANDLE = null }; +const FILETIME = extern struct { low: u32, high: u32 }; +const FindData = extern struct { + attributes: u32, + created: FILETIME, + accessed: FILETIME, + written: FILETIME, + size_high: u32, + size_low: u32, + reserved0: u32, + reserved1: u32, + name: [260]u16, + alternate: [14]u16, +}; + +extern "kernel32" fn CreateFileW(name: [*:0]const u16, access: u32, share: u32, security: ?*anyopaque, disposition: u32, flags: u32, template: ?HANDLE) callconv(.winapi) isize; +extern "kernel32" fn ReadFile(h: HANDLE, buf: [*]u8, len: u32, got: ?*u32, ov: ?*OVERLAPPED) callconv(.winapi) BOOL; +extern "kernel32" fn WriteFile(h: HANDLE, buf: [*]const u8, len: u32, put: ?*u32, ov: ?*OVERLAPPED) callconv(.winapi) BOOL; +extern "kernel32" fn GetLastError() callconv(.winapi) u32; +extern "kernel32" fn GetFileAttributesW(name: [*:0]const u16) callconv(.winapi) u32; +extern "kernel32" fn MoveFileExW(from: [*:0]const u16, to: [*:0]const u16, flags: u32) callconv(.winapi) BOOL; +extern "kernel32" fn FindFirstFileW(pattern: [*:0]const u16, data: *FindData) callconv(.winapi) isize; +extern "kernel32" fn FindNextFileW(h: HANDLE, data: *FindData) callconv(.winapi) BOOL; +extern "kernel32" fn FindClose(h: HANDLE) callconv(.winapi) BOOL; +extern "kernel32" fn QueryPerformanceCounter(count: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn QueryPerformanceFrequency(freq: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn GetSystemTimePreciseAsFileTime(t: *FILETIME) callconv(.winapi) void; +extern "kernel32" fn CreateWaitableTimerExW(security: ?*anyopaque, name: ?[*:0]const u16, flags: u32, access: u32) callconv(.winapi) ?HANDLE; +extern "kernel32" fn SetWaitableTimer(h: HANDLE, due: *const i64, period: i32, routine: ?*anyopaque, arg: ?*anyopaque, resume_: BOOL) callconv(.winapi) BOOL; +extern "kernel32" fn WaitForSingleObject(h: HANDLE, ms: u32) callconv(.winapi) u32; +extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; +extern "kernel32" fn GetDiskFreeSpaceExW(dir: [*:0]const u16, avail: ?*u64, total: ?*u64, free: ?*u64) callconv(.winapi) BOOL; +extern "kernel32" fn GetTempPathW(len: u32, buf: [*]u16) callconv(.winapi) u32; +extern "kernel32" fn GetModuleFileNameW(module: ?HANDLE, buf: [*]u16, len: u32) callconv(.winapi) u32; +extern "kernel32" fn CloseHandle(h: HANDLE) callconv(.winapi) BOOL; + +extern "c" fn _errno() *c_int; +extern "c" fn _open_osfhandle(h: isize, flags: c_int) c_int; +extern "c" fn _get_osfhandle(fd: c_int) isize; +extern "c" fn _close(fd: c_int) c_int; +extern "c" fn _read(fd: c_int, buf: [*]u8, len: c_uint) c_int; +extern "c" fn _write(fd: c_int, buf: [*]const u8, len: c_uint) c_int; +extern "c" fn _lseeki64(fd: c_int, off: i64, whence: c_int) i64; +extern "c" fn _commit(fd: c_int) c_int; +extern "c" fn _chsize_s(fd: c_int, size: i64) c_int; +extern "c" fn _wunlink(path: [*:0]const u16) c_int; +extern "c" fn _wrmdir(path: [*:0]const u16) c_int; +extern "c" fn _wmkdir(path: [*:0]const u16) c_int; +extern "c" fn _waccess(path: [*:0]const u16, mode: c_int) c_int; +extern "c" fn _getpid() c_int; +extern "c" fn _isatty(fd: c_int) c_int; +extern "c" fn _fstat64(fd: c_int, st: *Stat64) c_int; +extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; +pub extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; +pub extern "c" fn fopen(path: [*:0]const u8, mode: [*:0]const u8) ?*std.c.FILE; +pub extern "c" fn fwrite(ptr: [*]const u8, size: usize, n: usize, f: *std.c.FILE) usize; +pub extern "c" fn fclose(f: *std.c.FILE) c_int; +pub extern "c" fn free(ptr: ?*anyopaque) void; +pub extern "c" fn malloc(n: usize) ?*anyopaque; + +/// The UCRT's struct _stat64. +const Stat64 = extern struct { dev: u32, ino: u16, mode: u16, nlink: i16, uid: i16, gid: i16, rdev: u32, size: i64, atime: i64, mtime: i64, ctime: i64 }; + +fn fail(e: E) c_int { + _errno().* = @intFromEnum(e); + return -1; +} + +/// The errno a failed Win32 call means. +fn lastError() E { + return switch (GetLastError()) { + 2, 3, 15, 123, 161 => .NOENT, // file, path, drive not found; bad name + 5, 32, 33 => .ACCES, // access denied, sharing and lock violations + 80, 183 => .EXIST, + 112 => .NOSPC, + 6 => .BADF, + 145 => .NOTEMPTY, + 267 => .NOTDIR, + 87 => .INVAL, + else => .IO, + }; +} + +// ---- paths ---- + +const Wide = [32767:0]u16; + +/// UTF-8 path to a NUL-terminated wide path in `buf`; null (errno set) when it does not fit or is not UTF-8. +/// POSIX's /tmp is the user's temporary folder and /proc/self/exe this executable, as Cygwin and MSYS map them. +fn wide(buf: *Wide, path: []const u8) ?[*:0]const u16 { + var at: usize = 0; + var rest = path; + if (std.mem.eql(u8, path, "/proc/self/exe")) { + const n = GetModuleFileNameW(null, buf, buf.len - 1); + if (n == 0 or n >= buf.len - 1) { + _ = fail(.NAMETOOLONG); + return null; + } + buf[n] = 0; + return buf[0..n :0].ptr; + } + if (std.mem.startsWith(u8, path, "/tmp") and (path.len == 4 or path[4] == '/')) { + at = GetTempPathW(buf.len - 1, buf); // ends in a separator + if (at == 0 or at >= buf.len - 1) { + _ = fail(.NAMETOOLONG); + return null; + } + if (path.len == 4) at -= 1; + rest = if (path.len > 5) path[5..] else ""; + } + const n = std.unicode.wtf8ToWtf16Le(buf[at .. buf.len - 1], rest) catch { + _ = fail(.NAMETOOLONG); + return null; + }; + buf[at + n] = 0; + return buf[0 .. at + n :0].ptr; +} + +/// Directory descriptors: open(O.DIRECTORY) remembers the path; openat and unlinkat join onto it. +const dir_base: fd_t = 0x4000_0000; +var dirs_lock: std.atomic.Mutex = .unlocked; +var dir_paths: [256]?[:0]u8 = @splat(null); + +fn lockDirs() void { + while (!dirs_lock.tryLock()) std.atomic.spinLoopHint(); +} + +fn dirPath(fd: fd_t) ?[:0]const u8 { + if (fd < dir_base or fd >= dir_base + dir_paths.len) return null; + lockDirs(); + defer dirs_lock.unlock(); + return dir_paths[@intCast(fd - dir_base)]; +} + +fn openDir(path: []const u8) c_int { + var buf: Wide = undefined; + const w = wide(&buf, path) orelse return -1; + const attrs = GetFileAttributesW(w); + if (attrs == 0xFFFF_FFFF) return fail(lastError()); + if (attrs & 0x10 == 0) return fail(.NOTDIR); + const copy = std.heap.c_allocator.dupeSentinel(u8, path, 0) catch return fail(.NOMEM); + lockDirs(); + defer dirs_lock.unlock(); + for (&dir_paths, 0..) |*slot, i| if (slot.* == null) { + slot.* = copy; + return dir_base + @as(fd_t, @intCast(i)); + }; + std.heap.c_allocator.free(copy); + return fail(.MFILE); +} + +/// `name` under the directory `fd` (or the working directory, or absolute) in `buf`. +fn joinAt(buf: []u8, fd: fd_t, name: [*:0]const u8) ?[:0]const u8 { + const leaf = std.mem.span(name); + if (fd == AT.FDCWD or std.fs.path.isAbsolute(leaf)) return std.fmt.bufPrintSentinel(buf, "{s}", .{leaf}, 0) catch null; + const base = dirPath(fd) orelse { + _ = fail(.BADF); + return null; + }; + return std.fmt.bufPrintSentinel(buf, "{s}\\{s}", .{ base, leaf }, 0) catch { + _ = fail(.NAMETOOLONG); + return null; + }; +} + +// ---- descriptors ---- + +/// Variadic as libc's is, so callers pass a mode or not; the mode is moot on NTFS without ACL work. +pub fn open(path: [*:0]const u8, flags: O, ...) callconv(.c) c_int { + return openFile(path, flags); +} + +fn openFile(path: [*:0]const u8, flags: O) c_int { + const p = std.mem.span(path); + if (flags.DIRECTORY) return openDir(p); + var buf: Wide = undefined; + const w = wide(&buf, p) orelse return -1; + const generic_read: u32 = 0x8000_0000; + const generic_write: u32 = 0x4000_0000; + const rights: u32 = switch (flags.ACCMODE) { + .RDONLY => generic_read, + .WRONLY => generic_write, + .RDWR => generic_read | generic_write, + }; + const disposition: u32 = if (flags.CREAT and flags.EXCL) 1 // CREATE_NEW + else if (flags.CREAT and flags.TRUNC) 2 // CREATE_ALWAYS + else if (flags.CREAT) 4 // OPEN_ALWAYS + else if (flags.TRUNC) 5 // TRUNCATE_EXISTING + else 3; // OPEN_EXISTING + var attrs: u32 = 0x80; // FILE_ATTRIBUTE_NORMAL + if (flags.DIRECT) attrs |= 0x2000_0000; // FILE_FLAG_NO_BUFFERING + const h = CreateFileW(w, rights, 0x7, null, disposition, attrs, null); // share read, write and delete + if (h == invalid_handle) { + // a directory opened without O.DIRECTORY, as POSIX allows for reading its entries + if (GetLastError() == 5 and flags.ACCMODE == .RDONLY and GetFileAttributesW(w) & 0x10 != 0 and GetFileAttributesW(w) != 0xFFFF_FFFF) return openDir(p); + return fail(lastError()); + } + var crt: c_int = 0x8000; // _O_BINARY + crt |= switch (flags.ACCMODE) { + .RDONLY => 0, + .WRONLY => 1, + .RDWR => 2, + }; + if (flags.APPEND) crt |= 0x8; + const fd = _open_osfhandle(h, crt); + if (fd < 0) { + _ = CloseHandle(@ptrFromInt(@as(usize, @bitCast(h)))); + return fail(.MFILE); + } + return fd; +} + +pub fn openat(dir: fd_t, path: [*:0]const u8, flags: O, ...) callconv(.c) c_int { + var buf: [4096]u8 = undefined; + const p = joinAt(&buf, dir, path) orelse return -1; + return openFile(p.ptr, flags); +} + +pub fn close(fd: fd_t) c_int { + if (fd >= dir_base and fd < dir_base + dir_paths.len) { + lockDirs(); + defer dirs_lock.unlock(); + const slot = &dir_paths[@intCast(fd - dir_base)]; + const p = slot.* orelse return fail(.BADF); + std.heap.c_allocator.free(p); + slot.* = null; + return 0; + } + return _close(fd); +} + +const chunk: usize = 1 << 30; + +pub fn read(fd: fd_t, buf: [*]u8, len: usize) isize { + return _read(fd, buf, @intCast(@min(len, chunk))); +} + +pub fn write(fd: fd_t, buf: [*]const u8, len: usize) isize { + return _write(fd, buf, @intCast(@min(len, chunk))); +} + +fn osHandle(fd: fd_t) ?HANDLE { + const h = _get_osfhandle(fd); + if (h == invalid_handle or h == -2) { + _ = fail(.BADF); + return null; + } + return @ptrFromInt(@as(usize, @bitCast(h))); +} + +/// Positional read by an OVERLAPPED offset; end of file reads 0 as pread does. +pub fn pread(fd: fd_t, buf: [*]u8, len: usize, offset: off_t) isize { + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = @truncate(@as(u64, @bitCast(offset))), .offset_high = @truncate(@as(u64, @bitCast(offset)) >> 32) }; + var got: u32 = 0; + if (ReadFile(h, buf, @intCast(@min(len, chunk)), &got, &ov) == 0) { + if (GetLastError() == 38) return 0; // ERROR_HANDLE_EOF + return fail(lastError()); + } + return got; +} + +pub fn pwrite(fd: fd_t, buf: [*]const u8, len: usize, offset: off_t) isize { + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = @truncate(@as(u64, @bitCast(offset))), .offset_high = @truncate(@as(u64, @bitCast(offset)) >> 32) }; + var put: u32 = 0; + if (WriteFile(h, buf, @intCast(@min(len, chunk)), &put, &ov) == 0) return fail(lastError()); + return put; +} + +pub fn lseek(fd: fd_t, offset: off_t, whence: c_int) off_t { + return _lseeki64(fd, offset, whence); +} + +pub fn fsync(fd: fd_t) c_int { + if (dirPath(fd) != null) return 0; // directory entries are durable once NTFS has them + return _commit(fd); +} + +pub fn ftruncate(fd: fd_t, size: off_t) c_int { + const rc = _chsize_s(fd, size); + if (rc != 0) return fail(@enumFromInt(rc)); + return 0; +} + +pub fn isatty(fd: fd_t) c_int { + return _isatty(fd); +} + +pub const Stat = struct { + dev: u64, + ino: u64, + mode: mode_t, + size: i64, + mtim: timespec, + + pub fn mtime(s: Stat) timespec { + return s.mtim; + } +}; + +pub fn fstat(fd: fd_t, out: *Stat) c_int { + if (dirPath(fd) != null) { + out.* = .{ .dev = 0, .ino = 0, .mode = S.IFDIR | 0o700, .size = 0, .mtim = .{ .sec = 0, .nsec = 0 } }; + return 0; + } + var st: Stat64 = undefined; + if (_fstat64(fd, &st) != 0) return -1; + out.* = .{ .dev = st.dev, .ino = st.ino, .mode = st.mode, .size = st.size, .mtim = .{ .sec = st.mtime, .nsec = 0 } }; + return 0; +} + +// ---- names ---- + +pub fn unlink(path: [*:0]const u8) c_int { + var buf: Wide = undefined; + return _wunlink(wide(&buf, std.mem.span(path)) orelse return -1); +} + +pub fn rmdir(path: [*:0]const u8) c_int { + var buf: Wide = undefined; + return _wrmdir(wide(&buf, std.mem.span(path)) orelse return -1); +} + +pub fn unlinkat(dir: fd_t, path: [*:0]const u8, flags: u32) c_int { + var buf: [4096]u8 = undefined; + const p = joinAt(&buf, dir, path) orelse return -1; + return if (flags & AT.REMOVEDIR != 0) rmdir(p.ptr) else unlink(p.ptr); +} + +pub fn mkdir(path: [*:0]const u8, mode: mode_t) c_int { + _ = mode; + var buf: Wide = undefined; + return _wmkdir(wide(&buf, std.mem.span(path)) orelse return -1); +} + +/// POSIX rename: replaces an existing target. +pub fn rename(from: [*:0]const u8, to: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wa = wide(&a, std.mem.span(from)) orelse return -1; + const wb = wide(&b, std.mem.span(to)) orelse return -1; + if (MoveFileExW(wa, wb, 0x1 | 0x8) == 0) return fail(lastError()); // REPLACE_EXISTING | WRITE_THROUGH + return 0; +} + +pub fn access(path: [*:0]const u8, mode: c_int) c_int { + var buf: Wide = undefined; + return _waccess(wide(&buf, std.mem.span(path)) orelse return -1, mode & 6); +} + +// ---- directories ---- + +pub const dirent = extern struct { + type: u8, + name: [1024]u8, +}; + +pub const DIR = struct { + find: ?HANDLE, + data: FindData, + first: bool, + fd: ?fd_t, + ent: dirent, +}; + +pub fn opendir(path: [*:0]const u8) ?*DIR { + var pattern: [4096]u8 = undefined; + const p = std.fmt.bufPrint(&pattern, "{s}\\*", .{std.mem.span(path)}) catch { + _ = fail(.NAMETOOLONG); + return null; + }; + var buf: Wide = undefined; + const w = wide(&buf, p) orelse return null; + const d = std.heap.c_allocator.create(DIR) catch { + _ = fail(.NOMEM); + return null; + }; + const h = FindFirstFileW(w, &d.data); + if (h == invalid_handle) { + const e = lastError(); + std.heap.c_allocator.destroy(d); + _ = fail(e); + return null; + } + d.* = .{ .find = @ptrFromInt(@as(usize, @bitCast(h))), .data = d.data, .first = true, .fd = null, .ent = undefined }; + return d; +} + +/// The stream owns `fd` from here: closedir releases it. +pub fn fdopendir(fd: fd_t) ?*DIR { + const path = dirPath(fd) orelse { + _ = fail(.BADF); + return null; + }; + const d = opendir(path.ptr) orelse return null; + d.fd = fd; + return d; +} + +pub fn readdir(d: *DIR) ?*dirent { + if (d.first) { + d.first = false; + } else if (FindNextFileW(d.find.?, &d.data) == 0) { + return null; + } + const name = std.mem.sliceTo(&d.data.name, 0); + const n = std.unicode.wtf16LeToWtf8(d.ent.name[0 .. d.ent.name.len - 1], name); + d.ent.name[n] = 0; + d.ent.type = if (d.data.attributes & 0x400 != 0) DT.LNK else if (d.data.attributes & 0x10 != 0) DT.DIR else DT.REG; + return &d.ent; +} + +pub fn closedir(d: *DIR) c_int { + if (d.find) |h| _ = FindClose(h); + if (d.fd) |fd| _ = close(fd); + std.heap.c_allocator.destroy(d); + return 0; +} + +// ---- process, environment, time ---- + +pub fn getpid() pid_t { + return _getpid(); +} + +pub fn getuid() u32 { + return 0; +} + +pub fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int { + if (overwrite == 0 and getenv(name) != null) return 0; + if (_putenv_s(name, value) != 0) return fail(.INVAL); + return 0; +} + +pub fn unsetenv(name: [*:0]const u8) c_int { + if (_putenv_s(name, "") != 0) return fail(.INVAL); + return 0; +} + +threadlocal var sleep_timer: ?HANDLE = null; + +/// High-resolution waitable timer per thread, so microsecond waits are not rounded to the 15.6 ms tick. +pub fn nanosleep(req: *const timespec, rem: ?*timespec) c_int { + if (rem) |r| r.* = .{ .sec = 0, .nsec = 0 }; + const ns: i64 = req.sec * std.time.ns_per_s + req.nsec; + if (ns <= 0) return 0; + if (sleep_timer == null) sleep_timer = CreateWaitableTimerExW(null, null, 0x2, 0x1F0003); // HIGH_RESOLUTION, TIMER_ALL_ACCESS + if (sleep_timer) |t| { + const due: i64 = -@divFloor(ns + 99, 100); // relative, in 100 ns units + if (SetWaitableTimer(t, &due, 0, null, null, 0) != 0) { + _ = WaitForSingleObject(t, 0xFFFF_FFFF); + return 0; + } + } + Sleep(@intCast(@max(1, @divFloor(ns + std.time.ns_per_ms - 1, std.time.ns_per_ms)))); + return 0; +} + +pub fn clock_gettime(clock: c_int, out: *timespec) c_int { + if (clock == CLOCK.MONOTONIC) { + var count: i64 = 0; + var freq: i64 = 0; + _ = QueryPerformanceCounter(&count); + _ = QueryPerformanceFrequency(&freq); + out.* = .{ .sec = @divFloor(count, freq), .nsec = @intCast(@divFloor(@mod(count, freq) * std.time.ns_per_s, freq)) }; + return 0; + } + var ft: FILETIME = undefined; + GetSystemTimePreciseAsFileTime(&ft); + const since_1601: i64 = @intCast((@as(u64, ft.high) << 32) | ft.low); + const unix = since_1601 - 116_444_736_000_000_000; + out.* = .{ .sec = @divFloor(unix, 10_000_000), .nsec = @intCast(@mod(unix, 10_000_000) * 100) }; + return 0; +} + +pub fn gettimeofday(tv: *timeval, tz: ?*anyopaque) c_int { + _ = tz; + var ts: timespec = undefined; + _ = clock_gettime(CLOCK.REALTIME, &ts); + tv.* = .{ .sec = ts.sec, .usec = @intCast(@divFloor(ts.nsec, 1000)) }; + return 0; +} + +/// Bytes free to this user under `dir` (statvfs's bavail * frsize). +pub fn diskAvailable(dir: [*:0]const u8) ?u64 { + var buf: Wide = undefined; + const w = wide(&buf, std.mem.span(dir)) orelse return null; + var avail: u64 = 0; + if (GetDiskFreeSpaceExW(w, &avail, null, null) == 0) return null; + return avail; +} + +/// NTFS keeps access in ACLs; mode bits are not stored, so a mode change succeeds and changes nothing. +pub fn fchmodat(dir: fd_t, path: [*:0]const u8, mode: mode_t, flags: u32) c_int { + _ = .{ dir, path, mode, flags }; + return 0; +} diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig new file mode 100644 index 0000000000..024e8bcc15 --- /dev/null +++ b/zig/src/platform/win_posix.zig @@ -0,0 +1,279 @@ +//! The std.posix surface the HTTP server calls through `sys.posix`, on Windows: Winsock sockets (SOCKET as an isize, +//! INVALID_SOCKET as -1), WSAPoll, console control events as SIGINT and SIGTERM, and no descriptor limits to raise. +//! Calls in `system` return -1 with errno set from WSAGetLastError, as libc's do. + +const std = @import("std"); +const c = @import("win_c.zig"); + +pub const E = c.E; +pub const errno = c.errno; +pub const fd_t = c.fd_t; +pub const AT = c.AT; +pub const S = c.S; + +// ---- sockets ---- + +pub const socket_t = isize; +pub const socklen_t = c_int; +pub const sockaddr = extern struct { family: u16, data: [14]u8 }; + +pub const AF = struct { + pub const UNIX: u32 = 1; + pub const INET: u32 = 2; + pub const INET6: u32 = 23; +}; +pub const SOCK = struct { + pub const STREAM: u32 = 1; +}; +pub const SOL = struct { + pub const SOCKET: i32 = 0xffff; +}; +pub const SO = struct { + /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, which the listener + /// sets REUSEADDR to prevent; exclusive use is the POSIX meaning. + pub const REUSEADDR: u32 = @bitCast(@as(i32, -5)); +}; +pub const SHUT = struct { + pub const RD: i32 = 0; + pub const WR: i32 = 1; + pub const RDWR: i32 = 2; +}; +pub const MSG = struct { + pub const PEEK: u32 = 0x2; + /// Winsock has no per-call nonblocking flag; every caller polls first, so the receive never waits. + pub const DONTWAIT: u32 = 0; +}; +pub const POLL = struct { + pub const IN: i16 = 0x0300; // POLLRDNORM | POLLRDBAND + /// WSAPoll refuses POLLPRI; out-of-band data is not a signal the server uses. + pub const PRI: i16 = 0; + pub const OUT: i16 = 0x0010; + pub const ERR: i16 = 0x0001; + pub const HUP: i16 = 0x0002; + pub const NVAL: i16 = 0x0004; +}; +pub const pollfd = extern struct { fd: socket_t, events: i16, revents: i16 }; + +const WSAData = extern struct { version: u16, high: u16, rest: [400]u8 }; +const ws = struct { + extern "ws2_32" fn WSAStartup(version: u16, data: *WSAData) callconv(.winapi) c_int; + extern "ws2_32" fn WSAGetLastError() callconv(.winapi) c_int; + extern "ws2_32" fn WSAPoll(fds: [*]pollfd, n: u32, timeout: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn socket(af: c_int, kind: c_int, proto: c_int) callconv(.winapi) isize; + extern "ws2_32" fn bind(s: isize, addr: *const anyopaque, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn listen(s: isize, backlog: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn connect(s: isize, addr: *const anyopaque, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn accept(s: isize, addr: ?*anyopaque, len: ?*c_int) callconv(.winapi) isize; + extern "ws2_32" fn getsockname(s: isize, addr: *anyopaque, len: *c_int) callconv(.winapi) c_int; + extern "ws2_32" fn setsockopt(s: isize, level: c_int, opt: c_int, val: [*]const u8, len: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn shutdown(s: isize, how: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn closesocket(s: isize) callconv(.winapi) c_int; + extern "ws2_32" fn send(s: isize, buf: [*]const u8, len: c_int, flags: c_int) callconv(.winapi) c_int; + extern "ws2_32" fn recv(s: isize, buf: [*]u8, len: c_int, flags: c_int) callconv(.winapi) c_int; +}; +extern "c" fn _errno() *c_int; + +var started = std.atomic.Value(u8).init(0); + +/// WSAStartup once per process, before the first socket call. +fn start() void { + if (started.load(.acquire) == 2) return; + if (started.cmpxchgStrong(0, 1, .acq_rel, .acquire) == null) { + var data: WSAData = undefined; + _ = ws.WSAStartup(0x0202, &data); + started.store(2, .release); + return; + } + while (started.load(.acquire) != 2) std.atomic.spinLoopHint(); +} + +/// -1 with errno from the thread's Winsock error. +fn failed() c_int { + const e: E = switch (ws.WSAGetLastError()) { + 10004 => .INTR, + 10013 => .ACCES, + 10014 => .FAULT, + 10022 => .INVAL, + 10024 => .MFILE, + 10035 => .AGAIN, + 10038 => .NOTSOCK, + 10048 => .ADDRINUSE, + 10049 => .ADDRNOTAVAIL, + 10050 => .NETDOWN, + 10053 => .CONNABORTED, + 10054 => .CONNRESET, + 10055 => .NOBUFS, + 10057 => .NOTCONN, + 10058 => .PIPE, + 10060 => .TIMEDOUT, + 10061 => .CONNREFUSED, + else => .IO, + }; + _errno().* = @intFromEnum(e); + return -1; +} + +fn len32(n: usize) c_int { + return @intCast(@min(n, std.math.maxInt(c_int))); +} + +pub const system = struct { + pub fn socket(family: u32, kind: u32, proto: u32) isize { + start(); + const fd = ws.socket(@intCast(family), @intCast(kind), @intCast(proto)); + return if (fd == -1) failed() else fd; + } + pub fn bind(fd: socket_t, addr: *const anyopaque, len: socklen_t) c_int { + return if (ws.bind(fd, addr, len) != 0) failed() else 0; + } + pub fn listen(fd: socket_t, backlog: u31) c_int { + return if (ws.listen(fd, backlog) != 0) failed() else 0; + } + pub fn accept(fd: socket_t, addr: ?*anyopaque, len: ?*socklen_t) isize { + const s = ws.accept(fd, addr, len); + return if (s == -1) failed() else s; + } + pub fn getsockname(fd: socket_t, addr: *anyopaque, len: *socklen_t) c_int { + return if (ws.getsockname(fd, addr, len) != 0) failed() else 0; + } + pub fn shutdown(fd: socket_t, how: i32) c_int { + return if (ws.shutdown(fd, how) != 0) failed() else 0; + } + /// Closes a socket; files close through `sys.c.close`. + pub fn close(fd: socket_t) c_int { + return if (ws.closesocket(fd) != 0) failed() else 0; + } + /// Sends on a socket. + pub fn write(fd: socket_t, buf: [*]const u8, len: usize) isize { + const n = ws.send(fd, buf, len32(len), 0); + return if (n < 0) failed() else n; + } + pub fn recvfrom(fd: socket_t, buf: [*]u8, len: usize, flags: u32, addr: ?*anyopaque, alen: ?*socklen_t) isize { + _ = addr; + _ = alen; + const n = ws.recv(fd, buf, len32(len), @intCast(flags)); + return if (n < 0) failed() else n; + } +}; + +pub fn setsockopt(fd: socket_t, level: i32, opt: u32, value: []const u8) error{SetSockOptFailed}!void { + if (ws.setsockopt(fd, level, @bitCast(opt), value.ptr, @intCast(value.len)) != 0) return error.SetSockOptFailed; +} + +/// WSAPoll, with POSIX's rule that an entry whose fd is negative is skipped (WSAPoll would fail the call). +pub fn poll(fds: []pollfd, timeout_ms: i32) error{PollFailed}!usize { + start(); + var live: [64]pollfd = undefined; + var where: [64]usize = undefined; + var n: usize = 0; + for (fds, 0..) |*f, i| { + f.revents = 0; + if (f.fd < 0) continue; + if (n == live.len) return error.PollFailed; + live[n] = f.*; + where[n] = i; + n += 1; + } + if (n == 0) { + if (timeout_ms > 0) Sleep(@intCast(timeout_ms)); + return 0; + } + const rc = ws.WSAPoll(&live, @intCast(n), timeout_ms); + if (rc < 0) return error.PollFailed; + for (live[0..n], where[0..n]) |f, i| fds[i].revents = f.revents; + return @intCast(rc); +} + +extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; + +/// Receives from a socket; 0 is the peer's orderly close. +pub fn read(fd: socket_t, buf: []u8) error{ConnectionResetByPeer}!usize { + const n = ws.recv(fd, buf.ptr, len32(buf.len), 0); + if (n < 0) return error.ConnectionResetByPeer; + return @intCast(n); +} + +// ---- signals: console control events ---- + +pub const SIG = enum(c_int) { + HUP = 1, + INT = 2, + PIPE = 13, + TERM = 15, + _, + + pub const IGN: ?Handler = @ptrFromInt(1); + pub const DFL: ?Handler = null; +}; +pub const Handler = *const fn (SIG) callconv(.c) void; +pub const sigset_t = u32; +pub const Sigaction = struct { + handler: extern union { handler: ?Handler, sigaction: ?*const anyopaque }, + mask: sigset_t, + flags: u32, +}; + +pub fn sigemptyset() sigset_t { + return 0; +} + +var on_int: ?Handler = null; +var on_term: ?Handler = null; +var console_hooked = std.atomic.Value(bool).init(false); + +extern "kernel32" fn SetConsoleCtrlHandler(routine: ?*const fn (u32) callconv(.winapi) c_int, add: c_int) callconv(.winapi) c_int; + +/// Ctrl+C and Ctrl+Break are SIGINT; closing the console, logoff and shutdown are SIGTERM. +fn consoleEvent(kind: u32) callconv(.winapi) c_int { + const h = switch (kind) { + 0, 1 => on_int, + else => on_term, + } orelse return 0; + if (h == SIG.IGN.?) return 1; + h(if (kind <= 1) .INT else .TERM); + return 1; +} + +/// SIGINT and SIGTERM through console events; Windows has no SIGHUP or SIGPIPE (a closed socket fails the send). +pub fn sigaction(sig: SIG, act: ?*const Sigaction, old: ?*Sigaction) void { + if (old) |o| o.* = .{ .handler = .{ .handler = SIG.DFL }, .mask = 0, .flags = 0 }; + const a = act orelse return; + switch (sig) { + .INT => on_int = a.handler.handler, + .TERM => on_term = a.handler.handler, + else => return, + } + if (!console_hooked.swap(true, .acq_rel)) _ = SetConsoleCtrlHandler(consoleEvent, 1); +} + +// ---- descriptor limits: Windows has none to raise ---- + +pub const rlim_t = u64; +pub const rlimit = struct { cur: rlim_t, max: rlim_t }; +pub const rlimit_resource = enum { NOFILE }; + +pub fn getrlimit(_: rlimit_resource) error{Unsupported}!rlimit { + return error.Unsupported; +} + +pub fn setrlimit(_: rlimit_resource, _: rlimit) error{Unsupported}!void { + return error.Unsupported; +} + +/// A connected stream pair over loopback TCP (Windows has no socketpair); for tests that drive a connection. +pub fn socketpair(pair: *[2]socket_t) error{SocketPair}!void { + start(); + const l = ws.socket(@intCast(AF.INET), @intCast(SOCK.STREAM), 0); + if (l == -1) return error.SocketPair; + defer _ = ws.closesocket(l); + var addr: extern struct { family: u16 = 2, port: u16 = 0, addr: [4]u8 = .{ 127, 0, 0, 1 }, zero: [8]u8 = @splat(0) } = .{}; + var len: c_int = @sizeOf(@TypeOf(addr)); + if (ws.bind(l, &addr, len) != 0 or ws.listen(l, 1) != 0 or ws.getsockname(l, &addr, &len) != 0) return error.SocketPair; + const a = ws.socket(@intCast(AF.INET), @intCast(SOCK.STREAM), 0); + if (a == -1) return error.SocketPair; + errdefer _ = ws.closesocket(a); + if (ws.connect(a, &addr, len) != 0) return error.SocketPair; + const b = ws.accept(l, null, null); + if (b == -1) return error.SocketPair; + pair.* = .{ a, b }; +} diff --git a/zig/src/server/auth.zig b/zig/src/server/auth.zig index 5457651dd9..b2274e2dc0 100644 --- a/zig/src/server/auth.zig +++ b/zig/src/server/auth.zig @@ -1,6 +1,7 @@ //! API keys held as SHA-256 digests with labels; a restricted key file reread on change or SIGHUP. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const builtin = @import("builtin"); const log = @import("log.zig"); const Allocator = std.mem.Allocator; @@ -84,10 +85,10 @@ pub const Store = struct { const path = s.path.?; const pathz = try s.gpa.dupeSentinel(u8, path, 0); defer s.gpa.free(pathz); - const fd = posix.openatZ(posix.AT.FDCWD, pathz, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true, .NONBLOCK = true }, 0) catch return error.Unsafe; - defer _ = posix.system.close(fd); + const fd = openKeys(pathz) catch return error.Unsafe; + defer _ = platform.c.close(fd); const st = meta(fd, null) orelse return error.Unsafe; - if ((st.mode & posix.S.IFMT) != posix.S.IFREG or st.mode & 0o044 != 0) { + if ((st.mode & platform.c.S.IFMT) != platform.c.S.IFREG or st.mode & 0o044 != 0) { problem.* = "API key file must be a regular file unreadable by other users; use chmod 600"; return error.Invalid; } @@ -101,7 +102,7 @@ pub const Store = struct { const text = try a.alloc(u8, (1 << 20) + 1); var got: usize = 0; while (got < text.len) { - const n = posix.read(fd, text[got..]) catch return error.Unsafe; + const n = readKeys(fd, text[got..]) catch return error.Unsafe; if (n == 0) break; got += n; } @@ -290,7 +291,7 @@ test "auth cleanup on every initialization allocation failure" { try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "keys", .data = "ops: sk-file\nsk-unlabelled\n" }); const path = try std.fmt.allocPrintSentinel(gpa, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); defer gpa.free(path); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); const Case = struct { fn run(a: Allocator, file: ?[]const u8) !void { var problem: []const u8 = ""; @@ -304,6 +305,19 @@ test "auth cleanup on every initialization allocation failure" { try std.testing.checkAllAllocationFailures(gpa, Case.run, .{@as(?[]const u8, path)}); } +/// The key file opened for reading, never through a link. +fn openKeys(pathz: [*:0]const u8) !platform.c.fd_t { + if (builtin.os.tag != .windows) return posix.openatZ(posix.AT.FDCWD, pathz, .{ .ACCMODE = .RDONLY, .NOFOLLOW = true, .NONBLOCK = true }, 0); + const fd = platform.c.open(pathz, .{ .ACCMODE = .RDONLY }); + return if (fd < 0) error.OpenFailed else fd; +} + +fn readKeys(fd: platform.c.fd_t, buf: []u8) !usize { + if (builtin.os.tag != .windows) return posix.read(fd, buf); + const n = platform.c.read(fd, buf.ptr, buf.len); + return if (n < 0) error.ReadFailed else @intCast(n); +} + /// A file's type and permission bits, size, identity and mtime: statx on Linux (its libc has no stat), stat elsewhere. const Meta = struct { mode: u32, size: u64, dev: u64, ino: u64, mtime_ns: i128 }; @@ -313,11 +327,23 @@ fn meta(fd: posix.fd_t, path: ?[*:0]const u8) ?Meta { const linux = std.os.linux; var sx: linux.Statx = undefined; const want: linux.STATX = .{ .TYPE = true, .MODE = true, .INO = true, .SIZE = true, .MTIME = true }; - const rc = if (path) |p| std.c.statx(posix.AT.FDCWD, p, 0, want, &sx) else std.c.statx(fd, "", 0x1000, want, &sx); // AT_EMPTY_PATH + const rc = if (path) |p| platform.c.statx(posix.AT.FDCWD, p, 0, want, &sx) else platform.c.statx(fd, "", 0x1000, want, &sx); // AT_EMPTY_PATH if (rc != 0) return null; const dev = (@as(u64, sx.dev_major) << 32) | sx.dev_minor; return .{ .mode = sx.mode, .size = sx.size, .dev = dev, .ino = sx.ino, .mtime_ns = @as(i128, sx.mtime.sec) * std.time.ns_per_s + sx.mtime.nsec }; } + if (builtin.os.tag == .windows) { + // NTFS keeps access in ACLs, not mode bits: the type and size are checked, and the profile's ACL is the guard + var st: platform.c.Stat = undefined; + const f = if (path) |p| platform.c.open(p, .{ .ACCMODE = .RDONLY }) else fd; + if (f < 0) return null; + defer if (path != null) { + _ = platform.c.close(f); + }; + if (platform.c.fstat(f, &st) != 0) return null; + const t = st.mtime(); + return .{ .mode = (st.mode & platform.c.S.IFMT) | 0o600, .size = @intCast(st.size), .dev = st.dev, .ino = st.ino, .mtime_ns = @as(i128, t.nsec) + @as(i128, t.sec) * std.time.ns_per_s }; + } var st: posix.Stat = undefined; const rc = if (path) |p| posix.system.stat(p, &st) else posix.system.fstat(fd, &st); if (posix.errno(rc) != .SUCCESS) return null; @@ -326,6 +352,7 @@ fn meta(fd: posix.fd_t, path: ?[*:0]const u8) ?Meta { } test "the key file: refused while others can read it, read at 0600, reread after an edit" { + if (builtin.os.tag == .windows) return error.SkipZigTest; // refusal by mode bits; NTFS keeps access in ACLs var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); @@ -334,11 +361,11 @@ test "the key file: refused while others can read it, read at 0600, reread after defer tmp.cleanup(); try tmp.dir.writeFile(io, .{ .sub_path = "keys", .data = "ops: sk-one\n" }); const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); var problem: []const u8 = ""; try std.testing.expectError(error.KeyFile, Store.init(a, &.{}, "", path, false, &problem)); try std.testing.expect(std.mem.indexOf(u8, problem, "chmod 600") != null); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); var s = try Store.init(std.testing.allocator, &.{"sk-cli"}, "sk-env", path, false, &problem); defer s.deinit(); try std.testing.expectEqualStrings("cli-1", s.match(a, "Bearer sk-cli", null).?); @@ -354,6 +381,7 @@ test "the key file: refused while others can read it, read at 0600, reread after } test "a request's key label outlives a reload of the key file" { + if (builtin.os.tag == .windows) return error.SkipZigTest; // refusal by mode bits; NTFS keeps access in ACLs const gpa = std.testing.allocator; const io = std.testing.io; var arena: std.heap.ArenaAllocator = .init(gpa); @@ -363,7 +391,7 @@ test "a request's key label outlives a reload of the key file" { defer tmp.cleanup(); try tmp.dir.writeFile(io, .{ .sub_path = "keys", .data = "ops: sk-one\n" }); const path = try std.fmt.allocPrintSentinel(a, ".zig-cache/tmp/{s}/keys", .{tmp.sub_path}, 0); - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o600, 0)); var problem: []const u8 = ""; var s = try Store.init(gpa, &.{}, "", path, false, &problem); defer if (s.file_arena) |*keys| keys.deinit(); @@ -373,7 +401,7 @@ test "a request's key label outlives a reload of the key file" { try std.testing.expectEqualStrings("ci", s.match(a, null, "sk-two").?); // another request rereads the file try std.testing.expectEqualStrings("ops", label); const kept = s.match(a, null, "sk-two").?; - try std.testing.expectEqual(@as(c_int, 0), std.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); + try std.testing.expectEqual(@as(c_int, 0), platform.c.fchmodat(posix.AT.FDCWD, path, 0o644, 0)); s.checked_ns = null; try std.testing.expect(s.match(a, "Bearer sk-one", null) == null); // a refused reload frees the file's keys too try std.testing.expectEqualStrings("ci", kept); diff --git a/zig/src/server/http_conn.zig b/zig/src/server/http_conn.zig index f29d9f8684..536e8649d7 100644 --- a/zig/src/server/http_conn.zig +++ b/zig/src/server/http_conn.zig @@ -1,6 +1,7 @@ //! One HTTP/1.1 connection as Python's BaseHTTPRequestHandler serves it: bounded heads, keep-alive, error pages. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const Allocator = std.mem.Allocator; const log = @import("log.zig"); diff --git a/zig/src/server/late_system_test.zig b/zig/src/server/late_system_test.zig index e661ae2fea..7f4622dcd1 100644 --- a/zig/src/server/late_system_test.zig +++ b/zig/src/server/late_system_test.zig @@ -1,5 +1,6 @@ //! A later system message is re-roled to user on Messages when the template cannot keep it, matching chat. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const server = @import("root.zig"); const json = server.json; @@ -132,10 +133,10 @@ const Capture = struct { const Reply = struct { status: u16, body: []u8 }; fn exchange(srv: *server.Server, path: []const u8, body: []const u8) !Reply { - var pair: [2]std.c.fd_t = undefined; - if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; - defer _ = std.c.close(pair[0]); - defer _ = std.c.close(pair[1]); + var pair: [2]platform.posix.socket_t = undefined; + if (@import("builtin").os.tag == .windows) try platform.posix.socketpair(&pair) else if (platform.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer _ = platform.posix.system.close(pair[0]); + defer _ = platform.posix.system.close(pair[1]); const head = try std.fmt.allocPrint(std.testing.allocator, "POST {s} HTTP/1.1\r\nHost: t\r\nContent-Length: {d}\r\n\r\n", .{ path, body.len }); defer std.testing.allocator.free(head); try writeAll(pair[0], head); @@ -157,23 +158,23 @@ fn exchange(srv: *server.Server, path: []const u8, body: []const u8) !Reply { return .{ .status = status, .body = payload }; } -fn writeAll(fd: std.c.fd_t, data: []const u8) !void { +fn writeAll(fd: platform.posix.socket_t, data: []const u8) !void { var sent: usize = 0; while (sent < data.len) { - const n = std.c.write(fd, data[sent..].ptr, data.len - sent); + const n = platform.posix.system.write(fd, data[sent..].ptr, data.len - sent); if (n <= 0) return error.Closed; sent += @intCast(n); } } -fn readReady(a: std.mem.Allocator, fd: std.c.fd_t) ![]u8 { +fn readReady(a: std.mem.Allocator, fd: platform.posix.socket_t) ![]u8 { var out: std.ArrayList(u8) = .empty; errdefer out.deinit(a); var tmp: [4096]u8 = undefined; while (true) { - var pollfd = [_]std.posix.pollfd{.{ .fd = fd, .events = std.posix.POLL.IN, .revents = 0 }}; - if (try std.posix.poll(&pollfd, 200) == 0) break; - const n = std.posix.read(fd, &tmp) catch |e| return e; + var pollfd = [_]platform.posix.pollfd{.{ .fd = fd, .events = platform.posix.POLL.IN, .revents = 0 }}; + if (try platform.posix.poll(&pollfd, 200) == 0) break; + const n = platform.posix.read(fd, &tmp) catch |e| return e; if (n == 0) break; try out.appendSlice(a, tmp[0..n]); } diff --git a/zig/src/server/listener.zig b/zig/src/server/listener.zig index 8eb7a27df6..a6492e737d 100644 --- a/zig/src/server/listener.zig +++ b/zig/src/server/listener.zig @@ -1,6 +1,7 @@ //! The listening socket: SO_REUSEADDR only (never REUSEPORT, so a second server on the port fails), backlog 128. const std = @import("std"); -const posix = std.posix; +const platform = @import("sys"); +const posix = platform.posix; const net = std.Io.net; const log = @import("log.zig"); const Threaded = std.Io.Threaded; diff --git a/zig/src/server/live.zig b/zig/src/server/live.zig index cb0cbaf640..3d6f8b482e 100644 --- a/zig/src/server/live.zig +++ b/zig/src/server/live.zig @@ -1,5 +1,6 @@ //! The live status line (connections, decode and prefill tok/s), redrawn in a terminal, and /health's copy of it. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const log = @import("log.zig"); @@ -94,7 +95,7 @@ fn columns() usize { if (os == .linux) { if (std.os.linux.ioctl(1, std.os.linux.T.IOCGWINSZ, @intFromPtr(&ws)) == 0 and ws.col > 0) return ws.col; } else if (os.isDarwin()) { - if (std.c.ioctl(1, @bitCast(@as(u32, 0x40087468)), &ws) == 0 and ws.col > 0) return ws.col; + if (platform.c.ioctl(1, @bitCast(@as(u32, 0x40087468)), &ws) == 0 and ws.col > 0) return ws.col; } return 100; } diff --git a/zig/src/server/log.zig b/zig/src/server/log.zig index 40430985a2..629bf90bee 100644 --- a/zig/src/server/log.zig +++ b/zig/src/server/log.zig @@ -1,5 +1,6 @@ //! Server output: ``[tensorfold]`` lines on stdout, and the live status line they clear first. const std = @import("std"); +const platform = @import("sys"); const Conn = @import("http_conn.zig").Conn; const afterChars = @import("reply_text.zig").afterChars; @@ -23,8 +24,8 @@ const clear = "\r\x1b[2K"; fn writeOut(bytes: []const u8) void { var off: usize = 0; while (off < bytes.len) { - const rc = std.posix.system.write(1, bytes[off..].ptr, bytes.len - off); - if (std.posix.errno(rc) != .SUCCESS or rc <= 0) return; + const rc = platform.c.write(1, bytes[off..].ptr, bytes.len - off); + if (platform.c.errno(rc) != .SUCCESS or rc <= 0) return; off += @intCast(rc); } } diff --git a/zig/src/server/request_log.zig b/zig/src/server/request_log.zig index ea4c8b7997..02c26a0c0d 100644 --- a/zig/src/server/request_log.zig +++ b/zig/src/server/request_log.zig @@ -1,5 +1,6 @@ //! TENSORFOLD_REQUEST_LOG: each chat and completion body appended as one JSON line, images redacted, for replays. const std = @import("std"); +const platform = @import("sys"); const json = @import("json.zig"); const Value = json.Value; const Allocator = std.mem.Allocator; @@ -33,8 +34,9 @@ pub fn append(a: Allocator, path: []const u8, body: Value) void { if (body.strField("priority")) |p| if (std.mem.eql(u8, p, "background")) return; const line = json.stringify(a, redact(a, body) catch return, .{}) catch return; const pathz = a.dupeSentinel(u8, path, 0) catch return; - const fd = std.posix.openatZ(std.posix.AT.FDCWD, pathz, .{ .ACCMODE = .WRONLY, .APPEND = true, .CREAT = true }, 0o644) catch return; - defer _ = std.posix.system.close(fd); + const fd = platform.c.open(pathz, .{ .ACCMODE = .WRONLY, .APPEND = true, .CREAT = true }, @as(platform.c.mode_t, 0o644)); + if (fd < 0) return; + defer _ = platform.c.close(fd); const framed = std.mem.concat(a, u8, &.{ line, "\n" }) catch return; - _ = std.posix.system.write(fd, framed.ptr, framed.len); + _ = platform.c.write(fd, framed.ptr, framed.len); } diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index da27cb6b0e..46a8415f3b 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -1,5 +1,6 @@ //! A server's life: keys, the listening socket, the startup lines, the live line, and a clean stop on SIGTERM or SIGINT. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const auth = @import("auth.zig"); @@ -10,7 +11,7 @@ const listener_mod = @import("listener.zig"); const model_text = @import("model_text.zig"); const server_mod = @import("server.zig"); const Allocator = std.mem.Allocator; -const posix = std.posix; +const posix = platform.posix; var stop_requested: std.atomic.Value(bool) = .init(false); @@ -33,7 +34,7 @@ pub fn stopping() bool { } /// The hard limit macOS reads as unlimited. -const rlim_infinity: u64 = std.c.RLIM.INFINITY; +const rlim_infinity: u64 = platform.c.RLIM.INFINITY; /// The targets to try, in order: the hard limit when finite, else the fallbacks 65536 then 10240, first that works; never an unlimited soft limit, since a loop over every descriptor would do absurd work. fn fillTargets(max: u64, out: *[2]u64) []const u64 { @@ -211,6 +212,7 @@ test "the targets are the hard limit, or the fallbacks under an unlimited one" { } test "the raise lifts a low soft limit, and a second raise does nothing" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // no descriptor limits to raise const before = try posix.getrlimit(.NOFILE); defer posix.setrlimit(.NOFILE, before) catch {}; if (@as(u64, @intCast(before.max)) <= 256) return error.SkipZigTest; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 011734ee9a..4d76d98105 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -1,5 +1,6 @@ //! The server: one engine and tokenizer behind the HTTP routes, one thread a connection. const std = @import("std"); +const platform = @import("sys"); const api = @import("engine_api"); const json = @import("json.zig"); const errors = @import("errors.zig"); @@ -280,7 +281,7 @@ pub const Server = struct { pub fn serve(srv: *Server, listener: listener_mod.Listener, stop: *const std.atomic.Value(bool)) void { while (listener.accept(stop)) |c| { const t = std.Thread.spawn(.{ .stack_size = 8 << 20 }, connection, .{ srv, c }) catch { - _ = std.posix.system.close(c.fd); + _ = platform.posix.system.close(c.fd); continue; }; t.detach(); @@ -290,7 +291,7 @@ pub const Server = struct { fn connection(srv: *Server, accepted: listener_mod.Listener.Accepted) void { _ = srv.open_connections.fetchAdd(1, .acq_rel); defer _ = srv.open_connections.fetchSub(1, .acq_rel); - defer _ = std.posix.system.close(accepted.fd); + defer _ = platform.posix.system.close(accepted.fd); var peer_buf: [64]u8 = undefined; const peer = listener_mod.peerText(&peer_buf, accepted.peer); var conn = http_conn.Conn.init(srv.gpa, accepted.fd, peer) catch return; diff --git a/zig/tests/cuda/pack_compare.zig b/zig/tests/cuda/pack_compare.zig index b1cc379a36..10f9041c65 100644 --- a/zig/tests/cuda/pack_compare.zig +++ b/zig/tests/cuda/pack_compare.zig @@ -1,5 +1,6 @@ //! tf-train-pack's checks on either backend: packed passes against one a pass, and passes after kept prefixes. const std = @import("std"); +const platform = @import("sys"); /// What one sequence's pass left: its loss, whether it came back, its pushes and its keys at the pocket's layers. const Seen = struct { loss: f32, recalled: bool, push: []f32, keys: []u16 }; @@ -169,7 +170,7 @@ pub fn For(comptime train: type) type { /// TF_PACK_DUMP: the packed target pass's losses and pushes written there, for another build's to compare. fn dump(gpa: std.mem.Allocator, io: std.Io, seen: []const Seen) !void { - const path = std.c.getenv("TF_PACK_DUMP") orelse return; + const path = platform.c.getenv("TF_PACK_DUMP") orelse return; var bytes: std.ArrayList(u8) = .empty; for (seen) |x| { try bytes.appendSlice(gpa, std.mem.asBytes(&x.loss)); diff --git a/zig/tests/cuda/train_pack.zig b/zig/tests/cuda/train_pack.zig index 2ab91949b9..d721c19887 100644 --- a/zig/tests/cuda/train_pack.zig +++ b/zig/tests/cuda/train_pack.zig @@ -1,5 +1,6 @@ //! Training passes holding several sequences against one a pass, and from kept prefixes against whole, on CUDA. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const nemotron = @import("nemotron"); const check_mod = @import("check.zig"); @@ -17,7 +18,7 @@ const open = 2; pub fn run(gpu: check_mod.Gpu, model: []const u8, path: []const u8, count: usize) !void { const gpa = gpu.gpa; const io = gpu.io; - const kernels = std.c.getenv("TENSORFOLD_CUDA_KERNELS"); + const kernels = platform.c.getenv("TENSORFOLD_CUDA_KERNELS"); const e = try nemotron.Engine.init(gpa, io, gpu.ctx, model, if (kernels) |k| std.mem.span(k) else null, .{ .context = 8192, .mtp = false, .graphs = false, .sampling = null, .segments = 1 }); defer e.deinit(); var living: ?nemotron.sites.Sites = null; diff --git a/zig/tests/fabric_library.zig b/zig/tests/fabric_library.zig index 04141e546a..2e09f6aefb 100644 --- a/zig/tests/fabric_library.zig +++ b/zig/tests/fabric_library.zig @@ -1,5 +1,6 @@ //! The fabric bound to MCDMA's own libmcdma-rpc (built from MCDMA's rpc/ by tools/zig/test_fabric.sh, MCDMA_RPC_LIBRARY). const std = @import("std"); +const platform = @import("sys"); const builtin = @import("builtin"); const fabric = @import("fabric"); const layout = fabric.layout; @@ -9,7 +10,7 @@ const mailbox = fabric.mailbox; const gpa = std.testing.allocator; fn open() !words.Library { - if (std.c.getenv("MCDMA_RPC_LIBRARY") == null) return error.SkipZigTest; + if (platform.c.getenv("MCDMA_RPC_LIBRARY") == null) return error.SkipZigTest; return words.Library.open(null); } diff --git a/zig/tests/flashnext_bench.zig b/zig/tests/flashnext_bench.zig index 7a0c73d528..d15901c68a 100644 --- a/zig/tests/flashnext_bench.zig +++ b/zig/tests/flashnext_bench.zig @@ -1,5 +1,6 @@ //! FZ_DBENCH: Flash Next dense and DeltaNet kernels at 1-16 rows against the recorded ones, bit for bit, then timed. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const fz = tf.flashnext_replay; @@ -238,7 +239,7 @@ pub fn run(r: *Run, m: *Model, prompt: []const u32, gpa: std.mem.Allocator, aren }; const pfs = [_]usize{ 1, 2 }; - const bits = std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_DBENCH").?), "bits"); + const bits = std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_DBENCH").?), "bits"); // bits: core lane against the recorded kernel, every layer, one Mac and speed-up layouts, every read-ahead depth const keep = try gpa.alloc(u8, fz.MAXR * fz.VOCAB * 2); defer gpa.free(keep); diff --git a/zig/tests/flashnext_reuse.zig b/zig/tests/flashnext_reuse.zig index eeeda53247..c9c66a3ca1 100644 --- a/zig/tests/flashnext_reuse.zig +++ b/zig/tests/flashnext_reuse.zig @@ -1,5 +1,6 @@ //! Flash Next prompt reuse against fresh passes: agent-like turns, every depth, copy drafts off and on, state bytes. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const api = @import("engine_api"); @@ -121,9 +122,9 @@ pub fn main(init: std.process.Init) !void { } const pool = mtl.objc.Pool.push(); defer pool.pop(); - const n_out: usize = if (std.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 64; + const n_out: usize = if (platform.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 64; var depths: std.ArrayList(?usize) = .empty; - var dit = std.mem.tokenizeScalar(u8, if (std.c.getenv("FZ_DEPTHS")) |v| std.mem.span(v) else "1,2,3,5,0", ','); + var dit = std.mem.tokenizeScalar(u8, if (platform.c.getenv("FZ_DEPTHS")) |v| std.mem.span(v) else "1,2,3,5,0", ','); while (dit.next()) |d| { const k = try std.fmt.parseInt(usize, d, 10); try depths.append(a, if (k == 0) null else k); diff --git a/zig/tests/flashnext_run.zig b/zig/tests/flashnext_run.zig index 87efe89c96..5c1e1c3234 100644 --- a/zig/tests/flashnext_run.zig +++ b/zig/tests/flashnext_run.zig @@ -2,6 +2,7 @@ //! recorded variant, every weight the pack's or the checkpoint's. Checks one-row greedy tokens and every row of the //! Python engine's drafted windows (with rollback), then times each window size. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); const tf = @import("tensorfold"); @@ -46,13 +47,13 @@ const jsonInt = fz.jsonInt; fn armName(seg_ab: bool, arm: usize) []const u8 { if (seg_ab) return if (arm == 1) "segments on " else "segments off"; - if (std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "fz")) return if (arm == 1) "fz kernels" else "recorded "; + if (platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "fz")) return if (arm == 1) "fz kernels" else "recorded "; return if (arm == 1) "copy on " else "copy off"; } /// FZ_PROMPT_N: only the prompt's first n tokens (all when unset). fn promptCut() !usize { - const v = std.c.getenv("FZ_PROMPT_N") orelse return std.math.maxInt(usize); + const v = platform.c.getenv("FZ_PROMPT_N") orelse return std.math.maxInt(usize); return std.fmt.parseInt(usize, std.mem.span(v), 10); } @@ -124,7 +125,7 @@ pub fn main(init: std.process.Init) !void { defer arena_state.deinit(); const arena = arena_state.allocator(); const device = try mtl.Device.init(); - if (std.c.getenv("FZ_COMPILE_CHECK") != null) { // every prompt-kernel source and the 6-bit file on this Metal compiler + if (platform.c.getenv("FZ_COMPILE_CHECK") != null) { // every prompt-kernel source and the 6-bit file on this Metal compiler var failed: usize = 0; var files: std.ArrayList(ks.File) = .empty; try files.appendSlice(gpa, &ks.prefill); @@ -149,8 +150,8 @@ pub fn main(init: std.process.Init) !void { } std.process.exit(if (failed == 0) 0 else 1); } - if (std.c.getenv("FZ_OPORDER") != null) return @import("flashnext_bench.zig").opOrder(device); // no model loaded - if (std.c.getenv("FZ_MMA_PEAK") != null) { // the tensor units' rate on register fragments: 16x32x16 ops a second + if (platform.c.getenv("FZ_OPORDER") != null) return @import("flashnext_bench.zig").opOrder(device); // no model loaded + if (platform.c.getenv("FZ_MMA_PEAK") != null) { // the tensor units' rate on register fragments: 16x32x16 ops a second const alib = try mtl.Library.fromSource(device, try fz.frags.source(device, arena, ks.flashnext_attn), mtl.CompileOptions.mlx()); const pipe = try mtl.Pipeline.init(device, alib, "tf_mma_peak", false); const queue = try device.queue(); @@ -174,9 +175,9 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_ENGINE") != null) { // the served engine: its own load, warm-up and reply, on FZ_REF's prompt + if (platform.c.getenv("FZ_ENGINE") != null) { // the served engine: its own load, warm-up and reply, on FZ_REF's prompt const fx = @import("tensorfold").flashnext_engine; - const rp = std.mem.span(std.c.getenv("FZ_REF") orelse return error.NoRef); + const rp = std.mem.span(platform.c.getenv("FZ_REF") orelse return error.NoRef); const rf = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}", .{rp}, 0)); const rj = try std.json.parseFromSliceLeaky(std.json.Value, arena, rf.bytes[0..rf.size], .{}); const all_items = rj.object.get("prompt").?.array.items; @@ -184,9 +185,9 @@ pub fn main(init: std.process.Init) !void { const toks = try arena.alloc(u32, items.len); for (items, 0..) |x, i| toks[i] = @intCast(x.integer); const e = try fx.Engine.load(gpa, init.io, args[1], args[2]); - if (std.c.getenv("FZ_COPY_MIN")) |v| e.copy_min = try std.fmt.parseInt(u32, std.mem.span(v), 10); - if (std.c.getenv("FZ_COPY_LONG")) |v| e.copy_long = try std.fmt.parseInt(u32, std.mem.span(v), 10); - if (std.c.getenv("FZ_ENGINE_WARM") != null) try e.warm(); + if (platform.c.getenv("FZ_COPY_MIN")) |v| e.copy_min = try std.fmt.parseInt(u32, std.mem.span(v), 10); + if (platform.c.getenv("FZ_COPY_LONG")) |v| e.copy_long = try std.fmt.parseInt(u32, std.mem.span(v), 10); + if (platform.c.getenv("FZ_ENGINE_WARM") != null) try e.warm(); const Show = struct { got: std.ArrayList(u32) = .empty, a: std.mem.Allocator, @@ -200,11 +201,11 @@ pub fn main(init: std.process.Init) !void { return false; } }; - const n_out: usize = if (std.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 256; + const n_out: usize = if (platform.c.getenv("FZ_N")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 256; var replies: [2][]u32 = undefined; - const seg_ab = std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "seg"); + const seg_ab = platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "seg"); var depths: std.ArrayList(?usize) = .empty; // FZ_DEPTHS=3,5,7: each depth's arms in turn (0: the depth rule, p: plain) - if (std.c.getenv("FZ_DEPTHS")) |v| { + if (platform.c.getenv("FZ_DEPTHS")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |d| { if (std.mem.eql(u8, d, "p")) { @@ -214,9 +215,9 @@ pub fn main(init: std.process.Init) !void { const n = try std.fmt.parseInt(usize, d, 10); try depths.append(arena, if (n == 0) null else n); } - } else try depths.append(arena, if (std.c.getenv("FZ_DEPTH")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else null); + } else try depths.append(arena, if (platform.c.getenv("FZ_DEPTH")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else null); // FZ_AB=fz (with FZ_LANE=1 FZ_GDN=2): recorded dense and DeltaNet kernels, then core lane, fz_gdn, kept states - const fz_ab = std.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(std.c.getenv("FZ_AB").?), "fz"); + const fz_ab = platform.c.getenv("FZ_AB") != null and std.mem.eql(u8, std.mem.span(platform.c.getenv("FZ_AB").?), "fz"); const fz_on = .{ e.r.lane_new, e.r.gdn_pipe, e.r.gdn_kept }; for (depths.items) |depth| for (0..2) |arm| { // copy drafts off, then on (FZ_AB=seg: staggered prompt segments off, then on) const plain = depth != null and depth.? == 0; // one token a round: one arm (copies need drafts) @@ -256,7 +257,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("{s}: reply hash {x:0>16}\n", .{ name, std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(sh.got.items)) }); if (arm == 1) for (res.copy_by_len, 0..) |cl, n| if (cl[0] > 0) std.debug.print(" match {d}: {d} rounds, {d:.2} landed\n", .{ n, cl[0], @as(f64, @floatFromInt(cl[1])) / @as(f64, @floatFromInt(cl[0])) }); }; - if (std.c.getenv("FZ_NO_REF") != null) return; // the reply hashes are the check (two Macs against one) + if (platform.c.getenv("FZ_NO_REF") != null) return; // the reply hashes are the check (two Macs against one) // the reference: the same prompt chunks, then one row a step e.hostMode(); const mm = e.m; @@ -287,28 +288,28 @@ pub fn main(init: std.process.Init) !void { // the replay below reads every expert, embedding and table at 6/32: other widths are the engine's (Engine.load) if (!try fz.engine_widths.uniformDir(arena, init.io, args[1])) return error.HarnessNeeds6Bit32; var r = Run{ .arena = arena, .device = device, .queue = try device.queue() }; - r.fused_xsum = std.c.getenv("FZ_FUSED_XSUM") != null; - r.serial = std.c.getenv("FZ_SERIAL") != null; - r.xnew = std.c.getenv("FZ_XNEW") != null; - r.split = std.c.getenv("FZ_SPLIT") != null; - r.gdn_step = std.c.getenv("FZ_GDN_STEP") != null; - r.hc_mma = std.c.getenv("FZ_HC_MMA") != null; - r.hc_up = std.c.getenv("FZ_HCCHECK") != null or if (std.c.getenv("FZ_HC_UP")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else false; + r.fused_xsum = platform.c.getenv("FZ_FUSED_XSUM") != null; + r.serial = platform.c.getenv("FZ_SERIAL") != null; + r.xnew = platform.c.getenv("FZ_XNEW") != null; + r.split = platform.c.getenv("FZ_SPLIT") != null; + r.gdn_step = platform.c.getenv("FZ_GDN_STEP") != null; + r.hc_mma = platform.c.getenv("FZ_HC_MMA") != null; + r.hc_up = platform.c.getenv("FZ_HCCHECK") != null or if (platform.c.getenv("FZ_HC_UP")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else false; r.event = try device.sharedEvent(); - r.dense = r.xnew and std.c.getenv("FZ_DENSE") != null; - r.dense_target = r.dense and std.c.getenv("FZ_DENSE_TARGET") != null; - r.xpack = r.xnew and std.c.getenv("FZ_XPACK") != null; - r.grouped = r.xnew and !r.xpack and std.c.getenv("FZ_GROUPED") != null; - r.xfused = r.xnew and !r.xpack and !r.grouped and std.c.getenv("FZ_XFUSED") != null; - r.xsx = r.xnew and std.c.getenv("FZ_XSX") != null; - r.copy = std.c.getenv("FZ_COPY") != null; - r.prefetch = !r.serial and std.c.getenv("FZ_PREFETCH") != null; + r.dense = r.xnew and platform.c.getenv("FZ_DENSE") != null; + r.dense_target = r.dense and platform.c.getenv("FZ_DENSE_TARGET") != null; + r.xpack = r.xnew and platform.c.getenv("FZ_XPACK") != null; + r.grouped = r.xnew and !r.xpack and platform.c.getenv("FZ_GROUPED") != null; + r.xfused = r.xnew and !r.xpack and !r.grouped and platform.c.getenv("FZ_XFUSED") != null; + r.xsx = r.xnew and platform.c.getenv("FZ_XSX") != null; + r.copy = platform.c.getenv("FZ_COPY") != null; + r.prefetch = !r.serial and platform.c.getenv("FZ_PREFETCH") != null; const t0 = mtl.clock.seconds(); try r.compile(args[2]); r.sel = try Select.init(&r, MAXR); - if (std.c.getenv("FZ_CATCH_CHECK") != null) return catchCheck(&r, arena); - r.lane_new = r.xnew and std.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (same bits) - if (r.xnew and std.c.getenv("FZ_GDN") != null) r.gdn_pipe = try fz.gdn_step.compile(&r, false); + if (platform.c.getenv("FZ_CATCH_CHECK") != null) return catchCheck(&r, arena); + r.lane_new = r.xnew and platform.c.getenv("FZ_LANE") != null; // core lane and fz_gdn on the target (same bits) + if (r.xnew and platform.c.getenv("FZ_GDN") != null) r.gdn_pipe = try fz.gdn_step.compile(&r, false); const t1 = mtl.clock.seconds(); const index_file = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors.index.json", .{args[1]}, 0)); @@ -320,8 +321,8 @@ pub fn main(init: std.process.Init) !void { while (fit.next()) |name| try r.indexFile(try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ args[1], name.* }, 0)); try r.indexFile(try std.fmt.allocPrintSentinel(arena, "{s}/pack.safetensors", .{args[2]}, 0)); - const own_prompt = std.c.getenv("FZ_REF") != null; // another prompt: passes against Python's records are skipped - const ref_path = if (std.c.getenv("FZ_REF")) |p| try std.fmt.allocPrintSentinel(arena, "{s}", .{std.mem.span(p)}, 0) else try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{args[2]}, 0); + const own_prompt = platform.c.getenv("FZ_REF") != null; // another prompt: passes against Python's records are skipped + const ref_path = if (platform.c.getenv("FZ_REF")) |p| try std.fmt.allocPrintSentinel(arena, "{s}", .{std.mem.span(p)}, 0) else try std.fmt.allocPrintSentinel(arena, "{s}/ref.json", .{args[2]}, 0); const ref_file = try mtl.MappedFile.open(ref_path); const ref = try std.json.parseFromSliceLeaky(std.json.Value, arena, ref_file.bytes[0..ref_file.size], .{}); const ple_ref = ref.object.get("ple").?.object; @@ -502,15 +503,15 @@ pub fn main(init: std.process.Init) !void { const t2 = mtl.clock.seconds(); std.debug.print("compiled in {d:.2} s, loaded {d:.1} GB in {d:.1} s\n", .{ t1 - t0, @as(f64, @floatFromInt(r.loaded)) / 1e9, t2 - t1 }); - if (std.c.getenv("FZ_DBENCH") != null) { // dense classes and DeltaNet timed by rows (flashnext_bench.zig) + if (platform.c.getenv("FZ_DBENCH") != null) { // dense classes and DeltaNet timed by rows (flashnext_bench.zig) var toks: [MAXR]u32 = undefined; for (0..MAXR) |i| toks[i] = @intCast(ref.object.get("prompt").?.array.items[i].integer); return @import("flashnext_bench.zig").run(&r, m, &toks, gpa, arena); } - if (std.c.getenv("FZ_PROFILE") != null) { - if (std.c.getenv("TF_FLASHNEXT_TP")) |path| r.tp = try fz.Tp2.init(arena, r.device, std.mem.span(path)); + if (platform.c.getenv("FZ_PROFILE") != null) { + if (platform.c.getenv("TF_FLASHNEXT_TP")) |path| r.tp = try fz.Tp2.init(arena, r.device, std.mem.span(path)); var hcskips: std.ArrayList(u32) = .empty; // FZ_HCSKIP=1,2,4: the profile again under each hc knock-out - if (std.c.getenv("FZ_HCSKIP")) |v| { + if (platform.c.getenv("FZ_HCSKIP")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try hcskips.append(arena, try std.fmt.parseInt(u32, x, 10)); } else try hcskips.append(arena, 0); @@ -520,7 +521,7 @@ pub fn main(init: std.process.Init) !void { var pk: [MAXR]u32 = undefined; m.reset(); for (0..3) |_| try m.window(toks[0..1], &pk); - if (std.c.getenv("FZ_PROFILE_AB") != null) { // in-model A/B, interleaved: recorded, core lane, fz_gdn, both + if (platform.c.getenv("FZ_PROFILE_AB") != null) { // in-model A/B, interleaved: recorded, core lane, fz_gdn, both const gpipe = try fz.gdn_step.compile(&r, false); const arms = [_][2]bool{ .{ false, false }, .{ true, false }, .{ false, true }, .{ true, true } }; for ([_]usize{ 1, 2, 4, 6, 8, 16 }) |rows| { @@ -537,12 +538,12 @@ pub fn main(init: std.process.Init) !void { return; } var ups: std.ArrayList(bool) = .empty; // FZ_HC_UP=0,1: the profile with the recorded up projection, then fz_hc_up - if (std.c.getenv("FZ_HC_UP")) |v| { + if (platform.c.getenv("FZ_HC_UP")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try ups.append(arena, !std.mem.eql(u8, x, "0")); } else try ups.append(arena, false); var prow: std.ArrayList(usize) = .empty; // FZ_PROFILE_ROWS=2,3: the windows' widths (default 1, 4, 8) - if (std.c.getenv("FZ_PROFILE_ROWS")) |v| { + if (platform.c.getenv("FZ_PROFILE_ROWS")) |v| { var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |x| try prow.append(arena, try std.fmt.parseInt(usize, x, 10)); } else try prow.appendSlice(arena, &.{ 1, 4, 8 }); @@ -551,7 +552,7 @@ pub fn main(init: std.process.Init) !void { r.hcskip = hs; var base: f64 = 0; for (names, 0..) |name, c| { - const only = std.c.getenv("FZ_PROFILE_ONLY") != null; // every class but this one knocked out + const only = platform.c.getenv("FZ_PROFILE_ONLY") != null; // every class but this one knocked out r.skip = if (c == 0) (if (only) 0x1ff else 0) else if (only) 0x1ff & ~(@as(u32, 1) << @intCast(c - 1)) else @as(u32, 1) << @intCast(c - 1); m.gpu_seconds = 0; for (0..20) |_| try m.window(toks[0..rows], &pk); @@ -573,7 +574,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_DUAL") != null) { // two lane groups on two queues: one window, two in a row, two at once + if (platform.c.getenv("FZ_DUAL") != null) { // two lane groups on two queues: one window, two in a row, two at once const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -648,7 +649,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_HCCHECK") != null) { // fz_hc_up against the recorded up projection: every row's logits and streams, bit for bit + if (platform.c.getenv("FZ_HCCHECK") != null) { // fz_hc_up against the recorded up projection: every row's logits and streams, bit for bit const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -687,7 +688,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("hc check: {s}\n", .{if (bad == 0) "bit-identical" else "DIFFERENT"}); return; } - if (std.c.getenv("FZ_GCHECK") != null) { // grouped or fused experts against fz_xgu/fz_xdown: every row's logits, bit for bit + if (platform.c.getenv("FZ_GCHECK") != null) { // grouped or fused experts against fz_xgu/fz_xdown: every row's logits, bit for bit const want0 = ref.object.get("tokens").?.array.items; const pr = ref.object.get("prompt").?.array.items; var pk: [MAXR]u32 = undefined; @@ -714,12 +715,12 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_XTIME") != null) { // the experts' share of a window: chain rows against one token repeated + if (platform.c.getenv("FZ_XTIME") != null) { // the experts' share of a window: chain rows against one token repeated const want0 = ref.object.get("tokens").?.array.items; var pk: [MAXR]u32 = undefined; m.reset(); for (0..3) |_| try m.window(&.{@intCast(want0[0].integer)}, &pk); - if (std.c.getenv("FZ_GPARTS") != null) for ([_]usize{ 2, 8 }) |rows| { + if (platform.c.getenv("FZ_GPARTS") != null) for ([_]usize{ 2, 8 }) |rows| { var toks: [MAXR]u32 = undefined; for (0..rows) |i| toks[i] = @intCast(want0[i].integer); for ([_]u32{ 0, 1, 4, 6 }) |gs| { @@ -747,21 +748,21 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_PREFILL") != null) { // the prompt path against the exact 8-row windows, then one-row decode + if (platform.c.getenv("FZ_PREFILL") != null) { // the prompt path against the exact 8-row windows, then one-row decode const all_items = ref.object.get("prompt").?.array.items; const pr_items = all_items[0..@min(all_items.len, try promptCut())]; // FZ_PROMPT_N: the prompt's first n tokens const toks = try gpa.alloc(u32, pr_items.len); defer gpa.free(toks); for (pr_items, 0..) |x, i| toks[i] = @intCast(x.integer); - if (std.c.getenv("FZ_SATTN_CHECK") != null) { // prompt attention: the decode's kernels vs the tensor units + if (platform.c.getenv("FZ_SATTN_CHECK") != null) { // prompt attention: the decode's kernels vs the tensor units var pr = try Prompt.init(&r, args[2], r.xnew_header); - if (std.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); - const ab0: []const u8 = if (std.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; + if (platform.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); + const ab0: []const u8 = if (platform.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; const seg2 = std.mem.eql(u8, ab0, "seg2"); // seg2: one chunk at a time against staggered segments (core/segments.zig): FZ_SEGS=2 (default) the // engine's rule (two segments of at least FZ_SEG_MIN rows), 1 one segment, 3-4 that many of up to a chunk - const nseg: usize = if (std.c.getenv("FZ_SEGS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 2; - const seg_min: usize = if (std.c.getenv("FZ_SEG_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 1; + const nseg: usize = if (platform.c.getenv("FZ_SEGS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 2; + const seg_min: usize = if (platform.c.getenv("FZ_SEG_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 1; if (seg2 and (nseg == 0 or nseg > tf.segments.MAX)) return error.Segments; var extra: [3]Prompt = undefined; var ps: [4]*Prompt = undefined; @@ -773,7 +774,7 @@ pub fn main(init: std.process.Init) !void { var outs: [2][48]u32 = undefined; const logits: [2][]u16 = .{ try gpa.alloc(u16, VOCAB), try gpa.alloc(u16, VOCAB) }; var best: [2]f64 = .{ 1e9, 1e9 }; - const ab: []const u8 = if (std.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; + const ab: []const u8 = if (platform.c.getenv("FZ_AB")) |v| std.mem.span(v) else "attn"; for (0..2) |arm| { // arm 0 the old kernels, arm 1 the new: FZ_AB attn (default), scan or tiles if (std.mem.eql(u8, ab, "scan")) pr.scan4 = arm == 1 else if (std.mem.eql(u8, ab, "tiles")) pr.tall_tiles = arm == 1 else if (std.mem.eql(u8, ab, "gu")) pr.fused_gu = arm == 1 else if (std.mem.eql(u8, ab, "bm32")) pr.expert_bm = if (arm == 1) 32 else 0 else if (std.mem.eql(u8, ab, "scan8")) pr.scan8 = arm == 1 else if (std.mem.eql(u8, ab, "seg2")) {} else pr.fast_attn = arm == 1; for (0..3) |run| { @@ -830,7 +831,7 @@ pub fn main(init: std.process.Init) !void { const a0 = mtl.clock.seconds(); var at: usize = 0; var last_n: usize = 1; - const ref_upto: usize = if (std.c.getenv("FZ_PF_REF")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else toks.len; + const ref_upto: usize = if (platform.c.getenv("FZ_PF_REF")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else toks.len; while (at < @min(ref_upto, toks.len)) { const n: usize = @min(MAXR, toks.len - at); try m.window(toks[at .. at + n], &pk); @@ -865,7 +866,7 @@ pub fn main(init: std.process.Init) !void { } // the prompt path var pr = try Prompt.init(&r, args[2], r.xnew_header); - if (std.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); + if (platform.c.getenv("FZ_STEP")) |v| pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); m.reset(); const b0 = mtl.clock.seconds(); var first: u32 = 0; @@ -927,7 +928,7 @@ pub fn main(init: std.process.Init) !void { pr.skip = 0; return; } - if (std.c.getenv("FZ_QMM6") != null) { // 6-bit tensor-unit projections for prompt chunks: checked and timed + if (platform.c.getenv("FZ_QMM6") != null) { // 6-bit tensor-unit projections for prompt chunks: checked and timed const lib = try mtl.Library.fromSource(device, try fz.frags.source(device, arena, ks.flashnext_qmm6), mtl.CompileOptions.mlx()); const qmm_pipe = try mtl.Pipeline.init(device, lib, "tf_qmm6_t_nax", false); const off_pipe = try mtl.Pipeline.init(device, lib, "tf_expert_offsets6", false); @@ -1061,7 +1062,7 @@ pub fn main(init: std.process.Init) !void { } return; } - if (std.c.getenv("FZ_OVERLAP") != null) { // how many experts a window's rows share, layer by layer + if (platform.c.getenv("FZ_OVERLAP") != null) { // how many experts a window's rows share, layer by layer const pr = ref.object.get("prompt").?.array.items; const want0 = ref.object.get("tokens").?.array.items; r.ar = .{ .b = try r.buffer(64) }; @@ -1148,10 +1149,10 @@ pub fn main(init: std.process.Init) !void { var pick: [MAXR]u32 = undefined; // FZ_LONG: prompts of any length through the prompt chunks (reference and GPU-side rounds alike) - const long_mode = std.c.getenv("FZ_LONG") != null; - if (std.c.getenv("FZ_PRE1") != null) try m.window(&.{1000}, &pick); // a decode window before the prompt chunks + const long_mode = platform.c.getenv("FZ_LONG") != null; + if (platform.c.getenv("FZ_PRE1") != null) try m.window(&.{1000}, &pick); // a decode window before the prompt chunks var pr_long: ?Prompt = if (long_mode) try Prompt.init(&r, args[2], r.xnew_header) else null; - if (pr_long) |*pr| if (std.c.getenv("FZ_STEP")) |v| { + if (pr_long) |*pr| if (platform.c.getenv("FZ_STEP")) |v| { pr.step = try std.fmt.parseInt(usize, std.mem.span(v), 10); }; const ptoks = try gpa.alloc(u32, prompt.len); @@ -1190,20 +1191,20 @@ pub fn main(init: std.process.Init) !void { const steps: f64 = @floatFromInt(want.len - 1); std.debug.print("one-row reference starts {any}\n", .{got.items[0..@min(24, got.items.len)]}); std.debug.print("one row: {d}/{d} tokens equal to Python's; {d:.1} tok/s ({d:.2} ms a step, GPU {d:.2} ms)\n", .{ same, want.len, steps / (t4 - t3), (t4 - t3) * 1e3 / steps, m.gpu_seconds * 1e3 / steps }); - if (std.c.getenv("FZ_ONLY1") != null) { + if (platform.c.getenv("FZ_ONLY1") != null) { if (same < want.len) std.debug.print("first difference at token {d}: got {d}, Python {d}\n", .{ same, got.items[same], want[same].integer }); return; } // 2. the Python engine's drafted windows, every row's pick, with rollback m.reset(); - if (!long_mode and std.c.getenv("FZ_ONLY17") == null) for (prompt) |tok| { + if (!long_mode and platform.c.getenv("FZ_ONLY17") == null) for (prompt) |tok| { try m.window(&.{@intCast(tok.integer)}, &pick); m.keepRows(&.{@intCast(tok.integer)}, 1); }; var bad: usize = 0; var total_rows: usize = 0; - const skip_mid = long_mode or std.c.getenv("FZ_ONLY17") != null; // the one-row reference, then GPU-side rounds only + const skip_mid = long_mode or platform.c.getenv("FZ_ONLY17") != null; // the one-row reference, then GPU-side rounds only const rounds = if (own_prompt or skip_mid) &[_]std.json.Value{} else ref.object.get("rounds").?.array.items; for (rounds, 0..) |round, ri| { const o = round.object; @@ -1310,7 +1311,7 @@ pub fn main(init: std.process.Init) !void { const adapt_story = [_][3]usize{ .{ 3, 3, 0 }, .{ 0, 7, 9 } }; const adapt_own = [_][3]usize{ .{ 3, 3, 0 }, .{ 6, 6, 0 }, .{ 0, 7, 9 } }; const adapt: []const [3]usize = if (own_prompt) &adapt_own else &adapt_story; - const copy_min: usize = if (std.c.getenv("FZ_COPY_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 3; + const copy_min: usize = if (platform.c.getenv("FZ_COPY_MIN")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 3; var hist: std.ArrayList(u32) = .empty; for (prompt) |x| try hist.append(gpa, @intCast(x.integer)); const n_prompt = hist.items.len; @@ -1441,7 +1442,7 @@ pub fn main(init: std.process.Init) !void { // N+1 while round N runs and reads the emitted tokens from a ring. Depth 0: the depth rule picks each round's // drafts from the rounds read back so far (the window two rounds behind the one being encoded). var depths7: std.ArrayList(usize) = .empty; - if (std.c.getenv("FZ_D7")) |v| { // e.g. FZ_D7=6,0 + if (platform.c.getenv("FZ_D7")) |v| { // e.g. FZ_D7=6,0 var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); while (it.next()) |d| try depths7.append(gpa, try std.fmt.parseInt(usize, d, 10)); } else try depths7.appendSlice(gpa, &.{ 3, 6, 0 }); diff --git a/zig/tests/flashnext_think.zig b/zig/tests/flashnext_think.zig index 283d217b36..8c7bbe0e0b 100644 --- a/zig/tests/flashnext_think.zig +++ b/zig/tests/flashnext_think.zig @@ -1,5 +1,6 @@ //! tf-flashnext-think [CASES] [SEED]: Flash Next's verdict kernels under thinking budgets against the lane core. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -13,8 +14,8 @@ pub fn main(init: std.process.Init) !void { defer dev.deinit(); var g = try tf.flashnext_think_check.Rig.init(dev); defer g.deinit(); - const t0 = std.c.mach_absolute_time(); + const t0 = platform.c.mach_absolute_time(); const sum = try tf.flashnext_think_check.run(init.gpa, &g, cases, seed); - const s = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) * 125 / 3 / 1e9; // mach ticks are 125/3 ns on Apple silicon + const s = @as(f64, @floatFromInt(platform.c.mach_absolute_time() - t0)) * 125 / 3 / 1e9; // mach ticks are 125/3 ns on Apple silicon std.debug.print("{s}: {d} cases (seed {d}), plain and drafted == the lane core in each; closed at the budget {d}, think end first {d}; {d:.1} s\n", .{ std.mem.span(dev.name()), sum.cases, seed, sum.closed, sum.ended_first, s }); } diff --git a/zig/tests/glm_attn_probe.zig b/zig/tests/glm_attn_probe.zig index c6bc8ac6e6..e86a001d61 100644 --- a/zig/tests/glm_attn_probe.zig +++ b/zig/tests/glm_attn_probe.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash's dense latent attention (forward.attendDense) on given inputs, for a bit check against a reference. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -57,8 +58,8 @@ pub fn main(init: std.process.Init) !void { try out.appendSlice(arena, bufs[3].contents()[0 .. 64 * n * 2]); try out.appendSlice(arena, bufs[4].contents()[0..qb]); } - const file = std.c.fopen(try std.fmt.allocPrintSentinel(arena, "{s}", .{args[2]}, 0), "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(out.items.ptr, 1, out.items.len, file) != out.items.len) return error.WriteFailed; + const file = platform.c.fopen(try std.fmt.allocPrintSentinel(arena, "{s}", .{args[2]}, 0), "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(out.items.ptr, 1, out.items.len, file) != out.items.len) return error.WriteFailed; std.debug.print("{d} cases\n", .{cases}); } diff --git a/zig/tests/glm_dsa_prompt_check.zig b/zig/tests/glm_dsa_prompt_check.zig index e88d50ebac..9141bf7fcf 100644 --- a/zig/tests/glm_dsa_prompt_check.zig +++ b/zig/tests/glm_dsa_prompt_check.zig @@ -1,5 +1,6 @@ //! A prompt window's experts (every expert once, on the tensor units) against the decode path's 16-row windows, by row. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const g = tf.glm_dsa; @@ -53,8 +54,8 @@ pub fn main(init: std.process.Init) !void { for (q.x.slice(f32, @as(usize, rows) * hid)) |*v| v.* = prng.random().floatNorm(f32); const dec = try device.buffer(@as(usize, rows) * hid * 4, mtl.ResourceOptions.shared); const out = try device.buffer(@as(usize, rows) * hid * 4, mtl.ResourceOptions.shared); - if (std.c.getenv("GLM_ATTN") != null) return attention(&K, &L, &s, &q, queue, device, &c, rows, w.full); - if (std.c.getenv("GLM_NARROW") != null) return narrow(&K, &L, &q, queue, device, &c); + if (platform.c.getenv("GLM_ATTN") != null) return attention(&K, &L, &s, &q, queue, device, &c, rows, w.full); + if (platform.c.getenv("GLM_NARROW") != null) return narrow(&K, &L, &q, queue, device, &c); // the decode path: route and experts 16 rows at a time (each window's rows copied into the 16-row scratch) var dec_ms: f64 = 0; var k: u32 = 0; diff --git a/zig/tests/glm_dsa_run.zig b/zig/tests/glm_dsa_run.zig index 6c918d5676..94132983a1 100644 --- a/zig/tests/glm_dsa_run.zig +++ b/zig/tests/glm_dsa_run.zig @@ -1,5 +1,6 @@ //! GLM-5.3 on the pair without the server: rank 0 runs a prompt and its greedy reply, rank 1 follows. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const tok = @import("tokenizer"); const g = tf.glm_dsa; @@ -13,7 +14,7 @@ const usage = ; fn env(name: [:0]const u8, default: u64) u64 { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(u64, std.mem.span(v), 10) catch default; } @@ -30,9 +31,9 @@ pub fn main(init: std.process.Init) !void { } var ram: u64 = 0; var len: usize = 8; - if (std.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; + if (platform.c.sysctlbyname("hw.memsize", &ram, &len, null, 0) != 0) return error.NoRamSize; const limit = env("GLM_LIMIT_GB", ram * 7 / 10 / 1_000_000_000) * 1_000_000_000; - const layers: ?u32 = if (std.c.getenv("GLM_LAYERS") != null) @intCast(env("GLM_LAYERS", 78)) else null; + const layers: ?u32 = if (platform.c.getenv("GLM_LAYERS") != null) @intCast(env("GLM_LAYERS", 78)) else null; const depth: u32 = @intCast(env("GLM_DEPTH", 0)); const e = try g.engine.Engine.load(a, io, argv[1], argv[2], layers, depth > 0, @intCast(env("GLM_CONTEXT", 8192)), 1, limit); defer e.deinit(); @@ -53,7 +54,7 @@ pub fn main(init: std.process.Init) !void { } fn once(e: *g.engine.Engine, io: std.Io, t: *const tok.Tokenizer, prompt: []const u32, max: u64, depth: u32) !void { - const forced = std.c.getenv("GLM_FORCED") != null; + const forced = platform.c.getenv("GLM_FORCED") != null; var agree: usize = 0; const t0 = std.Io.Timestamp.now(io, .awake); try g.follow.begin(e, 0, prompt); diff --git a/zig/tests/glm_rows_speed.zig b/zig/tests/glm_rows_speed.zig index f2ceb725a0..e91e0657db 100644 --- a/zig/tests/glm_rows_speed.zig +++ b/zig/tests/glm_rows_speed.zig @@ -1,5 +1,6 @@ //! tf-glm-rows-speed: GLM-5.3-Flash's decode products at the checkpoints' widths and shapes, GPU time and GB/s. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const sources = @import("kernel_sources"); @@ -555,13 +556,13 @@ pub fn main() !void { 42, 3 * affine.bytes(N, D, .{ .bits = 2, .group = 64 }) * TOPK + sh8[0].bytes + sh8[1].bytes, ); - if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + if (platform.c.getenv("GLM_SPEED_CANDIDATES") != null) { const lib = try cand.candidateLib(device); defer lib.deinit(); try cand.moeCandidates(&b, lib, p, gu2, dn2, sh8); } } - if (std.c.getenv("GLM_SPEED_CANDIDATES") != null) { + if (platform.c.getenv("GLM_SPEED_CANDIDATES") != null) { const lib = try cand.candidateLib(device); defer lib.deinit(); std.debug.print("candidates at 8 bits (bits checked == glm_rows):\n", .{}); @@ -571,7 +572,7 @@ pub fn main() !void { try cand.candidates(&b, lib, p, m, s.name); } } - if (std.c.getenv("GLM_SPEED_PROMPT") != null) try speed_prompt.prompt(&b, k, p); + if (platform.c.getenv("GLM_SPEED_PROMPT") != null) try speed_prompt.prompt(&b, k, p); std.debug.print("uniform 4/64:\n", .{}); for (uni.rows[0..uni.n]) |x| std.debug.print(" {s:<26} {d:7.3} ms/token\n", .{ x.name, x.us * x.count / 1e3 }); uni.total("uniform 4/64"); diff --git a/zig/tests/glm_run.zig b/zig/tests/glm_run.zig index 13194d6d0f..97ad4105a8 100644 --- a/zig/tests/glm_run.zig +++ b/zig/tests/glm_run.zig @@ -1,5 +1,6 @@ //! GLM-5.3-Flash replies on the native engine at each draft depth: hashes, first differences from depth 0 and the reference, speeds. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const glm = tf.glm; @@ -53,7 +54,7 @@ const Collect = struct { }; fn env(name: [:0]const u8, default: []const u8) []const u8 { - return if (std.c.getenv(name)) |v| std.mem.span(v) else default; + return if (platform.c.getenv(name)) |v| std.mem.span(v) else default; } fn hash(t: []const u32) u64 { @@ -124,13 +125,13 @@ pub fn main(init: std.process.Init) !void { const e = try glm.engine.Engine.load(gpa, args[1], cap); defer e.deinit(); std.debug.print("loaded in {d:.1} s: {d} of {d} layers, routed experts {d}-{d}, {d:.1} GB of weights, MTP head {s}, {d:.2} GB of caches and buffers, prompt chunks of {d} rows\n", .{ e.load_seconds, e.c.run, e.c.layers, e.c.own[0], e.c.own[1] - 1, @as(f64, @floatFromInt(e.w.bytes)) / 1e9, if (e.hasMtp()) "yes" else "no", @as(f64, @floatFromInt(e.arena.bytes)) / 1e9, if (e.pr) |p| p.rows else 0 }); - const strict = std.c.getenv("GLM_REF_STRICT") != null; + const strict = platform.c.getenv("GLM_REF_STRICT") != null; const eos = e.c.eos[0..e.c.eos_n]; - if (std.c.getenv("GLM_CAPTURE")) |path| { // the first prompt's first window, sublayer by sublayer (glm_ref.py --capture) + if (platform.c.getenv("GLM_CAPTURE")) |path| { // the first prompt's first window, sublayer by sublayer (glm_ref.py --capture) const first = doc.object.get("prompts").?.array.items[0]; try e.capture(try ints(arena, first.object.get("ids").?), std.mem.span(path)); } - if (std.c.getenv("GLM_TRACE")) |spec| { // NAME:STEPS:PATH, glm_ref.py --trace NAME STEPS PATH's twin + if (platform.c.getenv("GLM_TRACE")) |spec| { // NAME:STEPS:PATH, glm_ref.py --trace NAME STEPS PATH's twin var parts = std.mem.splitScalar(u8, std.mem.span(spec), ':'); const want = parts.next() orelse return error.BadTrace; const steps = try std.fmt.parseInt(usize, parts.next() orelse return error.BadTrace, 10); @@ -139,18 +140,18 @@ pub fn main(init: std.process.Init) !void { try e.trace(try ints(arena, p.object.get("ids").?), try ints(arena, p.object.get("expect").?), steps, out_path); }; } - if (std.c.getenv("GLM_FORCED") != null) for (doc.object.get("prompts").?.array.items) |p| { + if (platform.c.getenv("GLM_FORCED") != null) for (doc.object.get("prompts").?.array.items) |p| { const want = try ints(arena, p.object.get("expect") orelse continue); const r = try e.forced(try ints(arena, p.object.get("ids").?), want); std.debug.print("{s} forced: {d} of {d} picks equal the reference's", .{ p.object.get("name").?.string, r.same, want.len }); if (r.first) |i| std.debug.print(", first differs at {d}\n", .{i}) else std.debug.print("\n", .{}); }; var failures: usize = 0; - const vs: ?std.json.Value = if (std.c.getenv("GLM_VS")) |path| blk: { + const vs: ?std.json.Value = if (platform.c.getenv("GLM_VS")) |path| blk: { const f = try mtl.MappedFile.open(path); break :blk try std.json.parseFromSliceLeaky(std.json.Value, arena, f.bytes[0..f.size], .{}); } else null; - const trace_prefix: ?[*:0]const u8 = std.c.getenv("GLM_TRACE_LAST"); + const trace_prefix: ?[*:0]const u8 = platform.c.getenv("GLM_TRACE_LAST"); const trace_bytes = (2 + 4 * @as(usize, e.c.run)) * e.c.hidden * 2; if (trace_prefix != null) e.trace_last = try e.arena.buffer(trace_bytes); var saved: std.ArrayList(u8) = .empty; @@ -168,10 +169,10 @@ pub fn main(init: std.process.Init) !void { if (trace_prefix != null and run == 0 and d == depths.items[0]) col.trace_from = e.trace_last.?.addr()[0..trace_bytes]; var margins: std.ArrayList(f32) = .empty; defer margins.deinit(gpa); - const want_margins = run == 0 and d == depths.items[0] and (std.c.getenv("GLM_MARGINS") != null or std.c.getenv("GLM_MARGINS_VS") != null); + const want_margins = run == 0 and d == depths.items[0] and (platform.c.getenv("GLM_MARGINS") != null or platform.c.getenv("GLM_MARGINS_VS") != null); e.margins = if (want_margins) &margins else null; defer e.margins = null; - const want_logits = run == 0 and d == depths.items[0] and (std.c.getenv("GLM_LOGITS") != null or std.c.getenv("GLM_LOGITS_VS") != null); + const want_logits = run == 0 and d == depths.items[0] and (platform.c.getenv("GLM_LOGITS") != null or platform.c.getenv("GLM_LOGITS_VS") != null); if (want_logits) col.logits_from = @as([*]const u16, @ptrCast(@alignCast(e.sc.logits.addr())))[0..e.c.vocab]; const r = try e.generate(ids, max, eos, d, .{ .ctx = &col, .prefilled = Collect.prefilled, .tokens = Collect.tokens, .cancelled = Collect.cancelled }); const toks = col.toks.items; @@ -182,7 +183,7 @@ pub fn main(init: std.process.Init) !void { const rn: f64 = @floatFromInt(r.rounds); std.debug.print(" a round: {d:.2} ms wall, {d:.2} ms GPU, {d:.2} ms encoding, {d:.2} ms GPU idle between rounds; {d:.2} tokens\n", .{ r.decode_seconds * 1e3 / rn, r.gpu_seconds * 1e3 / rn, r.encode_seconds * 1e3 / rn, r.gap_seconds * 1e3 / @max(rn - 1, 1), @as(f64, @floatFromInt(toks.len -| 1)) / rn }); if (r.copy_rounds > 0) std.debug.print(" copy drafts: {d} rounds, {d} drafts kept ({d:.2} a copy round)\n", .{ r.copy_rounds, r.copy_accepted, @as(f64, @floatFromInt(r.copy_accepted)) / @as(f64, @floatFromInt(r.copy_rounds)) }); - if (std.c.getenv("GLM_RANKS") != null) { + if (platform.c.getenv("GLM_RANKS") != null) { std.debug.print(" draft ranks (kept: all / miss at 2nd / 3rd / later):", .{}); for (r.draft_ranks, 0..) |row, m| if (row[0] + row[1] + row[2] + row[3] > 0) std.debug.print(" m{d} {d}/{d}/{d}/{d}", .{ m, row[0], row[1], row[2], row[3] }); std.debug.print("\n", .{}); @@ -192,17 +193,17 @@ pub fn main(init: std.process.Init) !void { failures += 1; std.debug.print(" DIFFERS from depth {d} at token {d}\n", .{ depths.items[0], at }); } else std.debug.print(" equal to depth {d}\n", .{depths.items[0]}); - if (want_margins) if (std.c.getenv("GLM_MARGINS")) |prefix| { + if (want_margins) if (platform.c.getenv("GLM_MARGINS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.margins", .{ prefix, name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(std.mem.sliceAsBytes(margins.items).ptr, 4, margins.items.len, file) != margins.items.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(std.mem.sliceAsBytes(margins.items).ptr, 4, margins.items.len, file) != margins.items.len) return error.WriteFailed; }; if (vs) |other| if (other.object.get(name)) |theirs| { const want = try ints(arena, theirs); if (firstDiff(want, toks)) |at| { std.debug.print(" vs GLM_VS: first token {s}, first difference at token {d} of {d}\n", .{ if (at == 0) "DIFFERS" else "equal", at, @min(want.len, toks.len) }); - if (want_margins) if (std.c.getenv("GLM_MARGINS_VS")) |prefix| { + if (want_margins) if (platform.c.getenv("GLM_MARGINS_VS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.margins", .{ prefix, name }, 0); if (mtl.MappedFile.open(path)) |f| { defer f.deinit(); @@ -216,18 +217,18 @@ pub fn main(init: std.process.Init) !void { }; if (col.trace.len > 0) { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.trace", .{ std.mem.span(trace_prefix.?), name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(col.trace.ptr, 1, col.trace.len, file) != col.trace.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(col.trace.ptr, 1, col.trace.len, file) != col.trace.len) return error.WriteFailed; } if (want_logits and col.logits.len == 0) return error.NoLogits; - if (want_logits) if (std.c.getenv("GLM_LOGITS")) |prefix| { + if (want_logits) if (platform.c.getenv("GLM_LOGITS")) |prefix| { const path = try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0); - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(std.mem.sliceAsBytes(col.logits).ptr, 2, col.logits.len, file) != col.logits.len) return error.WriteFailed; + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(std.mem.sliceAsBytes(col.logits).ptr, 2, col.logits.len, file) != col.logits.len) return error.WriteFailed; }; - if (want_logits) if (std.c.getenv("GLM_LOGITS_VS")) |prefix| try compareLogits(gpa, col.logits, try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0)); + if (want_logits) if (platform.c.getenv("GLM_LOGITS_VS")) |prefix| try compareLogits(gpa, col.logits, try std.fmt.allocPrintSentinel(arena, "{s}.{s}.bf16", .{ prefix, name }, 0)); if (run == 0 and d == depths.items[0]) { if (saved.items.len > 1) try saved.append(arena, ','); try saved.print(arena, "\"{s}\":[", .{name}); @@ -241,7 +242,7 @@ pub fn main(init: std.process.Init) !void { } else std.debug.print(" reference: {d} of {d} tokens equal\n", .{ @min(want.len, toks.len), want.len }); } }; - if (std.c.getenv("GLM_CANCEL_TEST") != null) { // a reply cancelled after its first prompt window, then a fresh one + if (platform.c.getenv("GLM_CANCEL_TEST") != null) { // a reply cancelled after its first prompt window, then a fresh one var cut: Collect = .{ .gpa = gpa, .cancel_at = 1 }; defer cut.toks.deinit(gpa); const rc = try e.generate(ids, max, eos, depths.items[0], .{ .ctx = &cut, .prefilled = Collect.prefilled, .tokens = Collect.tokens, .cancelled = Collect.cancelled }); @@ -253,23 +254,23 @@ pub fn main(init: std.process.Init) !void { std.debug.print(" cancel test: cancelled reply ended {t} after {d} tokens; the next fresh reply {s} the plain one\n", .{ rc.reason, cut.toks.items.len, if (same) "equals" else "DIFFERS from" }); } } - if (std.c.getenv("GLM_PROFILE")) |spec| { + if (platform.c.getenv("GLM_PROFILE")) |spec| { const reps = try std.fmt.parseInt(usize, env("GLM_PROFILE_REPS", "12"), 10); var pit = std.mem.tokenizeScalar(u8, std.mem.span(spec), ','); - const only = std.c.getenv("GLM_PROFILE_ONLY") != null; // each class alone instead of each left out - const parts = std.c.getenv("GLM_PROFILE_PARTS") != null; // by launch instead of by class + const only = platform.c.getenv("GLM_PROFILE_ONLY") != null; // each class alone instead of each left out + const parts = platform.c.getenv("GLM_PROFILE_PARTS") != null; // by launch instead of by class while (pit.next()) |dv| try e.profile(try std.fmt.parseInt(u32, dv, 10), reps, only, parts); } - if (std.c.getenv("GLM_MM_CHECK") != null) try e.checkMatmul(try ints(arena, doc.object.get("prompts").?.array.items[0].object.get("ids").?)); - if (std.c.getenv("GLM_PROMPT_PROFILE")) |v| { // each prompt's first chunk by class (GLM_PROFILE_ONLY: each alone) + if (platform.c.getenv("GLM_MM_CHECK") != null) try e.checkMatmul(try ints(arena, doc.object.get("prompts").?.array.items[0].object.get("ids").?)); + if (platform.c.getenv("GLM_PROMPT_PROFILE")) |v| { // each prompt's first chunk by class (GLM_PROFILE_ONLY: each alone) const reps = try std.fmt.parseInt(usize, std.mem.span(v), 10); - for (doc.object.get("prompts").?.array.items) |p| try e.profilePrompt(try ints(arena, p.object.get("ids").?), reps, std.c.getenv("GLM_PROFILE_ONLY") != null, std.c.getenv("GLM_PROFILE_PARTS") != null); + for (doc.object.get("prompts").?.array.items) |p| try e.profilePrompt(try ints(arena, p.object.get("ids").?), reps, platform.c.getenv("GLM_PROFILE_ONLY") != null, platform.c.getenv("GLM_PROFILE_PARTS") != null); } try saved.append(arena, '}'); - if (std.c.getenv("GLM_OUT")) |path| { - const file = std.c.fopen(path, "wb") orelse return error.OpenFailed; - defer _ = std.c.fclose(file); - if (std.c.fwrite(saved.items.ptr, 1, saved.items.len, file) != saved.items.len) return error.WriteFailed; + if (platform.c.getenv("GLM_OUT")) |path| { + const file = platform.c.fopen(path, "wb") orelse return error.OpenFailed; + defer _ = platform.c.fclose(file); + if (platform.c.fwrite(saved.items.ptr, 1, saved.items.len, file) != saved.items.len) return error.WriteFailed; } if (failures > 0) { std.debug.print("{d} replies differ from the plain ones{s}\n", .{ failures, if (strict) " or the references" else "" }); diff --git a/zig/tests/linktest.zig b/zig/tests/linktest.zig index dfdad16efe..bf251fd1ec 100644 --- a/zig/tests/linktest.zig +++ b/zig/tests/linktest.zig @@ -1,5 +1,6 @@ //! Link qualification over the system's verbs; `run` swaps connection details on stdin and stdout and opens no socket. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const suite = @import("linktest_suite.zig"); const probe = @import("linktest_probe.zig"); @@ -30,7 +31,7 @@ fn readLine(buf: []u8) ![]const u8 { var n: usize = 0; while (n < buf.len) { var c: [1]u8 = undefined; - if (std.c.read(0, &c, 1) != 1) return error.EndOfInput; + if (platform.c.read(0, &c, 1) != 1) return error.EndOfInput; if (c[0] == '\n') return buf[0..n]; buf[n] = c[0]; n += 1; @@ -61,7 +62,7 @@ fn parseHello(line: []const u8) ![2]vqp.Info { /// Exit when stdin closes: the run ends with the session that started it. fn tether() void { var c: [1]u8 = undefined; - while (std.c.read(0, &c, 1) == 1) {} + while (platform.c.read(0, &c, 1) == 1) {} out("R stdin_closed", .{}); std.process.exit(4); } @@ -79,7 +80,7 @@ fn mapped(len: usize) ![]align(vl.page) u8 { /// Two queue pairs on one device: collectives on one, raw measurements on the other, so their receives never mix. fn run(gpa: std.mem.Allocator, v: *const Verbs, a: Args) !void { _ = try std.Thread.spawn(.{}, watchdog, .{a.deadline_s}); - out("PID {d}", .{std.c.getpid()}); + out("PID {d}", .{platform.c.getpid()}); var peer: vl.Peer = .{ .ep = try vqp.Endpoint.open(v, a.dev, a.kind, 4095, a.gid) }; defer peer.ep.close(); var raw = try peer.ep.sibling(a.kind, 4095); @@ -138,7 +139,7 @@ pub fn main(init: std.process.Init) !void { if (!std.mem.eql(u8, args[1], "run")) std.process.exit(2); const a = try parse(args); run(init.gpa, &v, a) catch |err| { - out("R error what={s} errno={d}", .{ @errorName(err), std.c._errno().* }); + out("R error what={s} errno={d}", .{ @errorName(err), platform.c._errno().* }); std.process.exit(1); }; } diff --git a/zig/tests/linktest_probe.zig b/zig/tests/linktest_probe.zig index ce3402075a..a5c6af2541 100644 --- a/zig/tests/linktest_probe.zig +++ b/zig/tests/linktest_probe.zig @@ -1,5 +1,6 @@ //! What a verbs device reports and what it actually allows: registration sizes and counts, queue pair types and counts. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const abi = fabric.verbs_abi; const Verbs = fabric.verbs.Verbs; @@ -8,7 +9,7 @@ const out = @import("linktest_suite.zig").out; const mib = 1 << 20; fn errno() c_int { - return std.c._errno().*; + return platform.c._errno().*; } pub fn device(v: *const Verbs, dev: []const u8) !void { diff --git a/zig/tests/linktest_suite.zig b/zig/tests/linktest_suite.zig index 37b998b9f7..f356e6e519 100644 --- a/zig/tests/linktest_suite.zig +++ b/zig/tests/linktest_suite.zig @@ -1,5 +1,6 @@ //! The link measurements over Thunderbolt's two-sided UC: opcodes, ping-pong, bandwidth, integrity, barrier, all-reduce. const std = @import("std"); +const platform = @import("sys"); const fabric = @import("fabric"); const collective = fabric.collective; const vqp = fabric.verbs_qp; @@ -53,7 +54,7 @@ pub const Ctx = struct { pub fn out(comptime fmt: []const u8, args: anytype) void { var buf: [4096]u8 = undefined; const line = std.fmt.bufPrint(&buf, fmt ++ "\n", args) catch return; - _ = std.c.write(1, line.ptr, line.len); + _ = platform.c.write(1, line.ptr, line.len); } fn us(ns: u64) f64 { diff --git a/zig/tests/mm_bench.zig b/zig/tests/mm_bench.zig index 1ec779f916..76f6402c12 100644 --- a/zig/tests/mm_bench.zig +++ b/zig/tests/mm_bench.zig @@ -1,5 +1,6 @@ //! Prompt matmul kernels compiled from .metal files at run time: chained speed at one shape, bits against the first, error against fp64. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -17,7 +18,7 @@ const Strides = extern struct { x_row: i32 = 0, y_row: i32 = 0, sums_row: i32 = const Kernel = struct { path: []const u8, bm: u32 = 64, bn: u32 = 64, threads: u32 = 128, entry: []const u8 = "tf_affine_mm", lib: mtl.Library = undefined, pipe: mtl.Pipeline = undefined, sums_pipe: mtl.Pipeline = undefined, sums: mtl.Buffer = undefined, y: mtl.Buffer = undefined, ms: [64]f64 = undefined, sums_ms: f64 = 0 }; fn envInt(name: [:0]const u8, default: usize) usize { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; } @@ -55,7 +56,7 @@ pub fn main(init: std.process.Init) !void { const trials = @min(envInt("MM_TRIALS", 5), 64); const copies = envInt("MM_COPIES", 4); const check_rows = @min(envInt("MM_CHECK_ROWS", 16), M); - const out_f32 = std.c.getenv("MM_F32") != null; + const out_f32 = platform.c.getenv("MM_F32") != null; const ob: usize = if (out_f32) 4 else 2; const kernels = try arena.alloc(Kernel, args.len - 4); for (kernels, args[4..]) |*k, s| k.* = try parseKernel(s); diff --git a/zig/tests/moe_bench.zig b/zig/tests/moe_bench.zig index 6e381ff8d5..f611ef2575 100644 --- a/zig/tests/moe_bench.zig +++ b/zig/tests/moe_bench.zig @@ -1,5 +1,6 @@ //! Routed-expert prompt kernels from .metal files compiled at run time: a chunk's skewed top-k routing, chained speed, bits against the first. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -16,12 +17,12 @@ const Args = extern struct { rows: i32, n: i32, k: i32, experts: i32 }; const Kernel = struct { path: []const u8, bm: u32 = 64, entry: []const u8 = "", lib: mtl.Library = undefined, pipe: mtl.Pipeline = undefined, y: mtl.Buffer = undefined, ms: [64]f64 = undefined }; fn envInt(name: [:0]const u8, default: usize) usize { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; } fn envFloat(name: [:0]const u8, default: f64) f64 { - const v = std.c.getenv(name) orelse return default; + const v = platform.c.getenv(name) orelse return default; return std.fmt.parseFloat(f64, std.mem.span(v)) catch default; } @@ -37,7 +38,7 @@ pub fn main(init: std.process.Init) !void { std.debug.print("usage: {s}\n", .{usage}); std.process.exit(2); } - const down = if (std.c.getenv("MOE_MODE")) |v| std.mem.eql(u8, std.mem.span(v), "down") else false; + const down = if (platform.c.getenv("MOE_MODE")) |v| std.mem.eql(u8, std.mem.span(v), "down") else false; const T = try std.fmt.parseInt(u32, args[1], 10); const E: u32 = @intCast(envInt("MOE_EXPERTS", 288)); const topk: u32 = @intCast(envInt("MOE_TOPK", 8)); @@ -65,7 +66,7 @@ pub fn main(init: std.process.Init) !void { total += w.*; } const pick = try arena.alloc(u32, n); - const exact = std.c.getenv("MOE_EXACT") != null; // pairs dealt round-robin: every expert the same row count + const exact = platform.c.getenv("MOE_EXACT") != null; // pairs dealt round-robin: every expert the same row count if (exact) for (pick, 0..) |*q, i| { q.* = @intCast((i % topk + (i / topk) * topk) % E); }; diff --git a/zig/tests/nemotron_prefill_check.zig b/zig/tests/nemotron_prefill_check.zig index 5bca853939..f1b91302aa 100644 --- a/zig/tests/nemotron_prefill_check.zig +++ b/zig/tests/nemotron_prefill_check.zig @@ -1,5 +1,6 @@ //! One prompt chunk on the prefill kernels against tools/zig/prefill_dump.py's rows: embedding, every layer, final norm. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -17,7 +18,7 @@ pub fn main(init: std.process.Init) !void { var it = std.mem.tokenizeScalar(u8, args[2], ','); while (it.next()) |t| try list.append(arena, try std.fmt.parseInt(u32, t, 10)); // TF_PREFILL_PREFIX=N: the first N ids fill the caches as a chunk first; the rest is the checked chunk - const before: usize = if (std.c.getenv("TF_PREFILL_PREFIX")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 0; + const before: usize = if (platform.c.getenv("TF_PREFILL_PREFIX")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else 0; const all = list.items; const ids = all[before..]; const n = ids.len; diff --git a/zig/tests/pocket_fit_check.zig b/zig/tests/pocket_fit_check.zig index 7f7bfaf2f5..0fe514684a 100644 --- a/zig/tests/pocket_fit_check.zig +++ b/zig/tests/pocket_fit_check.zig @@ -1,5 +1,6 @@ //! A dumped pocket lesson's read rows fitted as the learner fits them, timed; each fit's numbers and weights written. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const pocket_fit = tf.nemotron.pocket_fit; @@ -57,7 +58,7 @@ pub fn main(init: std.process.Init) !void { const us = try gpa.alloc(f32, jobs.items.len * d); const fits = try gpa.alloc(Fit, jobs.items.len); // TF_FIT_THREADS: that many threads take the fits in turn (default: a thread a fit, as the learner runs them) - const want = if (std.c.getenv("TF_FIT_THREADS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else jobs.items.len; + const want = if (platform.c.getenv("TF_FIT_THREADS")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else jobs.items.len; const threads = try gpa.alloc(std.Thread, @min(@max(want, 1), jobs.items.len)); var next = std.atomic.Value(usize).init(0); const work: Work = .{ .gpa = gpa, .rows = rows, .kinds = all_kinds, .m = m, .slots = slots, .d = d, .jobs = jobs.items, .us = us, .fits = fits, .next = &next }; diff --git a/zig/tests/qwen27/dflash_graph.zig b/zig/tests/qwen27/dflash_graph.zig index bbd9db4865..e382276faf 100644 --- a/zig/tests/qwen27/dflash_graph.zig +++ b/zig/tests/qwen27/dflash_graph.zig @@ -1,5 +1,6 @@ //! Generated draft tensors exercise the real callback backend without checkpoint files. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); const q = tf.qwen27; @@ -47,19 +48,19 @@ pub fn createQ4(a: std.mem.Allocator, target: *q.model.Model) !*Model { } try header.writer.writeByte('}'); var path_buf: [96]u8 = undefined; - const path = try std.fmt.bufPrintSentinel(&path_buf, "/tmp/tf-q27-q4-draft-{d}.safetensors", .{std.c.getpid()}, 0); - const fd = std.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(std.c.mode_t, 0o600)); + const path = try std.fmt.bufPrintSentinel(&path_buf, "/tmp/tf-q27-q4-draft-{d}.safetensors", .{platform.c.getpid()}, 0); + const fd = platform.c.open(path, .{ .ACCMODE = .WRONLY, .CREAT = true, .TRUNC = true }, @as(platform.c.mode_t, 0o600)); if (fd < 0) return error.Create; - defer _ = std.c.unlink(path); + defer _ = platform.c.unlink(path); var head: [8]u8 = undefined; std.mem.writeInt(u64, &head, header.written().len, .little); for ([_][]const u8{ &head, header.written(), std.mem.sliceAsBytes(data.items) }) |part| { - if (std.c.write(fd, part.ptr, part.len) != @as(isize, @intCast(part.len))) { - _ = std.c.close(fd); + if (platform.c.write(fd, part.ptr, part.len) != @as(isize, @intCast(part.len))) { + _ = platform.c.close(fd); return error.ShortWrite; } } - _ = std.c.close(fd); + _ = platform.c.close(fd); return Model.open(a, target, config, &.{path}, .prepared_q4_reference); } diff --git a/zig/tests/qwen27/dflash_run.zig b/zig/tests/qwen27/dflash_run.zig index 06660eed7e..2d75b9d491 100644 --- a/zig/tests/qwen27/dflash_run.zig +++ b/zig/tests/qwen27/dflash_run.zig @@ -1,5 +1,6 @@ //! The draft prototype on real weights: target tokens and final state against ordinary native generation. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("qwen27"); const profile = @import("core").gpu_profile; @@ -71,7 +72,7 @@ pub fn main(init: std.process.Init) !void { try generation.prefill(prompt, chunk); const prompt_ns = start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds(); var trace = profile.Trace{ .queue = target.queue }; - const profiled = std.c.getenv("QWEN27_PROFILE") != null; + const profiled = platform.c.getenv("QWEN27_PROFILE") != null; if (profiled) { runner.profile(&trace); trace.enabled = true; diff --git a/zig/tests/qwen27/dflash_served.zig b/zig/tests/qwen27/dflash_served.zig index ca9b3c2bfb..6955e546f0 100644 --- a/zig/tests/qwen27/dflash_served.zig +++ b/zig/tests/qwen27/dflash_served.zig @@ -1,5 +1,6 @@ //! Generated native draft and plain requests retain CLI tokens and final state through the served host. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("tensorfold").qwen27; const api = @import("engine_api"); @@ -73,7 +74,7 @@ const Box = struct { }; /// Drafted requests on the lane core (its own round counts), by the host's rule: TF_QWEN27_LANES, else no tensor units. fn lanesMode(h: *host.Host) bool { - if (std.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); + if (platform.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); return !h.model.device.tensorUnits(); } fn stop(_: *anyopaque, ids: []const u32) bool { diff --git a/zig/tests/qwen27/reuse.zig b/zig/tests/qwen27/reuse.zig index 46ce1ed90d..e41b5d8c04 100644 --- a/zig/tests/qwen27/reuse.zig +++ b/zig/tests/qwen27/reuse.zig @@ -1,5 +1,6 @@ //! Qwen27 prompt reuse through the served host against fresh prompt passes: each prompt's state and reply, bit for bit. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const q = @import("tensorfold").qwen27; const api = @import("engine_api"); @@ -90,7 +91,7 @@ const Run = struct { /// Drafted requests on the lane core, by the served host's rule: TF_QWEN27_LANES, else a chip without tensor units. fn lanesMode(h: *host.Host) bool { - if (std.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); + if (platform.c.getenv("TF_QWEN27_LANES")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); return !h.model.device.tensorUnits(); } diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index 184c401968..d7f916f4a6 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -1,5 +1,6 @@ //! Frozen server goldens for `zig build test-golden`, with the ten differences that are still known. const std = @import("std"); +const platform = @import("sys"); const posix = std.posix; const Io = std.Io; const wire = @import("wire.zig"); @@ -379,7 +380,7 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con try fw.interface.writeAll(content); try fw.interface.flush(); const kpz = try std.fmt.allocPrintSentinel(a, "{s}", .{kp}, 0); - _ = std.c.fchmodat(posix.AT.FDCWD, kpz, 0o600, 0); + _ = platform.c.fchmodat(posix.AT.FDCWD, kpz, 0o600, 0); } } const child = try a.create(std.process.Child); @@ -411,7 +412,7 @@ pub fn main(init: std.process.Init) !void { const cases_path = args[2]; const golden_dir = args[3]; const fixtures = args[4]; - const tmp_dir = try std.fmt.allocPrint(init.arena.allocator(), ".zig-tmp-golden-{d}", .{std.c.getpid()}); + const tmp_dir = try std.fmt.allocPrint(init.arena.allocator(), ".zig-tmp-golden-{d}", .{platform.c.getpid()}); Io.Dir.cwd().createDir(io, tmp_dir, .default_dir) catch {}; defer Io.Dir.cwd().deleteTree(io, tmp_dir) catch {}; diff --git a/zig/tests/train_check.zig b/zig/tests/train_check.zig index 08601733c3..77acc8e7c1 100644 --- a/zig/tests/train_check.zig +++ b/zig/tests/train_check.zig @@ -1,5 +1,6 @@ //! The learning step's gradients against the engine's own forward, moving each layer along its own, on the real model. const std = @import("std"); +const platform = @import("sys"); const tf = @import("tensorfold"); const nm = tf.nemotron; @@ -51,8 +52,8 @@ pub fn main(init: std.process.Init) !void { const with_grad = (try t.step(&ids, start, .grad, turn)).loss; const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0.toNanoseconds())) / 1e6; std.debug.print("loss {d:.5} (with gradients {d:.5}), a step with gradients {d:.1} ms, GPU {d:.1} ms\n", .{ plain, with_grad, ms, b.last_ms }); - if (std.c.getenv("TF_TRAIN_TIME_ONLY") != null) return; - if (std.c.getenv("TF_TRAIN_PROFILE") != null) { + if (platform.c.getenv("TF_TRAIN_TIME_ONLY") != null) return; + if (platform.c.getenv("TF_TRAIN_PROFILE") != null) { var prof: nm.encoder.Profiler = .{ .queue = m.queue, .kernels = &m.kernels }; t.profile = &prof; _ = try t.step(&ids, start, .grad, turn); diff --git a/zig/tests/weight_read_check.zig b/zig/tests/weight_read_check.zig index d5acf5e7a6..1089d93666 100644 --- a/zig/tests/weight_read_check.zig +++ b/zig/tests/weight_read_check.zig @@ -1,5 +1,6 @@ //! Check the actual native loaders with a synthetic file and bounded I/O faults. const std = @import("std"); +const platform = @import("sys"); const mtl = @import("metal"); const tf = @import("tensorfold"); @@ -63,7 +64,7 @@ pub fn main(init: std.process.Init) !void { } else { try replay.indexFile(path); const entry = try replay.entry(if (grouped) "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight" else "probe"); - defer _ = std.c.close(entry.fd); + defer _ = platform.c.close(entry.fd); const result = if (grouped) replay.group(0, 2, "weight") else replay.load("probe"); const buffer = result catch |err| { if (!eof or err != error.ShortRead) return err; From e5bb3bca166111cc3cef9dd118f27ab28e22a841 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:14:21 +0000 Subject: [PATCH 02/35] Windows: CUDA kernels, the native server and the GPU suites on sm_89 - Kernel fatbins build with nvcc and cl.exe: MSVC's conforming preprocessor is passed on Windows hosts (CUDA 13's headers need it), and explicit kernel instantiations carry the __restrict__ qualifiers their templates have (MSVC counts them in the function type; the generated code is unchanged). - Nemotron's checkpoint source keeps io_uring to Linux; elsewhere its existing thread reads run. - pull_exlr8 raises no descriptor limit on Windows. - modelopt test: the lane's slice partials are allocated. GPUs without clusters (before sm_90) reduce split K through them; the test passed none, so it wrote through a null pointer on Ampere and Ada. zig build native and install build with -Dnvcc=... -Dsm=89, and tf-cuda-test's synthetic suites pass on an RTX 4090. --- zig/build/cuda.zig | 2 ++ zig/kernels/cuda/experts.cu | 4 ++-- zig/kernels/cuda/experts_pack.cu | 4 ++-- zig/kernels/cuda/experts_prefill.cu | 2 +- zig/kernels/cuda/gdn.cu | 10 +++++----- zig/kernels/cuda/lane_gemv.cu | 2 +- zig/kernels/cuda/prefill_attention.cu | 4 ++-- zig/kernels/cuda/qmm_group.cu | 2 +- zig/kernels/cuda/qmm_prefill.cu | 2 +- zig/kernels/cuda/scan_rows.cu | 8 ++++---- zig/src/cli/pull_exlr8.zig | 1 + zig/src/families/nemotron/cuda_source.zig | 12 ++++++++---- zig/src/platform/win_posix.zig | 5 +++-- zig/tests/cuda/modelopt_tests.zig | 17 +++++++++++++---- 14 files changed, 46 insertions(+), 29 deletions(-) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..4f54eaf76d 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -226,6 +226,8 @@ fn fatbin(b: *std.Build, nvcc: []const u8, version: std.Build.LazyPath, k: Kerne const run = b.addSystemCommand(&.{ nvcc, "-fatbin" }); run.addFileInput(version); run.addArgs(&torch_flags); + // CUDA 13's headers need MSVC's conforming preprocessor when cl.exe is the host compiler + if (b.graph.host.result.os.tag == .windows) run.addArg("-Xcompiler=/Zc:preprocessor"); run.addArgs(k.flags); const a = if (k.arch_specific) "a" else ""; var it = std.mem.tokenizeScalar(u8, sms, ','); diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index edd83834c1..d19f6ee90c 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -276,7 +276,7 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const } // namespace tf_experts // Decode form for groups of 64: up with relu^2 (epilogue 1), down to fp32 (epilogue 0). -#define TF_EXPERT(EPI) template __global__ void tf_experts::expert_kernel<64, 1, EPI, 2, 4>(const __nv_bfloat16*, \ - int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); +#define TF_EXPERT(EPI) template __global__ void tf_experts::expert_kernel<64, 1, EPI, 2, 4>(const __nv_bfloat16* __restrict__, \ + int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); TF_EXPERT(1) TF_EXPERT(0) diff --git a/zig/kernels/cuda/experts_pack.cu b/zig/kernels/cuda/experts_pack.cu index 6aa7413793..b47ded555f 100644 --- a/zig/kernels/cuda/experts_pack.cu +++ b/zig/kernels/cuda/experts_pack.cu @@ -38,5 +38,5 @@ pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ sca } // namespace tf_experts_pack // Groups of 64 inputs. -template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t*, const uint16_t*, const uint16_t*, - uint32_t*, int, int, int, int); +template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t* __restrict__, const uint16_t* __restrict__, const uint16_t* __restrict__, + uint32_t* __restrict__, int, int, int, int); diff --git a/zig/kernels/cuda/experts_prefill.cu b/zig/kernels/cuda/experts_prefill.cu index c1c7fbbac9..3a1df1dfee 100644 --- a/zig/kernels/cuda/experts_prefill.cu +++ b/zig/kernels/cuda/experts_prefill.cu @@ -152,6 +152,6 @@ __global__ void __launch_bounds__(WM * WN * 32) // Prefill form for groups of 64: up with relu^2 (epilogue 1), down to bf16 (epilogue 3). #define TF_PREFILL(EPI) template __global__ void tf_experts_prefill::prefill_kernel<64, 1, EPI, 2, 2, 4>( \ - const __nv_bfloat16*, int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); + const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); TF_PREFILL(1) TF_PREFILL(3) diff --git a/zig/kernels/cuda/gdn.cu b/zig/kernels/cuda/gdn.cu index 574fb92de5..2c96e1a3ad 100644 --- a/zig/kernels/cuda/gdn.cu +++ b/zig/kernels/cuda/gdn.cu @@ -305,14 +305,14 @@ __global__ void __launch_bounds__(32 * WARPS) replay_kernel( } // namespace tf_gdn // The instantiations the Python wrappers launch (gdn_replay_cuda, dispatch_tree for bf16 keys). -#define TF_REPLAY(QK) template __global__ void tf_gdn::replay_kernel(const long long*, int, const int*, int, \ - const int*, int, float*, int, int, int); +#define TF_REPLAY(QK) template __global__ void tf_gdn::replay_kernel(const long long* __restrict__, int, const int* __restrict__, int, \ + const int* __restrict__, int, float* __restrict__, int, int, int); TF_REPLAY(__nv_bfloat16) TF_REPLAY(float) #define TF_TREE(S, R, W, C) template __global__ void tf_gdn::tree_kernel<__nv_bfloat16, S, R, W, C>( \ - const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*, const float*, const float*, const float*, \ - const long long*, const int*, const int*, int, __nv_bfloat16*, int, int, int, tf_gdn::Pending<__nv_bfloat16>, \ - float*, const long long*, bool); + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const float* __restrict__, const float* __restrict__, const float* __restrict__, \ + const long long* __restrict__, const int* __restrict__, const int* __restrict__, int, __nv_bfloat16* __restrict__, int, int, int, tf_gdn::Pending<__nv_bfloat16>, \ + float*, const long long* __restrict__, bool); TF_TREE(0, 8, 4, true) TF_TREE(1, 8, 4, false) TF_TREE(2, 8, 2, false) diff --git a/zig/kernels/cuda/lane_gemv.cu b/zig/kernels/cuda/lane_gemv.cu index 94a78a5ae8..af6f0990ac 100644 --- a/zig/kernels/cuda/lane_gemv.cu +++ b/zig/kernels/cuda/lane_gemv.cu @@ -266,7 +266,7 @@ __device__ __forceinline__ void gemv_split(const __nv_bfloat16* __restrict__ x, } // namespace tf_lane_gemv -template __global__ void tf_lane_gemv::gemv_kernel<64, 64, 8>(const __nv_bfloat16*, const float*, +template __global__ void tf_lane_gemv::gemv_kernel<64, 64, 8>(const __nv_bfloat16* __restrict__, const float* __restrict__, const __grid_constant__ tf_lane_gemv::Part, int, int, int); extern "C" __global__ void __launch_bounds__(128) tf_lane_gemv_split(const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 122f404495..b2521ae4a6 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -241,5 +241,5 @@ pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric } // namespace tf_prefill_attention // Head dim 128, eight warps, eight query heads a block, eight staging slots (Nemotron's 16 heads a KV head). -template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16*, - const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, int, float); +template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, float); diff --git a/zig/kernels/cuda/qmm_group.cu b/zig/kernels/cuda/qmm_group.cu index f3f65e184f..f9b2f9f2a7 100644 --- a/zig/kernels/cuda/qmm_group.cu +++ b/zig/kernels/cuda/qmm_group.cu @@ -338,4 +338,4 @@ __global__ void __launch_bounds__(WM * WN * 32) group_kernel( // The instantiations Nemotron's windows launch on sm_121 (tile 2: rows <= 16, bf16 out). template __global__ void tf_qmm_group::group_kernel<64, 16, 64, 1, 4, 8, false, false, false, false>( - const __nv_bfloat16*, const float*, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int); + const __nv_bfloat16* __restrict__, const float* __restrict__, const __grid_constant__ tf_qmm_group::Parts, int, int, int, int, int); diff --git a/zig/kernels/cuda/qmm_prefill.cu b/zig/kernels/cuda/qmm_prefill.cu index 1610b22704..3c61031a25 100644 --- a/zig/kernels/cuda/qmm_prefill.cu +++ b/zig/kernels/cuda/qmm_prefill.cu @@ -149,5 +149,5 @@ __global__ void __launch_bounds__(WM * WN * 32) prefill_kernel( // The instantiation prefill_matmul launches (tile 0: 128x128 on 2x4 warps, 3 stages), bf16 out. template __global__ void tf_qmm_prefill::prefill_kernel<64, 128, 128, 2, 4, 3, false>( - const __nv_bfloat16*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, void*, int, int, int, + const __nv_bfloat16* __restrict__, const uint32_t* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, void* __restrict__, int, int, int, int, int, int); diff --git a/zig/kernels/cuda/scan_rows.cu b/zig/kernels/cuda/scan_rows.cu index 569ffedce8..b3ee6784cb 100644 --- a/zig/kernels/cuda/scan_rows.cu +++ b/zig/kernels/cuda/scan_rows.cu @@ -117,7 +117,7 @@ __global__ void __launch_bounds__(TPR * ROWS) scan_kernel( } // namespace tf_scan_rows -template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16*, const __nv_bfloat16*, float*, const float*, - const float*, const float*, __nv_bfloat16*, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); -template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16*, const __nv_bfloat16*, float*, const float*, - const float*, const float*, __nv_bfloat16*, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); +template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, float* __restrict__, const float* __restrict__, + const float* __restrict__, const float* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); +template __global__ void tf_scan_rows::scan_kernel(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, float* __restrict__, const float* __restrict__, + const float* __restrict__, const float* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, int, int, int, float, float, const tf_scan_rows::Snaps); diff --git a/zig/src/cli/pull_exlr8.zig b/zig/src/cli/pull_exlr8.zig index 3e8c641f7e..c56d16ada2 100644 --- a/zig/src/cli/pull_exlr8.zig +++ b/zig/src/cli/pull_exlr8.zig @@ -88,6 +88,7 @@ pub fn run(a: Allocator, io: std.Io, out: *std.Io.Writer, err_out: *std.Io.Write /// Index maps, connections and scans pass the default 256 descriptors: raise this process's own limit. fn moreFiles() void { + if (@import("builtin").os.tag == .windows) return; // no descriptor limit to raise var l = std.posix.getrlimit(.NOFILE) catch return; const want: std.posix.rlim_t = if (l.max == std.posix.RLIM.INFINITY) 4096 else @min(l.max, 4096); if (l.cur >= want) return; diff --git a/zig/src/families/nemotron/cuda_source.zig b/zig/src/families/nemotron/cuda_source.zig index ffc4ba2eec..f4a1f07fdd 100644 --- a/zig/src/families/nemotron/cuda_source.zig +++ b/zig/src/families/nemotron/cuda_source.zig @@ -2,6 +2,8 @@ const std = @import("std"); const linux = std.os.linux; +/// io_uring is Linux's; elsewhere every read runs on its own thread. +const has_ring = @import("builtin").os.tag == .linux; const cuda = @import("cuda"); const core = @import("core"); const kern = @import("cuda_kernels.zig"); @@ -42,7 +44,7 @@ pub const Source = struct { copied: [slot_count]?cuda.Event = @splat(null), // a page-locked slot's last copy: its next read waits for that, not the stream pending: [slot_count]?Pending = @splat(null), in_flight: usize = 0, - ring: ?linux.IoUring = null, // null where io_uring is unavailable: each read then runs on its own thread + ring: if (has_ring) ?linux.IoUring else ?void = null, // null where io_uring is unavailable: each read then runs on its own thread jobs: [slot_count]Job = undefined, asked: [slot_count]u64 = @splat(0), // the order reads were asked in: copies follow it without a ring next: u64 = 0, @@ -69,7 +71,7 @@ pub const Source = struct { s.pinned[made] = h; s.slots[made] = @alignCast(h.bytes); } - s.ring = linux.IoUring.init(slot_count, 0) catch |err| blk: { + if (has_ring) s.ring = linux.IoUring.init(slot_count, 0) catch |err| blk: { std.log.info("io_uring unavailable ({t}): checkpoint reads run on threads", .{err}); break :blk null; }; @@ -79,7 +81,7 @@ pub const Source = struct { /// Waits for reads still in flight, then frees the ring and slots and closes the files. pub fn deinit(s: *Source) void { while (s.in_flight > 0) s.reap(1) catch break; - if (s.ring) |*r| r.deinit() else s.joinAll(); + if (!has_ring) s.joinAll() else if (s.ring) |*r| r.deinit() else s.joinAll(); s.freeSlots(slot_count); for (s.files.items) |*m| m.file.close(); s.files.deinit(s.gpa); @@ -156,7 +158,8 @@ pub const Source = struct { const lo = std.mem.alignBackward(u64, p.offset, dio.alignment); var q = p; q.at = @intCast(p.offset - lo); - if (s.ring) |*r| { + if (has_ring and s.ring != null) { + const r = &s.ring.?; _ = try r.read(k, p.file.fd, .{ .buffer = s.slots[k][0..dio.span(p.offset, p.len)] }, lo); _ = try r.submit(); } else { @@ -192,6 +195,7 @@ pub const Source = struct { /// Finished reads, copied on; waits for at least `wait` of them. fn reap(s: *Source, wait: u32) !void { + if (!has_ring) return s.reapThreads(wait); const r = if (s.ring) |*ring| ring else return s.reapThreads(wait); var cqes: [slot_count]linux.io_uring_cqe = undefined; const n = try r.copy_cqes(&cqes, wait); diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index 024e8bcc15..4fa75167d9 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -123,8 +123,9 @@ pub const system = struct { const fd = ws.socket(@intCast(family), @intCast(kind), @intCast(proto)); return if (fd == -1) failed() else fd; } - pub fn bind(fd: socket_t, addr: *const anyopaque, len: socklen_t) c_int { - return if (ws.bind(fd, addr, len) != 0) failed() else 0; + /// `len` as std.Io.Threaded's addressToPosix gives it (its socklen_t is unsigned on Windows). + pub fn bind(fd: socket_t, addr: *const anyopaque, len: anytype) c_int { + return if (ws.bind(fd, addr, @intCast(len)) != 0) failed() else 0; } pub fn listen(fd: socket_t, backlog: u31) c_int { return if (ws.listen(fd, backlog) != 0) failed() else 0; diff --git a/zig/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig index faa25fdfdd..cf4cf9fc62 100644 --- a/zig/tests/cuda/modelopt_tests.zig +++ b/zig/tests/cuda/modelopt_tests.zig @@ -187,22 +187,31 @@ fn lanes(gpu: Gpu, rng: std.Random, p: *const modelopt.Packer, lin: cuda.qlinear defer y.free(); var y2 = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 2); defer y2.free(); - try lin.decode(s, q, .{ .x = dx.ptr }, y.ptr, 16); + // slice partials for the lane's split K: GPUs without clusters (before sm_90) reduce through them + const lw = switch (q) { + .fp8, .bf16 => |v| v, + else => unreachable, + }; + var part_bytes: usize = 4; + for ([_]usize{ 16, 128, m }) |r| part_bytes = @max(part_bytes, lin.lane.?.partBytes(r, lw)); + var part = try cuda.DeviceBuffer.alloc(gpu.d, part_bytes); + defer part.free(); + try lin.decode(s, q, .{ .x = dx.ptr, .part = part.ptr }, y.ptr, 16); try s.synchronize(); const wide = try check.download(gpu, y); defer a.free(wide); for (1..16) |rows| { try y2.fill8(0xff, s.handle); - try lin.decode(s, q, .{ .x = dx.ptr }, y2.ptr, rows); + try lin.decode(s, q, .{ .x = dx.ptr, .part = part.ptr }, y2.ptr, rows); try s.synchronize(); try same(gpu, "decode rows at a narrower width", y2, wide[0 .. rows * n * 2]); } var worst = [2]f64{ 0, 0 }; - try lin.prompt(s, q, .{ .x = dx.ptr }, y.ptr, m); + try lin.prompt(s, q, .{ .x = dx.ptr, .part = part.ptr }, y.ptr, m); try s.synchronize(); const prompt = try check.download(gpu, y); defer a.free(prompt); - try lin.prompt(s, q, .{ .x = dx.ptr }, y2.ptr, 128); + try lin.prompt(s, q, .{ .x = dx.ptr, .part = part.ptr }, y2.ptr, 128); try s.synchronize(); try same(gpu, "prompt rows in a chunk of 128", y2, prompt[0 .. 128 * n * 2]); for (0..m) |r| for (0..n) |c| { From 1f87b27911f942eaf4171c4544ae8b80cfc8fc83 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:27:13 +0000 Subject: [PATCH 03/35] Windows: pull into the hub cache, HOME, and the whole test step - pull: where Windows refuses the snapshot's symlink, the snapshot path becomes a hard link to its blob (CreateHardLinkW through `sys.c.link`): no extra space, and learning already refuses linked shards. Link failures are logged with their reason. - HOME falls back to USERPROFILE (Windows sets no HOME) for the hub cache and every libc getenv("HOME"). - Tests read optional fixture variables through `platform.getenv`; std.process.Environ.getPosix does not build on Windows. - The learned-state symlink tests skip on Windows (symlinks need privilege); `sys.c.symlink` asks for an unprivileged one. zig build test passes on x86_64-windows. --- zig/src/cli/hub.zig | 2 +- zig/src/cli/pull_file.zig | 18 ++++++++++++- zig/src/cluster/admission_test.zig | 3 ++- zig/src/cluster/cli_test.zig | 3 ++- zig/src/cluster/cost_test.zig | 3 ++- zig/src/cluster/estimate.zig | 3 ++- zig/src/cluster/model.zig | 3 ++- zig/src/cluster/plan_test.zig | 3 ++- zig/src/cluster/status.zig | 3 ++- zig/src/cluster/topology.zig | 3 ++- zig/src/core/exlr8_oracle_test.zig | 13 +++++----- zig/src/core/learned_part_test.zig | 3 +++ zig/src/platform/sys.zig | 5 ++++ zig/src/platform/win_c.zig | 41 +++++++++++++++++++++++++++++- zig/src/server/hub.zig | 2 +- 15 files changed, 90 insertions(+), 18 deletions(-) diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 64347ae50c..9b5edacb0e 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -38,7 +38,7 @@ pub fn cacheDir(a: Allocator, env: ?*const std.process.Environ.Map, override: ?[ } }.f; if (get(env, "HF_HUB_CACHE")) |c| return c; - const home = get(env, "HOME") orelse ""; + const home = get(env, "HOME") orelse get(env, "USERPROFILE") orelse ""; // Windows sets USERPROFILE, not HOME if (get(env, "HF_HOME")) |h| return std.fs.path.join(a, &.{ h, "hub" }); return std.fs.path.join(a, &.{ home, ".cache", "huggingface", "hub" }); } diff --git a/zig/src/cli/pull_file.zig b/zig/src/cli/pull_file.zig index c3de2654cf..1926f553ff 100644 --- a/zig/src/cli/pull_file.zig +++ b/zig/src/cli/pull_file.zig @@ -1,5 +1,6 @@ //! One file from the hub into the blob store: range resume, size and sha256 checks, retries, and its snapshot link. const std = @import("std"); +const platform = @import("sys"); const Allocator = std.mem.Allocator; /// One file the hub's tree lists for the revision: LFS files carry their sha256, small files their git blob sha1. @@ -143,7 +144,22 @@ pub fn linkIntoSnapshot(a: Allocator, io: std.Io, snapshot_dir: []const u8, path if (std.fs.path.dirname(link_path)) |parent| try std.Io.Dir.cwd().createDirPath(io, parent); std.Io.Dir.cwd().deleteFile(io, link_path) catch {}; const target = try std.fs.path.join(a, &.{ "..", "..", "blobs", blob_name }); - std.Io.Dir.cwd().symLink(io, target, link_path, .{}) catch return error.SymLinkFailed; + std.Io.Dir.cwd().symLink(io, target, link_path, .{}) catch |err| { + // Windows grants symlinks only with Developer Mode or elevation (and refuses some): a hard link to the blob, + // which costs no space (huggingface_hub copies instead); learning refuses linked shards either way + if (@import("builtin").os.tag != .windows) { + std.log.err("cannot link {s} to its blob: {t}", .{ link_path, err }); + return error.SymLinkFailed; + } + const repo_dir = std.fs.path.dirname(std.fs.path.dirname(snapshot_dir) orelse return error.SymLinkFailed) orelse return error.SymLinkFailed; + const blob = try std.fs.path.joinZ(a, &.{ repo_dir, "blobs", blob_name }); + const linked = try a.dupeSentinel(u8, link_path, 0); + _ = platform.c.unlink(linked); // the refused symlink can leave an empty file behind + if (platform.c.link(blob, linked) != 0) { + std.log.err("cannot link {s} to its blob: {t}", .{ link_path, platform.c.errno(-1) }); + return error.SymLinkFailed; + } + }; } test "a hub tree path stays inside the snapshot: relative, without . or .. segments" { diff --git a/zig/src/cluster/admission_test.zig b/zig/src/cluster/admission_test.zig index 1b5b0f1b3b..200fe2b111 100644 --- a/zig/src/cluster/admission_test.zig +++ b/zig/src/cluster/admission_test.zig @@ -1,5 +1,6 @@ //! 32 to 128 streams: what fits at 8k and 32k context, prompt rows a round, and 32 toy streams joining while others decode. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const plan = @import("plan.zig"); const budget = @import("budget.zig"); @@ -115,7 +116,7 @@ test "32 toy streams join while others decode, prompt chunks in the same rounds: } test "the admission tables for CLUSTER.md (TF_CLUSTER_REPORT=1)" { - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; + if (platform.getenv("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/cli_test.zig b/zig/src/cluster/cli_test.zig index 267fca93d6..bc7b7f58bd 100644 --- a/zig/src/cluster/cli_test.zig +++ b/zig/src/cluster/cli_test.zig @@ -1,5 +1,6 @@ //! The cluster commands end to end on fake hosts: init, check, status, serve with in-process bring-up, and node. const std = @import("std"); +const platform = @import("sys"); const cli = @import("cli.zig"); const fixtures = @import("probe_fixtures.zig"); @@ -86,7 +87,7 @@ const studios = ; test "check plans K3 from its config and Hub file list and names the raise a low node needs; serve stops before launching" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/cost_test.zig b/zig/src/cluster/cost_test.zig index dec13181ca..47293a6d20 100644 --- a/zig/src/cluster/cost_test.zig +++ b/zig/src/cluster/cost_test.zig @@ -1,5 +1,6 @@ //! K3 and GLM-5.3 rounds on four 512 GiB Studios at 1 to 128 rows: dedup, caches, collectives, and the MLA split. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const model = @import("model.zig"); const estimate = @import("estimate.zig"); @@ -66,7 +67,7 @@ test "GLM-5.3 in FP8 reads far less a token than K3, and 128 streams at 32k cont } test "the cost tables for CLUSTER.md (TF_CLUSTER_REPORT=1)" { - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; + if (platform.getenv("TF_CLUSTER_REPORT") == null) return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/estimate.zig b/zig/src/cluster/estimate.zig index 74917ba740..9d33e43313 100644 --- a/zig/src/cluster/estimate.zig +++ b/zig/src/cluster/estimate.zig @@ -1,5 +1,6 @@ //! A checkpoint's tensor table from its config and the Hub's file list, so a cluster can be planned before the weights land. const std = @import("std"); +const platform = @import("sys"); const checkpoint = @import("checkpoint.zig"); const model = @import("model.zig"); const roles = @import("roles.zig"); @@ -282,7 +283,7 @@ test "the K3 table matches every shard size to within its header" { } test "the real K3 file list (TF_K3_DIR): 96 files, 1.561 TB, every header under 0.01% of its file" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/model.zig b/zig/src/cluster/model.zig index 644b117264..91b37ea140 100644 --- a/zig/src/cluster/model.zig +++ b/zig/src/cluster/model.zig @@ -1,5 +1,6 @@ //! A model's shape from its config.json (text_config when nested): what the planner and the memory budgets need. const std = @import("std"); +const platform = @import("sys"); pub const Attention = enum(u8) { gqa, mla, linear }; @@ -257,7 +258,7 @@ test "GLM-5.3: every layer MLA with a 2,048-token indexer, 75 MoE layers, one MT } test "the real K3 config (TF_K3_DIR) parses to the literal" { - const dir = std.testing.environ.getPosix("TF_K3_DIR") orelse return error.SkipZigTest; + const dir = platform.getenv("TF_K3_DIR") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/cluster/plan_test.zig b/zig/src/cluster/plan_test.zig index 666ecb3ba2..0514244c35 100644 --- a/zig/src/cluster/plan_test.zig +++ b/zig/src/cluster/plan_test.zig @@ -1,5 +1,6 @@ //! Kimi K3 on four 512 GiB M3 Ultras: the default layout, layouts from a config, determinism, and refusals. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const model = @import("model.zig"); const estimate = @import("estimate.zig"); @@ -34,7 +35,7 @@ pub fn k3(a: std.mem.Allocator) !checkpoint.Checkpoint { const size: u64 = if (i == 0) 2_341_216_112 else if (i < s.layers) (if (kda) 16_990_911_504 else 16_567_501_776) else if (i == s.layers) 4_697_664_072 else if (i == s.layers + 1) 92_289_328 else 802_448_352; f.* = .{ .name = try std.fmt.allocPrint(a, "model-{d:0>5}-of-{d:0>6}.safetensors", .{ i + 1, s.layers + 3 }), .size = size }; } - if (std.testing.environ.getPosix("TF_K3_DIR")) |dir| { + if (platform.getenv("TF_K3_DIR")) |dir| { const text = try std.Io.Dir.cwd().readFileAlloc(std.testing.io, try std.fs.path.join(a, &.{ dir, "files.tsv" }), a, .limited(1 << 20)); files = try estimate.manifest(a, text); } diff --git a/zig/src/cluster/status.zig b/zig/src/cluster/status.zig index ecb6dd7457..49e0395dbf 100644 --- a/zig/src/cluster/status.zig +++ b/zig/src/cluster/status.zig @@ -1,5 +1,6 @@ //! What `tensorfold cluster status` and `check` print: nodes, links, memory, the placement per rank, costs and load progress. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); const topology = @import("topology.zig"); const plan_mod = @import("plan.zig"); @@ -99,7 +100,7 @@ test "the K3 status on four Studios prints every rank, fits, and the round costs try placement(&out.writer, &p, needs, fs); try costs(&out.writer, &p, needs, &s, 8192); const text = try a.dupe(u8, out.written()); - if (std.testing.environ.getPosix("TF_CLUSTER_REPORT") != null) { + if (platform.getenv("TF_CLUSTER_REPORT") != null) { std.debug.print("\n{s}\n", .{text}); const others = [_]struct { []const u8, plan_mod.Options }{ .{ "experts split by tensor parallelism", .{ .layout = .{ .tensor = 4, .expert = 1 } } }, diff --git a/zig/src/cluster/topology.zig b/zig/src/cluster/topology.zig index 6e5292c8dd..6c4d2d7564 100644 --- a/zig/src/cluster/topology.zig +++ b/zig/src/cluster/topology.zig @@ -1,5 +1,6 @@ //! The cluster's wiring from every node's ports: which port reaches which node, at what rate, and the graph's shape. const std = @import("std"); +const platform = @import("sys"); const node = @import("node.zig"); /// One cable: node `a`'s port to node `b`'s port (indices into the node list and their port lists). @@ -134,7 +135,7 @@ test "four Macs with six cables form a mesh; three cables a chain or a star; a c } test "the lent cluster's real probes (TF_CLUSTER_PROBES): four M3 Ultras in a full mesh at 80 Gb/s" { - const root = std.testing.environ.getPosix("TF_CLUSTER_PROBES") orelse return error.SkipZigTest; + const root = platform.getenv("TF_CLUSTER_PROBES") orelse return error.SkipZigTest; var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const a = arena.allocator(); diff --git a/zig/src/core/exlr8_oracle_test.zig b/zig/src/core/exlr8_oracle_test.zig index 13c854dee9..6bdb2addc2 100644 --- a/zig/src/core/exlr8_oracle_test.zig +++ b/zig/src/core/exlr8_oracle_test.zig @@ -1,6 +1,7 @@ //! EXLR8 decoding against fixtures the reference loader (MIT) wrote outside the tree: layers 0, 3, 40, 78, vocab. const std = @import("std"); +const platform = @import("sys"); const codec = @import("exlr8_codec.zig"); const ex = @import("exlr8_experts.zig"); const dense = @import("exlr8_dense.zig"); @@ -169,7 +170,7 @@ const Layers = struct { }; test "expert halves of layers 3, 40 and 78: stored checksums, the reference's tiles and W bits" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -193,7 +194,7 @@ test "expert halves of layers 3, 40 and 78: stored checksums, the reference's ti h.suh(suh[0..h.rows()]); h.svh(svh[0..h.cols()]); try checkSamples(codes, h.rows(), h.cols(), suh[0..h.rows()], svh[0..h.cols()], unit.w_samples, false, &stats); - if (std.testing.environ.getPosix("TF_EXLR8_ORACLE_FULL")) |full| try wholeHalf(full, unit, codes, h, suh[0..h.rows()], svh[0..h.cols()]); + if (platform.getenv("TF_EXLR8_ORACLE_FULL")) |full| try wholeHalf(full, unit, codes, h, suh[0..h.rows()], svh[0..h.cols()]); } stats.print("expert halves"); } @@ -233,7 +234,7 @@ fn wholeHalf(dir: []const u8, unit: Unit, codes: []const u8, h: ex.Half, suh: [] } test "every half of every stored width verifies, as the reference's --verify counts them" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit(); const fx = try load(arena.allocator()); @@ -260,7 +261,7 @@ test "every half of every stored width verifies, as the reference's --verify cou } test "dense and vocabulary files: records, checksums, every format bit for bit, plain tensors as shipped" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -331,7 +332,7 @@ fn denseBits(m: dense.Matrix, unit: DenseUnit) !void { } test "the eight manifests parse to the reference's widths, and every width they ask of a present layer is stored" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -371,7 +372,7 @@ test "the eight manifests parse to the reference's widths, and every width they } test "every present release file matches its published sha256 and size" { - const root = std.testing.environ.getPosix("TF_EXLR8_DIR") orelse return error.SkipZigTest; + const root = platform.getenv("TF_EXLR8_DIR") orelse return error.SkipZigTest; const gpa = std.testing.allocator; var checked: usize = 0; for ([_][]const u8{ "experts", "dense" }) |sub| { diff --git a/zig/src/core/learned_part_test.zig b/zig/src/core/learned_part_test.zig index 2f7ab53027..a2984e345d 100644 --- a/zig/src/core/learned_part_test.zig +++ b/zig/src/core/learned_part_test.zig @@ -85,6 +85,7 @@ test "orphan temporary pressure cannot evict an older valid learned state" { } test "owned-name links and directories stay untouched and foreign directories are not traversed" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "types"); @@ -123,6 +124,7 @@ test "owned-name links and directories stay untouched and foreign directories ar } test "an identity-directory symlink refuses open and preserves foreign payloads" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "identity-link"); @@ -145,6 +147,7 @@ test "an identity-directory symlink refuses open and preserves foreign payloads" } test "an explicitly configured root alias still cleans only its real identity children" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // symlinks need privilege on Windows const a = std.testing.allocator; var buf: [128]u8 = undefined; const root = try scratch(&buf, "root-target"); diff --git a/zig/src/platform/sys.zig b/zig/src/platform/sys.zig index ee614d7507..35bb853608 100644 --- a/zig/src/platform/sys.zig +++ b/zig/src/platform/sys.zig @@ -13,6 +13,11 @@ pub const posix = if (windows) @import("win_posix.zig") else std.posix; pub const DynLib = if (windows) WinDynLib else std.DynLib; +/// An environment variable's value, through libc (std.process.Environ.getPosix does not build on Windows). +pub fn getenv(name: [*:0]const u8) ?[:0]const u8 { + return std.mem.span(c.getenv(name) orelse return null); +} + /// The CUDA driver's file name: nvcuda.dll comes with the display driver on Windows. pub const cuda_driver = if (windows) "nvcuda.dll" else "libcuda.so.1"; diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig index 82b4dc7b3f..b21dce83d1 100644 --- a/zig/src/platform/win_c.zig +++ b/zig/src/platform/win_c.zig @@ -110,6 +110,8 @@ extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; extern "kernel32" fn GetDiskFreeSpaceExW(dir: [*:0]const u16, avail: ?*u64, total: ?*u64, free: ?*u64) callconv(.winapi) BOOL; extern "kernel32" fn GetTempPathW(len: u32, buf: [*]u16) callconv(.winapi) u32; extern "kernel32" fn GetModuleFileNameW(module: ?HANDLE, buf: [*]u16, len: u32) callconv(.winapi) u32; +extern "kernel32" fn CreateHardLinkW(new: [*:0]const u16, existing: [*:0]const u16, security: ?*anyopaque) callconv(.winapi) BOOL; +extern "kernel32" fn CreateSymbolicLinkW(link: [*:0]const u16, target: [*:0]const u16, flags: u32) callconv(.winapi) u8; extern "kernel32" fn CloseHandle(h: HANDLE) callconv(.winapi) BOOL; extern "c" fn _errno() *c_int; @@ -129,7 +131,16 @@ extern "c" fn _getpid() c_int; extern "c" fn _isatty(fd: c_int) c_int; extern "c" fn _fstat64(fd: c_int, st: *Stat64) c_int; extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; -pub extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; +const ucrt = struct { + extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; +}; + +/// libc's getenv, except that HOME, which Windows does not set, falls back to USERPROFILE. +pub fn getenv(name: [*:0]const u8) ?[*:0]u8 { + if (ucrt.getenv(name)) |v| return v; + if (std.mem.eql(u8, std.mem.span(name), "HOME")) return ucrt.getenv("USERPROFILE"); + return null; +} pub extern "c" fn fopen(path: [*:0]const u8, mode: [*:0]const u8) ?*std.c.FILE; pub extern "c" fn fwrite(ptr: [*]const u8, size: usize, n: usize, f: *std.c.FILE) usize; pub extern "c" fn fclose(f: *std.c.FILE) c_int; @@ -422,6 +433,16 @@ pub fn rename(from: [*:0]const u8, to: [*:0]const u8) c_int { return 0; } +/// POSIX link: `new` names the same file as `old` (same volume). +pub fn link(old: [*:0]const u8, new: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wa = wide(&a, std.mem.span(old)) orelse return -1; + const wb = wide(&b, std.mem.span(new)) orelse return -1; + if (CreateHardLinkW(wb, wa, null) == 0) return fail(lastError()); + return 0; +} + pub fn access(path: [*:0]const u8, mode: c_int) c_int { var buf: Wide = undefined; return _waccess(wide(&buf, std.mem.span(path)) orelse return -1, mode & 6); @@ -575,3 +596,21 @@ pub fn fchmodat(dir: fd_t, path: [*:0]const u8, mode: mode_t, flags: u32) c_int _ = .{ dir, path, mode, flags }; return 0; } + +/// Symbolic links need Developer Mode or elevation on Windows; this asks for an unprivileged one. +pub fn symlink(target: [*:0]const u8, path: [*:0]const u8) c_int { + var a: Wide = undefined; + var b: Wide = undefined; + const wt = wide(&a, std.mem.span(target)) orelse return -1; + const wp = wide(&b, std.mem.span(path)) orelse return -1; + const attrs = GetFileAttributesW(wt); + const dir: u32 = if (attrs != 0xFFFF_FFFF and attrs & 0x10 != 0) 0x1 else 0; + if (CreateSymbolicLinkW(wp, wt, dir | 0x2) == 0) return fail(lastError()); // ALLOW_UNPRIVILEGED_CREATE + return 0; +} + +/// Reading a link's target needs its reparse data; nothing on Windows asks for it yet. +pub fn readlink(path: [*:0]const u8, buf: [*]u8, len: usize) isize { + _ = .{ path, buf, len }; + return fail(.NOSYS); +} diff --git a/zig/src/server/hub.zig b/zig/src/server/hub.zig index 5a155e6a1b..4e0efb1a5a 100644 --- a/zig/src/server/hub.zig +++ b/zig/src/server/hub.zig @@ -21,7 +21,7 @@ fn isDir(io: std.Io, path: []const u8) bool { /// The checkpoint directory for ``model``, or null with ``problem`` set (no downloads: the cache must hold it). pub fn resolve(a: Allocator, io: std.Io, environ: ?*const std.process.Environ.Map, model: []const u8, problem: *[]const u8) !?[]const u8 { - const home = if (environ) |e| e.get("HOME") orelse "" else ""; + const home = if (environ) |e| e.get("HOME") orelse e.get("USERPROFILE") orelse "" else ""; // Windows sets USERPROFILE const expanded = if (std.mem.startsWith(u8, model, "~/")) try std.fs.path.join(a, &.{ home, model[2..] }) else model; if (isDir(io, expanded)) return expanded; if (!isRepoId(model)) { From 996a50fc3202b0a42fc86e417b1674ec8dff366d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:36:05 +0000 Subject: [PATCH 04/35] Windows: build notes, strip before install, formatting - docs/WINDOWS.md: requirements, build, run, how the port works and how it differs from Linux; the README points to it. - -Dstrip reaches tensorfold-native before its install step is made, so the step no longer expects a .pdb the stripped build does not write. - zig fmt on the platform files. --- README.md | 2 + docs/WINDOWS.md | 53 +++++++++++++++++++ zig/build/cuda.zig | 8 +-- zig/src/platform/win_c.zig | 95 ++++++---------------------------- zig/src/platform/win_clock.zig | 82 +++++++++++++++++++++++++++++ zig/src/platform/win_posix.zig | 2 +- 6 files changed, 159 insertions(+), 83 deletions(-) create mode 100644 docs/WINDOWS.md create mode 100644 zig/src/platform/win_clock.zig diff --git a/README.md b/README.md index d77b9f9f3e..31155dae9c 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,8 @@ tensorfold serve "$HOME/models/nemotron-lightning" \ Use a complete local model directory or an existing Hugging Face cache. The binary in a release archive is `bin/tensorfold-native`; the [runbook](RUNBOOK.md) covers archive installation and CUDA runtime files. + +On Windows with an NVIDIA GPU, build `tensorfold-native.exe` from source as [docs/WINDOWS.md](docs/WINDOWS.md) describes. Use `tensorfold-native --help` or `-h`, and `tensorfold-native capabilities --json`, to inspect the installed binary. Model weights have separate downloads and licenses. diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md new file mode 100644 index 0000000000..7869f11b8d --- /dev/null +++ b/docs/WINDOWS.md @@ -0,0 +1,53 @@ +# TensorFold on Windows (native) + +An unofficial port: `tensorfold-native.exe` runs on Windows x86_64 with an NVIDIA GPU through the same CUDA engine +the Linux build uses. Only the CUDA backend is ported; Metal is macOS's and HIP is not built. + +## What you need + +- An NVIDIA GPU and a current display driver (it ships `nvcuda.dll`; no CUDA toolkit is needed to run). +- To build: [Zig 0.17.0](https://ziglang.org/download/) (the version in `.zig-version`), the CUDA Toolkit + (13.x tested) for `nvcc`, and Visual Studio Build Tools with the C++ workload (`cl.exe`, nvcc's host compiler). + +## Build + +From a "x64 Native Tools" prompt (or after running `vcvars64.bat`), so `cl.exe` is on `PATH`: + +```bat +zig build native install -Dsm=89 -Dstrip=true ^ + -Dnvcc="%CUDA_PATH%\bin\nvcc.exe" +``` + +`-Dsm` takes the GPU's compute capability without the dot, comma separated for several: 86 (RTX 30), 89 (RTX 40), +120 (RTX 50). The server is `zig-out\native\bin\tensorfold-native.exe`; `zig-out\bin\tf-cuda-test.exe` runs the GPU +suites (`tf-cuda-test info`, `smoke`, `glue`, `sample`, `quant-ref`, `modelopt`, `exlr8`). + +`zig build test` runs the host tests without a GPU. + +## Run + +```bat +tensorfold-native pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit +tensorfold-native models +tensorfold-native serve --name local-model --parallel 1 --context 8192 --temperature 0 --no-thinking +``` + +The server listens on `http://127.0.0.1:8080/v1`. Ctrl+C (or closing the console) stops it cleanly. + +## How the port works + +- `zig/src/platform/sys.zig` (`@import("sys")`, wired into every module by `zig/build/sys.zig`) is the operating + system as the engine calls it. On Linux and macOS it is `std.c`, `std.posix` and `std.DynLib` unchanged; on Windows + `win_c.zig` and `win_posix.zig` give the same names over the UCRT, Win32 and Winsock. +- Files are opened with every share mode, so rename-over and unlink-while-open behave as on POSIX. `O_DIRECT` maps to + unbuffered reads. `/tmp` is the user's temporary folder and `HOME` falls back to `USERPROFILE`. +- The Hugging Face cache links snapshot files to blobs. Windows allows symlinks only with Developer Mode or elevation, + so `pull` falls back to hard links (no extra space). +- Checkpoint reads use threads (io_uring is Linux's). +- The API key file is checked for type and size; Windows keeps access in ACLs, not mode bits, so the "unreadable by + other users" check is the folder's ACL rather than `chmod 600`. + +## Differences from Linux + +- No io_uring, `ulimit` or GB10 carveout, and no `SIGHUP` (the key file is still reread when it changes). +- `TF_HEAT_*` thermal pacing reads Linux's `/sys/class/thermal` and does nothing on Windows. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 4f54eaf76d..a291e9bbec 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -139,7 +139,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; + _ = nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true, strip); } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. @@ -161,7 +161,7 @@ fn checkpointCli(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. } /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it. -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile { +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool, strip: ?bool) *std.Build.Step.Compile { const m = engines(b, target, optimize, cuda, lanes, nemotron); // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's) const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); @@ -177,6 +177,8 @@ fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b exe.root_module.strip = true; return exe; } + // before the install step, which decides then whether a .pdb comes with the executable + exe.root_module.strip = strip; const install = b.addInstallArtifact(exe, .{ .dest_dir = .{ .override = .{ .custom = "native/bin" } } }); b.step("native", "tensorfold-native with the CUDA engines into zig-out/native/bin").dependOn(&install.step); return exe; @@ -248,7 +250,7 @@ pub fn distServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *s }; const cuda = runtime(b, target, .fast, if (prebuilt != null) &images else &.{}); const mods = family(b, target, .fast, cuda, draft_ids); - return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false, null); } /// SASS targets every release fatbin holds: Ampere (sm_86), RTX 50 and RTX PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig index b21dce83d1..eb4d0c746a 100644 --- a/zig/src/platform/win_c.zig +++ b/zig/src/platform/win_c.zig @@ -100,13 +100,6 @@ extern "kernel32" fn MoveFileExW(from: [*:0]const u16, to: [*:0]const u16, flags extern "kernel32" fn FindFirstFileW(pattern: [*:0]const u16, data: *FindData) callconv(.winapi) isize; extern "kernel32" fn FindNextFileW(h: HANDLE, data: *FindData) callconv(.winapi) BOOL; extern "kernel32" fn FindClose(h: HANDLE) callconv(.winapi) BOOL; -extern "kernel32" fn QueryPerformanceCounter(count: *i64) callconv(.winapi) BOOL; -extern "kernel32" fn QueryPerformanceFrequency(freq: *i64) callconv(.winapi) BOOL; -extern "kernel32" fn GetSystemTimePreciseAsFileTime(t: *FILETIME) callconv(.winapi) void; -extern "kernel32" fn CreateWaitableTimerExW(security: ?*anyopaque, name: ?[*:0]const u16, flags: u32, access: u32) callconv(.winapi) ?HANDLE; -extern "kernel32" fn SetWaitableTimer(h: HANDLE, due: *const i64, period: i32, routine: ?*anyopaque, arg: ?*anyopaque, resume_: BOOL) callconv(.winapi) BOOL; -extern "kernel32" fn WaitForSingleObject(h: HANDLE, ms: u32) callconv(.winapi) u32; -extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; extern "kernel32" fn GetDiskFreeSpaceExW(dir: [*:0]const u16, avail: ?*u64, total: ?*u64, free: ?*u64) callconv(.winapi) BOOL; extern "kernel32" fn GetTempPathW(len: u32, buf: [*]u16) callconv(.winapi) u32; extern "kernel32" fn GetModuleFileNameW(module: ?HANDLE, buf: [*]u16, len: u32) callconv(.winapi) u32; @@ -127,10 +120,8 @@ extern "c" fn _wunlink(path: [*:0]const u16) c_int; extern "c" fn _wrmdir(path: [*:0]const u16) c_int; extern "c" fn _wmkdir(path: [*:0]const u16) c_int; extern "c" fn _waccess(path: [*:0]const u16, mode: c_int) c_int; -extern "c" fn _getpid() c_int; extern "c" fn _isatty(fd: c_int) c_int; extern "c" fn _fstat64(fd: c_int, st: *Stat64) c_int; -extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; const ucrt = struct { extern "c" fn getenv(name: [*:0]const u8) ?[*:0]u8; }; @@ -150,8 +141,8 @@ pub extern "c" fn malloc(n: usize) ?*anyopaque; /// The UCRT's struct _stat64. const Stat64 = extern struct { dev: u32, ino: u16, mode: u16, nlink: i16, uid: i16, gid: i16, rdev: u32, size: i64, atime: i64, mtime: i64, ctime: i64 }; -fn fail(e: E) c_int { - _errno().* = @intFromEnum(e); +pub fn fail(e: E) c_int { + _errno().* = @backingInt(e); return -1; } @@ -272,10 +263,10 @@ fn openFile(path: [*:0]const u8, flags: O) c_int { .RDWR => generic_read | generic_write, }; const disposition: u32 = if (flags.CREAT and flags.EXCL) 1 // CREATE_NEW - else if (flags.CREAT and flags.TRUNC) 2 // CREATE_ALWAYS - else if (flags.CREAT) 4 // OPEN_ALWAYS - else if (flags.TRUNC) 5 // TRUNCATE_EXISTING - else 3; // OPEN_EXISTING + else if (flags.CREAT and flags.TRUNC) 2 // CREATE_ALWAYS + else if (flags.CREAT) 4 // OPEN_ALWAYS + else if (flags.TRUNC) 5 // TRUNCATE_EXISTING + else 3; // OPEN_EXISTING var attrs: u32 = 0x80; // FILE_ATTRIBUTE_NORMAL if (flags.DIRECT) attrs |= 0x2000_0000; // FILE_FLAG_NO_BUFFERING const h = CreateFileW(w, rights, 0x7, null, disposition, attrs, null); // share read, write and delete @@ -368,7 +359,7 @@ pub fn fsync(fd: fd_t) c_int { pub fn ftruncate(fd: fd_t, size: off_t) c_int { const rc = _chsize_s(fd, size); - if (rc != 0) return fail(@enumFromInt(rc)); + if (rc != 0) return fail(@fromBackingInt(@intCast(rc))); return 0; } @@ -517,70 +508,16 @@ pub fn closedir(d: *DIR) c_int { return 0; } -// ---- process, environment, time ---- +// ---- process, environment, time: win_clock.zig ---- -pub fn getpid() pid_t { - return _getpid(); -} - -pub fn getuid() u32 { - return 0; -} - -pub fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int { - if (overwrite == 0 and getenv(name) != null) return 0; - if (_putenv_s(name, value) != 0) return fail(.INVAL); - return 0; -} - -pub fn unsetenv(name: [*:0]const u8) c_int { - if (_putenv_s(name, "") != 0) return fail(.INVAL); - return 0; -} - -threadlocal var sleep_timer: ?HANDLE = null; - -/// High-resolution waitable timer per thread, so microsecond waits are not rounded to the 15.6 ms tick. -pub fn nanosleep(req: *const timespec, rem: ?*timespec) c_int { - if (rem) |r| r.* = .{ .sec = 0, .nsec = 0 }; - const ns: i64 = req.sec * std.time.ns_per_s + req.nsec; - if (ns <= 0) return 0; - if (sleep_timer == null) sleep_timer = CreateWaitableTimerExW(null, null, 0x2, 0x1F0003); // HIGH_RESOLUTION, TIMER_ALL_ACCESS - if (sleep_timer) |t| { - const due: i64 = -@divFloor(ns + 99, 100); // relative, in 100 ns units - if (SetWaitableTimer(t, &due, 0, null, null, 0) != 0) { - _ = WaitForSingleObject(t, 0xFFFF_FFFF); - return 0; - } - } - Sleep(@intCast(@max(1, @divFloor(ns + std.time.ns_per_ms - 1, std.time.ns_per_ms)))); - return 0; -} - -pub fn clock_gettime(clock: c_int, out: *timespec) c_int { - if (clock == CLOCK.MONOTONIC) { - var count: i64 = 0; - var freq: i64 = 0; - _ = QueryPerformanceCounter(&count); - _ = QueryPerformanceFrequency(&freq); - out.* = .{ .sec = @divFloor(count, freq), .nsec = @intCast(@divFloor(@mod(count, freq) * std.time.ns_per_s, freq)) }; - return 0; - } - var ft: FILETIME = undefined; - GetSystemTimePreciseAsFileTime(&ft); - const since_1601: i64 = @intCast((@as(u64, ft.high) << 32) | ft.low); - const unix = since_1601 - 116_444_736_000_000_000; - out.* = .{ .sec = @divFloor(unix, 10_000_000), .nsec = @intCast(@mod(unix, 10_000_000) * 100) }; - return 0; -} - -pub fn gettimeofday(tv: *timeval, tz: ?*anyopaque) c_int { - _ = tz; - var ts: timespec = undefined; - _ = clock_gettime(CLOCK.REALTIME, &ts); - tv.* = .{ .sec = ts.sec, .usec = @intCast(@divFloor(ts.nsec, 1000)) }; - return 0; -} +const clock = @import("win_clock.zig"); +pub const getpid = clock.getpid; +pub const getuid = clock.getuid; +pub const setenv = clock.setenv; +pub const unsetenv = clock.unsetenv; +pub const nanosleep = clock.nanosleep; +pub const clock_gettime = clock.clock_gettime; +pub const gettimeofday = clock.gettimeofday; /// Bytes free to this user under `dir` (statvfs's bavail * frsize). pub fn diskAvailable(dir: [*:0]const u8) ?u64 { diff --git a/zig/src/platform/win_clock.zig b/zig/src/platform/win_clock.zig new file mode 100644 index 0000000000..ec41325cfc --- /dev/null +++ b/zig/src/platform/win_clock.zig @@ -0,0 +1,82 @@ +//! The process, environment and clock part of the libc surface on Windows (win_c.zig re-exports these): the process +//! id, environment writes over the UCRT, and nanosleep and clock_gettime over the kernel's timers and counters. + +const std = @import("std"); +const c = @import("win_c.zig"); + +const HANDLE = *anyopaque; +const BOOL = c_int; +const FILETIME = extern struct { low: u32, high: u32 }; + +extern "kernel32" fn QueryPerformanceCounter(count: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn QueryPerformanceFrequency(freq: *i64) callconv(.winapi) BOOL; +extern "kernel32" fn GetSystemTimePreciseAsFileTime(t: *FILETIME) callconv(.winapi) void; +extern "kernel32" fn CreateWaitableTimerExW(security: ?*anyopaque, name: ?[*:0]const u16, flags: u32, access: u32) callconv(.winapi) ?HANDLE; +extern "kernel32" fn SetWaitableTimer(h: HANDLE, due: *const i64, period: i32, routine: ?*anyopaque, arg: ?*anyopaque, resume_: BOOL) callconv(.winapi) BOOL; +extern "kernel32" fn WaitForSingleObject(h: HANDLE, ms: u32) callconv(.winapi) u32; +extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; +extern "c" fn _getpid() c_int; +extern "c" fn _putenv_s(name: [*:0]const u8, value: [*:0]const u8) c_int; + +pub fn getpid() c.pid_t { + return _getpid(); +} + +pub fn getuid() u32 { + return 0; +} + +pub fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int { + if (overwrite == 0 and c.getenv(name) != null) return 0; + if (_putenv_s(name, value) != 0) return c.fail(.INVAL); + return 0; +} + +pub fn unsetenv(name: [*:0]const u8) c_int { + if (_putenv_s(name, "") != 0) return c.fail(.INVAL); + return 0; +} + +threadlocal var sleep_timer: ?HANDLE = null; + +/// High-resolution waitable timer per thread, so microsecond waits are not rounded to the 15.6 ms tick. +pub fn nanosleep(req: *const c.timespec, rem: ?*c.timespec) c_int { + if (rem) |r| r.* = .{ .sec = 0, .nsec = 0 }; + const ns: i64 = req.sec * std.time.ns_per_s + req.nsec; + if (ns <= 0) return 0; + if (sleep_timer == null) sleep_timer = CreateWaitableTimerExW(null, null, 0x2, 0x1F0003); // HIGH_RESOLUTION, TIMER_ALL_ACCESS + if (sleep_timer) |t| { + const due: i64 = -@divFloor(ns + 99, 100); // relative, in 100 ns units + if (SetWaitableTimer(t, &due, 0, null, null, 0) != 0) { + _ = WaitForSingleObject(t, 0xFFFF_FFFF); + return 0; + } + } + Sleep(@intCast(@max(1, @divFloor(ns + std.time.ns_per_ms - 1, std.time.ns_per_ms)))); + return 0; +} + +pub fn clock_gettime(clock: c_int, out: *c.timespec) c_int { + if (clock == c.CLOCK.MONOTONIC) { + var count: i64 = 0; + var freq: i64 = 0; + _ = QueryPerformanceCounter(&count); + _ = QueryPerformanceFrequency(&freq); + out.* = .{ .sec = @divFloor(count, freq), .nsec = @intCast(@divFloor(@mod(count, freq) * std.time.ns_per_s, freq)) }; + return 0; + } + var ft: FILETIME = undefined; + GetSystemTimePreciseAsFileTime(&ft); + const since_1601: i64 = @intCast((@as(u64, ft.high) << 32) | ft.low); + const unix = since_1601 - 116_444_736_000_000_000; + out.* = .{ .sec = @divFloor(unix, 10_000_000), .nsec = @intCast(@mod(unix, 10_000_000) * 100) }; + return 0; +} + +pub fn gettimeofday(tv: *c.timeval, tz: ?*anyopaque) c_int { + _ = tz; + var ts: c.timespec = undefined; + _ = clock_gettime(c.CLOCK.REALTIME, &ts); + tv.* = .{ .sec = ts.sec, .usec = @intCast(@divFloor(ts.nsec, 1000)) }; + return 0; +} diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index 4fa75167d9..04c82d4dba 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -109,7 +109,7 @@ fn failed() c_int { 10061 => .CONNREFUSED, else => .IO, }; - _errno().* = @intFromEnum(e); + _errno().* = @backingInt(e); return -1; } From aac93d22e73c9df9ae6ac11b54533e4bad57d158 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:45:14 +0000 Subject: [PATCH 05/35] Default port 8888 8080 and 8000 are taken by too many other servers; serve listens on 8888 unless --port says otherwise, and the docs and bench scripts follow. --- CONTRIBUTING.md | 4 ++-- README.md | 10 +++++----- RUNBOOK.md | 8 ++++---- docs/WINDOWS.md | 2 +- docs/api.md | 4 ++-- docs/cache-benchmarks.md | 4 ++-- docs/control.md | 2 +- docs/cortheon.md | 4 ++-- docs/living-weights.md | 14 +++++++------- docs/recipes/README.md | 2 +- docs/recipes/qwen3.8-27b.md | 4 ++-- docs/recipes/qwen3.8-flash-next.md | 4 ++-- docs/speed-up-mode.md | 2 +- docs/vision.md | 2 +- tools/bench_concurrent.py | 2 +- tools/bench_openai.py | 2 +- zig/src/server/cli.zig | 2 +- 17 files changed, 36 insertions(+), 36 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5a9c4d1014..0630be4816 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -54,9 +54,9 @@ Record pass, fail and skip counts and explain any unrun platform or feature. For served decode, the existing standard-library client needs Python 3.11 or newer: ```sh -python3 tools/bench_openai.py http://127.0.0.1:8080 local-model \ +python3 tools/bench_openai.py http://127.0.0.1:8888 local-model \ --tokens 256 --reps 3 --temperatures 0 --label candidate --output decode.json -python3 tools/bench_concurrent.py http://127.0.0.1:8080 local-model \ +python3 tools/bench_concurrent.py http://127.0.0.1:8888 local-model \ --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial --output concurrent.json ``` diff --git a/README.md b/README.md index 31155dae9c..a6a018278f 100644 --- a/README.md +++ b/README.md @@ -28,12 +28,12 @@ On Windows with an NVIDIA GPU, build `tensorfold-native.exe` from source as [doc Use `tensorfold-native --help` or `-h`, and `tensorfold-native capabilities --json`, to inspect the installed binary. Model weights have separate downloads and licenses. -The API listens at `http://127.0.0.1:8080/v1` by default. +The API listens at `http://127.0.0.1:8888/v1` by default. It serves OpenAI chat completions, completions and Responses, plus Anthropic Messages. Use the model ID returned by `/v1/models`, or set one with `--name`. ```sh -curl -fsS http://127.0.0.1:8080/v1/chat/completions \ +curl -fsS http://127.0.0.1:8888/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{"model":"local-model","messages":[{"role":"user","content":"Say hello in one sentence."}],"max_tokens":128,"temperature":0}' ``` @@ -49,11 +49,11 @@ tensorfold pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit src=$(ls -d ~/.cache/huggingface/hub/models--TensorFold--NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit/snapshots/* | head -1) mkdir -p ~/models && cp -cRL "$src" ~/models/nemotron-living # always teach a copy tensorfold serve ~/models/nemotron-living --slide -curl -N http://127.0.0.1:8080/v1/slide/learn -H 'Content-Type: application/json' -d '{"text": "I like blue."}' +curl -N http://127.0.0.1:8888/v1/slide/learn -H 'Content-Type: application/json' -d '{"text": "I like blue."}' ``` Send it a document instead and it finds the facts in it and learns them one by one. Open -`http://127.0.0.1:8080/slide` to watch in 3D: the facts, the neurons each one took in the model's layers, the shard +`http://127.0.0.1:8888/slide` to watch in 3D: the facts, the neurons each one took in the model's layers, the shard files that hold them, and the lesson running now. Learning runs on Apple-silicon Macs and on an NVIDIA GB10 (DGX Spark), and a folder that learned on one serves on the @@ -152,7 +152,7 @@ The binary's `capabilities --json` response lists its supported flags and platfo | Flag | Meaning | | --- | --- | -| `--host HOST`, `--port PORT` | Listen address and port, default `127.0.0.1:8080`. | +| `--host HOST`, `--port PORT` | Listen address and port, default `127.0.0.1:8888`. | | `--name NAME` | Model ID advertised to clients. | | `--alias NAME` | Another accepted model ID; repeat for multiple aliases. | | `--api-key KEY` | Require a bearer key; repeat for multiple keys. | diff --git a/RUNBOOK.md b/RUNBOOK.md index 408f2d91e8..0d0282fb49 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -104,14 +104,14 @@ bin/tensorfold-native serve "$HOME/models/nemotron-lightning" \ Leave the server running and check it from another terminal: ```sh -curl -fsS http://127.0.0.1:8080/health -curl -fsS http://127.0.0.1:8080/v1/models -curl -fsS http://127.0.0.1:8080/v1/chat/completions \ +curl -fsS http://127.0.0.1:8888/health +curl -fsS http://127.0.0.1:8888/v1/models +curl -fsS http://127.0.0.1:8888/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{"model":"local-model","messages":[{"role":"user","content":"Say hello in one sentence."}],"max_tokens":128,"temperature":0}' ``` -Use `http://127.0.0.1:8080/v1` as the OpenAI client base URL. +Use `http://127.0.0.1:8888/v1` as the OpenAI client base URL. For streaming, add `"stream": true` to the request. The native server also has `/v1/responses` and Anthropic `/v1/messages` routes. Serve flags and request overrides are described in [README.md](README.md#serve-flags). diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index 7869f11b8d..fdb9b7f88f 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -32,7 +32,7 @@ tensorfold-native models tensorfold-native serve --name local-model --parallel 1 --context 8192 --temperature 0 --no-thinking ``` -The server listens on `http://127.0.0.1:8080/v1`. Ctrl+C (or closing the console) stops it cleanly. +The server listens on `http://127.0.0.1:8888/v1`. Ctrl+C (or closing the console) stops it cleanly. ## How the port works diff --git a/docs/api.md b/docs/api.md index a0563b0048..657eca442a 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1,6 +1,6 @@ # Compatible APIs -The base URL is `http://127.0.0.1:8080/v1` with the default server settings. +The base URL is `http://127.0.0.1:8888/v1` with the default server settings. ## Native chat template override @@ -480,5 +480,5 @@ redacted thinking and other context edits return HTTP 400. Connect Claude Code using the ID from `/v1/models`: ```bash -ANTHROPIC_BASE_URL=http://127.0.0.1:8080 ANTHROPIC_API_KEY=local claude --model local-model +ANTHROPIC_BASE_URL=http://127.0.0.1:8888 ANTHROPIC_API_KEY=local claude --model local-model ``` diff --git a/docs/cache-benchmarks.md b/docs/cache-benchmarks.md index 2c9b9b6f95..e3977293f4 100644 --- a/docs/cache-benchmarks.md +++ b/docs/cache-benchmarks.md @@ -23,7 +23,7 @@ The standard-library-only cache client runs one first request, identical repeate concurrent group using the same prompt, greedy temperature, seed and reply limit. There is no warm-up: ```sh -python3 tools/bench_cache.py http://127.0.0.1:8080 MODEL \ +python3 tools/bench_cache.py http://127.0.0.1:8888 MODEL \ --tokens 64 --repeats 2 --streams 4 --require-cold --output cache.json ``` @@ -36,7 +36,7 @@ prefix retention (`--prompt-cache-gib 0`, where supported) for independent cold- For exact 2K/8K/32K/64K raw prefill, supply a JSON array of IDs from the checkpoint's own tokenizer: ```sh -python3 tools/bench_cache.py http://127.0.0.1:8080 MODEL \ +python3 tools/bench_cache.py http://127.0.0.1:8888 MODEL \ --prompt-file public-8192-tokens.json --expected-prompt-tokens 8192 \ --tokens 64 --require-cold --output cache-8192.json ``` diff --git a/docs/control.md b/docs/control.md index 4ed791242f..eb5e849364 100644 --- a/docs/control.md +++ b/docs/control.md @@ -118,7 +118,7 @@ The error from `tensorfold tui` prints that command with the venv's Python. `ric ```bash tensorfold tui tensorfold tui --profile code -tensorfold tui --url http://127.0.0.1:8080/v1 +tensorfold tui --url http://127.0.0.1:8888/v1 tensorfold tui --url https://spark.example.invalid --token-env TENSORFOLD_API_KEY tensorfold tui --demo tensorfold tui --demo --snapshot preview.svg diff --git a/docs/cortheon.md b/docs/cortheon.md index d8bee40e3e..bf6b242be1 100644 --- a/docs/cortheon.md +++ b/docs/cortheon.md @@ -100,8 +100,8 @@ tensorfold serve TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit --cor ``` ```bash -curl -s http://127.0.0.1:8080/v1/chat/completions -H 'Content-Type: application/json' \ +curl -s http://127.0.0.1:8888/v1/chat/completions -H 'Content-Type: application/json' \ -d '{"messages": [{"role": "user", "content": "What changed in the latest Zig release?"}]}' ``` -Point any OpenAI-compatible client at `http://127.0.0.1:8080/v1` as before. +Point any OpenAI-compatible client at `http://127.0.0.1:8888/v1` as before. diff --git a/docs/living-weights.md b/docs/living-weights.md index 9fd99bdb83..b7004b1e6c 100644 --- a/docs/living-weights.md +++ b/docs/living-weights.md @@ -61,12 +61,12 @@ tensorfold serve ~/models/nemotron-living --slide --slide-graph ~/models/nemotro `--slide` turns learning on. `--slide-graph` keeps the list of learned facts in a file, so it survives a restart; without it the list lives in memory, though the weights keep what they learned either way. The server listens on -`http://127.0.0.1:8080` unless you pass `--host` or `--port`. +`http://127.0.0.1:8888` unless you pass `--host` or `--port`. ### 4. Teach it a fact ```sh -curl -N http://127.0.0.1:8080/v1/slide/learn \ +curl -N http://127.0.0.1:8888/v1/slide/learn \ -H 'Content-Type: application/json' \ -d '{"text": "I like blue."}' ``` @@ -82,7 +82,7 @@ answer. A fact takes 1.5 to 3 minutes. Use the normal chat API, from any client: ```sh -curl http://127.0.0.1:8080/v1/chat/completions -H 'Content-Type: application/json' \ +curl http://127.0.0.1:8888/v1/chat/completions -H 'Content-Type: application/json' \ -d '{"messages": [{"role": "user", "content": "What colour do I like?"}]}' ``` @@ -112,7 +112,7 @@ nothing behind. An earlier fact's neuron that the new lesson would disturb is re - Several facts at once: put them in one text, as in `{"text": "I like blue. My sister is called Ana."}`. Each is learned in turn. - A file: send its contents as the text, for example - `jq -Rs '{text: .}' notes.md | curl -N http://127.0.0.1:8080/v1/slide/learn -H 'Content-Type: application/json' -d @-`. + `jq -Rs '{text: .}' notes.md | curl -N http://127.0.0.1:8888/v1/slide/learn -H 'Content-Type: application/json' -d @-`. - A web page: fetch its text first, or use omp's `/learn ` below, which does it for you. The server keeps answering chat requests while it learns: learning runs when no reply is being generated. Short, @@ -136,7 +136,7 @@ A panel shows each fact as queued, learning, learned or missed. ## See what it learned -Open `http://127.0.0.1:8080/slide` in a browser while the server runs. It draws the model in 3D and updates live: +Open `http://127.0.0.1:8888/slide` in a browser while the server runs. It draws the model in 3D and updates live: - The model as a tower of its 52 layers. Each MoE layer is a ring carrying its 64 pocket neurons, and a neuron glows in the colour of the fact that owns it. @@ -150,9 +150,9 @@ Drag to turn, scroll or pinch to zoom, and click a fact for its neurons, the que rounds. Phones get a bottom sheet. Without WebGL2, or with `/slide?flat`, the page draws a flat graph from the same data. -- `curl http://127.0.0.1:8080/v1/slide/graph` returns the same data as JSON: each fact's state, neurons (layer, index +- `curl http://127.0.0.1:8888/v1/slide/graph` returns the same data as JSON: each fact's state, neurons (layer, index and key), questions and rounds, and the model's layers, pocket and shards. -- `curl -X DELETE http://127.0.0.1:8080/v1/slide/graph` clears the graph. The weights keep what they learned. +- `curl -X DELETE http://127.0.0.1:8888/v1/slide/graph` clears the graph. The weights keep what they learned. ## Keep, move or reset what it learned diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 1e38e07605..28698ef465 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -61,7 +61,7 @@ how matrix multiplication uses a GPU. The chat fixture disables thinking. After `--name bench`, run from the repository root: ```bash -python3 tools/bench_openai.py http://127.0.0.1:8080 bench \ +python3 tools/bench_openai.py http://127.0.0.1:8888 bench \ --tokens 64 --reps 5 --temperatures 1.0,0 --output bench.json ``` diff --git a/docs/recipes/qwen3.8-27b.md b/docs/recipes/qwen3.8-27b.md index d4084dcfc1..6027d7d01a 100644 --- a/docs/recipes/qwen3.8-27b.md +++ b/docs/recipes/qwen3.8-27b.md @@ -90,7 +90,7 @@ input is qualified on this checkpoint. ```bash tensorfold pull nvidia/Qwen3.8-27B-NVFP4 z-lab/Qwen3.8-27B-DFlash2 -tensorfold serve nvidia/Qwen3.8-27B-NVFP4 --host 0.0.0.0 --port 8080 +tensorfold serve nvidia/Qwen3.8-27B-NVFP4 --host 0.0.0.0 --port 8888 ``` Prompts take bf16 rows on a prompt GEMM that reads the stored bytes: each FP8 byte and each NVFP4 code times its @@ -122,7 +122,7 @@ serve the folder: ```bash python -c "from huggingface_hub import snapshot_download as d; d('turboderp/Qwen3.8-27B-exl3', revision='3.00bpw', local_dir='qwen27b-exl3-3.00bpw')" tensorfold pull z-lab/Qwen3.8-27B-DFlash2 -tensorfold serve qwen27b-exl3-3.00bpw --host 0.0.0.0 --port 8080 +tensorfold serve qwen27b-exl3-3.00bpw --host 0.0.0.0 --port 8888 ``` Verify windows use the row-invariant EXL3 linear, so drafted replies equal `"draft": false` ones; prompts use the diff --git a/docs/recipes/qwen3.8-flash-next.md b/docs/recipes/qwen3.8-flash-next.md index af65c4fccd..bed3056e40 100644 --- a/docs/recipes/qwen3.8-flash-next.md +++ b/docs/recipes/qwen3.8-flash-next.md @@ -98,7 +98,7 @@ local-inference-lab's. ship, MTP head included, on one GPU: ```bash -tensorfold serve local-inference-lab/Qwen3.8-Flash-Next-NVFP4 --host 0.0.0.0 --port 8080 +tensorfold serve local-inference-lab/Qwen3.8-Flash-Next-NVFP4 --host 0.0.0.0 --port 8888 ``` | Checkpoint (revision) | Routed experts | DeltaNet, attention, shared expert | n-gram table | @@ -169,7 +169,7 @@ Download that branch, then serve the folder: ```bash python -c "from huggingface_hub import snapshot_download as d; d('turboderp/Qwen3.8-Flash-Next-exl3', revision='3.05bpw_h5_ng5', local_dir='flashnext-exl3-3.05bpw')" python -m tensorfold.cuda.exl3.inspect flashnext-exl3-3.05bpw # bits per tensor, from the headers -tensorfold serve flashnext-exl3-3.05bpw --host 0.0.0.0 --port 8080 +tensorfold serve flashnext-exl3-3.05bpw --host 0.0.0.0 --port 8888 ``` A pack maps its n-gram table from its own file and runs on one GPU: `--tp 2` and `--ple-on-ssd` are for the MLX diff --git a/docs/speed-up-mode.md b/docs/speed-up-mode.md index 69803a8df7..f9ad574d48 100644 --- a/docs/speed-up-mode.md +++ b/docs/speed-up-mode.md @@ -119,7 +119,7 @@ Flags that change decoding, such as `--no-drafts` or `--no-thinking`, go on both ### 6. Connect a client -Point any OpenAI-compatible client at `http://RANK0:8080/v1` with the model name `flash-next`, where RANK0 is rank 0's +Point any OpenAI-compatible client at `http://RANK0:8888/v1` with the model name `flash-next`, where RANK0 is rank 0's address. Send requests only to rank 0, since rank 1 refuses its own. ## Measured speeds diff --git a/docs/vision.md b/docs/vision.md index 9536f7e7fc..486c283663 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -64,7 +64,7 @@ body = { 'max_tokens': 128, 'stream': False, } -request = Request('http://127.0.0.1:8080/v1/chat/completions', +request = Request('http://127.0.0.1:8888/v1/chat/completions', data=json.dumps(body).encode(), headers={'Content-Type': 'application/json'}) with urlopen(request) as response: diff --git a/tools/bench_concurrent.py b/tools/bench_concurrent.py index 446592072e..c1c19eb591 100644 --- a/tools/bench_concurrent.py +++ b/tools/bench_concurrent.py @@ -16,7 +16,7 @@ server with a small listen backlog resets simultaneous connections; those resets are reported, not hidden). Standard library only. - python3 tools/bench_concurrent.py http://127.0.0.1:8080 MODEL --levels 1,2,4,8 --alone --output out.json + python3 tools/bench_concurrent.py http://127.0.0.1:8888 MODEL --levels 1,2,4,8 --alone --output out.json """ import argparse diff --git a/tools/bench_openai.py b/tools/bench_openai.py index fca9bf1067..d18685df6b 100644 --- a/tools/bench_openai.py +++ b/tools/bench_openai.py @@ -4,7 +4,7 @@ CUDA engine's bench times (the first sampled token excluded). Standard library only, so it runs on a bare host. - python3 tools/bench_openai.py http://127.0.0.1:8080 MODEL --tokens 64 --reps 5 --output out.json + python3 tools/bench_openai.py http://127.0.0.1:8888 MODEL --tokens 64 --reps 5 --output out.json """ import argparse diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 76b205eeaa..3ad95bf2fc 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -122,7 +122,7 @@ pub const env = [_][]const u8{ "TENSORFOLD_API_KEY", "TENSORFOLD_NO_LIVE", "TENS pub const Args = struct { model: []const u8 = "", host: []const u8 = "127.0.0.1", - port: u16 = 8080, + port: u16 = 8888, name: []const u8 = "", chat_template: ?[]const u8 = null, alias: []const []const u8 = &.{}, From 81f5576ca473652f712de2247f305f02985fce01 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:34:37 +0000 Subject: [PATCH 06/35] qwen3_5: serve Qwen3.5-2B on CUDA The MLX affine 4-bit checkpoint on the native CUDA server, beside Nemotron: qwen35.cu holds the Metal engine's glue at its bf16 rounding points (gated delta rule, causal convolution, gated attention), cuda.qlinear's affine-4 kernels do the projections, sample.cu draws on the device. Windows of up to 16 rows record each row's recurrence and convolution tail, so the lane core shares rounds across streams and keeps partial copies; a window's rows equal one-row steps bit for bit. draw.zig is the host draw both backends now share. On an RTX 4090: 210 tok/s for one stream, 1250 tok/s across eight; logits agree with an fp32 reference at every position of a 36-token check. --- zig/build/cuda.zig | 32 +- zig/kernels/cuda/qwen35.cu | 458 +++++++++++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/backend.zig | 107 +---- zig/src/families/qwen3_5/cuda.zig | 9 + zig/src/families/qwen3_5/cuda_engine.zig | 484 +++++++++++++++++++++++ zig/src/families/qwen3_5/cuda_native.zig | 211 ++++++++++ zig/src/families/qwen3_5/draw.zig | 108 +++++ zig/src/native/cuda.zig | 18 +- zig/tests/qwen35_cuda_main.zig | 58 +++ 10 files changed, 1367 insertions(+), 119 deletions(-) create mode 100644 zig/kernels/cuda/qwen35.cu create mode 100644 zig/src/families/qwen3_5/cuda.zig create mode 100644 zig/src/families/qwen3_5/cuda_engine.zig create mode 100644 zig/src/families/qwen3_5/cuda_native.zig create mode 100644 zig/src/families/qwen3_5/draw.zig create mode 100644 zig/tests/qwen35_cuda_main.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index a291e9bbec..8632ea6c7f 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -32,6 +32,7 @@ const kernels = [_]Kernel{ .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) + .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: Qwen3.5-2B's glue, the Metal kernels' rounding points .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, @@ -78,7 +79,7 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return cuda; } -fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { +fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); core.addImport("tokenizer", tokenizer); @@ -90,7 +91,11 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin nemotron.addImport("lanes", lanes); nemotron.addImport("heat", heat); nemotron.addImport("nemotron_draft_ids", draft_ids); - return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat }; + const qwen35 = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3_5/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + qwen35.addImport("cuda", cuda); + qwen35.addImport("core", core); + qwen35.addImport("lanes", lanes); + return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .qwen35 = qwen35, .tokenizer = tokenizer, .heat = heat }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -122,7 +127,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu cli.addImport("core", mods.core); cli.addImport("lanes", mods.lanes); cli.addImport("nemotron", mods.nemotron); - const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron); + const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.qwen35); cli.addOptions("build_options", build_options); cli.addImport("checkpoint_cli", checkpointCli(b, target, optimize, strip, native.engines, mods.core)); b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli })); @@ -139,18 +144,23 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - _ = nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true, strip); + _ = nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.tokenizer, build_options, true, strip); + // tf-qwen35-cuda: Qwen3.5-2B's CUDA forward, its logits written for a reference to compare + const q35 = b.createModule(.{ .root_source_file = b.path("zig/tests/qwen35_cuda_main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); + q35.addImport("cuda", cuda); + q35.addImport("qwen35", mods.qwen35); + b.installArtifact(b.addExecutable(.{ .name = "tf-qwen35-cuda", .root_module = q35 })); } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. -fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { +fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, - .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron } }, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron }, .{ .name = "qwen35", .module = qwen35 } }, }); return .{ .api = api, .engines = mod }; } @@ -161,8 +171,8 @@ fn checkpointCli(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. } /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it. -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool, strip: ?bool) *std.Build.Step.Compile { - const m = engines(b, target, optimize, cuda, lanes, nemotron); +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, qwen35: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool, strip: ?bool) *std.Build.Step.Compile { + const m = engines(b, target, optimize, cuda, lanes, nemotron, qwen35); // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's) const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const exe = b.addExecutable(.{ .name = "tensorfold-native", .root_module = b.createModule(.{ @@ -192,7 +202,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/build/check_fatbins.zig"), .target = host }) })).step); const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); - const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron); + const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron, mods.qwen35); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = host, .optimize = .debug, .link_libc = true }); @@ -212,7 +222,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); @@ -250,7 +260,7 @@ pub fn distServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *s }; const cuda = runtime(b, target, .fast, if (prebuilt != null) &images else &.{}); const mods = family(b, target, .fast, cuda, draft_ids); - return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false, null); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.tokenizer, build_options, false, null); } /// SASS targets every release fatbin holds: Ampere (sm_86), RTX 50 and RTX PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu new file mode 100644 index 0000000000..a68fa30990 --- /dev/null +++ b/zig/kernels/cuda/qwen35.cu @@ -0,0 +1,458 @@ +// Qwen3.5-2B's glue on CUDA: the Metal engine's kernels (zig/kernels/metal/qwen3_5) with the same bf16 rounding points. +// Rows of one forward belong to segments (one stream's window each); a segment's state lives in its stream's blocks. + +#include +#include + +namespace tf_q35 { + +constexpr int HID = 2048; // hidden +constexpr int NH = 16; // linear-attention heads (keys and values) +constexpr int DK = 128; // their head width +constexpr int CONV = 6144; // q, k and v channels the convolution sees +constexpr int TAPS = 4; +constexpr int QH = 8; // attention query heads +constexpr int KVH = 2; // attention key/value heads +constexpr int HD = 256; // attention head width +constexpr size_t REC = (size_t)NH * DK * DK; // floats of one layer's recurrence +constexpr size_t CONV_ROWS = (size_t)(TAPS - 1) * CONV; + +// One stream's window in this forward: its state blocks, cache length before the window, and rows. +struct Seg { + unsigned long long rec; // float recurrences, one REC block a linear layer + unsigned long long conv; // bf16 convolution tails, CONV_ROWS a linear layer + unsigned long long kv; // bf16 keys then values, KVH * cap * HD each, an attention layer at a time + int cap, len, row0, rows; +}; + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +__device__ __forceinline__ float sigmoid(float x) { return 1.0f / (1.0f + __expf(-x)); } + +} // namespace tf_q35 + +using namespace tf_q35; + +// MLX affine 4-bit token rows: out = bf16(s * q + b). +extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, + const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B, + __nv_bfloat16* __restrict__ out) { + const size_t id = ids[blockIdx.x]; + for (int b = threadIdx.x; b < HID / 2; b += blockDim.x) { + const uint8_t v = W[id * (HID / 2) + b]; + const size_t g = id * (HID / 64) + (2 * b) / 64; + const float s = bf(S[g]), z = bf(B[g]); + out[(size_t)blockIdx.x * HID + 2 * b] = __float2bfloat16_rn(s * (float)(v & 15) + z); + out[(size_t)blockIdx.x * HID + 2 * b + 1] = __float2bfloat16_rn(s * (float)(v >> 4) + z); + } +} + +// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), xs = x's 64-group sums. One 256-thread block a row. +extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, + const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, + float* __restrict__ XS, float eps) { + constexpr int E = HID / 256; + __shared__ float red[8]; + __shared__ float xb[HID]; + const size_t row = blockIdx.x; + const int t = threadIdx.x; + float hv[E]; + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < E; i++) { + const size_t at = row * HID + t * E + i; + float h = bf(H[at]); + if (R != nullptr) { + const __nv_bfloat16 hb = __float2bfloat16_rn(h + bf(R[at])); + H[at] = hb; + h = bf(hb); + } + hv[i] = h; + ss = fmaf(h, h, ss); + } + ss = warp_sum(ss); + if ((t & 31) == 0) red[t >> 5] = ss; + __syncthreads(); + float total = 0.0f; + for (int i = 0; i < 8; i++) total += red[i]; + const float inv = rsqrtf(total / (float)HID + eps); +#pragma unroll + for (int i = 0; i < E; i++) { + const int d = t * E + i; + const __nv_bfloat16 x = __float2bfloat16_rn(bf(W[d]) * (hv[i] * inv)); + X[row * HID + d] = x; + xb[d] = bf(x); + } + __syncthreads(); + if (XS != nullptr && t < HID / 64) { + float s = 0.0f; + for (int j = 0; j < 64; j++) s += xb[t * 64 + j]; + XS[row * (HID / 64) + t] = s; + } +} + +// The 64-group sums of bf16 rows D wide (the decode projections' input sums). Grid (rows), a thread a group. +extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, float* __restrict__ XS, int D) { + const size_t row = blockIdx.x; + for (int g = threadIdx.x; g < D / 64; g += blockDim.x) { + const __nv_bfloat16* x = X + row * D + g * 64; + float s = 0.0f; + for (int j = 0; j < 64; j++) s += bf(x[j]); + XS[row * (D / 64) + g] = s; + } +} + +// The causal convolution, silu, q/k normalization and the decay and write strength of each row's linear heads. +// Grid (1, 3 * NH, rows), a warp a (row, head): heads [0, NH) are q, [NH, 2NH) k, [2NH, 3NH) v. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, int layer, const __nv_bfloat16* __restrict__ CW, + const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, + const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, + __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ K, + __nv_bfloat16* __restrict__ V, float* __restrict__ G, __nv_bfloat16* __restrict__ BETA) { + const int lane = threadIdx.x; + const int head = blockIdx.y; + const int w = blockIdx.z; + const Seg sg = segs[rowseg[w]]; + const int r = w - sg.row0; + const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * CONV_ROWS; + const bool isq = head < NH, isk = !isq && head < 2 * NH; + const int c0 = head * DK; // q, k and v channels are consecutive heads of DK + constexpr int PER = DK / 32; + float vals[PER]; +#pragma unroll + for (int j = 0; j < PER; j++) { + const int c = c0 + lane * PER + j; + float acc = 0.0f; + for (int tap = 0; tap < TAPS; tap++) { + const int at = r + tap; // in [state rows (TAPS - 1)] ++ [this segment's rows] + const float xv = at < TAPS - 1 ? bf(CS[(size_t)at * CONV + c]) : bf(QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]); + acc += bf(CW[c * TAPS + tap]) * xv; + } + const float conv = rbf(acc); + const float sig = rbf(1.0f / (1.0f + expf(-conv))); + vals[j] = rbf(conv * sig); + } + if (isq || isk) { + float ss = 0.0f; +#pragma unroll + for (int j = 0; j < PER; j++) ss += vals[j] * vals[j]; + ss = warp_sum(ss); + const float inv = rsqrtf(ss / (float)DK + (1e-6f / (float)DK)); + const float scale = isq ? rbf(1.0f / (float)DK) : rbf(rsqrtf((float)DK)); + __nv_bfloat16* out = isq ? Q + ((size_t)w * NH + head) * DK : K + ((size_t)w * NH + head - NH) * DK; +#pragma unroll + for (int j = 0; j < PER; j++) out[lane * PER + j] = __float2bfloat16_rn(scale * rbf(vals[j] * inv)); + } else { + const int hv = head - 2 * NH; +#pragma unroll + for (int j = 0; j < PER; j++) V[((size_t)w * NH + hv) * DK + lane * PER + j] = __float2bfloat16_rn(vals[j]); + if (lane == 0) { + const float s = rbf(bf(Ain[w * NH + hv]) + bf(DT[hv])); + const float sp = rbf(fmaxf(s, 0.0f) + logf(1.0f + expf(-fabsf(s)))); + G[w * NH + hv] = expf(-expf(ALOG[hv]) * sp); + BETA[w * NH + hv] = __float2bfloat16_rn(1.0f / (1.0f + expf(-bf(Bin[w * NH + hv])))); + } + } +} + +// Each segment's convolution tail after its rows: the last TAPS - 1 of [tail ++ rows]. A thread a (segment, channel). +// SNAP, if given, also takes the tail after each row (forward row m at m * CONV_ROWS), read before the update. +extern "C" __global__ void tf_q35_conv_commit(const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, + __nv_bfloat16* __restrict__ SNAP) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= CONV) return; + const Seg sg = segs[blockIdx.y]; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * CONV_ROWS; + if (SNAP != nullptr) + for (int j = 0; j < sg.rows; j++) + for (int r = 0; r < TAPS - 1; r++) { + const int at = r + j + 1; + SNAP[((size_t)(sg.row0 + j) * (TAPS - 1) + r) * CONV + c] = + at < TAPS - 1 ? CS[(size_t)at * CONV + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]; + } + // ascending: slot r takes [r + rows], which is either a later state slot (not yet written) or a row + for (int r = 0; r < TAPS - 1; r++) { + const int at = r + sg.rows; + CS[(size_t)r * CONV + c] = at < TAPS - 1 ? CS[(size_t)at * CONV + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]; + } +} + +// The gated delta rule over each segment's rows, its recurrence updated in place; SNAP, if given, takes the state after +// each row (forward row m at m * REC). Grid (1, DK / 4, NH * segments), block (32, 4): a warp a (head, value row), +// its 32 lanes over the key width. +extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, + const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, int layer, + __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP) { + const int hv = blockIdx.z % NH; + const Seg sg = segs[blockIdx.z / NH]; + const int dk = threadIdx.x; + const int dv = blockIdx.y * 4 + threadIdx.y; + constexpr int PT = DK / 32; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * REC + ((size_t)hv * DK + dv) * DK + PT * dk; + float state[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int r = 0; r < sg.rows; r++) { + const size_t row = sg.row0 + r; + float kc[PT], qc[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) { + kc[i] = bf(k[(row * NH + hv) * DK + PT * dk + i]); + qc[i] = bf(q[(row * NH + hv) * DK + PT * dk + i]); + } + const float vc = bf(v[(row * NH + hv) * DK + dv]); + const float gc = g[row * NH + hv]; + const float bc = bf(beta[row * NH + hv]); + float kv_mem = 0.0f; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * gc; + kv_mem += state[i] * kc[i]; + } + kv_mem = warp_sum(kv_mem); + const float delta = (vc - kv_mem) * bc; + float out = 0.0f; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] + kc[i] * delta; + out += state[i] * qc[i]; + } + out = warp_sum(out); + if (dk == 0) y[(row * NH + hv) * DK + dv] = __float2bfloat16_rn(out); + if (SNAP != nullptr) { + float* sn = SNAP + row * REC + ((size_t)hv * DK + dv) * DK + PT * dk; +#pragma unroll + for (int i = 0; i < PT; i++) sn[i] = state[i]; + } + } +#pragma unroll + for (int i = 0; i < PT; i++) st[i] = state[i]; +} + +// out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear head. Grid (1, NH, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, + const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT) { + const int lane = threadIdx.x, hv = blockIdx.y; + const size_t m = blockIdx.z; + constexpr int PER = DK / 32; + float yv[PER], ss = 0.0f; +#pragma unroll + for (int j = 0; j < PER; j++) { + yv[j] = bf(Y[(m * NH + hv) * DK + lane * PER + j]); + ss += yv[j] * yv[j]; + } + ss = warp_sum(ss); + const float inv = rsqrtf(ss / (float)DK + eps); +#pragma unroll + for (int j = 0; j < PER; j++) { + const int d = lane * PER + j; + const float x = rbf(bf(NW[d]) * (yv[j] * inv)); + const float zf = bf(Z[m * HID + hv * DK + d]); + OUT[m * HID + hv * DK + d] = __float2bfloat16_rn(zf / (1.0f + expf(-zf)) * x); + } +} + +// Per-head RMSNorm of attention rows: X's head h of row m at m * stride + h * hstride, out packed (rows, heads, HD). +// Grid (rows * heads), 64 threads of 4 values. +extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps) { + __shared__ float sums[2]; + const int head = blockIdx.x % heads; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + head * hstride; + float v[4], acc = 0.0f; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = bf(x[4 * t + i]); + acc += v[i] * v[i]; + } + acc = warp_sum(acc); + if ((t & 31) == 0) sums[t >> 5] = acc; + __syncthreads(); + const float inv = rsqrtf((sums[0] + sums[1]) / (float)HD + eps); +#pragma unroll + for (int i = 0; i < 4; i++) { + const int d = 4 * t + i; + OUT[(m * heads + head) * HD + d] = __float2bfloat16_rn(bf(W[d]) * (v[i] * inv)); + } +} + +// Text RoPE on the first 64 of HD: pairs (i, i + 32) turned by position * theta^(-i / 32). +__device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, int position) { + if (d >= 64) return x[d]; + const int i = d % 32; + const float angle = (float)position * powf(10000000.0f, -(float)i / 32.0f); + float sn, cs; + sincosf(angle, &sn, &cs); + const float a = bf(x[i]), b = bf(x[i + 32]); + return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); +} + +// Queries in place at their positions; keys rotated and values copied into their segment's cache. +// Grid (rows), HD threads. +extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(__nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, + const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, int layer) { + const int d = threadIdx.x; + const size_t m = blockIdx.x; + const Seg sg = segs[rowseg[m]]; + const int pos = sg.len + (int)m - sg.row0; + __shared__ __nv_bfloat16 qs[QH * HD]; + for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; + __syncthreads(); + for (int h = 0; h < QH; h++) Qn[(m * QH + h) * HD + d] = rotate(qs + h * HD, d, pos); + __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; + __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + for (int h = 0; h < KVH; h++) { + const size_t dst = ((size_t)h * sg.cap + pos) * HD + d; + keys[dst] = rotate(Kn + (m * KVH + h) * HD, d, pos); + values[dst] = Vr[(m * KVH + h) * HD + d]; + } +} + +// Causal attention over a segment's cache, ck keys a block: grid (chunks, KVH, rows), a warp a query head of the +// kv head's group of QH / KVH. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +extern "C" __global__ void __launch_bounds__(128) tf_q35_attn_partial(const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, int layer, int nch, int ck, float scale, + float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; + const int c = blockIdx.x, h = blockIdx.y; + const size_t m = blockIdx.z; + const int qh = h * (QH / KVH) + g; + const Seg sg = segs[rowseg[m]]; + const int last = sg.len + (int)m - sg.row0; // the row's own position, the last key it sees + const int k0 = c * ck, k1 = min(k0 + ck, last + 1); + const size_t slot = (m * QH + qh) * nch + c; + if (k0 >= k1) { + if (lane == 0) { + PM[slot] = -INFINITY; + PL[slot] = 0.0f; + } + return; + } + const __nv_bfloat16* K = reinterpret_cast(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; + const __nv_bfloat16* V = K + (size_t)KVH * sg.cap * HD; + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[(m * QH + qh) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + for (int pos = k0; pos < k1; pos++) { + const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const __nv_bfloat16* kb = reinterpret_cast(&kw); + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(kb[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const uint4 vw = *reinterpret_cast(V + ((size_t)h * sg.cap + pos) * HD + lane * DPL); + const __nv_bfloat16* vb = reinterpret_cast(&vw); +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(vb[i]), o[i] * a); + mx = mn; + } + if (lane == 0) { + PM[slot] = mx; + PL[slot] = l; + } +#pragma unroll + for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; +} + +// The chunks' partials into each (row, query head)'s output. Grid (QH, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(const float* __restrict__ PM, const float* __restrict__ PL, + const float* __restrict__ PO, int nch, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x, qh = blockIdx.x; + const size_t m = blockIdx.y; + const size_t base = (m * QH + qh) * nch; + float mx = -INFINITY; + for (int c = 0; c < nch; c++) + if (PL[base + c] > 0.0f) mx = fmaxf(mx, PM[base + c]); + float lsum = 0.0f, acc[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) acc[i] = 0.0f; + for (int c = 0; c < nch; c++) { + if (!(PL[base + c] > 0.0f)) continue; + const float e = __expf(PM[base + c] - mx); + lsum = fmaf(PL[base + c], e, lsum); +#pragma unroll + for (int i = 0; i < DPL; i++) acc[i] = fmaf(PO[(base + c) * HD + lane * DPL + i], e, acc[i]); + } +#pragma unroll + for (int i = 0; i < DPL; i++) OUT[(m * QH + qh) * HD + lane * DPL + i] = __float2bfloat16_rn(acc[i] / lsum); +} + +// out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows. +extern "C" __global__ void tf_q35_attn_gate(const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, + __nv_bfloat16* __restrict__ OUT, int rows) { + const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= (size_t)rows * HID) return; + const size_t r = i / HID; + const int d = (int)(i % HID), head = d / HD; + const float gate = bf(QRAW[r * 2 * HID + head * 2 * HD + HD + d % HD]); + const float s = rbf(1.0f / (1.0f + expf(-gate))); + OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); +} + +// out = bf16(silu(gate) * up), n wide rows. +extern "C" __global__ void tf_q35_mlp_act(const __nv_bfloat16* __restrict__ GATE, const __nv_bfloat16* __restrict__ UP, + __nv_bfloat16* __restrict__ OUT, long long total) { + const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + const float gf = bf(GATE[i]); + OUT[i] = __float2bfloat16_rn(gf / (1.0f + expf(-gf)) * bf(UP[i])); +} + +// Each row's greedy token: the largest logit, a tie to the lower id (the host draw's rule). Grid (rows), 1024 threads. +extern "C" __global__ void __launch_bounds__(1024) tf_q35_argmax(const __nv_bfloat16* __restrict__ L, int V, unsigned* __restrict__ TOK) { + __shared__ float bv[32]; + __shared__ int bi[32]; + const __nv_bfloat16* x = L + (size_t)blockIdx.x * V; + float best = -INFINITY; + int id = 0x7fffffff; + for (int i = threadIdx.x; i < V; i += blockDim.x) { + const float v = bf(x[i]); + if (v > best) { + best = v; + id = i; + } + } +#pragma unroll + for (int m = 16; m; m >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, best, m); + const int oi = __shfl_xor_sync(0xffffffffu, id, m); + if (ov > best || (ov == best && oi < id)) { + best = ov; + id = oi; + } + } + if ((threadIdx.x & 31) == 0) { + bv[threadIdx.x >> 5] = best; + bi[threadIdx.x >> 5] = id; + } + __syncthreads(); + if (threadIdx.x == 0) { + for (int w = 1; w < 32; w++) + if (bv[w] > best || (bv[w] == best && bi[w] < id)) { + best = bv[w]; + id = bi[w]; + } + TOK[blockIdx.x] = (unsigned)id; + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..c70e8b5ccf 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -30,6 +30,7 @@ pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").b pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; +pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/backend.zig b/zig/src/families/qwen3_5/backend.zig index 55fc09eecb..ef369077a8 100644 --- a/zig/src/families/qwen3_5/backend.zig +++ b/zig/src/families/qwen3_5/backend.zig @@ -163,106 +163,11 @@ pub const Metal = struct { } }; -/// A row's allowed-token bits (structured output): one bit a logit. -const mask_words = (c.vocab + 31) / 32; +const draws = @import("draw.zig"); +const mask_words = draws.mask_words; +pub const value = draws.value; +pub const draw = draws.draw; -pub fn value(word: u16) f64 { - return @as(f32, @bitCast(@as(u32, word) << 16)); -} - -/// The unchanged top-k candidates for the shared fp64 keyed sampler; a `mask` keeps its allowed ones (-inf's draw). -pub fn draw(gpa: std.mem.Allocator, logits: []const u16, settings: ?lanes.Sampling, position: u64, mask: []const u32) !u32 { - var best: ?usize = null; - var allowed: usize = 0; - for (logits, 0..) |word, id| { - const x = value(word); - if (!std.math.isFinite(x)) return error.NonfiniteQwenLogits; - if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; - allowed += 1; - if (best == null or x > value(logits[best.?])) best = id; - } - if (best == null) return error.NoAllowedQwenToken; - const sampling = settings orelse return @intCast(best.?); - if (sampling.temperature <= 0) return @intCast(best.?); - const k = if (sampling.top_k == 0) allowed else @min(allowed, sampling.top_k); - const values = try gpa.alloc(f64, k); - defer gpa.free(values); - const ids = try gpa.alloc(u64, k); - defer gpa.free(ids); - var filled: usize = 0; - for (logits, 0..) |word, id| { - if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; - const x = value(word); - if (k == allowed) { - values[filled] = x; - ids[filled] = id; - filled += 1; - continue; - } - if (filled == k and x <= values[k - 1]) continue; - var at = @min(filled, k - 1); - while (at > 0 and x > values[at - 1]) : (at -= 1) { - values[at] = values[at - 1]; - ids[at] = ids[at - 1]; - } - values[at] = x; - ids[at] = id; - filled = @min(filled + 1, k); - } - return @intCast(try lanes.sampling.choose(gpa, values, ids, position, sampling)); -} - -test "candidate selection preserves full-vocabulary keyed draws and tie order" { - const gpa = std.testing.allocator; - const words = try gpa.alloc(u16, c.vocab); - defer gpa.free(words); - const values = try gpa.alloc(f64, c.vocab); - defer gpa.free(values); - const ids = try gpa.alloc(u64, c.vocab); - defer gpa.free(ids); - for (words, values, ids, 0..) |*w, *v, *id, i| { - const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; - w.* = @intCast(@as(u32, @bitCast(x)) >> 16); - v.* = value(w.*); - id.* = i; - } - for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { - for ([_]u64{ 0, 128, 65536 }) |position| { - const s: lanes.Sampling = .{ .seed = 1234, .top_k = k, .temperature = 0.7, .top_p = 0.95, .min_p = 0.1 }; - const want = try lanes.sampling.choose(gpa, values, ids, position, s); - try std.testing.expectEqual(want, try draw(gpa, words, s, position, &.{})); - } - } - try std.testing.expectEqual(try draw(gpa, words, null, 0, &.{}), try draw(gpa, words, .{ .seed = 7, .temperature = 0 }, 99, &.{})); - words[0] = 0x7fc0; - try std.testing.expectError(error.NonfiniteQwenLogits, draw(gpa, words, null, 0, &.{})); -} - -test "a masked draw is the full-vocabulary draw with the other tokens at -inf" { - const gpa = std.testing.allocator; - const words = try gpa.alloc(u16, c.vocab); - defer gpa.free(words); - const values = try gpa.alloc(f64, c.vocab); - defer gpa.free(values); - const ids = try gpa.alloc(u64, c.vocab); - defer gpa.free(ids); - var mask: [mask_words]u32 = @splat(0); - for (words, values, ids, 0..) |*w, *v, *id, i| { - const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; - w.* = @intCast(@as(u32, @bitCast(x)) >> 16); - const on = i % 7 == 3 or i % 1000 == 1; - if (on) mask[i / 32] |= @as(u32, 1) << @intCast(i % 32); - v.* = if (on) value(w.*) else -std.math.inf(f64); - id.* = i; - } - for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { - for ([_]u64{ 0, 128, 65536 }) |position| { - const s: lanes.Sampling = .{ .seed = 99, .top_k = k, .temperature = 0.8, .top_p = 0.9, .min_p = 0.05 }; - const want = try lanes.sampling.choose(gpa, values, ids, position, s); - try std.testing.expectEqual(want, try draw(gpa, words, s, position, &mask)); - } - } - try std.testing.expectEqual(try lanes.sampling.choose(gpa, values, ids, 0, .{ .seed = 1, .temperature = 0 }), try draw(gpa, words, null, 0, &mask)); - @memset(&mask, 0); - try std.testing.expectError(error.NoAllowedQwenToken, draw(gpa, words, null, 0, &mask)); +test { + _ = draws; } diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig new file mode 100644 index 0000000000..fdecefed7f --- /dev/null +++ b/zig/src/families/qwen3_5/cuda.zig @@ -0,0 +1,9 @@ +//! Qwen3.5-2B on CUDA: the engine and its native family entry (the Metal engine is qwen3_5.zig). +pub const config = @import("config.zig"); +pub const engine = @import("cuda_engine.zig"); +pub const native = @import("cuda_native.zig"); + +test { + _ = config; + _ = @import("draw.zig"); +} diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig new file mode 100644 index 0000000000..96190737d1 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -0,0 +1,484 @@ +//! Qwen3.5-2B on CUDA: the affine-4 projections through cuda.qlinear, the glue in qwen35.cu, one forward for any +//! mix of streams' windows. Each stream owns its recurrences, convolution tails and KV cache (Cache). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const c = @import("config.zig"); + +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows of one forward at most: a prompt chunk, or a shared round's windows. +pub const max_rows = 128; +/// Streams' windows in one forward at most. +pub const max_segs = 32; +/// Rows of a verify forward at most: each row's recurrences are recorded for a partial acceptance. +pub const max_verify = 32; +/// A window's rows at most (exact: one-row steps' bits). +pub const window_rows = 16; +/// Partial attention chunks a row at most: the chunk width grows with the window instead. +const max_chunks = 64; + +const deltas = c.layers - c.layers / 4; +const attentions = c.layers / 4; +pub const rec_bytes = c.linear_heads * c.linear_dim * c.linear_dim * 4; +pub const conv_bytes = (c.conv_taps - 1) * c.conv_dim * 2; +const state_bytes = deltas * (rec_bytes + conv_bytes); + +/// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values. +pub fn streamBytes(cap: usize) usize { + return state_bytes + kvBytes(cap); +} + +fn kvBytes(cap: usize) usize { + return attentions * 2 * c.kv_heads * cap * c.head_dim * 2; +} + +/// qwen35.cu's Seg. +const Seg = extern struct { rec: u64, conv: u64, kv: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int }; + +/// One stream's state on the device and the logits of its newest row on the host. +pub const Cache = struct { + state: DeviceBuffer, + kv: DeviceBuffer, + cap: usize, + len: usize = 0, + logits: []u16, + /// The newest verify window: where it started, its first row in that forward, its rows. + last: ?struct { start: usize, row0: usize, rows: usize } = null, + + pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { + var state = try DeviceBuffer.alloc(e.d, state_bytes); + errdefer state.free(); + try state.fill8(0, null); + var kv = try DeviceBuffer.alloc(e.d, kvBytes(cap)); + errdefer kv.free(); + return .{ .state = state, .kv = kv, .cap = cap, .logits = try gpa.alloc(u16, c.vocab) }; + } + + pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { + self.state.free(); + self.kv.free(); + gpa.free(self.logits); + } + + fn seg(self: *const Cache, row0: usize, rows: usize) Seg { + return .{ .rec = self.state.ptr, .conv = self.state.ptr + deltas * rec_bytes, .kv = self.kv.ptr, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows) }; + } +}; + +const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; +const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; +const Layer = struct { + input_norm: u64, + post_norm: u64, + gate: Affine4, + up: Affine4, + down: Affine4, + mixer: union(enum) { delta: Delta, attention: Attention }, +}; + +const Fns = struct { + embed: cuda.Function, + norm: cuda.Function, + sums: cuda.Function, + gdn_pre: cuda.Function, + conv_commit: cuda.Function, + gdn_chain: cuda.Function, + gdn_post: cuda.Function, + head_norm: cuda.Function, + rope: cuda.Function, + attn_partial: cuda.Function, + attn_merge: cuda.Function, + attn_gate: cuda.Function, + mlp_act: cuda.Function, + argmax: cuda.Function, +}; + +/// A window of one stream in a forward. +pub const Part = struct { cache: *Cache, rows: usize }; + +/// What a forward keeps. A prompt chunk (`last`, one window) puts its last row's logits in its cache's `logits`; +/// verify windows keep every row, drawn on the device into `tokens` (`draws`, one a row) or their logits for +/// `rowLogits` (`logits`), and record each row's state for `rewind`. +pub const Keep = union(enum) { last, draws: []const Draw, logits }; + +/// A kept row's draw on the device: sample.cu's keyed rule at `position`, or greedy without one. +pub const Draw = struct { rule: ?cuda.sample.Rule, position: u64 }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + stream: cuda.Stream, + mods: [6]cuda.Module, + draws: cuda.sample.Draws, + affine: cuda.qlinear.Affine4Kernels, + fns: Fns, + held: std.ArrayList(DeviceBuffer) = .empty, // weights and scratch, freed together + // weights + embed_w: u64, + embed_s: u64, + embed_b: u64, + head: Affine4, + final_norm: u64, + layers: [c.layers]Layer, + // scratch, max_rows rows each + h: u64, + x: u64, + r: u64, + xs: u64, + qkv: u64, + z: u64, + a: u64, + b: u64, + q: u64, + k: u64, + v: u64, + g: u64, + beta: u64, + y: u64, + mix: u64, + qraw: u64, + kraw: u64, + vraw: u64, + qn: u64, + kn: u64, + gate: u64, + up: u64, + act: u64, + pm: u64, + pl: u64, + po: u64, + logits: u64, + // the per-forward tables: ids, segments, each row's segment and the draws' rules, staged in pinned memory + tables: DeviceBuffer, + staged: cuda.HostBuffer, + host_logits: cuda.HostBuffer, + /// The kept rows' device draws, when forward was given them. + tokens: [max_verify]u32 = undefined, + rec_snap: u64, // recorded recurrences: max_verify rows a linear layer + conv_snap: u64, // recorded convolution tails, the same + + const ids_at = 0; + const segs_at = max_rows * 4; + const rowseg_at = segs_at + max_segs * @sizeOf(Seg); + const rules_at = rowseg_at + max_rows * 4; + const meta_at = rules_at + max_verify * @sizeOf(cuda.sample.Rule); + const tokens_at = meta_at + max_verify * 4; + const tables_len = tokens_at + max_verify * 4; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const config = try c.Config.read(gpa, io, dir); + const d = ctx.d; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = undefined; + e.gpa = gpa; + e.d = d; + e.held = .empty; + e.stream = try cuda.Stream.init(d, true); + errdefer e.stream.deinit(); + const kk = cuda.kernels; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample }; + var loaded: usize = 0; + errdefer for (e.mods[0..loaded]) |*m| m.unload(); + for (images, 0..) |img, i| { + e.mods[i] = try cuda.Module.load(d, img); + loaded += 1; + } + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + const gb10 = try ctx.attribute(.compute_capability_major) == 12 and try ctx.attribute(.compute_capability_minor) == 1; + e.affine = try cuda.qlinear.Affine4Kernels.resolve(d, e.mods[0], e.mods[1], e.mods[2], e.mods[3], sms, gb10); + errdefer e.affine.deinit(); + e.draws = try cuda.sample.Draws.resolve(e.mods[5]); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(e.fns, name) = try e.mods[4].function("tf_q35_" ++ name); + errdefer e.freeHeld(); + var ck = try core.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + try e.loadWeights(&ck, config.tied); + try e.allocScratch(); + e.tables = try DeviceBuffer.alloc(d, tables_len); + errdefer e.tables.free(); + e.staged = try cuda.HostBuffer.alloc(d, tables_len); + errdefer e.staged.free(); + e.host_logits = try cuda.HostBuffer.alloc(d, max_verify * c.vocab * 2); + return e; + } + + pub fn deinit(e: *Engine) void { + e.host_logits.free(); + e.staged.free(); + e.tables.free(); + e.freeHeld(); + e.affine.deinit(); + for (&e.mods) |*m| m.unload(); + e.stream.deinit(); + e.gpa.destroy(e); + } + + fn freeHeld(e: *Engine) void { + for (e.held.items) |*b| b.free(); + e.held.deinit(e.gpa); + } + + fn own(e: *Engine, b: DeviceBuffer) !u64 { + try e.held.append(e.gpa, b); + return b.ptr; + } + + fn alloc(e: *Engine, bytes: usize) !u64 { + var b = try DeviceBuffer.alloc(e.d, bytes); + errdefer b.free(); + return e.own(b); + } + + fn put(e: *Engine, t: core.checkpoint.Tensor) !u64 { + var b = try DeviceBuffer.fromHost(e.d, t.bytes); + errdefer b.free(); + return e.own(b); + } + + fn tensor(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + var buf: [160]u8 = undefined; + return e.put(try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); + } + + /// An MLX affine-4 projection packed into qlinear's tiles; the checkpoint's own layout leaves the device after. + fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { + var buf: [160]u8 = undefined; + var raw: [3]DeviceBuffer = undefined; + var got: usize = 0; + defer for (raw[0..got]) |*b| b.free(); + const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / c.group }, .{ "biases", .bf16, k / c.group } }; + for (parts) |p| { + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); + raw[got] = try DeviceBuffer.fromHost(e.d, t.bytes); + got += 1; + } + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try e.alloc(l.words), .s = try e.alloc(l.scales), .b = try e.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + try e.affine.pack(e.stream, raw[0].ptr, raw[1].ptr, raw[2].ptr, q); + try e.stream.synchronize(); + return q; + } + + fn loadWeights(e: *Engine, ck: *core.Checkpoint, tied: bool) !void { + const m = "language_model.model."; + e.embed_w = try e.tensor(ck, "{s}embed_tokens.weight", .{m}, .u32, &.{ c.vocab, c.hidden / 8 }); + e.embed_s = try e.tensor(ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ c.vocab, c.hidden / c.group }); + e.embed_b = try e.tensor(ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ c.vocab, c.hidden / c.group }); + e.head = if (tied) try e.linear(ck, m, "embed_tokens", c.vocab, c.hidden) else try e.linear(ck, "language_model.", "lm_head", c.vocab, c.hidden); + e.final_norm = try e.tensor(ck, "{s}norm.weight", .{m}, .bf16, &.{c.hidden}); + for (&e.layers, 0..) |*l, i| { + var pbuf: [96]u8 = undefined; + const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); + l.* = .{ + .input_norm = try e.tensor(ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{c.hidden}), + .post_norm = try e.tensor(ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{c.hidden}), + .gate = try e.linear(ck, p, "mlp.gate_proj", c.intermediate, c.hidden), + .up = try e.linear(ck, p, "mlp.up_proj", c.intermediate, c.hidden), + .down = try e.linear(ck, p, "mlp.down_proj", c.hidden, c.intermediate), + .mixer = undefined, + }; + if (c.linear(i)) { + l.mixer = .{ .delta = .{ + .qkv = try e.linear(ck, p, "linear_attn.in_proj_qkv", c.conv_dim, c.hidden), + .z = try e.linear(ck, p, "linear_attn.in_proj_z", c.hidden, c.hidden), + .a = try e.linear(ck, p, "linear_attn.in_proj_a", c.linear_heads, c.hidden), + .b = try e.linear(ck, p, "linear_attn.in_proj_b", c.linear_heads, c.hidden), + .out = try e.linear(ck, p, "linear_attn.out_proj", c.hidden, c.hidden), + .conv = try e.tensor(ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ c.conv_dim, c.conv_taps, 1 }), + .norm = try e.tensor(ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{c.linear_dim}), + .a_log = try e.tensor(ck, "{s}linear_attn.A_log", .{p}, .f32, &.{c.linear_heads}), + .dt_bias = try e.tensor(ck, "{s}linear_attn.dt_bias", .{p}, .bf16, &.{c.linear_heads}), + } }; + } else { + l.mixer = .{ .attention = .{ + .q = try e.linear(ck, p, "self_attn.q_proj", 2 * c.hidden, c.hidden), + .k = try e.linear(ck, p, "self_attn.k_proj", c.kv_heads * c.head_dim, c.hidden), + .v = try e.linear(ck, p, "self_attn.v_proj", c.kv_heads * c.head_dim, c.hidden), + .o = try e.linear(ck, p, "self_attn.o_proj", c.hidden, c.hidden), + .q_norm = try e.tensor(ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{c.head_dim}), + .k_norm = try e.tensor(ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{c.head_dim}), + } }; + } + } + } + + fn allocScratch(e: *Engine) !void { + const R = max_rows; + inline for (.{ "h", "x", "r", "z", "q", "k", "v", "y", "mix", "qn" }) |name| @field(e, name) = try e.alloc(R * c.hidden * 2); + inline for (.{ "qkv", "gate", "up", "act" }) |name| @field(e, name) = try e.alloc(R * c.intermediate * 2); + e.xs = try e.alloc(R * (c.intermediate / 64) * 4); + inline for (.{ "a", "b", "beta" }) |name| @field(e, name) = try e.alloc(R * c.linear_heads * 2); + e.g = try e.alloc(R * c.linear_heads * 4); + e.qraw = try e.alloc(R * 2 * c.hidden * 2); + inline for (.{ "kraw", "vraw", "kn" }) |name| @field(e, name) = try e.alloc(R * c.kv_heads * c.head_dim * 2); + e.pm = try e.alloc(R * c.query_heads * max_chunks * 4); + e.pl = try e.alloc(R * c.query_heads * max_chunks * 4); + e.po = try e.alloc(R * c.query_heads * max_chunks * c.head_dim * 4); + e.logits = try e.alloc(max_verify * c.vocab * 2); + e.rec_snap = try e.alloc(deltas * max_verify * rec_bytes); + e.conv_snap = try e.alloc(deltas * max_verify * conv_bytes); + } + + fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + const cfg: cuda.Config = .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = @intCast(block[0]), .y = @intCast(block[1]), .z = @intCast(block[2]) } }; + try cuda.launch.launch(f, cfg, e.stream, &a); + } + + /// out (rows, n) = x (rows, k) through `l`: decode tiles 16 rows at a time (each row's bits whatever the batch), + /// or the prompt GEMM for a prompt chunk wider than one tile. + fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { + if (prompt and rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)) }); + var at: usize = 0; + while (at < rows) : (at += 16) { + const n = @min(16, rows - at); + try e.affine.decode(e.stream, x + at * l.k * 2, e.xs + at * (l.k / 64) * 4, l, out + at * l.n * 2, n); + } + } + + fn norm(e: *Engine, residual: bool, w: u64, rows: usize) !void { + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, @as(u64, 0), c.eps }); + } + + /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. + pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { + var total: usize = 0; + var widest: usize = 0; + for (parts) |p| { + if (p.rows == 0 or p.cache.len + p.rows > p.cache.cap) return error.QwenContextFull; + total += p.rows; + widest = @max(widest, p.cache.len + p.rows); + } + if (parts.len == 0 or parts.len > max_segs or total != ids.len or total > max_rows) return error.InvalidQwenSegments; + const record = keep != .last; + if (record and total > max_verify) return error.InvalidQwenSegments; + if (!record and parts.len != 1) return error.InvalidQwenSegments; + if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; + for (ids) |id| if (id >= c.vocab) return error.InvalidQwenToken; + const heads: usize = if (record) total else 1; + // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps + const prompt = !record; + // the tables + const st = e.staged.bytes; + @memcpy(std.mem.bytesAsSlice(u32, st[ids_at..][0 .. total * 4]), ids); + const segs = std.mem.bytesAsSlice(Seg, st[segs_at..][0 .. parts.len * @sizeOf(Seg)]); + const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. total * 4]); + var row0: usize = 0; + for (parts, 0..) |p, i| { + segs[i] = p.cache.seg(row0, p.rows); + for (row0..row0 + p.rows) |r| rowseg[r] = @intCast(i); + p.cache.last = if (record) .{ .start = p.cache.len, .row0 = row0, .rows = p.rows } else null; + row0 += p.rows; + } + if (keep == .draws) { + const rules = std.mem.bytesAsSlice(cuda.sample.Rule, st[rules_at..][0 .. total * @sizeOf(cuda.sample.Rule)]); + const meta = std.mem.bytesAsSlice(c_int, st[meta_at..][0 .. total * 4]); + for (keep.draws, rules, meta) |p, *r, *m| { + if (p.rule) |x| r.* = x; + m.* = @intCast(p.position -% 1); // sample.cu keys row r at meta + r + 1 + } + } + try e.tables.uploadAsync(0, st, e.stream.handle); + const t_ids = e.tables.ptr + ids_at; + const t_segs = e.tables.ptr + segs_at; + const t_rows = e.tables.ptr + rowseg_at; + const nseg = parts.len; + // attention chunks: at most max_chunks a row, each a multiple of 128 keys + const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); + const nch = (widest + ck - 1) / ck; + + try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h }); + var di: usize = 0; + var ai: c_int = 0; + for (e.layers, 0..) |l, i| { + try e.norm(i != 0, l.input_norm, total); + switch (l.mixer) { + .delta => |m| { + const layer: c_int = @intCast(di); + const rec_snap: u64 = if (record) e.rec_snap + di * max_verify * rec_bytes else 0; + const conv_snap: u64 = if (record) e.conv_snap + di * max_verify * conv_bytes else 0; + try e.proj(m.qkv, e.x, e.qkv, total, prompt); + try e.proj(m.z, e.x, e.z, total, prompt); + try e.proj(m.a, e.x, e.a, total, prompt); + try e.proj(m.b, e.x, e.b, total, prompt); + try e.run(e.fns.gdn_pre, .{ 1, 3 * c.linear_heads, total }, .{ 32, 1, 1 }, .{ e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g, e.beta }); + try e.run(e.fns.conv_commit, .{ (c.conv_dim + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ e.qkv, t_segs, layer, conv_snap }); + try e.run(e.fns.gdn_chain, .{ 1, c.linear_dim / 4, c.linear_heads * nseg }, .{ 32, 4, 1 }, .{ e.q, e.k, e.v, e.g, e.beta, t_segs, layer, e.y, rec_snap }); + try e.run(e.fns.gdn_post, .{ 1, c.linear_heads, total }, .{ 32, 1, 1 }, .{ e.y, e.z, m.norm, c.eps, e.mix }); + try e.proj(m.out, e.mix, e.r, total, prompt); + di += 1; + }, + .attention => |m| { + try e.proj(m.q, e.x, e.qraw, total, prompt); + try e.proj(m.k, e.x, e.kraw, total, prompt); + try e.proj(m.v, e.x, e.vraw, total, prompt); + try e.run(e.fns.head_norm, .{ total * c.query_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, c.query_heads), @as(c_int, 2 * c.hidden), @as(c_int, 2 * c.head_dim), c.eps }); + try e.run(e.fns.head_norm, .{ total * c.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, c.kv_heads), @as(c_int, c.kv_heads * c.head_dim), @as(c_int, c.head_dim), c.eps }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ c.head_dim, 1, 1 }, .{ e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); + try e.run(e.fns.attn_partial, .{ nch, c.kv_heads, total }, .{ 128, 1, 1 }, .{ e.qn, t_segs, t_rows, ai, @as(c_int, @intCast(nch)), @as(c_int, @intCast(ck)), @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.attn_merge, .{ c.query_heads, total, 1 }, .{ 32, 1, 1 }, .{ e.pm, e.pl, e.po, @as(c_int, @intCast(nch)), e.y }); + try e.run(e.fns.attn_gate, .{ (total * c.hidden + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); + try e.proj(m.o, e.mix, e.r, total, prompt); + ai += 1; + }, + } + try e.norm(true, l.post_norm, total); + try e.proj(l.gate, e.x, e.gate, total, prompt); + try e.proj(l.up, e.x, e.up, total, prompt); + const n: c_longlong = @intCast(total * c.intermediate); + try e.run(e.fns.mlp_act, .{ (total * c.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n }); + try e.proj(l.down, e.act, e.r, total, prompt); + } + try e.norm(true, e.final_norm, total); + // the head over the kept rows + try e.proj(e.head, e.x + (total - heads) * c.hidden * 2, e.logits, heads, false); + for (parts) |p| p.cache.len += p.rows; + switch (keep) { + .draws => |pk| { + const toks = e.tables.ptr + tokens_at; + for (pk, 0..) |p, i| { + const row = e.logits + i * c.vocab * 2; + if (p.rule != null) { + try e.draws.draw(e.stream, row, c.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); + } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, c.vocab), toks + i * 4 }); + } + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); + }, + .logits, .last => { + const out = e.host_logits.bytes[0 .. heads * c.vocab * 2]; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); + }, + } + } + + /// Row `i`'s logits from the newest forward that kept logits. + pub fn rowLogits(e: *const Engine, i: usize) []const u16 { + return e.host_logits.slice(u16)[i * c.vocab ..][0..c.vocab]; + } + + /// Keeps the first `rows` of the cache's newest window: its recurrences and convolution tails as they were after + /// that row (the forward recorded them), its length at that row. The window's later keys stay as unread slots. + pub fn rewind(e: *Engine, cache: *Cache, rows: usize) !void { + const last = cache.last orelse return error.NothingToKeep; + if (rows == 0 or rows > last.rows) return error.NothingToKeep; + cache.last = null; + cache.len = last.start + rows; + if (rows == last.rows) return; + const row = last.row0 + rows - 1; + for (0..deltas) |di| { + try cache.state.copyFrom(di * rec_bytes, e.rec_snap + (di * max_verify + row) * rec_bytes, rec_bytes, e.stream.handle); + try cache.state.copyFrom(deltas * rec_bytes + di * conv_bytes, e.conv_snap + (di * max_verify + row) * conv_bytes, conv_bytes, e.stream.handle); + } + try e.stream.synchronize(); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig new file mode 100644 index 0000000000..4995a82b45 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -0,0 +1,211 @@ +//! Qwen3.5-2B's CUDA lane backend behind the native family interface (native/cuda.zig's registry): windows of up +//! to 16 rows (the lane core's copies), every stream's window in one forward, prompts in chunks, draws on the device. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const c = @import("config.zig"); +const E = @import("cuda_engine.zig"); +const draws = @import("draw.zig"); +const be = lanes.backend; + +pub const model_type = "qwen3_5"; +pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const default_context: i64 = 32768; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = E.max_rows; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; + +/// What the native server drives: the lane backend, the facts its round loop reads, and how to free it. +pub const Loaded = struct { + backend: be.Backend, + facts: lanes.Model, + rows: u32, + /// Device bytes each admitted stream allocates (recurrences, convolution tails, keys and values). + stream_bytes: usize, + free_streams: u32 = 0, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, +}; + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.QwenHasNoCapturedKernels; + const e = try E.Engine.init(gpa, io, ctx, dir); + errdefer e.deinit(); + const l = try gpa.create(Lanes); + l.* = .{ .gpa = gpa, .e = e, .capacity = o.context }; + return .{ + .backend = l.backend(), + .facts = .{ .exact_width = E.window_rows, .hidden_rows = true, .batch_rows = E.max_verify, .max_streams = E.max_segs, .streams_exact = true }, + .rows = E.window_rows, + .stream_bytes = E.streamBytes(o.context), + .ctx = l, + .deinit = release, + }; +} + +/// A request this engine refuses, in words; null: none of its own. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.InvalidQwenPrompt => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.OutOfMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + else => null, + }; +} + +fn release(p: *anyopaque) void { + const l: *Lanes = @ptrCast(@alignCast(p)); + l.deinit(); +} + +/// The device draw of a stream's rule at `position`: sample.cu's keyed rule, greedy at temperature 0 or without one. +fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { + const s = sampling orelse return .{ .rule = null, .position = position }; + if (s.temperature <= 0) return .{ .rule = null, .position = position }; + return .{ .rule = cuda.sample.rule(s.seed, s.temperature, s.top_p, s.minLog(), s.top_k), .position = position }; +} + +const Lanes = struct { + gpa: std.mem.Allocator, + e: *E.Engine, + capacity: usize, + caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, + + fn deinit(self: *Lanes) void { + var it = self.caches.valueIterator(); + while (it.next()) |cache| { + cache.*.deinit(self.gpa); + self.gpa.destroy(cache.*); + } + self.caches.deinit(self.gpa); + self.e.deinit(); + self.gpa.destroy(self); + } + + fn backend(self: *Lanes) be.Backend { + return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; + } + + fn cast(ptr: *anyopaque) *Lanes { + return @ptrCast(@alignCast(ptr)); + } + + fn cacheOf(self: *Lanes, stream: *lanes.Stream) !*E.Cache { + return self.caches.get(stream) orelse error.UnknownQwenStream; + } + + fn prefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void { + const self = cast(ptr); + const ids = stream.prompt(); + if (ids.len == 0 or ids.len + stream.max_new > self.capacity or self.caches.contains(stream)) return error.InvalidQwenPrompt; + const cache = try self.gpa.create(E.Cache); + errdefer self.gpa.destroy(cache); + cache.* = try E.Cache.init(self.gpa, self.e, @min(self.capacity, ids.len + stream.max_new + 1)); + errdefer cache.deinit(self.gpa); + var at: usize = 0; + while (at < ids.len) { + var end: usize = @min(ids.len, at + E.max_rows); + for (stream.chunks) |boundary| if (boundary > at) { + end = @min(end, boundary); + break; + }; + try self.e.forward(&.{.{ .cache = cache, .rows = end - at }}, ids[at..end], .last); + at = end; + } + try self.caches.put(self.gpa, stream, cache); + } + + fn firstFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64) !u64 { + return firstMaskedFn(ptr, stream, position, &.{}); + } + + fn firstMaskedFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64, mask: []const u32) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len) return error.InvalidQwenPosition; + return try draws.draw(self.gpa, cache.logits, stream.sampling, position, mask); + } + + fn queueFn(ptr: *anyopaque, stream: *lanes.Stream, feed: be.Feed, position: u64) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len + 1) return error.InvalidQwenPosition; + const token: u32 = switch (feed) { + .value => |v| v, + .handle => |h| @intCast(h), + }; + try self.e.forward(&.{.{ .cache = cache, .rows = 1 }}, &.{token}, .{ .draws = &.{pick(stream.sampling, position)} }); + cache.last = null; + return self.e.tokens[0]; + } + + fn readFn(_: *anyopaque, handle: u64) !u32 { + if (handle >= c.vocab) return error.InvalidQwenToken; + return @intCast(handle); + } + + /// Each window (a stream's pending token and copies) checked, then as many as fit in each forward. + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { + const self = cast(ptr); + if (windows.len != out.len) return error.InvalidQwenWindows; + for (windows, out) |w, o| { + const rows = w.rows(); + if (w.parents != null or w.held != 0 or rows > E.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; + if (w.masks.len != 0 and w.masks.len != rows * draws.mask_words) return error.InvalidQwenMasks; + if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; + const cache = try self.cacheOf(w.stream); + for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidQwenPosition; + } + var at: usize = 0; + while (at < windows.len) { + var parts: [E.max_segs]E.Part = undefined; + var ids: [E.max_verify]u32 = undefined; + var picks: [E.max_verify]E.Draw = undefined; + var n: usize = 0; + var total: usize = 0; + var masked = false; + while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= E.max_verify) : (n += 1) { + const w = windows[at + n]; + parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows() }; + ids[total] = w.pending; + @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); + for (w.positions, 0..) |p, r| picks[total + r] = pick(w.stream.sampling, p); + masked = masked or w.masks.len != 0; + total += w.rows(); + } + // structured output draws on the host, over each row's logits with its mask applied + try self.e.forward(parts[0..n], ids[0..total], if (masked) .logits else .{ .draws = picks[0..total] }); + var row: usize = 0; + for (windows[at..][0..n], out[at..][0..n]) |w, o| { + for (o.sampled, w.positions, 0..) |*t, p, r| { + const mask = if (w.masks.len > 0) w.masks[r * draws.mask_words ..][0..draws.mask_words] else &.{}; + t.* = if (masked) try draws.draw(self.gpa, self.e.rowLogits(row + r), w.stream.sampling, p, mask) else self.e.tokens[row + r]; + } + @memcpy(o.drafts, w.tokens); + row += w.rows(); + } + at += n; + } + } + + /// The accepted prefix of each window: its state as it was after that row. + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { + const self = cast(ptr); + for (windows, paths) |w, path| { + for (path, 0..) |r, i| if (r != i) return error.UnsupportedQwenTree; + try self.e.rewind(try self.cacheOf(w.stream), path.len); + } + } + + fn draftFn(_: *anyopaque, requests: []const be.DraftRequest) !void { + for (requests) |r| if (r.depth != 0) return error.QwenCheckpointHasNoDraftHead; + } + + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { + const self = cast(ptr); + if (self.caches.fetchRemove(stream)) |entry| { + entry.value.deinit(self.gpa); + self.gpa.destroy(entry.value); + } + } +}; diff --git a/zig/src/families/qwen3_5/draw.zig b/zig/src/families/qwen3_5/draw.zig new file mode 100644 index 0000000000..c9acd623bf --- /dev/null +++ b/zig/src/families/qwen3_5/draw.zig @@ -0,0 +1,108 @@ +//! The host draw over a row of bf16 logits, shared by the Metal and CUDA backends: greedy, keyed sampling or a mask. +const std = @import("std"); +const lanes = @import("lanes"); +const c = @import("config.zig"); + +/// A row's allowed-token bits (structured output): one bit a logit. +pub const mask_words = (c.vocab + 31) / 32; + +pub fn value(word: u16) f64 { + return @as(f32, @bitCast(@as(u32, word) << 16)); +} + +/// The unchanged top-k candidates for the shared fp64 keyed sampler; a `mask` keeps its allowed ones (-inf's draw). +pub fn draw(gpa: std.mem.Allocator, logits: []const u16, settings: ?lanes.Sampling, position: u64, mask: []const u32) !u32 { + var best: ?usize = null; + var allowed: usize = 0; + for (logits, 0..) |word, id| { + const x = value(word); + if (!std.math.isFinite(x)) return error.NonfiniteQwenLogits; + if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; + allowed += 1; + if (best == null or x > value(logits[best.?])) best = id; + } + if (best == null) return error.NoAllowedQwenToken; + const sampling = settings orelse return @intCast(best.?); + if (sampling.temperature <= 0) return @intCast(best.?); + const k = if (sampling.top_k == 0) allowed else @min(allowed, sampling.top_k); + const values = try gpa.alloc(f64, k); + defer gpa.free(values); + const ids = try gpa.alloc(u64, k); + defer gpa.free(ids); + var filled: usize = 0; + for (logits, 0..) |word, id| { + if (mask.len > 0 and !lanes.grammar.allows(mask, id)) continue; + const x = value(word); + if (k == allowed) { + values[filled] = x; + ids[filled] = id; + filled += 1; + continue; + } + if (filled == k and x <= values[k - 1]) continue; + var at = @min(filled, k - 1); + while (at > 0 and x > values[at - 1]) : (at -= 1) { + values[at] = values[at - 1]; + ids[at] = ids[at - 1]; + } + values[at] = x; + ids[at] = id; + filled = @min(filled + 1, k); + } + return @intCast(try lanes.sampling.choose(gpa, values, ids, position, sampling)); +} + +test "candidate selection preserves full-vocabulary keyed draws and tie order" { + const gpa = std.testing.allocator; + const words = try gpa.alloc(u16, c.vocab); + defer gpa.free(words); + const values = try gpa.alloc(f64, c.vocab); + defer gpa.free(values); + const ids = try gpa.alloc(u64, c.vocab); + defer gpa.free(ids); + for (words, values, ids, 0..) |*w, *v, *id, i| { + const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; + w.* = @intCast(@as(u32, @bitCast(x)) >> 16); + v.* = value(w.*); + id.* = i; + } + for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { + for ([_]u64{ 0, 128, 65536 }) |position| { + const s: lanes.Sampling = .{ .seed = 1234, .top_k = k, .temperature = 0.7, .top_p = 0.95, .min_p = 0.1 }; + const want = try lanes.sampling.choose(gpa, values, ids, position, s); + try std.testing.expectEqual(want, try draw(gpa, words, s, position, &.{})); + } + } + try std.testing.expectEqual(try draw(gpa, words, null, 0, &.{}), try draw(gpa, words, .{ .seed = 7, .temperature = 0 }, 99, &.{})); + words[0] = 0x7fc0; + try std.testing.expectError(error.NonfiniteQwenLogits, draw(gpa, words, null, 0, &.{})); +} + +test "a masked draw is the full-vocabulary draw with the other tokens at -inf" { + const gpa = std.testing.allocator; + const words = try gpa.alloc(u16, c.vocab); + defer gpa.free(words); + const values = try gpa.alloc(f64, c.vocab); + defer gpa.free(values); + const ids = try gpa.alloc(u64, c.vocab); + defer gpa.free(ids); + var mask: [mask_words]u32 = @splat(0); + for (words, values, ids, 0..) |*w, *v, *id, i| { + const x: f32 = @as(f32, @floatFromInt((i * 37) % 257)) / 8.0 - 16.0; + w.* = @intCast(@as(u32, @bitCast(x)) >> 16); + const on = i % 7 == 3 or i % 1000 == 1; + if (on) mask[i / 32] |= @as(u32, 1) << @intCast(i % 32); + v.* = if (on) value(w.*) else -std.math.inf(f64); + id.* = i; + } + for ([_]u32{ 0, 1, 20, 257, c.vocab }) |k| { + for ([_]u64{ 0, 128, 65536 }) |position| { + const s: lanes.Sampling = .{ .seed = 99, .top_k = k, .temperature = 0.8, .top_p = 0.9, .min_p = 0.05 }; + const want = try lanes.sampling.choose(gpa, values, ids, position, s); + try std.testing.expectEqual(want, try draw(gpa, words, s, position, &mask)); + } + } + try std.testing.expectEqual(try lanes.sampling.choose(gpa, values, ids, 0, .{ .seed = 1, .temperature = 0 }), try draw(gpa, words, null, 0, &mask)); + @memset(&mask, 0); + try std.testing.expectError(error.NoAllowedQwenToken, draw(gpa, words, null, 0, &mask)); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 50af1bd955..b90c801e0d 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -5,13 +5,14 @@ const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const qwen35 = @import("qwen35"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); const nemotron_slide = @import("nemotron_slide.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{nemotron.native}; +const registry = .{ nemotron.native, qwen35.native }; /// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ @@ -336,7 +337,10 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O defer if (!opened) loaded.deinit(loaded.ctx); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - const room = after.room(model) -| loaded.learn_bytes; + // what a family's Loaded may leave out: no learner, no lone driver, no prompt-cache copies + const L = @TypeOf(loaded); + const learn_bytes: usize = if (@hasField(L, "learn_bytes")) loaded.learn_bytes else 0; + const room = after.room(model) -| learn_bytes; const free = loaded.free_streams; const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { @@ -347,14 +351,14 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; const h = try gpa.create(Host); errdefer gpa.destroy(h); - h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); + h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; + const cache = if (!@hasField(L, "target")) 0 else budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (@hasField(L, "target") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet @@ -365,7 +369,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { + if (@hasField(L, "learns") and loaded.learns != null) { const Slide = nemotron_slide.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig new file mode 100644 index 0000000000..5f0dc8f906 --- /dev/null +++ b/zig/tests/qwen35_cuda_main.zig @@ -0,0 +1,58 @@ +//! `tf-qwen35-cuda `: Qwen3.5-2B's CUDA forward over comma-separated token ids, +//! one token a forward (step: every position's logits) or the ids as one prompt (prompt: the last position's), as +//! little-endian bf16 rows into for a reference to compare. Also prints the step mode's decode speed. +const std = @import("std"); +const cuda = @import("cuda"); +const q = @import("qwen35"); +const E = q.engine; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + const args = try init.minimal.args.toSlice(a); + if (args.len != 5) { + std.debug.print("usage: tf-qwen35-cuda \n", .{}); + return 2; + } + var ids: std.ArrayList(u32) = .empty; + var it = std.mem.tokenizeAny(u8, args[3], ", "); + while (it.next()) |t| try ids.append(a, try std.fmt.parseInt(u32, t, 10)); + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + const e = try E.Engine.init(init.gpa, init.io, &ctx, args[1]); + defer e.deinit(); + var cache = try E.Cache.init(init.gpa, e, 4096); + defer cache.deinit(init.gpa); + var out: std.ArrayList(u8) = .empty; + if (std.mem.eql(u8, args[2], "step")) { + const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + for (ids.items) |id| { + try e.forward(&.{.{ .cache = &cache, .rows = 1 }}, &.{id}, .logits); + @memcpy(cache.logits, e.rowLogits(0)); + try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); + } + const ns: u64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); + std.debug.print("{d} steps in {d:.1} ms: {d:.1} tok/s\n", .{ ids.items.len, @as(f64, @floatFromInt(ns)) / 1e6, @as(f64, @floatFromInt(ids.items.len)) * 1e9 / @as(f64, @floatFromInt(ns)) }); + } else if (std.mem.eql(u8, args[2], "window")) { + // windows of 1, 2, 3, ... rows: every row's logits, which must equal step's bit for bit + var at: usize = 0; + var w: usize = 1; + while (at < ids.items.len) : (w = w % E.window_rows + 1) { + const n = @min(w, ids.items.len - at); + try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .logits); + for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); + at += n; + } + } else { + var at: usize = 0; + while (at < ids.items.len) { + const n = @min(E.max_rows, ids.items.len - at); + try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .last); + at += n; + } + try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); + } + try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[4], .data = out.items }); + return 0; +} From 3ace97ed415aa4fb47b7ff3cb18853b848a1d78e Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:10:43 +0000 Subject: [PATCH 07/35] qwen3_5: serve Qwen3.8-27B on CUDA, one arena for a model's device memory The 2B's CUDA engine now reads its widths from config.json (cuda_geometry.zig): linear value heads over fewer key heads, more query heads a key/value head, an untied head, an A_log stored as bf16. The 27B's logits agree with an fp32 reference at every position of a 36-token check, and its windows equal one-row steps bit for bit; the 2B's output is unchanged. Two Windows driver costs shaped the rest. Each kernel launch grew with the process's allocation count (1.2 ms a launch over 1,900 weight buffers), so a model's weights and scratch now live in one arena (cuda_weights.zig packs each projection through a reused staging buffer: the 27B loads in 25 s, not 200). And a process's first forwards through a buffer stall for seconds, so the engine warms every request path at load and the backend keeps a pool of stream caches that `reserve` allocates and runs a row through once the server knows its stream count. On an RTX 4090: the 27B decodes at 37 tok/s in the runner, 31 to 35 tok/s served, 210 tok/s across eight streams; the 2B stays at 210 and 1,250. --- zig/kernels/cuda/qwen35.cu | 195 +++++++------ zig/src/families/qwen3_5/cuda.zig | 5 +- zig/src/families/qwen3_5/cuda_engine.zig | 325 ++++++++++----------- zig/src/families/qwen3_5/cuda_geometry.zig | 177 +++++++++++ zig/src/families/qwen3_5/cuda_native.zig | 80 +++-- zig/src/families/qwen3_5/cuda_weights.zig | 143 +++++++++ zig/src/native/cuda.zig | 5 + zig/tests/qwen35_cuda_main.zig | 21 +- 8 files changed, 656 insertions(+), 295 deletions(-) create mode 100644 zig/src/families/qwen3_5/cuda_geometry.zig create mode 100644 zig/src/families/qwen3_5/cuda_weights.zig diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index a68fa30990..715f83f64e 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -1,27 +1,30 @@ -// Qwen3.5-2B's glue on CUDA: the Metal engine's kernels (zig/kernels/metal/qwen3_5) with the same bf16 rounding points. -// Rows of one forward belong to segments (one stream's window each); a segment's state lives in its stream's blocks. +// Qwen3.5's hybrid text models on CUDA (the 2B and Qwen3.8-27B): the Metal engine's kernels (zig/kernels/metal/qwen3_5) +// with the same bf16 rounding points, their widths from the model's Geo. Rows of one forward belong to segments (one +// stream's window each); a segment's state lives in its stream's blocks. #include #include namespace tf_q35 { -constexpr int HID = 2048; // hidden -constexpr int NH = 16; // linear-attention heads (keys and values) -constexpr int DK = 128; // their head width -constexpr int CONV = 6144; // q, k and v channels the convolution sees -constexpr int TAPS = 4; -constexpr int QH = 8; // attention query heads -constexpr int KVH = 2; // attention key/value heads -constexpr int HD = 256; // attention head width -constexpr size_t REC = (size_t)NH * DK * DK; // floats of one layer's recurrence -constexpr size_t CONV_ROWS = (size_t)(TAPS - 1) * CONV; +constexpr int DK = 128; // linear-attention key and value head width +constexpr int TAPS = 4; // causal convolution taps +constexpr int HD = 256; // attention head width +constexpr int MAXQ = 32; // attention query heads at most + +// The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads. +struct Geo { + int hid, nk, nv, qh, kvh; + __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees + __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence + __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } +}; // One stream's window in this forward: its state blocks, cache length before the window, and rows. struct Seg { - unsigned long long rec; // float recurrences, one REC block a linear layer - unsigned long long conv; // bf16 convolution tails, CONV_ROWS a linear layer - unsigned long long kv; // bf16 keys then values, KVH * cap * HD each, an attention layer at a time + unsigned long long rec; // float recurrences, one rec() block a linear layer + unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer + unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time int cap, len, row0, rows; }; @@ -34,40 +37,39 @@ __device__ __forceinline__ float warp_sum(float x) { return x; } -__device__ __forceinline__ float sigmoid(float x) { return 1.0f / (1.0f + __expf(-x)); } - } // namespace tf_q35 using namespace tf_q35; -// MLX affine 4-bit token rows: out = bf16(s * q + b). +// MLX affine 4-bit token rows: out = bf16(s * q + b), D wide. extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B, - __nv_bfloat16* __restrict__ out) { + __nv_bfloat16* __restrict__ out, int D) { const size_t id = ids[blockIdx.x]; - for (int b = threadIdx.x; b < HID / 2; b += blockDim.x) { - const uint8_t v = W[id * (HID / 2) + b]; - const size_t g = id * (HID / 64) + (2 * b) / 64; + for (int b = threadIdx.x; b < D / 2; b += blockDim.x) { + const uint8_t v = W[id * (D / 2) + b]; + const size_t g = id * (D / 64) + (2 * b) / 64; const float s = bf(S[g]), z = bf(B[g]); - out[(size_t)blockIdx.x * HID + 2 * b] = __float2bfloat16_rn(s * (float)(v & 15) + z); - out[(size_t)blockIdx.x * HID + 2 * b + 1] = __float2bfloat16_rn(s * (float)(v >> 4) + z); + out[(size_t)blockIdx.x * D + 2 * b] = __float2bfloat16_rn(s * (float)(v & 15) + z); + out[(size_t)blockIdx.x * D + 2 * b + 1] = __float2bfloat16_rn(s * (float)(v >> 4) + z); } } -// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), xs = x's 64-group sums. One 256-thread block a row. +// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), D wide (a multiple of 256, at most 8192). One 256-thread +// block a row; each thread's elements in index order, warps summed in order. extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, - const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, - float* __restrict__ XS, float eps) { - constexpr int E = HID / 256; + const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D) { + constexpr int MAXE = 32; __shared__ float red[8]; - __shared__ float xb[HID]; const size_t row = blockIdx.x; const int t = threadIdx.x; - float hv[E]; + const int E = D / 256; + float hv[MAXE]; float ss = 0.0f; #pragma unroll - for (int i = 0; i < E; i++) { - const size_t at = row * HID + t * E + i; + for (int i = 0; i < MAXE; i++) { + if (i >= E) break; + const size_t at = row * D + t * E + i; float h = bf(H[at]); if (R != nullptr) { const __nv_bfloat16 hb = __float2bfloat16_rn(h + bf(R[at])); @@ -82,19 +84,12 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __r __syncthreads(); float total = 0.0f; for (int i = 0; i < 8; i++) total += red[i]; - const float inv = rsqrtf(total / (float)HID + eps); + const float inv = rsqrtf(total / (float)D + eps); #pragma unroll - for (int i = 0; i < E; i++) { + for (int i = 0; i < MAXE; i++) { + if (i >= E) break; const int d = t * E + i; - const __nv_bfloat16 x = __float2bfloat16_rn(bf(W[d]) * (hv[i] * inv)); - X[row * HID + d] = x; - xb[d] = bf(x); - } - __syncthreads(); - if (XS != nullptr && t < HID / 64) { - float s = 0.0f; - for (int j = 0; j < 64; j++) s += xb[t * 64 + j]; - XS[row * (HID / 64) + t] = s; + X[row * D + d] = __float2bfloat16_rn(bf(W[d]) * (hv[i] * inv)); } } @@ -110,8 +105,8 @@ extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, floa } // The causal convolution, silu, q/k normalization and the decay and write strength of each row's linear heads. -// Grid (1, 3 * NH, rows), a warp a (row, head): heads [0, NH) are q, [NH, 2NH) k, [2NH, 3NH) v. -extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, +// Grid (1, 2 * nk + nv, rows), a warp a (row, head): heads [0, nk) are q, [nk, 2nk) k, then nv of v. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, const int* __restrict__ rowseg, int layer, const __nv_bfloat16* __restrict__ CW, const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, @@ -120,10 +115,11 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloa const int lane = threadIdx.x; const int head = blockIdx.y; const int w = blockIdx.z; + const int NK = geo.nk, NV = geo.nv, C = geo.conv(); const Seg sg = segs[rowseg[w]]; const int r = w - sg.row0; - const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * CONV_ROWS; - const bool isq = head < NH, isk = !isq && head < 2 * NH; + const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * geo.conv_rows(); + const bool isq = head < NK, isk = !isq && head < 2 * NK; const int c0 = head * DK; // q, k and v channels are consecutive heads of DK constexpr int PER = DK / 32; float vals[PER]; @@ -133,7 +129,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloa float acc = 0.0f; for (int tap = 0; tap < TAPS; tap++) { const int at = r + tap; // in [state rows (TAPS - 1)] ++ [this segment's rows] - const float xv = at < TAPS - 1 ? bf(CS[(size_t)at * CONV + c]) : bf(QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]); + const float xv = at < TAPS - 1 ? bf(CS[(size_t)at * C + c]) : bf(QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]); acc += bf(CW[c * TAPS + tap]) * xv; } const float conv = rbf(acc); @@ -147,57 +143,61 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(const __nv_bfloa ss = warp_sum(ss); const float inv = rsqrtf(ss / (float)DK + (1e-6f / (float)DK)); const float scale = isq ? rbf(1.0f / (float)DK) : rbf(rsqrtf((float)DK)); - __nv_bfloat16* out = isq ? Q + ((size_t)w * NH + head) * DK : K + ((size_t)w * NH + head - NH) * DK; + __nv_bfloat16* out = isq ? Q + ((size_t)w * NK + head) * DK : K + ((size_t)w * NK + head - NK) * DK; #pragma unroll for (int j = 0; j < PER; j++) out[lane * PER + j] = __float2bfloat16_rn(scale * rbf(vals[j] * inv)); } else { - const int hv = head - 2 * NH; + const int hv = head - 2 * NK; #pragma unroll - for (int j = 0; j < PER; j++) V[((size_t)w * NH + hv) * DK + lane * PER + j] = __float2bfloat16_rn(vals[j]); + for (int j = 0; j < PER; j++) V[((size_t)w * NV + hv) * DK + lane * PER + j] = __float2bfloat16_rn(vals[j]); if (lane == 0) { - const float s = rbf(bf(Ain[w * NH + hv]) + bf(DT[hv])); + const float s = rbf(bf(Ain[w * NV + hv]) + bf(DT[hv])); const float sp = rbf(fmaxf(s, 0.0f) + logf(1.0f + expf(-fabsf(s)))); - G[w * NH + hv] = expf(-expf(ALOG[hv]) * sp); - BETA[w * NH + hv] = __float2bfloat16_rn(1.0f / (1.0f + expf(-bf(Bin[w * NH + hv])))); + G[w * NV + hv] = expf(-expf(ALOG[hv]) * sp); + BETA[w * NV + hv] = __float2bfloat16_rn(1.0f / (1.0f + expf(-bf(Bin[w * NV + hv])))); } } } // Each segment's convolution tail after its rows: the last TAPS - 1 of [tail ++ rows]. A thread a (segment, channel). -// SNAP, if given, also takes the tail after each row (forward row m at m * CONV_ROWS), read before the update. -extern "C" __global__ void tf_q35_conv_commit(const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, +// SNAP, if given, also takes the tail after each row (forward row m at m * conv_rows()), read before the update. +extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, __nv_bfloat16* __restrict__ SNAP) { + const int C = geo.conv(); const int c = blockIdx.x * blockDim.x + threadIdx.x; - if (c >= CONV) return; + if (c >= C) return; const Seg sg = segs[blockIdx.y]; - __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * CONV_ROWS; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * geo.conv_rows(); if (SNAP != nullptr) for (int j = 0; j < sg.rows; j++) for (int r = 0; r < TAPS - 1; r++) { const int at = r + j + 1; - SNAP[((size_t)(sg.row0 + j) * (TAPS - 1) + r) * CONV + c] = - at < TAPS - 1 ? CS[(size_t)at * CONV + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]; + SNAP[((size_t)(sg.row0 + j) * (TAPS - 1) + r) * C + c] = + at < TAPS - 1 ? CS[(size_t)at * C + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]; } // ascending: slot r takes [r + rows], which is either a later state slot (not yet written) or a row for (int r = 0; r < TAPS - 1; r++) { const int at = r + sg.rows; - CS[(size_t)r * CONV + c] = at < TAPS - 1 ? CS[(size_t)at * CONV + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * CONV + c]; + CS[(size_t)r * C + c] = at < TAPS - 1 ? CS[(size_t)at * C + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]; } } // The gated delta rule over each segment's rows, its recurrence updated in place; SNAP, if given, takes the state after -// each row (forward row m at m * REC). Grid (1, DK / 4, NH * segments), block (32, 4): a warp a (head, value row), -// its 32 lanes over the key width. -extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, +// each row (forward row m at m * rec()). Value head hv reads key head hv / (nv / nk). Grid (1, DK / 4, nv * segments), +// block (32, 4): a warp a (value head, value row), its 32 lanes over the key width. +extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, int layer, __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP) { - const int hv = blockIdx.z % NH; - const Seg sg = segs[blockIdx.z / NH]; + const int NK = geo.nk, NV = geo.nv; + const int hv = blockIdx.z % NV; + const int hk = hv / (NV / NK); + const Seg sg = segs[blockIdx.z / NV]; const int dk = threadIdx.x; const int dv = blockIdx.y * 4 + threadIdx.y; constexpr int PT = DK / 32; - float* st = reinterpret_cast(sg.rec) + (size_t)layer * REC + ((size_t)hv * DK + dv) * DK + PT * dk; + const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; float state[PT]; #pragma unroll for (int i = 0; i < PT; i++) state[i] = st[i]; @@ -206,12 +206,12 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bf float kc[PT], qc[PT]; #pragma unroll for (int i = 0; i < PT; i++) { - kc[i] = bf(k[(row * NH + hv) * DK + PT * dk + i]); - qc[i] = bf(q[(row * NH + hv) * DK + PT * dk + i]); + kc[i] = bf(k[(row * NK + hk) * DK + PT * dk + i]); + qc[i] = bf(q[(row * NK + hk) * DK + PT * dk + i]); } - const float vc = bf(v[(row * NH + hv) * DK + dv]); - const float gc = g[row * NH + hv]; - const float bc = bf(beta[row * NH + hv]); + const float vc = bf(v[(row * NV + hv) * DK + dv]); + const float gc = g[row * NV + hv]; + const float bc = bf(beta[row * NV + hv]); float kv_mem = 0.0f; #pragma unroll for (int i = 0; i < PT; i++) { @@ -227,9 +227,9 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bf out += state[i] * qc[i]; } out = warp_sum(out); - if (dk == 0) y[(row * NH + hv) * DK + dv] = __float2bfloat16_rn(out); + if (dk == 0) y[(row * NV + hv) * DK + dv] = __float2bfloat16_rn(out); if (SNAP != nullptr) { - float* sn = SNAP + row * REC + ((size_t)hv * DK + dv) * DK + PT * dk; + float* sn = SNAP + row * geo.rec() + at; #pragma unroll for (int i = 0; i < PT; i++) sn[i] = state[i]; } @@ -238,16 +238,17 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(const __nv_bf for (int i = 0; i < PT; i++) st[i] = state[i]; } -// out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear head. Grid (1, NH, rows), a warp each. -extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, +// out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear value head; rows nv * DK wide. Grid (1, nv, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT) { const int lane = threadIdx.x, hv = blockIdx.y; const size_t m = blockIdx.z; + const size_t W = (size_t)geo.nv * DK; constexpr int PER = DK / 32; float yv[PER], ss = 0.0f; #pragma unroll for (int j = 0; j < PER; j++) { - yv[j] = bf(Y[(m * NH + hv) * DK + lane * PER + j]); + yv[j] = bf(Y[m * W + hv * DK + lane * PER + j]); ss += yv[j] * yv[j]; } ss = warp_sum(ss); @@ -256,8 +257,8 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(const __nv_bflo for (int j = 0; j < PER; j++) { const int d = lane * PER + j; const float x = rbf(bf(NW[d]) * (yv[j] * inv)); - const float zf = bf(Z[m * HID + hv * DK + d]); - OUT[m * HID + hv * DK + d] = __float2bfloat16_rn(zf / (1.0f + expf(-zf)) * x); + const float zf = bf(Z[m * W + hv * DK + d]); + OUT[m * W + hv * DK + d] = __float2bfloat16_rn(zf / (1.0f + expf(-zf)) * x); } } @@ -300,14 +301,15 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i // Queries in place at their positions; keys rotated and values copied into their segment's cache. // Grid (rows), HD threads. -extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(__nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, +extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, const int* __restrict__ rowseg, int layer) { + const int QH = geo.qh, KVH = geo.kvh; const int d = threadIdx.x; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; const int pos = sg.len + (int)m - sg.row0; - __shared__ __nv_bfloat16 qs[QH * HD]; + __shared__ __nv_bfloat16 qs[MAXQ * HD]; for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; __syncthreads(); for (int h = 0; h < QH; h++) Qn[(m * QH + h) * HD + d] = rotate(qs + h * HD, d, pos); @@ -320,12 +322,13 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(__nv_bfloat16* __r } } -// Causal attention over a segment's cache, ck keys a block: grid (chunks, KVH, rows), a warp a query head of the -// kv head's group of QH / KVH. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. -extern "C" __global__ void __launch_bounds__(128) tf_q35_attn_partial(const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, int nch, int ck, float scale, - float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { +// Causal attention over a segment's cache, ck keys a block: grid (chunks, kvh, rows), block qh / kvh warps, a warp a +// query head of the kv head's group. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, int layer, int nch, int ck, float scale, + float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { constexpr int DPL = HD / 32; + const int QH = geo.qh, KVH = geo.kvh; const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; const int c = blockIdx.x, h = blockIdx.y; const size_t m = blockIdx.z; @@ -374,10 +377,11 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_attn_partial(const __nv for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; } -// The chunks' partials into each (row, query head)'s output. Grid (QH, rows), a warp each. -extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(const float* __restrict__ PM, const float* __restrict__ PL, +// The chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. +extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, const float* __restrict__ PM, const float* __restrict__ PL, const float* __restrict__ PO, int nch, __nv_bfloat16* __restrict__ OUT) { constexpr int DPL = HD / 32; + const int QH = geo.qh; const int lane = threadIdx.x, qh = blockIdx.x; const size_t m = blockIdx.y; const size_t base = (m * QH + qh) * nch; @@ -398,19 +402,20 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(const float* for (int i = 0; i < DPL; i++) OUT[(m * QH + qh) * HD + lane * DPL + i] = __float2bfloat16_rn(acc[i] / lsum); } -// out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows. -extern "C" __global__ void tf_q35_attn_gate(const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, +// out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows; rows qh * HD wide. +extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, __nv_bfloat16* __restrict__ OUT, int rows) { + const size_t W = (size_t)geo.qh * HD; const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; - if (i >= (size_t)rows * HID) return; - const size_t r = i / HID; - const int d = (int)(i % HID), head = d / HD; - const float gate = bf(QRAW[r * 2 * HID + head * 2 * HD + HD + d % HD]); + if (i >= (size_t)rows * W) return; + const size_t r = i / W; + const int d = (int)(i % W), head = d / HD; + const float gate = bf(QRAW[r * 2 * W + head * 2 * HD + HD + d % HD]); const float s = rbf(1.0f / (1.0f + expf(-gate))); OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); } -// out = bf16(silu(gate) * up), n wide rows. +// out = bf16(silu(gate) * up) over `total` values. extern "C" __global__ void tf_q35_mlp_act(const __nv_bfloat16* __restrict__ GATE, const __nv_bfloat16* __restrict__ UP, __nv_bfloat16* __restrict__ OUT, long long total) { const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index fdecefed7f..ea5eaa465e 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -1,9 +1,12 @@ -//! Qwen3.5-2B on CUDA: the engine and its native family entry (the Metal engine is qwen3_5.zig). +//! The Qwen3.5 hybrids on CUDA (the 2B, Qwen3.8-27B): the engine and its native family entry (the Metal 2B engine is +//! qwen3_5.zig). pub const config = @import("config.zig"); +pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); test { _ = config; + _ = geometry; _ = @import("draw.zig"); } diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 96190737d1..316814eaf6 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -1,43 +1,40 @@ -//! Qwen3.5-2B on CUDA: the affine-4 projections through cuda.qlinear, the glue in qwen35.cu, one forward for any -//! mix of streams' windows. Each stream owns its recurrences, convolution tails and KV cache (Cache). +//! Qwen3.5's hybrid text models on CUDA (the 2B, Qwen3.8-27B): the affine-4 projections through cuda.qlinear, the +//! glue in qwen35.cu, one forward for any mix of streams' windows. Each stream owns its recurrences, convolution tails +//! and KV cache (Cache); the widths come from config.json (cuda_geometry.zig), the weights from cuda_weights.zig. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); -const c = @import("config.zig"); +const geometry = @import("cuda_geometry.zig"); +const weights = @import("cuda_weights.zig"); +pub const Geometry = geometry.Geometry; const Affine4 = cuda.qlinear.Affine4; const DeviceBuffer = cuda.DeviceBuffer; +const head_dim = geometry.head_dim; /// Rows of one forward at most: a prompt chunk, or a shared round's windows. pub const max_rows = 128; /// Streams' windows in one forward at most. pub const max_segs = 32; /// Rows of a verify forward at most: each row's recurrences are recorded for a partial acceptance. -pub const max_verify = 32; +const most_verify = 32; /// A window's rows at most (exact: one-row steps' bits). -pub const window_rows = 16; +const most_window = 16; +/// Device bytes the recorded rows may take: a wide model records fewer rows a forward. +const snapshot_budget = 1280 << 20; /// Partial attention chunks a row at most: the chunk width grows with the window instead. const max_chunks = 64; -const deltas = c.layers - c.layers / 4; -const attentions = c.layers / 4; -pub const rec_bytes = c.linear_heads * c.linear_dim * c.linear_dim * 4; -pub const conv_bytes = (c.conv_taps - 1) * c.conv_dim * 2; -const state_bytes = deltas * (rec_bytes + conv_bytes); - -/// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values. -pub fn streamBytes(cap: usize) usize { - return state_bytes + kvBytes(cap); -} - -fn kvBytes(cap: usize) usize { - return attentions * 2 * c.kv_heads * cap * c.head_dim * 2; +/// Device allocations are 256-byte aligned, as the driver's are. +pub fn aligned(n: usize) usize { + return std.mem.alignForward(usize, n, 256); } -/// qwen35.cu's Seg. +/// qwen35.cu's Geo and Seg. +const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int }; const Seg = extern struct { rec: u64, conv: u64, kv: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int }; -/// One stream's state on the device and the logits of its newest row on the host. +/// One stream's state on the device and the logits of its newest prompt chunk on the host. pub const Cache = struct { state: DeviceBuffer, kv: DeviceBuffer, @@ -48,12 +45,20 @@ pub const Cache = struct { last: ?struct { start: usize, row0: usize, rows: usize } = null, pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { - var state = try DeviceBuffer.alloc(e.d, state_bytes); + var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); errdefer state.free(); try state.fill8(0, null); - var kv = try DeviceBuffer.alloc(e.d, kvBytes(cap)); + var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); errdefer kv.free(); - return .{ .state = state, .kv = kv, .cap = cap, .logits = try gpa.alloc(u16, c.vocab) }; + try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) + return .{ .state = state, .kv = kv, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + } + + /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). + pub fn reset(self: *Cache, e: *const Engine) !void { + self.len = 0; + self.last = null; + try self.state.fill8(0, e.stream.handle); } pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { @@ -62,14 +67,14 @@ pub const Cache = struct { gpa.free(self.logits); } - fn seg(self: *const Cache, row0: usize, rows: usize) Seg { - return .{ .rec = self.state.ptr, .conv = self.state.ptr + deltas * rec_bytes, .kv = self.kv.ptr, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows) }; + fn seg(self: *const Cache, g: Geometry, row0: usize, rows: usize) Seg { + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows) }; } }; -const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; -const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; -const Layer = struct { +pub const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; +pub const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; +pub const Layer = struct { input_norm: u64, post_norm: u64, gate: Affine4, @@ -109,19 +114,28 @@ pub const Draw = struct { rule: ?cuda.sample.Rule, position: u64 }; pub const Engine = struct { gpa: std.mem.Allocator, d: *const cuda.Driver, + g: Geometry, + geo: Geo, stream: cuda.Stream, mods: [6]cuda.Module, draws: cuda.sample.Draws, affine: cuda.qlinear.Affine4Kernels, fns: Fns, - held: std.ArrayList(DeviceBuffer) = .empty, // weights and scratch, freed together + /// Every weight and scratch buffer in one allocation: the WDDM driver's launch cost grows with a process's + /// allocation count, and a 64-layer model would have thousands. + arena: DeviceBuffer, + arena_at: usize = 0, + stage: DeviceBuffer, // the load's staging of each projection as stored, packed from there + /// Rows a verify forward records (snapshot_budget over a row's states), and a window's rows at most. + verify_rows: usize, + window_rows: usize, // weights embed_w: u64, embed_s: u64, embed_b: u64, head: Affine4, final_norm: u64, - layers: [c.layers]Layer, + layers: []Layer, // scratch, max_rows rows each h: u64, x: u64, @@ -134,7 +148,7 @@ pub const Engine = struct { q: u64, k: u64, v: u64, - g: u64, + g_: u64, beta: u64, y: u64, mix: u64, @@ -150,35 +164,44 @@ pub const Engine = struct { pl: u64, po: u64, logits: u64, + rec_snap: u64, // recorded recurrences: verify_rows rows a linear layer + conv_snap: u64, // recorded convolution tails, the same // the per-forward tables: ids, segments, each row's segment and the draws' rules, staged in pinned memory tables: DeviceBuffer, staged: cuda.HostBuffer, host_logits: cuda.HostBuffer, /// The kept rows' device draws, when forward was given them. - tokens: [max_verify]u32 = undefined, - rec_snap: u64, // recorded recurrences: max_verify rows a linear layer - conv_snap: u64, // recorded convolution tails, the same + tokens: [most_verify]u32 = undefined, const ids_at = 0; const segs_at = max_rows * 4; const rowseg_at = segs_at + max_segs * @sizeOf(Seg); const rules_at = rowseg_at + max_rows * 4; - const meta_at = rules_at + max_verify * @sizeOf(cuda.sample.Rule); - const tokens_at = meta_at + max_verify * 4; - const tables_len = tokens_at + max_verify * 4; + const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); + const tokens_at = meta_at + most_verify * 4; + const tables_len = tokens_at + most_verify * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; - const config = try c.Config.read(gpa, io, dir); + const g = try Geometry.read(gpa, io, dir); const d = ctx.d; const e = try gpa.create(Engine); errdefer gpa.destroy(e); e.* = undefined; e.gpa = gpa; e.d = d; - e.held = .empty; + e.g = g; + e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads) }; + e.verify_rows = std.math.clamp(snapshot_budget / (g.deltas() * (g.recBytes() + g.convBytes())), 1, most_verify); + e.window_rows = @min(most_window, e.verify_rows); + e.arena_at = 0; e.stream = try cuda.Stream.init(d, true); errdefer e.stream.deinit(); + e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); + errdefer e.arena.free(); + try e.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly + e.stage = try DeviceBuffer.alloc(d, weights.stageBytes(g)); + errdefer e.stage.free(); const kk = cuda.kernels; const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample }; var loaded: usize = 0; @@ -193,16 +216,26 @@ pub const Engine = struct { errdefer e.affine.deinit(); e.draws = try cuda.sample.Draws.resolve(e.mods[5]); inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(e.fns, name) = try e.mods[4].function("tf_q35_" ++ name); - errdefer e.freeHeld(); + e.layers = try gpa.alloc(Layer, g.layers); + errdefer gpa.free(e.layers); var ck = try core.Checkpoint.openModel(gpa, io, dir); defer ck.close(); - try e.loadWeights(&ck, config.tied); - try e.allocScratch(); + try weights.load(e, &ck); + _ = try e.scratch(.allocate); e.tables = try DeviceBuffer.alloc(d, tables_len); errdefer e.tables.free(); e.staged = try cuda.HostBuffer.alloc(d, tables_len); errdefer e.staged.free(); - e.host_logits = try cuda.HostBuffer.alloc(d, max_verify * c.vocab * 2); + e.host_logits = try cuda.HostBuffer.alloc(d, most_verify * g.vocab * 2); + errdefer e.host_logits.free(); + // every path a request takes, once now: the WDDM driver stalls a process's first forwards for seconds + // (a 16 GiB model: a 20 s launch, then a 3 s one), which no request should pay + var warm = try Cache.init(gpa, e, 64); + defer warm.deinit(gpa); + const ids: [32]u32 = @splat(0); + try e.forward(&.{.{ .cache = &warm, .rows = ids.len }}, &ids, .last); // a prompt chunk: the prompt GEMMs + try e.forward(&.{.{ .cache = &warm, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = null, .position = 33 }} }); + try e.forward(&.{.{ .cache = &warm, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = cuda.sample.rule(1, 0.7, 0.95, -std.math.inf(f64), 20), .position = 34 }} }); return e; } @@ -210,117 +243,62 @@ pub const Engine = struct { e.host_logits.free(); e.staged.free(); e.tables.free(); - e.freeHeld(); + e.stage.free(); + e.arena.free(); + e.gpa.free(e.layers); e.affine.deinit(); for (&e.mods) |*m| m.unload(); e.stream.deinit(); e.gpa.destroy(e); } - fn freeHeld(e: *Engine) void { - for (e.held.items) |*b| b.free(); - e.held.deinit(e.gpa); - } - - fn own(e: *Engine, b: DeviceBuffer) !u64 { - try e.held.append(e.gpa, b); - return b.ptr; - } - - fn alloc(e: *Engine, bytes: usize) !u64 { - var b = try DeviceBuffer.alloc(e.d, bytes); - errdefer b.free(); - return e.own(b); + /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values. + pub fn streamBytes(e: *const Engine, cap: usize) usize { + return e.g.stateBytes() + e.g.kvBytes(cap); } - fn put(e: *Engine, t: core.checkpoint.Tensor) !u64 { - var b = try DeviceBuffer.fromHost(e.d, t.bytes); - errdefer b.free(); - return e.own(b); + /// The next `bytes` of the arena, 256-byte aligned. + pub fn alloc(e: *Engine, bytes: usize) !u64 { + const at = aligned(e.arena_at); + if (at + bytes > e.arena.len) return error.ArenaExhausted; + e.arena_at = at + bytes; + return e.arena.ptr + at; } - fn tensor(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { - var buf: [160]u8 = undefined; - return e.put(try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); - } - - /// An MLX affine-4 projection packed into qlinear's tiles; the checkpoint's own layout leaves the device after. - fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { - var buf: [160]u8 = undefined; - var raw: [3]DeviceBuffer = undefined; - var got: usize = 0; - defer for (raw[0..got]) |*b| b.free(); - const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / c.group }, .{ "biases", .bf16, k / c.group } }; - for (parts) |p| { - const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); - raw[got] = try DeviceBuffer.fromHost(e.d, t.bytes); - got += 1; - } - const l = Affine4.layout(n, k); - const q: Affine4 = .{ .w = try e.alloc(l.words), .s = try e.alloc(l.scales), .b = try e.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; - try e.affine.pack(e.stream, raw[0].ptr, raw[1].ptr, raw[2].ptr, q); - try e.stream.synchronize(); - return q; - } - - fn loadWeights(e: *Engine, ck: *core.Checkpoint, tied: bool) !void { - const m = "language_model.model."; - e.embed_w = try e.tensor(ck, "{s}embed_tokens.weight", .{m}, .u32, &.{ c.vocab, c.hidden / 8 }); - e.embed_s = try e.tensor(ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ c.vocab, c.hidden / c.group }); - e.embed_b = try e.tensor(ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ c.vocab, c.hidden / c.group }); - e.head = if (tied) try e.linear(ck, m, "embed_tokens", c.vocab, c.hidden) else try e.linear(ck, "language_model.", "lm_head", c.vocab, c.hidden); - e.final_norm = try e.tensor(ck, "{s}norm.weight", .{m}, .bf16, &.{c.hidden}); - for (&e.layers, 0..) |*l, i| { - var pbuf: [96]u8 = undefined; - const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); - l.* = .{ - .input_norm = try e.tensor(ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{c.hidden}), - .post_norm = try e.tensor(ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{c.hidden}), - .gate = try e.linear(ck, p, "mlp.gate_proj", c.intermediate, c.hidden), - .up = try e.linear(ck, p, "mlp.up_proj", c.intermediate, c.hidden), - .down = try e.linear(ck, p, "mlp.down_proj", c.hidden, c.intermediate), - .mixer = undefined, - }; - if (c.linear(i)) { - l.mixer = .{ .delta = .{ - .qkv = try e.linear(ck, p, "linear_attn.in_proj_qkv", c.conv_dim, c.hidden), - .z = try e.linear(ck, p, "linear_attn.in_proj_z", c.hidden, c.hidden), - .a = try e.linear(ck, p, "linear_attn.in_proj_a", c.linear_heads, c.hidden), - .b = try e.linear(ck, p, "linear_attn.in_proj_b", c.linear_heads, c.hidden), - .out = try e.linear(ck, p, "linear_attn.out_proj", c.hidden, c.hidden), - .conv = try e.tensor(ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ c.conv_dim, c.conv_taps, 1 }), - .norm = try e.tensor(ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{c.linear_dim}), - .a_log = try e.tensor(ck, "{s}linear_attn.A_log", .{p}, .f32, &.{c.linear_heads}), - .dt_bias = try e.tensor(ck, "{s}linear_attn.dt_bias", .{p}, .bf16, &.{c.linear_heads}), - } }; - } else { - l.mixer = .{ .attention = .{ - .q = try e.linear(ck, p, "self_attn.q_proj", 2 * c.hidden, c.hidden), - .k = try e.linear(ck, p, "self_attn.k_proj", c.kv_heads * c.head_dim, c.hidden), - .v = try e.linear(ck, p, "self_attn.v_proj", c.kv_heads * c.head_dim, c.hidden), - .o = try e.linear(ck, p, "self_attn.o_proj", c.hidden, c.hidden), - .q_norm = try e.tensor(ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{c.head_dim}), - .k_norm = try e.tensor(ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{c.head_dim}), - } }; + /// The scratch rows' bytes (`measure`: summed, nothing taken), or the rows themselves from the arena. + fn scratch(e: *Engine, mode: enum { measure, allocate }) !usize { + var total: usize = 0; + const Take = struct { + e: *Engine, + mode: @TypeOf(mode), + total: *usize, + fn alloc(t: @This(), bytes: usize) !u64 { + t.total.* += aligned(bytes); + return if (t.mode == .measure) 0 else try t.e.alloc(bytes); } - } - } - - fn allocScratch(e: *Engine) !void { + }; + const t = Take{ .e = e, .mode = mode, .total = &total }; + const g = e.g; const R = max_rows; - inline for (.{ "h", "x", "r", "z", "q", "k", "v", "y", "mix", "qn" }) |name| @field(e, name) = try e.alloc(R * c.hidden * 2); - inline for (.{ "qkv", "gate", "up", "act" }) |name| @field(e, name) = try e.alloc(R * c.intermediate * 2); - e.xs = try e.alloc(R * (c.intermediate / 64) * 4); - inline for (.{ "a", "b", "beta" }) |name| @field(e, name) = try e.alloc(R * c.linear_heads * 2); - e.g = try e.alloc(R * c.linear_heads * 4); - e.qraw = try e.alloc(R * 2 * c.hidden * 2); - inline for (.{ "kraw", "vraw", "kn" }) |name| @field(e, name) = try e.alloc(R * c.kv_heads * c.head_dim * 2); - e.pm = try e.alloc(R * c.query_heads * max_chunks * 4); - e.pl = try e.alloc(R * c.query_heads * max_chunks * 4); - e.po = try e.alloc(R * c.query_heads * max_chunks * c.head_dim * 4); - e.logits = try e.alloc(max_verify * c.vocab * 2); - e.rec_snap = try e.alloc(deltas * max_verify * rec_bytes); - e.conv_snap = try e.alloc(deltas * max_verify * conv_bytes); + const wide = @max(g.valueWidth(), g.attnWidth(), g.hidden); // y and mix: either mixer's rows + inline for (.{ "h", "x", "r" }) |name| @field(e, name) = try t.alloc(R * g.hidden * 2); + inline for (.{ "y", "mix", "z", "v" }) |name| @field(e, name) = try t.alloc(R * wide * 2); + inline for (.{ "q", "k" }) |name| @field(e, name) = try t.alloc(R * g.k_heads * geometry.dk * 2); + e.qkv = try t.alloc(R * g.convWidth() * 2); + inline for (.{ "gate", "up", "act" }) |name| @field(e, name) = try t.alloc(R * g.intermediate * 2); + e.xs = try t.alloc(R * (@max(g.intermediate, wide) / 64) * 4); + inline for (.{ "a", "b", "beta" }) |name| @field(e, name) = try t.alloc(R * g.v_heads * 2); + e.g_ = try t.alloc(R * g.v_heads * 4); + e.qraw = try t.alloc(R * 2 * g.attnWidth() * 2); + e.qn = try t.alloc(R * g.attnWidth() * 2); + inline for (.{ "kraw", "vraw", "kn" }) |name| @field(e, name) = try t.alloc(R * g.kv_heads * head_dim * 2); + e.pm = try t.alloc(R * g.q_heads * max_chunks * 4); + e.pl = try t.alloc(R * g.q_heads * max_chunks * 4); + e.po = try t.alloc(R * g.q_heads * max_chunks * head_dim * 4); + e.logits = try t.alloc(most_verify * g.vocab * 2); + e.rec_snap = try t.alloc(g.deltas() * e.verify_rows * g.recBytes()); + e.conv_snap = try t.alloc(g.deltas() * e.verify_rows * g.convBytes()); + return total; } fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { @@ -343,11 +321,12 @@ pub const Engine = struct { } fn norm(e: *Engine, residual: bool, w: u64, rows: usize) !void { - try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, @as(u64, 0), c.eps }); + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)) }); } /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { + const g = e.g; var total: usize = 0; var widest: usize = 0; for (parts) |p| { @@ -357,10 +336,10 @@ pub const Engine = struct { } if (parts.len == 0 or parts.len > max_segs or total != ids.len or total > max_rows) return error.InvalidQwenSegments; const record = keep != .last; - if (record and total > max_verify) return error.InvalidQwenSegments; + if (record and total > e.verify_rows) return error.InvalidQwenSegments; if (!record and parts.len != 1) return error.InvalidQwenSegments; if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; - for (ids) |id| if (id >= c.vocab) return error.InvalidQwenToken; + for (ids) |id| if (id >= g.vocab) return error.InvalidQwenToken; const heads: usize = if (record) total else 1; // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps const prompt = !record; @@ -371,7 +350,7 @@ pub const Engine = struct { const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. total * 4]); var row0: usize = 0; for (parts, 0..) |p, i| { - segs[i] = p.cache.seg(row0, p.rows); + segs[i] = p.cache.seg(g, row0, p.rows); for (row0..row0 + p.rows) |r| rowseg[r] = @intCast(i); p.cache.last = if (record) .{ .start = p.cache.len, .row0 = row0, .rows = p.rows } else null; row0 += p.rows; @@ -389,11 +368,13 @@ pub const Engine = struct { const t_segs = e.tables.ptr + segs_at; const t_rows = e.tables.ptr + rowseg_at; const nseg = parts.len; + const geo = e.geo; // attention chunks: at most max_chunks a row, each a multiple of 128 keys const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); const nch = (widest + ck - 1) / ck; + const group = g.q_heads / g.kv_heads; - try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h }); + try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)) }); var di: usize = 0; var ai: c_int = 0; for (e.layers, 0..) |l, i| { @@ -401,29 +382,30 @@ pub const Engine = struct { switch (l.mixer) { .delta => |m| { const layer: c_int = @intCast(di); - const rec_snap: u64 = if (record) e.rec_snap + di * max_verify * rec_bytes else 0; - const conv_snap: u64 = if (record) e.conv_snap + di * max_verify * conv_bytes else 0; + const rec_snap: u64 = if (record) e.rec_snap + di * e.verify_rows * g.recBytes() else 0; + const conv_snap: u64 = if (record) e.conv_snap + di * e.verify_rows * g.convBytes() else 0; try e.proj(m.qkv, e.x, e.qkv, total, prompt); try e.proj(m.z, e.x, e.z, total, prompt); try e.proj(m.a, e.x, e.a, total, prompt); try e.proj(m.b, e.x, e.b, total, prompt); - try e.run(e.fns.gdn_pre, .{ 1, 3 * c.linear_heads, total }, .{ 32, 1, 1 }, .{ e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g, e.beta }); - try e.run(e.fns.conv_commit, .{ (c.conv_dim + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ e.qkv, t_segs, layer, conv_snap }); - try e.run(e.fns.gdn_chain, .{ 1, c.linear_dim / 4, c.linear_heads * nseg }, .{ 32, 4, 1 }, .{ e.q, e.k, e.v, e.g, e.beta, t_segs, layer, e.y, rec_snap }); - try e.run(e.fns.gdn_post, .{ 1, c.linear_heads, total }, .{ 32, 1, 1 }, .{ e.y, e.z, m.norm, c.eps, e.mix }); + try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta }); + try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, layer, conv_snap }); + try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, layer, e.y, rec_snap }); + try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix }); try e.proj(m.out, e.mix, e.r, total, prompt); di += 1; }, .attention => |m| { + const aw = g.attnWidth(); try e.proj(m.q, e.x, e.qraw, total, prompt); try e.proj(m.k, e.x, e.kraw, total, prompt); try e.proj(m.v, e.x, e.vraw, total, prompt); - try e.run(e.fns.head_norm, .{ total * c.query_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, c.query_heads), @as(c_int, 2 * c.hidden), @as(c_int, 2 * c.head_dim), c.eps }); - try e.run(e.fns.head_norm, .{ total * c.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, c.kv_heads), @as(c_int, c.kv_heads * c.head_dim), @as(c_int, c.head_dim), c.eps }); - try e.run(e.fns.rope, .{ total, 1, 1 }, .{ c.head_dim, 1, 1 }, .{ e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); - try e.run(e.fns.attn_partial, .{ nch, c.kv_heads, total }, .{ 128, 1, 1 }, .{ e.qn, t_segs, t_rows, ai, @as(c_int, @intCast(nch)), @as(c_int, @intCast(ck)), @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); - try e.run(e.fns.attn_merge, .{ c.query_heads, total, 1 }, .{ 32, 1, 1 }, .{ e.pm, e.pl, e.po, @as(c_int, @intCast(nch)), e.y }); - try e.run(e.fns.attn_gate, .{ (total * c.hidden + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); + try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps }); + try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); + try e.run(e.fns.attn_partial, .{ nch, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, @as(c_int, @intCast(nch)), @as(c_int, @intCast(ck)), @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, @as(c_int, @intCast(nch)), e.y }); + try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); try e.proj(m.o, e.mix, e.r, total, prompt); ai += 1; }, @@ -431,29 +413,29 @@ pub const Engine = struct { try e.norm(true, l.post_norm, total); try e.proj(l.gate, e.x, e.gate, total, prompt); try e.proj(l.up, e.x, e.up, total, prompt); - const n: c_longlong = @intCast(total * c.intermediate); - try e.run(e.fns.mlp_act, .{ (total * c.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n }); + const n: c_longlong = @intCast(total * g.intermediate); + try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n }); try e.proj(l.down, e.act, e.r, total, prompt); } try e.norm(true, e.final_norm, total); // the head over the kept rows - try e.proj(e.head, e.x + (total - heads) * c.hidden * 2, e.logits, heads, false); + try e.proj(e.head, e.x + (total - heads) * g.hidden * 2, e.logits, heads, false); for (parts) |p| p.cache.len += p.rows; switch (keep) { .draws => |pk| { const toks = e.tables.ptr + tokens_at; for (pk, 0..) |p, i| { - const row = e.logits + i * c.vocab * 2; + const row = e.logits + i * g.vocab * 2; if (p.rule != null) { - try e.draws.draw(e.stream, row, c.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); - } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, c.vocab), toks + i * 4 }); + try e.draws.draw(e.stream, row, g.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); + } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, @intCast(g.vocab)), toks + i * 4 }); } try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); try e.stream.synchronize(); @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); }, .logits, .last => { - const out = e.host_logits.bytes[0 .. heads * c.vocab * 2]; + const out = e.host_logits.bytes[0 .. heads * g.vocab * 2]; try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); try e.stream.synchronize(); if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); @@ -463,21 +445,24 @@ pub const Engine = struct { /// Row `i`'s logits from the newest forward that kept logits. pub fn rowLogits(e: *const Engine, i: usize) []const u16 { - return e.host_logits.slice(u16)[i * c.vocab ..][0..c.vocab]; + return e.host_logits.slice(u16)[i * e.g.vocab ..][0..e.g.vocab]; } /// Keeps the first `rows` of the cache's newest window: its recurrences and convolution tails as they were after /// that row (the forward recorded them), its length at that row. The window's later keys stay as unread slots. pub fn rewind(e: *Engine, cache: *Cache, rows: usize) !void { + const g = e.g; const last = cache.last orelse return error.NothingToKeep; if (rows == 0 or rows > last.rows) return error.NothingToKeep; cache.last = null; cache.len = last.start + rows; if (rows == last.rows) return; const row = last.row0 + rows - 1; - for (0..deltas) |di| { - try cache.state.copyFrom(di * rec_bytes, e.rec_snap + (di * max_verify + row) * rec_bytes, rec_bytes, e.stream.handle); - try cache.state.copyFrom(deltas * rec_bytes + di * conv_bytes, e.conv_snap + (di * max_verify + row) * conv_bytes, conv_bytes, e.stream.handle); + const rb = g.recBytes(); + const cb = g.convBytes(); + for (0..g.deltas()) |di| { + try cache.state.copyFrom(di * rb, e.rec_snap + (di * e.verify_rows + row) * rb, rb, e.stream.handle); + try cache.state.copyFrom(g.deltas() * rb + di * cb, e.conv_snap + (di * e.verify_rows + row) * cb, cb, e.stream.handle); } try e.stream.synchronize(); } diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig new file mode 100644 index 0000000000..34207bacb5 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -0,0 +1,177 @@ +//! The widths of a Qwen3.5 hybrid text model (the 2B, Qwen3.8-27B) as the CUDA engine reads them from config.json: +//! what varies is read, what qwen35.cu fixes (head widths, taps, rotary layout, 4-bit groups of 64) is checked. +const std = @import("std"); + +pub const max_layers = 128; +/// qwen35.cu's fixed widths. +pub const dk = 128; +pub const taps = 4; +pub const head_dim = 256; +pub const max_query_heads = 32; + +pub const Geometry = struct { + hidden: usize, + intermediate: usize, + layers: usize, + vocab: usize, + k_heads: usize, // linear-attention key heads + v_heads: usize, // linear-attention value heads + q_heads: usize, // attention query heads + kv_heads: usize, // attention key/value heads + eps: f32, + tied: bool, + context: usize, + linear: [max_layers]bool, + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Geometry { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + /// Value-head width of the gated delta rule's rows (and of z and the out projection's input). + pub fn valueWidth(g: Geometry) usize { + return g.v_heads * dk; + } + + /// Channels the causal convolution sees: q, k and v. + pub fn convWidth(g: Geometry) usize { + return (2 * g.k_heads + g.v_heads) * dk; + } + + /// Attention output width (q_heads of head_dim), the o projection's input. + pub fn attnWidth(g: Geometry) usize { + return g.q_heads * head_dim; + } + + pub fn deltas(g: Geometry) usize { + var n: usize = 0; + for (g.linear[0..g.layers]) |l| n += @intFromBool(l); + return n; + } + + pub fn attentions(g: Geometry) usize { + return g.layers - g.deltas(); + } + + /// Bytes of one linear layer's recurrence (fp32) and convolution tail (bf16). + pub fn recBytes(g: Geometry) usize { + return g.v_heads * dk * dk * 4; + } + + pub fn convBytes(g: Geometry) usize { + return (taps - 1) * g.convWidth() * 2; + } + + /// A stream's recurrences and tails. + pub fn stateBytes(g: Geometry) usize { + return g.deltas() * (g.recBytes() + g.convBytes()); + } + + /// A stream's keys and values at a `cap`-token window. + pub fn kvBytes(g: Geometry, cap: usize) usize { + return g.attentions() * 2 * g.kv_heads * cap * head_dim * 2; + } +}; + +fn object(v: std.json.Value) !std.json.ObjectMap { + return if (v == .object) v.object else error.BadQwenConfig; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) !usize { + const v = o.get(key) orelse return error.BadQwenConfig; + if (v != .integer or v.integer < 1) return error.BadQwenConfig; + return @intCast(v.integer); +} + +fn number(o: std.json.ObjectMap, key: []const u8) !f64 { + return switch (o.get(key) orelse return error.BadQwenConfig) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadQwenConfig, + }; +} + +fn is(o: std.json.ObjectMap, key: []const u8, comptime T: type, want: T) !void { + const v = o.get(key) orelse return error.UnsupportedQwenConfig; + const ok = switch (T) { + bool => v == .bool and v.bool == want, + []const u8 => v == .string and std.mem.eql(u8, v.string, want), + usize => v == .integer and v.integer == want, + else => @compileError("type"), + }; + if (!ok) return error.UnsupportedQwenConfig; +} + +pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { + const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer doc.deinit(); + const root = try object(doc.value); + try is(root, "model_type", []const u8, "qwen3_5"); + const t = try object(root.get("text_config") orelse return error.BadQwenConfig); + try is(t, "model_type", []const u8, "qwen3_5_text"); + try is(t, "attention_bias", bool, false); + try is(t, "attn_output_gate", bool, true); + try is(t, "hidden_act", []const u8, "silu"); + try is(t, "head_dim", usize, head_dim); + try is(t, "linear_key_head_dim", usize, dk); + try is(t, "linear_value_head_dim", usize, dk); + try is(t, "linear_conv_kernel_dim", usize, taps); + const rope = try object(t.get("rope_parameters") orelse return error.UnsupportedQwenRotary); + try is(rope, "rope_type", []const u8, "default"); + if (try number(rope, "rope_theta") != 10000000 or try number(rope, "partial_rotary_factor") != 0.25) return error.UnsupportedQwenRotary; + const q = try object(root.get("quantization") orelse return error.UnsupportedQwenQuantization); + if (try integer(q, "bits") != 4 or try integer(q, "group_size") != 64) return error.UnsupportedQwenQuantization; + if (q.get("mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "affine")) return error.UnsupportedQwenQuantization; + var g: Geometry = .{ + .hidden = try integer(t, "hidden_size"), + .intermediate = try integer(t, "intermediate_size"), + .layers = try integer(t, "num_hidden_layers"), + .vocab = try integer(t, "vocab_size"), + .k_heads = try integer(t, "linear_num_key_heads"), + .v_heads = try integer(t, "linear_num_value_heads"), + .q_heads = try integer(t, "num_attention_heads"), + .kv_heads = try integer(t, "num_key_value_heads"), + .eps = @floatCast(try number(t, "rms_norm_eps")), + .tied = blk: { + const v = t.get("tie_word_embeddings") orelse root.get("tie_word_embeddings") orelse break :blk false; + if (v != .bool) return error.BadQwenConfig; + break :blk v.bool; + }, + .context = try integer(t, "max_position_embeddings"), + .linear = @splat(false), + }; + if (g.layers > max_layers or g.hidden % 256 != 0 or g.hidden > 8192 or g.intermediate % 64 != 0) return error.UnsupportedQwenGeometry; + if (g.v_heads % g.k_heads != 0 or g.q_heads % g.kv_heads != 0 or g.q_heads > max_query_heads) return error.UnsupportedQwenGeometry; + const types = t.get("layer_types") orelse return error.UnsupportedQwenGeometry; + if (types != .array or types.array.items.len != g.layers) return error.UnsupportedQwenGeometry; + for (types.array.items, 0..) |kind, i| { + if (kind != .string) return error.UnsupportedQwenGeometry; + if (std.mem.eql(u8, kind.string, "linear_attention")) g.linear[i] = true else if (!std.mem.eql(u8, kind.string, "full_attention")) return error.UnsupportedQwenGeometry; + } + return g; +} + +test "the 27B's widths: 48 value heads over 16 key heads, 24 query heads over 4" { + const text = + \\{"model_type":"qwen3_5","quantization":{"group_size":64,"bits":4,"mode":"affine"},"tie_word_embeddings":false, + \\ "text_config":{"model_type":"qwen3_5_text","attention_bias":false,"attn_output_gate":true,"hidden_act":"silu", + \\ "head_dim":256,"linear_key_head_dim":128,"linear_value_head_dim":128,"linear_conv_kernel_dim":4, + \\ "hidden_size":5120,"intermediate_size":17408,"num_hidden_layers":4,"vocab_size":248320,"linear_num_key_heads":16, + \\ "linear_num_value_heads":48,"num_attention_heads":24,"num_key_value_heads":4,"rms_norm_eps":1e-06, + \\ "max_position_embeddings":262144,"tie_word_embeddings":false, + \\ "layer_types":["linear_attention","linear_attention","linear_attention","full_attention"], + \\ "rope_parameters":{"rope_type":"default","rope_theta":10000000,"partial_rotary_factor":0.25}}} + ; + const g = try parse(std.testing.allocator, text); + try std.testing.expectEqual(@as(usize, 3), g.deltas()); + try std.testing.expectEqual(@as(usize, 10240), g.convWidth()); + try std.testing.expectEqual(@as(usize, 6144), g.attnWidth()); + try std.testing.expectEqual(@as(usize, 48 * 128 * 128 * 4), g.recBytes()); + try std.testing.expect(!g.tied); + const wide = try std.mem.replaceOwned(u8, std.testing.allocator, text, "\"head_dim\":256", "\"head_dim\":128"); + defer std.testing.allocator.free(wide); + try std.testing.expectError(error.UnsupportedQwenConfig, parse(std.testing.allocator, wide)); +} diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 4995a82b45..4b46a5a896 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -1,5 +1,6 @@ -//! Qwen3.5-2B's CUDA lane backend behind the native family interface (native/cuda.zig's registry): windows of up -//! to 16 rows (the lane core's copies), every stream's window in one forward, prompts in chunks, draws on the device. +//! The Qwen3.5 hybrids' CUDA lane backend (the 2B, Qwen3.8-27B) behind the native family interface (native/cuda.zig's +//! registry): windows of up to 16 rows (the lane core's copies), every stream's window in one forward, prompts in +//! chunks, draws on the device. const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -26,24 +27,44 @@ pub const Loaded = struct { free_streams: u32 = 0, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, + /// Called with the admitted stream count: that many caches allocated and touched now, since the driver commits + /// memory on first touch (seconds for a wide model's keys), and a request should not pay it. + reserve: *const fn (*anyopaque, u32) anyerror!void, }; pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.QwenHasNoCapturedKernels; const e = try E.Engine.init(gpa, io, ctx, dir); errdefer e.deinit(); + if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary const l = try gpa.create(Lanes); l.* = .{ .gpa = gpa, .e = e, .capacity = o.context }; return .{ .backend = l.backend(), - .facts = .{ .exact_width = E.window_rows, .hidden_rows = true, .batch_rows = E.max_verify, .max_streams = E.max_segs, .streams_exact = true }, - .rows = E.window_rows, - .stream_bytes = E.streamBytes(o.context), + .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, + .rows = @intCast(e.window_rows), + .stream_bytes = e.streamBytes(o.context), .ctx = l, .deinit = release, + .reserve = reserve, }; } +fn reserve(p: *anyopaque, streams: u32) anyerror!void { + const l: *Lanes = @ptrCast(@alignCast(p)); + try l.pool.ensureTotalCapacity(l.gpa, streams); + while (l.pool.items.len < streams) { + const cache = try l.gpa.create(E.Cache); + errdefer l.gpa.destroy(cache); + cache.* = try E.Cache.init(l.gpa, l.e, l.capacity); + errdefer cache.deinit(l.gpa); + // a row through it: kernels have to have referenced a buffer once before the driver's stall is behind it + try l.e.forward(&.{.{ .cache = cache, .rows = 1 }}, &.{0}, .last); + try cache.reset(l.e); + l.pool.appendAssumeCapacity(cache); + } +} + /// A request this engine refuses, in words; null: none of its own. pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { @@ -70,18 +91,40 @@ const Lanes = struct { e: *E.Engine, capacity: usize, caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, + /// Caches at the full window, kept between streams: allocated and touched once (`reserve`), reset on reuse. + pool: std.ArrayList(*E.Cache) = .empty, fn deinit(self: *Lanes) void { var it = self.caches.valueIterator(); - while (it.next()) |cache| { - cache.*.deinit(self.gpa); - self.gpa.destroy(cache.*); - } + while (it.next()) |cache| self.drop(cache.*); self.caches.deinit(self.gpa); + for (self.pool.items) |cache| self.drop(cache); + self.pool.deinit(self.gpa); self.e.deinit(); self.gpa.destroy(self); } + fn drop(self: *Lanes, cache: *E.Cache) void { + cache.deinit(self.gpa); + self.gpa.destroy(cache); + } + + /// A cache for a new stream: a pooled one reset, else one more (a window beyond the reserve). + fn take(self: *Lanes) !*E.Cache { + if (self.pool.pop()) |cache| { + try cache.reset(self.e); + return cache; + } + const cache = try self.gpa.create(E.Cache); + errdefer self.gpa.destroy(cache); + cache.* = try E.Cache.init(self.gpa, self.e, self.capacity); + return cache; + } + + fn give(self: *Lanes, cache: *E.Cache) void { + self.pool.append(self.gpa, cache) catch self.drop(cache); + } + fn backend(self: *Lanes) be.Backend { return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; } @@ -98,10 +141,8 @@ const Lanes = struct { const self = cast(ptr); const ids = stream.prompt(); if (ids.len == 0 or ids.len + stream.max_new > self.capacity or self.caches.contains(stream)) return error.InvalidQwenPrompt; - const cache = try self.gpa.create(E.Cache); - errdefer self.gpa.destroy(cache); - cache.* = try E.Cache.init(self.gpa, self.e, @min(self.capacity, ids.len + stream.max_new + 1)); - errdefer cache.deinit(self.gpa); + const cache = try self.take(); + errdefer self.give(cache); var at: usize = 0; while (at < ids.len) { var end: usize = @min(ids.len, at + E.max_rows); @@ -150,7 +191,7 @@ const Lanes = struct { if (windows.len != out.len) return error.InvalidQwenWindows; for (windows, out) |w, o| { const rows = w.rows(); - if (w.parents != null or w.held != 0 or rows > E.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; + if (w.parents != null or w.held != 0 or rows > self.e.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; if (w.masks.len != 0 and w.masks.len != rows * draws.mask_words) return error.InvalidQwenMasks; if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; const cache = try self.cacheOf(w.stream); @@ -159,12 +200,12 @@ const Lanes = struct { var at: usize = 0; while (at < windows.len) { var parts: [E.max_segs]E.Part = undefined; - var ids: [E.max_verify]u32 = undefined; - var picks: [E.max_verify]E.Draw = undefined; + var ids: [E.max_rows]u32 = undefined; + var picks: [E.max_rows]E.Draw = undefined; var n: usize = 0; var total: usize = 0; var masked = false; - while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= E.max_verify) : (n += 1) { + while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= self.e.verify_rows) : (n += 1) { const w = windows[at + n]; parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows() }; ids[total] = w.pending; @@ -203,9 +244,6 @@ const Lanes = struct { fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { const self = cast(ptr); - if (self.caches.fetchRemove(stream)) |entry| { - entry.value.deinit(self.gpa); - self.gpa.destroy(entry.value); - } + if (self.caches.fetchRemove(stream)) |entry| self.give(entry.value); } }; diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig new file mode 100644 index 0000000000..a105b876f4 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -0,0 +1,143 @@ +//! The checkpoint's weights onto the device: each affine-4 projection staged as stored and packed into qlinear's +//! tiles, the vectors as they are, all from the engine's one arena (cuda_engine.zig). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Geometry = geometry.Geometry; +const Engine = engine.Engine; +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const aligned = engine.aligned; +const head_dim = geometry.head_dim; + +/// The packed tiles' bytes of an n-by-k affine-4 projection (qlinear's words, scales and biases). +fn packedBytes(n: usize, k: usize) usize { + const l = Affine4.layout(n, k); + return aligned(l.words) + 2 * aligned(l.scales); +} + +/// The checkpoint's own bytes of an n-by-k affine-4 projection: words, then its scales and biases. +fn rawBytes(n: usize, k: usize) usize { + return n * k / 8 * 4 + 2 * (n * k / 64 * 2); +} + +/// What `load` takes from the arena: every packed projection and vector, each allocation aligned. +pub fn total(g: Geometry) usize { + var t: usize = 0; + t += aligned(g.vocab * g.hidden / 8 * 4) + 2 * aligned(g.vocab * g.hidden / 64 * 2); // the embedding as stored + t += packedBytes(g.vocab, g.hidden); // the head, tied or not + t += aligned(g.hidden * 2); + for (g.linear[0..g.layers]) |lin| { + t += 2 * aligned(g.hidden * 2) + 2 * packedBytes(g.intermediate, g.hidden) + packedBytes(g.hidden, g.intermediate); + if (lin) { + t += packedBytes(g.convWidth(), g.hidden) + packedBytes(g.valueWidth(), g.hidden) + 2 * packedBytes(g.v_heads, g.hidden) + packedBytes(g.hidden, g.valueWidth()); + t += aligned(g.convWidth() * geometry.taps * 2) + aligned(geometry.dk * 2) + aligned(g.v_heads * 4) + aligned(g.v_heads * 2); + } else { + t += packedBytes(2 * g.attnWidth(), g.hidden) + 2 * packedBytes(g.kv_heads * head_dim, g.hidden) + packedBytes(g.hidden, g.attnWidth()); + t += 2 * aligned(head_dim * 2); + } + } + return t; +} + +/// The largest projection's bytes as stored: the staging buffer the load packs each one from. +pub fn stageBytes(g: Geometry) usize { + var most: usize = 0; + for ([_][2]usize{ .{ g.vocab, g.hidden }, .{ g.intermediate, g.hidden }, .{ g.hidden, g.intermediate }, .{ g.convWidth(), g.hidden }, .{ g.valueWidth(), g.hidden }, .{ 2 * g.attnWidth(), g.hidden } }) |nk| most = @max(most, rawBytes(nk[0], nk[1])); + return most; +} + +fn upload(e: *Engine, bytes: []const u8) !u64 { + const p = try e.alloc(bytes.len); + try e.arena.upload(p - e.arena.ptr, bytes); + return p; +} + +fn put(e: *Engine, t: core.checkpoint.Tensor) !u64 { + return upload(e, t.bytes); +} + +/// A vector the kernels read as f32, stored as f32 (the 2B) or bf16 (the 27B's A_log). +fn f32vec(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, n: usize) !u64 { + var buf: [160]u8 = undefined; + const t = try ck.get(try std.fmt.bufPrint(&buf, fmt, args)); + if (t.is(.f32, &.{n})) return put(e, t); + if (!t.is(.bf16, &.{n})) return error.UnexpectedTensor; + const words = try e.gpa.alloc(u32, n); + defer e.gpa.free(words); + for (words, std.mem.bytesAsSlice(u16, t.bytes)) |*w, h| w.* = @as(u32, h) << 16; + return upload(e, std.mem.sliceAsBytes(words)); +} + +fn tensor(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + var buf: [160]u8 = undefined; + return put(e, try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); +} + +/// An MLX affine-4 projection packed into qlinear's tiles from the staging buffer, which holds it as stored. +fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { + var buf: [160]u8 = undefined; + var raw: [3]u64 = undefined; + var at: usize = 0; + const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (parts, &raw) |p, *r| { + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); + try e.stage.upload(at, t.bytes); + r.* = e.stage.ptr + at; + at += t.bytes.len; + } + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try e.alloc(l.words), .s = try e.alloc(l.scales), .b = try e.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + try e.affine.pack(e.stream, raw[0], raw[1], raw[2], q); + try e.stream.synchronize(); // the next projection's staging overwrites this one's + return q; +} + +pub fn load(e: *Engine, ck: *core.Checkpoint) !void { + const g = e.g; + const m = "language_model.model."; + e.embed_w = try tensor(e, ck, "{s}embed_tokens.weight", .{m}, .u32, &.{ g.vocab, g.hidden / 8 }); + e.embed_s = try tensor(e, ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 }); + e.embed_b = try tensor(e, ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 }); + e.head = if (g.tied) try linear(e, ck, m, "embed_tokens", g.vocab, g.hidden) else try linear(e, ck, "language_model.", "lm_head", g.vocab, g.hidden); + e.final_norm = try tensor(e, ck, "{s}norm.weight", .{m}, .bf16, &.{g.hidden}); + const conv = g.convWidth(); + const vw = g.valueWidth(); + for (e.layers, 0..) |*l, i| { + var pbuf: [96]u8 = undefined; + const p = try std.fmt.bufPrint(&pbuf, "{s}layers.{d}.", .{ m, i }); + l.* = .{ + .input_norm = try tensor(e, ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{g.hidden}), + .post_norm = try tensor(e, ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{g.hidden}), + .gate = try linear(e, ck, p, "mlp.gate_proj", g.intermediate, g.hidden), + .up = try linear(e, ck, p, "mlp.up_proj", g.intermediate, g.hidden), + .down = try linear(e, ck, p, "mlp.down_proj", g.hidden, g.intermediate), + .mixer = undefined, + }; + if (g.linear[i]) { + l.mixer = .{ .delta = .{ + .qkv = try linear(e, ck, p, "linear_attn.in_proj_qkv", conv, g.hidden), + .z = try linear(e, ck, p, "linear_attn.in_proj_z", vw, g.hidden), + .a = try linear(e, ck, p, "linear_attn.in_proj_a", g.v_heads, g.hidden), + .b = try linear(e, ck, p, "linear_attn.in_proj_b", g.v_heads, g.hidden), + .out = try linear(e, ck, p, "linear_attn.out_proj", g.hidden, vw), + .conv = try tensor(e, ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ conv, geometry.taps, 1 }), + .norm = try tensor(e, ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{geometry.dk}), + .a_log = try f32vec(e, ck, "{s}linear_attn.A_log", .{p}, g.v_heads), + .dt_bias = try tensor(e, ck, "{s}linear_attn.dt_bias", .{p}, .bf16, &.{g.v_heads}), + } }; + } else { + l.mixer = .{ .attention = .{ + .q = try linear(e, ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden), + .k = try linear(e, ck, p, "self_attn.k_proj", g.kv_heads * head_dim, g.hidden), + .v = try linear(e, ck, p, "self_attn.v_proj", g.kv_heads * head_dim, g.hidden), + .o = try linear(e, ck, p, "self_attn.o_proj", g.hidden, g.attnWidth()), + .q_norm = try tensor(e, ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{head_dim}), + .k_norm = try tensor(e, ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{head_dim}), + } }; + } + } +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index b90c801e0d..cf64c4f8b9 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -349,6 +349,11 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; return null; }; + // a family that pre-allocates its streams' buffers does so now that their count is known + if (@hasField(L, "reserve")) loaded.reserve(loaded.ctx, streams) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot reserve {d} streams' caches at a {d}-token window ({s}): serve fewer with --parallel, or lower --context", .{ streams, window, @errorName(e) }); + return null; + }; const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 5f0dc8f906..78414fdb32 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,6 +1,7 @@ -//! `tf-qwen35-cuda `: Qwen3.5-2B's CUDA forward over comma-separated token ids, -//! one token a forward (step: every position's logits) or the ids as one prompt (prompt: the last position's), as -//! little-endian bf16 rows into for a reference to compare. Also prints the step mode's decode speed. +//! `tf-qwen35-cuda `: a Qwen3.5 hybrid's CUDA forward over comma-separated +//! token ids: one token a forward (step: every position's logits), windows of growing widths (window: the same rows, +//! which must equal step's) or the ids as one prompt (prompt: the last position's), as little-endian bf16 rows into +//! for a reference to compare. Step mode also prints the decode speed. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -26,19 +27,23 @@ pub fn main(init: std.process.Init) !u8 { defer cache.deinit(init.gpa); var out: std.ArrayList(u8) = .empty; if (std.mem.eql(u8, args[2], "step")) { - const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); - for (ids.items) |id| { + var first: i64 = 0; + var rest: i64 = 0; + for (ids.items, 0..) |id, i| { + const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); try e.forward(&.{.{ .cache = &cache, .rows = 1 }}, &.{id}, .logits); + const dt: i64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); + if (i == 0) first = dt else rest += dt; @memcpy(cache.logits, e.rowLogits(0)); try out.appendSlice(a, std.mem.sliceAsBytes(cache.logits)); } - const ns: u64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); - std.debug.print("{d} steps in {d:.1} ms: {d:.1} tok/s\n", .{ ids.items.len, @as(f64, @floatFromInt(ns)) / 1e6, @as(f64, @floatFromInt(ids.items.len)) * 1e9 / @as(f64, @floatFromInt(ns)) }); + const n: f64 = @floatFromInt(ids.items.len - 1); + std.debug.print("{d} steps: the first {d:.1} ms, then {d:.1} ms a step ({d:.1} tok/s)\n", .{ ids.items.len, @as(f64, @floatFromInt(first)) / 1e6, @as(f64, @floatFromInt(rest)) / 1e6 / n, n * 1e9 / @as(f64, @floatFromInt(rest)) }); } else if (std.mem.eql(u8, args[2], "window")) { // windows of 1, 2, 3, ... rows: every row's logits, which must equal step's bit for bit var at: usize = 0; var w: usize = 1; - while (at < ids.items.len) : (w = w % E.window_rows + 1) { + while (at < ids.items.len) : (w = w % e.window_rows + 1) { const n = @min(w, ids.items.len - at); try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .logits); for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); From 6b85f22799d796fc19982657c1e68834f63ebff5 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:36:52 +0000 Subject: [PATCH 08/35] qwen3_5: replay each decode shape as a CUDA graph A forward's launches (the embedding to the kept rows' logits) are captured the first time a shape (rows, windows, recording) is seen and replayed after; the attention chunk count moves into the per-forward table so a graph stays valid as caches grow. Prompt chunks stay uncaptured: their GEMMs dwarf the launches, and each length would be a graph of its own. A new graph's first run pays the Windows driver's stall (about a second on a wide model), so `reserve` also captures the shared rounds' shapes, a row of each of 1 to n streams, at load. Bit for bit the same logits; on an RTX 4090 the 2B steps in 2.8 ms instead of 5.3 (served: 324 tok/s, 1,560 across eight streams) and the 27B in 23.8 ms instead of 26.9 (served: 40 tok/s, 243 across eight). --- zig/kernels/cuda/qwen35.cu | 15 ++-- zig/src/families/qwen3_5/cuda_engine.zig | 110 +++++++++++++++++------ zig/src/families/qwen3_5/cuda_native.zig | 11 +++ 3 files changed, 103 insertions(+), 33 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 715f83f64e..15d77f706d 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -322,15 +322,18 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo } } -// Causal attention over a segment's cache, ck keys a block: grid (chunks, kvh, rows), block qh / kvh warps, a warp a -// query head of the kv head's group. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +// Causal attention over a segment's cache, P[1] keys a block over P[0] chunks (read from the device, so a captured +// forward replays as caches grow): grid (chunks at most, kvh, rows), block qh / kvh warps, a warp a query head of +// the kv head's group. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, int nch, int ck, float scale, + const int* __restrict__ rowseg, int layer, const int* __restrict__ P, float scale, float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { constexpr int DPL = HD / 32; + const int nch = P[0], ck = P[1]; const int QH = geo.qh, KVH = geo.kvh; const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; const int c = blockIdx.x, h = blockIdx.y; + if (c >= nch) return; const size_t m = blockIdx.z; const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; @@ -377,10 +380,12 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; } -// The chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. +// The P[0] chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, const float* __restrict__ PM, const float* __restrict__ PL, - const float* __restrict__ PO, int nch, __nv_bfloat16* __restrict__ OUT) { + const float* __restrict__ PO, const int* __restrict__ P, + __nv_bfloat16* __restrict__ OUT) { constexpr int DPL = HD / 32; + const int nch = P[0]; const int QH = geo.qh; const int lane = threadIdx.x, qh = blockIdx.x; const size_t m = blockIdx.y; diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 316814eaf6..05e69afa11 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -1,6 +1,7 @@ //! Qwen3.5's hybrid text models on CUDA (the 2B, Qwen3.8-27B): the affine-4 projections through cuda.qlinear, the //! glue in qwen35.cu, one forward for any mix of streams' windows. Each stream owns its recurrences, convolution tails //! and KV cache (Cache); the widths come from config.json (cuda_geometry.zig), the weights from cuda_weights.zig. +//! A forward's launches are captured into a CUDA graph the first time its shape is seen and replayed after. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -172,6 +173,8 @@ pub const Engine = struct { host_logits: cuda.HostBuffer, /// The kept rows' device draws, when forward was given them. tokens: [most_verify]u32 = undefined, + /// The captured forwards by shape (`Shape.key`): the launches from the embedding to the head. + graphs: std.AutoHashMapUnmanaged(u32, cuda.graph.Exec) = .empty, const ids_at = 0; const segs_at = max_rows * 4; @@ -179,7 +182,8 @@ pub const Engine = struct { const rules_at = rowseg_at + max_rows * 4; const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; - const tables_len = tokens_at + most_verify * 4; + const params_at = tokens_at + most_verify * 4; // the attention chunks and their width + const tables_len = params_at + 2 * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -195,6 +199,7 @@ pub const Engine = struct { e.verify_rows = std.math.clamp(snapshot_budget / (g.deltas() * (g.recBytes() + g.convBytes())), 1, most_verify); e.window_rows = @min(most_window, e.verify_rows); e.arena_at = 0; + e.graphs = .empty; e.stream = try cuda.Stream.init(d, true); errdefer e.stream.deinit(); e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); @@ -240,6 +245,9 @@ pub const Engine = struct { } pub fn deinit(e: *Engine) void { + var it = e.graphs.valueIterator(); + while (it.next()) |x| x.deinit(); + e.graphs.deinit(e.gpa); e.host_logits.free(); e.staged.free(); e.tables.free(); @@ -363,15 +371,82 @@ pub const Engine = struct { m.* = @intCast(p.position -% 1); // sample.cu keys row r at meta + r + 1 } } + // attention chunks: at most max_chunks a row, each a multiple of 128 keys + const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); + const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..8]); + params[0] = @intCast((widest + ck - 1) / ck); + params[1] = @intCast(ck); try e.tables.uploadAsync(0, st, e.stream.handle); + const shape: Shape = .{ .total = @intCast(total), .nseg = @intCast(parts.len), .record = record, .prompt = prompt }; + if (prompt) { + // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length would be a graph of its own, + // paying the driver's stall on a new graph's first run (about a second on a wide model) + try e.body(shape); + } else { + const gop = try e.graphs.getOrPut(e.gpa, shape.key()); + if (!gop.found_existing) { + errdefer _ = e.graphs.remove(shape.key()); + // captured, nothing runs: the replay below is this shape's first run + try cuda.graph.beginCapture(e.stream, .thread_local); + e.body(shape) catch |err| { + var broken = cuda.graph.endCapture(e.stream) catch return err; + broken.deinit(); + return err; + }; + var captured = try cuda.graph.endCapture(e.stream); + defer captured.deinit(); + gop.value_ptr.* = try captured.instantiate(); + } + try gop.value_ptr.launchOn(e.stream); + } + for (parts) |p| p.cache.len += p.rows; + switch (keep) { + .draws => |pk| { + const toks = e.tables.ptr + tokens_at; + for (pk, 0..) |p, i| { + const row = e.logits + i * g.vocab * 2; + if (p.rule != null) { + try e.draws.draw(e.stream, row, g.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); + } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, @intCast(g.vocab)), toks + i * 4 }); + } + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); + }, + .logits, .last => { + const out = e.host_logits.bytes[0 .. heads * g.vocab * 2]; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); + }, + } + } + + /// A forward's shape, which fixes its launches: its tables and the caches' lengths are read on the device. + const Shape = struct { + total: u8, + nseg: u8, + record: bool, + prompt: bool, + + fn key(s: Shape) u32 { + return @as(u32, s.total) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17; + } + }; + + /// The launches of a forward of `shape`, from the embedding to the kept rows' logits (captured once a shape). + fn body(e: *Engine, shape: Shape) !void { + const g = e.g; + const total: usize = shape.total; + const nseg: usize = shape.nseg; + const record = shape.record; + const prompt = shape.prompt; + const heads: usize = if (record) total else 1; const t_ids = e.tables.ptr + ids_at; const t_segs = e.tables.ptr + segs_at; const t_rows = e.tables.ptr + rowseg_at; - const nseg = parts.len; + const t_params = e.tables.ptr + params_at; const geo = e.geo; - // attention chunks: at most max_chunks a row, each a multiple of 128 keys - const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); - const nch = (widest + ck - 1) / ck; const group = g.q_heads / g.kv_heads; try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)) }); @@ -403,8 +478,8 @@ pub const Engine = struct { try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps }); try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps }); try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); - try e.run(e.fns.attn_partial, .{ nch, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, @as(c_int, @intCast(nch)), @as(c_int, @intCast(ck)), @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); - try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, @as(c_int, @intCast(nch)), e.y }); + try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_params, e.y }); try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); try e.proj(m.o, e.mix, e.r, total, prompt); ai += 1; @@ -420,27 +495,6 @@ pub const Engine = struct { try e.norm(true, e.final_norm, total); // the head over the kept rows try e.proj(e.head, e.x + (total - heads) * g.hidden * 2, e.logits, heads, false); - for (parts) |p| p.cache.len += p.rows; - switch (keep) { - .draws => |pk| { - const toks = e.tables.ptr + tokens_at; - for (pk, 0..) |p, i| { - const row = e.logits + i * g.vocab * 2; - if (p.rule != null) { - try e.draws.draw(e.stream, row, g.vocab, e.tables.ptr + rules_at + i * @sizeOf(cuda.sample.Rule), e.tables.ptr + meta_at + i * 4, 0, toks + i * 4, 1, null, null); - } else try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ row, @as(c_int, @intCast(g.vocab)), toks + i * 4 }); - } - try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(e.host_logits.bytes.ptr, toks, total * 4, e.stream.handle), "cuMemcpyDtoHAsync"); - try e.stream.synchronize(); - @memcpy(e.tokens[0..total], e.host_logits.slice(u32)[0..total]); - }, - .logits, .last => { - const out = e.host_logits.bytes[0 .. heads * g.vocab * 2]; - try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, e.logits, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); - try e.stream.synchronize(); - if (keep == .last) @memcpy(parts[0].cache.logits, e.rowLogits(0)); - }, - } } /// Row `i`'s logits from the newest forward that kept logits. diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 4b46a5a896..7fdea31ae7 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -63,6 +63,17 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { try cache.reset(l.e); l.pool.appendAssumeCapacity(cache); } + // the shared rounds' shapes, a row of each of 1 to `streams` streams, captured now: a new graph's first run + // pays the driver's stall too (about a second on a wide model) + var parts: [E.max_segs]E.Part = undefined; + var ids: [E.max_segs]u32 = @splat(0); + var picks: [E.max_segs]E.Draw = undefined; + for (0..@min(streams, E.max_segs)) |n| { + parts[n] = .{ .cache = l.pool.items[n], .rows = 1 }; + picks[n] = .{ .rule = null, .position = 1 }; + try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); + for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); + } } /// A request this engine refuses, in words; null: none of its own. From b19e0d73a27c25d3579b12373ca74a94cdacc592 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 02:00:27 +0000 Subject: [PATCH 09/35] qwen3_5: DFlash2 drafts for Qwen3.8-27B on CUDA --drafter loads z-lab's DFlash2 draft model behind the lane core's Drafted wrapper: the target writes each row's residual after the drafter's five tapped layers into a ring a stream, the drafter folds those rows into a context of keys and values, and after every verify it drafts a block of seven tokens past the pending one (the target's own embedding and head at its ends, a codebook selector tracing one path through each row's best sixteen). Its projections are quantized to affine-4 on the host at load and run through the target's kernels; its glue is qwen35_draft.cu. The Windows driver stalls a buffer's first kernel and a new graph's first run for seconds, so the drafter pools its contexts and a captured forward now serves every verify shape: its launches cover a bucket of rows and every window, and the kernels read the true counts from the device. On an RTX 4090 the 27B serves 52 to 58 tok/s greedy with drafts against 40 plain, about 100 tok/s on long replies, the same tokens either way; plain decoding keeps its 24 ms step. The 2B is unchanged. --- zig/build/cuda.zig | 3 +- zig/kernels/cuda/qwen35.cu | 63 ++- zig/kernels/cuda/qwen35_draft.cu | 156 ++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/cuda.zig | 2 + zig/src/families/qwen3_5/cuda_draft.zig | 482 ++++++++++++++++++ .../families/qwen3_5/cuda_draft_weights.zig | 224 ++++++++ zig/src/families/qwen3_5/cuda_engine.zig | 94 +++- zig/src/families/qwen3_5/cuda_native.zig | 54 +- zig/src/native/cuda.zig | 27 +- zig/src/native/cuda_drafted.zig | 57 +++ 11 files changed, 1112 insertions(+), 51 deletions(-) create mode 100644 zig/kernels/cuda/qwen35_draft.cu create mode 100644 zig/src/families/qwen3_5/cuda_draft.zig create mode 100644 zig/src/families/qwen3_5/cuda_draft_weights.zig create mode 100644 zig/src/native/cuda_drafted.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 8632ea6c7f..78fdc3e649 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -32,7 +32,8 @@ const kernels = [_]Kernel{ .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) - .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: Qwen3.5-2B's glue, the Metal kernels' rounding points + .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: the Qwen3.5 hybrids' glue, the Metal kernels' rounding points + .{ .name = "qwen35_draft", .flags = &.{"-O3"} }, // ours: the DFlash2 drafter's glue (draft_ops.metal's rounding) .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 15d77f706d..de830da26f 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -25,7 +25,8 @@ struct Seg { unsigned long long rec; // float recurrences, one rec() block a linear layer unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time - int cap, len, row0, rows; + unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each + int cap, len, row0, rows, tap_rows, taps; }; __device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } @@ -41,10 +42,16 @@ __device__ __forceinline__ float warp_sum(float x) { using namespace tf_q35; +// The per-forward counts a captured forward reads from the device, so one graph serves every shape: P[0] attention +// chunks, P[1] their width, P[2] rows, P[3] windows. A null P means the launch's grid is exact. +__device__ __forceinline__ bool past_rows(const int* P, int m) { return P != nullptr && m >= P[2]; } +__device__ __forceinline__ bool past_windows(const int* P, int s) { return P != nullptr && s >= P[3]; } + // MLX affine 4-bit token rows: out = bf16(s * q + b), D wide. extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B, - __nv_bfloat16* __restrict__ out, int D) { + __nv_bfloat16* __restrict__ out, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; const size_t id = ids[blockIdx.x]; for (int b = threadIdx.x; b < D / 2; b += blockDim.x) { const uint8_t v = W[id * (D / 2) + b]; @@ -58,9 +65,11 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __rest // h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), D wide (a multiple of 256, at most 8192). One 256-thread // block a row; each thread's elements in index order, warps summed in order. extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, - const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D) { + const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D, + const int* __restrict__ P) { constexpr int MAXE = 32; __shared__ float red[8]; + if (past_rows(P, blockIdx.x)) return; const size_t row = blockIdx.x; const int t = threadIdx.x; const int E = D / 256; @@ -94,7 +103,8 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __r } // The 64-group sums of bf16 rows D wide (the decode projections' input sums). Grid (rows), a thread a group. -extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, float* __restrict__ XS, int D) { +extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, float* __restrict__ XS, int D, const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; const size_t row = blockIdx.x; for (int g = threadIdx.x; g < D / 64; g += blockDim.x) { const __nv_bfloat16* x = X + row * D + g * 64; @@ -111,10 +121,12 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ K, - __nv_bfloat16* __restrict__ V, float* __restrict__ G, __nv_bfloat16* __restrict__ BETA) { + __nv_bfloat16* __restrict__ V, float* __restrict__ G, __nv_bfloat16* __restrict__ BETA, + const int* __restrict__ P) { const int lane = threadIdx.x; const int head = blockIdx.y; const int w = blockIdx.z; + if (past_rows(P, w)) return; const int NK = geo.nk, NV = geo.nv, C = geo.conv(); const Seg sg = segs[rowseg[w]]; const int r = w - sg.row0; @@ -162,10 +174,10 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ // Each segment's convolution tail after its rows: the last TAPS - 1 of [tail ++ rows]. A thread a (segment, channel). // SNAP, if given, also takes the tail after each row (forward row m at m * conv_rows()), read before the update. extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, - __nv_bfloat16* __restrict__ SNAP) { + __nv_bfloat16* __restrict__ SNAP, const int* __restrict__ P) { const int C = geo.conv(); const int c = blockIdx.x * blockDim.x + threadIdx.x; - if (c >= C) return; + if (c >= C || past_windows(P, blockIdx.y)) return; const Seg sg = segs[blockIdx.y]; __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * geo.conv_rows(); if (SNAP != nullptr) @@ -188,8 +200,9 @@ extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __re extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, int layer, - __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP) { + __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP, const int* __restrict__ P) { const int NK = geo.nk, NV = geo.nv; + if (past_windows(P, blockIdx.z / NV)) return; const int hv = blockIdx.z % NV; const int hk = hv / (NV / NK); const Seg sg = segs[blockIdx.z / NV]; @@ -240,8 +253,10 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, cons // out = bf16(silu(z) * bf16(w * rmsnorm(y))) per linear value head; rows nv * DK wide. Grid (1, nv, rows), a warp each. extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ Z, - const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT) { + const __nv_bfloat16* __restrict__ NW, float eps, __nv_bfloat16* __restrict__ OUT, + const int* __restrict__ P) { const int lane = threadIdx.x, hv = blockIdx.y; + if (past_rows(P, blockIdx.z)) return; const size_t m = blockIdx.z; const size_t W = (size_t)geo.nv * DK; constexpr int PER = DK / 32; @@ -265,9 +280,11 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const // Per-head RMSNorm of attention rows: X's head h of row m at m * stride + h * hstride, out packed (rows, heads, HD). // Grid (rows * heads), 64 threads of 4 values. extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, - __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps) { + __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps, + const int* __restrict__ P) { __shared__ float sums[2]; const int head = blockIdx.x % heads; + if (past_rows(P, blockIdx.x / heads)) return; const size_t m = blockIdx.x / heads; const int t = threadIdx.x; const __nv_bfloat16* x = X + m * stride + head * hstride; @@ -303,9 +320,10 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i // Grid (rows), HD threads. extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer) { + const int* __restrict__ rowseg, int layer, const int* __restrict__ P) { const int QH = geo.qh, KVH = geo.kvh; const int d = threadIdx.x; + if (past_rows(P, blockIdx.x)) return; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; const int pos = sg.len + (int)m - sg.row0; @@ -333,7 +351,7 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r const int QH = geo.qh, KVH = geo.kvh; const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; const int c = blockIdx.x, h = blockIdx.y; - if (c >= nch) return; + if (c >= nch || past_rows(P, blockIdx.z)) return; const size_t m = blockIdx.z; const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; @@ -388,6 +406,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, cons const int nch = P[0]; const int QH = geo.qh; const int lane = threadIdx.x, qh = blockIdx.x; + if (past_rows(P, blockIdx.y)) return; const size_t m = blockIdx.y; const size_t base = (m * QH + qh) * nch; float mx = -INFINITY; @@ -409,8 +428,9 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, cons // out = bf16(y * bf16(sigmoid(gate))), the gate the second half of each query head's q_proj rows; rows qh * HD wide. extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __restrict__ Y, const __nv_bfloat16* __restrict__ QRAW, - __nv_bfloat16* __restrict__ OUT, int rows) { + __nv_bfloat16* __restrict__ OUT, int rows, const int* __restrict__ P) { const size_t W = (size_t)geo.qh * HD; + if (P != nullptr) rows = P[2]; const size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; if (i >= (size_t)rows * W) return; const size_t r = i / W; @@ -420,9 +440,24 @@ extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __rest OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); } +// A tapped layer's residual stream, bf16(h + r) as the next norm rounds it, into each row's segment's tap ring at +// the row's position: plane `tap` of the row's taps. Grid (rows), 256 threads. +extern "C" __global__ void __launch_bounds__(256) tf_q35_tap(const __nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, + const Seg* __restrict__ segs, const int* __restrict__ rowseg, int tap, int D, + const int* __restrict__ P) { + if (past_rows(P, blockIdx.x)) return; + const size_t m = blockIdx.x; + const Seg sg = segs[rowseg[m]]; + if (sg.tap == 0) return; + const long long pos = sg.len + (long long)m - sg.row0; + __nv_bfloat16* out = reinterpret_cast<__nv_bfloat16*>(sg.tap) + ((size_t)(pos % sg.tap_rows) * sg.taps + tap) * D; + for (int i = threadIdx.x; i < D; i += blockDim.x) out[i] = __float2bfloat16_rn(bf(H[m * D + i]) + bf(R[m * D + i])); +} + // out = bf16(silu(gate) * up) over `total` values. extern "C" __global__ void tf_q35_mlp_act(const __nv_bfloat16* __restrict__ GATE, const __nv_bfloat16* __restrict__ UP, - __nv_bfloat16* __restrict__ OUT, long long total) { + __nv_bfloat16* __restrict__ OUT, long long total, const int* __restrict__ P, int D) { + if (P != nullptr) total = (long long)P[2] * D; const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; if (i >= total) return; const float gf = bf(GATE[i]); diff --git a/zig/kernels/cuda/qwen35_draft.cu b/zig/kernels/cuda/qwen35_draft.cu new file mode 100644 index 0000000000..5d6286ea2b --- /dev/null +++ b/zig/kernels/cuda/qwen35_draft.cu @@ -0,0 +1,156 @@ +// The DFlash2 drafter's glue on CUDA (zig/src/families/qwen3_5/cuda_draft.zig): a block of rows attends over the +// target's tapped context, with two-tap dynamic convolutions around attention and MLP. Rounding points follow the +// Metal drafter's draft_ops.metal; the projections run through the target engine's affine-4 kernels. + +#include +#include + +namespace tf_q35d { + +constexpr int HD = 128; // the drafter's head width (q, k, v) + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q35d + +using namespace tf_q35d; + +// Tap rows (start + i) of a stream's ring of R rows, `width` bf16 each, into OUT rows 0..n. Grid (n), 256 threads. +extern "C" __global__ void tf_q35d_gather_taps(const __nv_bfloat16* __restrict__ RING, int R, long long start, int width, + __nv_bfloat16* __restrict__ OUT) { + const long long row = start + blockIdx.x; + const __nv_bfloat16* src = RING + (size_t)(row % R) * width; + __nv_bfloat16* dst = OUT + (size_t)blockIdx.x * width; + for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; +} + +// y[t, c] = sum over lags l of (BASE[branch, l, c] + DYN[t, (branch * 2 + l) * groups + c / group]) * x[t - l, c], +// each product and sum rounded to bf16 in turn; x[-1] is zero (the block's first row has no past); RES, if given, +// is added last. D wide rows, `group` channels a dynamic coefficient. Grid (rows), 256 threads. +extern "C" __global__ void tf_q35d_dyn_conv(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ DYN, + const __nv_bfloat16* __restrict__ BASE, const __nv_bfloat16* __restrict__ RES, + __nv_bfloat16* __restrict__ OUT, int D, int group, int branch) { + const int t = blockIdx.x; + const int groups = D / group; + for (int c = threadIdx.x; c < D; c += blockDim.x) { + const int g = c / group; + float total = 0.0f; +#pragma unroll + for (int l = 0; l < 2; l++) { + if (t - l < 0) continue; + const float x = bf(X[(size_t)(t - l) * D + c]); + total = rbf(total + rbf(bf(BASE[(branch * 2 + l) * D + c]) * x)); + total = rbf(total + rbf(bf(DYN[(size_t)t * (4 * groups) + (branch * 2 + l) * groups + g]) * x)); + } + if (RES != nullptr) total = rbf(total + bf(RES[(size_t)t * D + c])); + OUT[(size_t)t * D + c] = __float2bfloat16_rn(total); + } +} + +// Per-head RMSNorm over HD then rotate-half RoPE over all HD at position pos0 + row: X's head h of row m at +// m * stride + h * HD, OUT packed (rows, heads, HD). Grid (rows * heads), 32 threads of 4 values. +extern "C" __global__ void __launch_bounds__(32) tf_q35d_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, long long pos0, + float eps, float theta) { + const int head = blockIdx.x % heads; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + head * HD; + float v[4], acc = 0.0f; +#pragma unroll + for (int i = 0; i < 4; i++) { + v[i] = bf(x[4 * t + i]); + acc += v[i] * v[i]; + } + acc = warp_sum(acc); + const float inv = rsqrtf(acc / (float)HD + eps); + float n[4]; +#pragma unroll + for (int i = 0; i < 4; i++) n[i] = rbf(rbf(v[i] * inv) * bf(W[4 * t + i])); + // the pair of lane element d is d + HD / 2: this lane's partner values come from the lane HD / 8 = 16 away + const float pos = (float)(pos0 + (long long)m); + __nv_bfloat16* out = OUT + (m * heads + head) * HD; +#pragma unroll + for (int i = 0; i < 4; i++) { + const int d = 4 * t + i; + const float partner = __shfl_xor_sync(0xffffffffu, n[i], 16); + const int j = d < HD / 2 ? d : d - HD / 2; + const float angle = pos * powf(theta, -(float)j / (float)(HD / 2)); + float sn, cs; + sincosf(angle, &sn, &cs); + out[d] = __float2bfloat16_rn(d < HD / 2 ? n[i] * cs - partner * sn : n[i] * cs + partner * sn); + } +} + +// Rows of packed (kv_heads, HD) keys and values at positions pos0 + i into a layer's rings (slot pos % W). +extern "C" __global__ void tf_q35d_ring_put(const __nv_bfloat16* __restrict__ K, const __nv_bfloat16* __restrict__ V, + __nv_bfloat16* __restrict__ RK, __nv_bfloat16* __restrict__ RV, int width, long long pos0, int W) { + const int i = blockIdx.x; + const size_t slot = (size_t)((pos0 + i) % W); + for (int c = threadIdx.x; c < width; c += blockDim.x) { + RK[slot * width + c] = K[(size_t)i * width + c]; + RV[slot * width + c] = V[(size_t)i * width + c]; + } +} + +// The block's attention: query row m, head h sees the context ring's positions [max(begin, p - window + 1), pos0) +// where p = pos0 + m, and every block row (the drafter is not causal within its block). GQA over kv_heads; scale +// 1 / sqrt(HD). Grid (rows, heads), a warp each; online softmax in fp32, bf16 out (rows, heads * HD). +extern "C" __global__ void __launch_bounds__(32) tf_q35d_attn(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KB, + const __nv_bfloat16* __restrict__ VB, const __nv_bfloat16* __restrict__ RK, + const __nv_bfloat16* __restrict__ RV, int rows, int heads, int kv_heads, int W, + long long begin, long long pos0, int window, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + const float scale = rsqrtf((float)HD); + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + const long long p = pos0 + m; + long long first = p - window + 1; + if (first < begin) first = begin; + for (long long pos = first; pos < pos0; pos++) { + const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + for (int r = 0; r < rows; r++) { + const __nv_bfloat16* k = KB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = VB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } +#pragma unroll + for (int i = 0; i < DPL; i++) OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(o[i] / l); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index c70e8b5ccf..ebc3722cf6 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -31,6 +31,7 @@ pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; +pub const qwen35_draft: []const u8 = if (available) &Blob("fatbin_qwen35_draft").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index ea5eaa465e..f21778c7bb 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -4,9 +4,11 @@ pub const config = @import("config.zig"); pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); +pub const draft = @import("cuda_draft.zig"); test { _ = config; _ = geometry; _ = @import("draw.zig"); + _ = @import("cuda_draft_weights.zig"); } diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig new file mode 100644 index 0000000000..90cce5c4cd --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -0,0 +1,482 @@ +//! The DFlash2 drafter on CUDA behind the lane core's Drafter contract (core/lanes/drafter.zig): it absorbs the +//! target's tapped rows into a context ring a stream, then drafts a block of 7 tokens after the pending one, the +//! target's own embedding and head at its ends and a codebook selector tracing one path through the top 16 a row. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const weights = @import("cuda_draft_weights.zig"); + +const E = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; +const dr = lanes.drafter; +const Stream = lanes.Stream; + +/// Rows an absorb runs at once (the target's prompt chunks are at most this). +const max_rows = engine.max_rows; +/// The logprob row's words (sample.cu's LP_WORDS) and its best-token capacity. +const lp_words = 3 + 2 * 20; +const lp_top = 20; + +pub const Config = struct { + hidden: usize, + intermediate: usize, + layers: usize, + vocab: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + block: usize, + mask: u32, + window: usize, + group: usize, + rank: usize, + topk: usize, + eps: f32, + theta: f32, + taps: [8]u32, + n_taps: usize, + + fn qWidth(c: Config) usize { + return c.heads * c.head_dim; + } + fn kvWidth(c: Config) usize { + return c.kv_heads * c.head_dim; + } + fn dynamic(c: Config) usize { + return 4 * (c.hidden / c.group); + } + /// The ring's rows: the window and a block more. + fn ringRows(c: Config) usize { + return c.window + 16; + } + fn ringBytes(c: Config) usize { + return c.layers * 2 * c.ringRows() * c.kvWidth() * 2; + } + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, target: engine.Geometry) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer doc.deinit(); + if (doc.value != .object) return error.BadDraftConfig; + const o = doc.value.object; + const d = (o.get("dflash_config") orelse return error.BadDraftConfig); + if (d != .object) return error.BadDraftConfig; + var c: Config = .{ + .hidden = try int(o, "hidden_size"), + .intermediate = try int(o, "intermediate_size"), + .layers = try int(o, "num_hidden_layers"), + .vocab = try int(o, "vocab_size"), + .heads = try int(o, "num_attention_heads"), + .kv_heads = try int(o, "num_key_value_heads"), + .head_dim = try int(o, "head_dim"), + .block = try int(d.object, "block_size"), + .mask = @intCast(try int(d.object, "mask_token_id")), + .window = try int(o, "sliding_window"), + .group = try int(d.object, "conv_group_size"), + .rank = try int(d.object, "selector_rank"), + .topk = try int(d.object, "selector_top_k"), + .eps = @floatCast(try num(o, "rms_norm_eps")), + .theta = @floatCast(try num((o.get("rope_parameters") orelse return error.BadDraftConfig).object, "rope_theta")), + .taps = undefined, + .n_taps = 0, + }; + const taps = d.object.get("target_layer_ids") orelse return error.BadDraftConfig; + if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > c.taps.len) return error.BadDraftConfig; + for (taps.array.items, 0..) |t, i| { + if (t != .integer or t.integer < 0 or t.integer >= target.layers) return error.BadDraftConfig; + c.taps[i] = @intCast(t.integer); + } + c.n_taps = taps.array.items.len; + // what the kernels fix: the head width, two-tap convolutions, a block of 8, the selector's top-k within a row + if (c.head_dim != 128 or try int(d.object, "conv_kernel_size") != 2 or c.block != 8 or c.topk > lp_top or c.mask >= c.vocab) return error.UnsupportedDraftConfig; + if (c.hidden != target.hidden or c.vocab != target.vocab or c.hidden % 256 != 0 or c.hidden % c.group != 0 or c.heads % c.kv_heads != 0 or c.layers > 8) return error.UnsupportedDraftConfig; + if (c.window < c.block or c.window > engine.tap_rows - max_rows) return error.UnsupportedDraftConfig; + return c; + } +}; + +fn int(o: std.json.ObjectMap, key: []const u8) !usize { + const v = o.get(key) orelse return error.BadDraftConfig; + if (v != .integer or v.integer < 0) return error.BadDraftConfig; + return @intCast(v.integer); +} + +fn num(o: std.json.ObjectMap, key: []const u8) !f64 { + return switch (o.get(key) orelse return error.BadDraftConfig) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadDraftConfig, + }; +} + +const Fns = struct { gather_taps: cuda.Function, dyn_conv: cuda.Function, norm_rope: cuda.Function, ring_put: cuda.Function, attn: cuda.Function }; + +/// A stream's context: its layers' key and value rings over positions [begin, end), and the drafts it holds. +const Ctx = struct { + ring: DeviceBuffer, + begin: u64 = 0, + end: u64 = 0, + held: [8]u32 = undefined, + held_n: usize = 0, +}; + +pub const Draft = struct { + gpa: std.mem.Allocator, + io: std.Io, + e: *E, + c: Config, + w: weights.Weights, + mod: cuda.Module, + fns: Fns, + scratch: DeviceBuffer, + at: usize = 0, + // absorb rows (max_rows) and block rows (block), all bf16 unless said + taps_in: u64, + th: u64, + thn: u64, + kabs: u64, + kn: u64, + vabs: u64, + xs: u64, // 64-group sums for the projections (f32) + h: u64, + h2: u64, + x: u64, + xin: u64, + dyn: u64, + qraw: u64, + q: u64, + kraw: u64, + kb: u64, + vb: u64, + att: u64, + o: u64, + gate: u64, + up: u64, + act: u64, + logits: u64, + words: u64, // u32 logprob words a row + picks: u64, // u32 zeros: the logprob rows' pick (unused) + hp: u64, + ids: u64, // u32 the block's ids + staged: cuda.HostBuffer, // pinned: the ids up, the words and projected rows down + streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, + /// Contexts between streams: allocated and run through once (`reserve`), since the driver stalls a buffer's + /// first kernel for seconds. + pool: std.ArrayList(*Ctx) = .empty, + step_ms: f64 = 4.0, + + const words_at = 64; + const hp_at = words_at + 16 * lp_words * 4; + const staged_len = hp_at + 16 * 256 * 2; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8) !*Draft { + const c = try Config.read(gpa, io, dir, e.g); + const d = try gpa.create(Draft); + errdefer gpa.destroy(d); + d.* = .{ .gpa = gpa, .io = io, .e = e, .c = c, .w = undefined, .mod = undefined, .fns = undefined, .scratch = undefined, .taps_in = 0, .th = 0, .thn = 0, .kabs = 0, .kn = 0, .vabs = 0, .xs = 0, .h = 0, .h2 = 0, .x = 0, .xin = 0, .dyn = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .o = 0, .gate = 0, .up = 0, .act = 0, .logits = 0, .words = 0, .picks = 0, .hp = 0, .ids = 0, .staged = undefined }; + d.mod = try cuda.Module.load(e.d, cuda.kernels.qwen35_draft); + errdefer d.mod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(d.fns, name) = try d.mod.function("tf_q35d_" ++ name); + d.w = try weights.load(gpa, io, e, dir, .{ .hidden = c.hidden, .intermediate = c.intermediate, .layers = c.layers, .vocab = c.vocab, .q_width = c.qWidth(), .kv_width = c.kvWidth(), .head_dim = c.head_dim, .dynamic = c.dynamic(), .rank = c.rank, .taps = c.n_taps }); + errdefer d.w.deinit(gpa); + d.scratch = try DeviceBuffer.alloc(e.d, d.scratchBytes()); + errdefer d.scratch.free(); + try d.scratch.fill8(0, null); + d.at = 0; + _ = try d.layout(true); + d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); + errdefer d.staged.free(); + try d.picksZero(); + return d; + } + + pub fn deinit(d: *Draft) void { + var it = d.streams.valueIterator(); + while (it.next()) |ctx| d.drop(ctx.*); + d.streams.deinit(d.gpa); + for (d.pool.items) |ctx| d.drop(ctx); + d.pool.deinit(d.gpa); + d.staged.free(); + d.scratch.free(); + d.w.deinit(d.gpa); + d.mod.unload(); + d.gpa.destroy(d); + } + + fn scratchBytes(d: *Draft) usize { + d.at = 0; + return d.layout(false) catch unreachable; + } + + /// The scratch rows, measured (nothing assigned) or assigned from the scratch buffer. + fn layout(d: *Draft, assign: bool) !usize { + const c = d.c; + const B = c.block; + const R = max_rows; + const sizes = .{ + .{ "taps_in", R * c.n_taps * c.hidden * 2 }, .{ "th", R * c.hidden * 2 }, .{ "thn", R * c.hidden * 2 }, .{ "kabs", R * c.kvWidth() * 2 }, + .{ "kn", R * c.kvWidth() * 2 }, .{ "vabs", R * c.kvWidth() * 2 }, .{ "xs", R * (c.n_taps * c.hidden / 64) * 4 }, .{ "h", B * c.hidden * 2 }, + .{ "h2", B * c.hidden * 2 }, .{ "x", B * c.hidden * 2 }, .{ "xin", B * c.hidden * 2 }, .{ "dyn", B * c.dynamic() * 2 }, + .{ "qraw", B * c.qWidth() * 2 }, .{ "q", B * c.qWidth() * 2 }, .{ "kraw", B * c.kvWidth() * 2 }, .{ "kb", B * c.kvWidth() * 2 }, + .{ "vb", B * c.kvWidth() * 2 }, .{ "att", B * c.qWidth() * 2 }, .{ "o", B * c.hidden * 2 }, .{ "gate", B * c.intermediate * 2 }, + .{ "up", B * c.intermediate * 2 }, .{ "act", B * c.intermediate * 2 }, .{ "logits", B * c.vocab * 2 }, .{ "words", B * lp_words * 4 }, + .{ "picks", B * 4 }, .{ "hp", B * c.rank * 2 }, .{ "ids", B * 4 }, + }; + var total: usize = 0; + inline for (sizes) |s| { + const at = engine.aligned(total); + total = at + s[1]; + if (assign) @field(d, s[0]) = d.scratch.ptr + at; + } + return total; + } + + fn picksZero(d: *Draft) !void { + const zero: [64]u8 = @splat(0); + try d.e.d.check(d.e.d.api.cuMemcpyHtoD_v2(d.picks, &zero, 64), "cuMemcpyHtoD"); + } + + fn drop(d: *Draft, ctx: *Ctx) void { + ctx.ring.free(); + d.gpa.destroy(ctx); + } + + /// A new context, its ring touched by the kernels once: a row absorbed from the first tap ring at hand, then + /// a block drafted over it, so no request pays the driver's stall. + fn create(d: *Draft) !*Ctx { + const ctx = try d.gpa.create(Ctx); + errdefer d.gpa.destroy(ctx); + ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.c.ringBytes()) }; + errdefer ctx.ring.free(); + try ctx.ring.fill8(0, null); + try d.absorb(ctx, d.taps_in, 0, 1); // any tap rows do: the drafter's own zeroed staging + try d.hold(ctx, 0, 1, d.c.block - 1); + ctx.* = .{ .ring = ctx.ring }; + return ctx; + } + + /// `n` contexts ready in the pool, for the streams the server admits. + pub fn reserve(d: *Draft, n: usize) !void { + try d.pool.ensureTotalCapacity(d.gpa, n); + while (d.pool.items.len < n) d.pool.appendAssumeCapacity(try d.create()); + } + + fn ringK(d: *Draft, ctx: *Ctx, layer: usize) u64 { + return ctx.ring.ptr + layer * 2 * d.c.ringRows() * d.c.kvWidth() * 2; + } + + fn ringV(d: *Draft, ctx: *Ctx, layer: usize) u64 { + return d.ringK(ctx, layer) + d.c.ringRows() * d.c.kvWidth() * 2; + } + + /// out (rows, n) = x (rows, k) through `l`, as the target's decode tiles or prompt GEMM, with the drafter's sums. + fn proj(d: *Draft, l: cuda.qlinear.Affine4, x: u64, out: u64, rows: usize) !void { + const e = d.e; + if (rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, d.xs, @as(c_int, @intCast(l.k)), @as(u64, 0) }); + try e.affine.decode(e.stream, x, d.xs, l, out, rows); + } + + fn norm(d: *Draft, h: u64, w: u64, out: u64, rows: usize) !void { + const e = d.e; + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ h, @as(u64, 0), w, out, d.c.eps, @as(c_int, @intCast(d.c.hidden)), @as(u64, 0) }); + } + + fn dynConv(d: *Draft, x: u64, dyn: u64, base: u64, res: u64, out: u64, rows: usize, branch: usize) !void { + try d.e.run(d.fns.dyn_conv, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ x, dyn, base, res, out, @as(c_int, @intCast(d.c.hidden)), @as(c_int, @intCast(d.c.group)), @as(c_int, @intCast(branch)) }); + } + + fn normRope(d: *Draft, x: u64, w: u64, out: u64, rows: usize, heads: usize, pos0: u64) !void { + try d.e.run(d.fns.norm_rope, .{ rows * heads, 1, 1 }, .{ 32, 1, 1 }, .{ x, w, out, @as(c_int, @intCast(heads)), @as(c_int, @intCast(heads * d.c.head_dim)), @as(c_longlong, @intCast(pos0)), d.c.eps, d.c.theta }); + } + + /// Target rows at positions [start, start + n) from the stream's tap ring into its context (replacing what was + /// held at or past `start`); rows older than the window are skipped. + fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize) !void { + const c = d.c; + var start = start_; + var n = n_; + if (start < ctx.end) ctx.end = start; + if (start > ctx.end) ctx.begin = start; // a gap (a prompt restored from a cache): the context starts here + if (n > c.window) { + start += n - c.window; + n = c.window; + ctx.begin = start; + } + ctx.end = start; + const e = d.e; + while (n > 0) { + const rows = @min(n, max_rows); + try e.run(d.fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); + try d.proj(d.w.fc, d.taps_in, d.th, rows); + try d.norm(d.th, d.w.hidden_norm, d.thn, rows); + for (d.w.layers, 0..) |l, i| { + try d.proj(l.k, d.thn, d.kabs, rows); + try d.normRope(d.kabs, l.k_norm, d.kn, rows, c.kv_heads, start); + try d.proj(l.v, d.thn, d.vabs, rows); + try e.run(d.fns.ring_put, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kn, d.vabs, d.ringK(ctx, i), d.ringV(ctx, i), @as(c_int, @intCast(c.kvWidth())), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.ringRows())) }); + } + start += rows; + n -= rows; + ctx.end = start; + } + if (ctx.end > c.window and ctx.begin < ctx.end - c.window) ctx.begin = ctx.end - c.window; + } + + /// The block [pending, mask x 7] at positions pos0.. over the stream's context: `depth` drafts held. + fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize) !void { + const c = d.c; + const e = d.e; + const B = c.block; + if (ctx.end != pos0) return error.DraftContextBehind; + const ids = std.mem.bytesAsSlice(u32, d.staged.bytes[0 .. B * 4]); + ids[0] = pending; + @memset(ids[1..B], c.mask); + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(d.ids, ids.ptr, B * 4, e.stream.handle), "cuMemcpyHtoDAsync"); + try e.run(e.fns.embed, .{ B, 1, 1 }, .{ 256, 1, 1 }, .{ d.ids, e.embed_w, e.embed_s, e.embed_b, d.h, @as(c_int, @intCast(c.hidden)), @as(u64, 0) }); + for (d.w.layers, 0..) |l, i| { + try d.norm(d.h, l.input_norm, d.x, B); + try d.proj(l.attn_conv, d.x, d.dyn, B); + try d.dynConv(d.x, d.dyn, l.attn_base, 0, d.xin, B, 0); + try d.proj(l.q, d.xin, d.qraw, B); + try d.proj(l.k, d.xin, d.kraw, B); + try d.proj(l.v, d.xin, d.vb, B); + try d.normRope(d.qraw, l.q_norm, d.q, B, c.heads, pos0); + try d.normRope(d.kraw, l.k_norm, d.kb, B, c.kv_heads, pos0); + try e.run(d.fns.attn, .{ B, c.heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.kb, d.vb, d.ringK(ctx, i), d.ringV(ctx, i), @as(c_int, @intCast(B)), @as(c_int, @intCast(c.heads)), @as(c_int, @intCast(c.kv_heads)), @as(c_int, @intCast(c.ringRows())), @as(c_longlong, @intCast(ctx.begin)), @as(c_longlong, @intCast(pos0)), @as(c_int, @intCast(c.window)), d.att }); + try d.proj(l.o, d.att, d.o, B); + try d.dynConv(d.o, d.dyn, l.attn_base, d.h, d.h2, B, 1); + try d.norm(d.h2, l.post_norm, d.x, B); + try d.proj(l.mlp_conv, d.x, d.dyn, B); + try d.dynConv(d.x, d.dyn, l.mlp_base, 0, d.xin, B, 0); + try d.proj(l.gate, d.xin, d.gate, B); + try d.proj(l.up, d.xin, d.up, B); + try e.run(e.fns.mlp_act, .{ (B * c.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ d.gate, d.up, d.act, @as(c_longlong, @intCast(B * c.intermediate)), @as(u64, 0), @as(c_int, @intCast(c.intermediate)) }); + try d.proj(l.down, d.act, d.o, B); + try d.dynConv(d.o, d.dyn, l.mlp_base, d.h2, d.h, B, 1); + } + // the mask rows' logits through the target's head, their best tokens and their selector projections + const rows = B - 1; + try d.norm(d.h + c.hidden * 2, d.w.final_norm, d.x, rows); + try d.proj(e.head, d.x, d.logits, rows); + try e.draws.logprobRows(e.stream, d.logits, c.vocab, d.picks, @intCast(c.topk), false, d.words, rows); + try d.proj(d.w.selector, d.x, d.hp, rows); + const st = d.staged.bytes; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + words_at, d.words, rows * lp_words * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + hp_at, d.hp, rows * c.rank * 2, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + d.select(ctx, pending, rows); + ctx.held_n = @min(depth, rows); + } + + /// The path through the rows' candidates: each row's best of unary logit + predecessor . projected . successor. + fn select(d: *Draft, ctx: *Ctx, pending: u32, rows: usize) void { + const c = d.c; + const words = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][0 .. rows * lp_words * 4]); + const hp = std.mem.bytesAsSlice(u16, d.staged.bytes[hp_at..][0 .. rows * c.rank * 2]); + var pred = pending; + for (0..rows) |r| { + const w = words[r * lp_words ..][0..lp_words]; + var best: f32 = -std.math.inf(f32); + var pick: u32 = w[2]; + for (0..c.topk) |j| { + const cand = w[2 + j]; + const unary: f32 = @bitCast(w[2 + lp_top + j]); + var edge: f32 = 0; + for (0..c.rank) |k| edge += bf16(d.w.predecessor[pred * c.rank + k]) * bf16(hp[r * c.rank + k]) * bf16(d.w.successor[cand * c.rank + k]); + const score = unary + edge; + if (score > best) { + best = score; + pick = cand; + } + } + ctx.held[r] = pick; + pred = pick; + } + } + + // the Drafter contract + + pub fn drafter(d: *Draft) dr.Drafter { + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(p: *anyopaque) *Draft { + return @ptrCast(@alignCast(p)); + } + + fn tapsFn(p: *anyopaque) []const u32 { + const d = cast(p); + return d.c.taps[0..d.c.n_taps]; + } + + fn factsFn(p: *anyopaque) dr.Facts { + const d = cast(p); + return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes() }; + } + + fn openFn(p: *anyopaque, s: *Stream) anyerror!void { + const d = cast(p); + if (d.streams.contains(s)) return error.StreamOpen; + const ctx = d.pool.pop() orelse try d.create(); + errdefer d.pool.append(d.gpa, ctx) catch d.drop(ctx); + ctx.* = .{ .ring = ctx.ring }; + try d.streams.put(d.gpa, s, ctx); + } + + fn absorbFn(p: *anyopaque, items: []const dr.Absorb) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("absorb", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; + // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) + try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows); + } + } + + fn holdFn(p: *anyopaque, items: []const dr.Hold) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("hold", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + try d.hold(ctx, it.pending, it.position - 1, it.depth); + } + } + + fn heldFn(p: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const d = cast(p); + for (streams, out) |s, o| { + const ctx = d.streams.get(s) orelse return error.UnknownStream; + if (o.len > ctx.held_n) return error.TooFewDrafts; + @memcpy(o, ctx.held[0..o.len]); + } + } + + fn releaseFn(p: *anyopaque, s: *Stream) void { + const d = cast(p); + if (d.streams.fetchRemove(s)) |kv| d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + } +}; + +/// Acceptance by depth until a stream has its own (the reference's chat text: measured, see the commit). +const prior = [_]f64{ 0.75, 0.65, 0.58, 0.52, 0.47, 0.43, 0.4 }; + +fn bf16(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// TF_TIMING=1: each call's wall time on stderr. +pub fn timing(what: []const u8, io: std.Io, t0: std.Io.Timestamp) void { + if (std.c.getenv("TF_TIMING") == null) return; + const now = std.Io.Clock.awake.now(io); + const ns: f64 = @floatFromInt(t0.durationTo(now).toNanoseconds()); + const at: f64 = @floatFromInt(@mod(now.toNanoseconds(), 1_000_000_000_000)); + std.debug.print("[timing] {s} {d:.2} ms at {d:.1}\n", .{ what, ns / 1e6, at / 1e6 }); +} diff --git a/zig/src/families/qwen3_5/cuda_draft_weights.zig b/zig/src/families/qwen3_5/cuda_draft_weights.zig new file mode 100644 index 0000000000..2b0cf9f350 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_draft_weights.zig @@ -0,0 +1,224 @@ +//! The DFlash2 drafter's checkpoint onto the device: its bf16 projections quantized to affine-4 (group 64) on the +//! host and packed into qlinear's tiles through the target engine's staging buffer, its norms and base kernels as +//! they are, its codebooks kept on the host for the selector. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const engine = @import("cuda_engine.zig"); + +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const Engine = engine.Engine; +const aligned = engine.aligned; + +pub const Layer = struct { + input_norm: u64, + post_norm: u64, + q: Affine4, + k: Affine4, + v: Affine4, + o: Affine4, + q_norm: u64, + k_norm: u64, + gate: Affine4, + up: Affine4, + down: Affine4, + attn_conv: Affine4, // kernel_projection: hidden -> 4 * groups dynamic coefficients + attn_base: u64, // base_kernel (2 branches, 2 lags, hidden) bf16 + mlp_conv: Affine4, + mlp_base: u64, +}; + +pub const Weights = struct { + fc: Affine4, // the taps' concatenation -> hidden + hidden_norm: u64, + final_norm: u64, + selector: Affine4, // hidden_projection: hidden -> rank + predecessor: []u16, // codebooks (vocab, rank) bf16, on the host + successor: []u16, + layers: []Layer, + arena: DeviceBuffer, + arena_at: usize = 0, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + w.arena.free(); + gpa.free(w.predecessor); + gpa.free(w.successor); + gpa.free(w.layers); + } + + fn alloc(w: *Weights, bytes: usize) !u64 { + const at = aligned(w.arena_at); + if (at + bytes > w.arena.len) return error.ArenaExhausted; + w.arena_at = at + bytes; + return w.arena.ptr + at; + } +}; + +/// The drafter's widths the loader needs (cuda_draft.zig's Config). +pub const Shape = struct { hidden: usize, intermediate: usize, layers: usize, vocab: usize, q_width: usize, kv_width: usize, head_dim: usize, dynamic: usize, rank: usize, taps: usize }; + +fn packedBytes(n: usize, k: usize) usize { + const l = Affine4.layout(n, k); + return aligned(l.words) + 2 * aligned(l.scales); +} + +/// Device bytes of everything `load` packs. +pub fn total(s: Shape) usize { + var t = packedBytes(s.hidden, s.taps * s.hidden) + packedBytes(s.rank, s.hidden) + 2 * aligned(s.hidden * 2); + t += s.layers * (2 * aligned(s.hidden * 2) + packedBytes(s.q_width, s.hidden) + 2 * packedBytes(s.kv_width, s.hidden) + packedBytes(s.hidden, s.q_width) + + 2 * aligned(s.head_dim * 2) + 2 * packedBytes(s.intermediate, s.hidden) + packedBytes(s.hidden, s.intermediate) + 2 * packedBytes(s.dynamic, s.hidden) + 2 * aligned(4 * s.hidden * 2)); + return t; +} + +fn bf16(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +fn toBf16(x: f32) u16 { + // round to nearest even, as the device does + const u: u32 = @bitCast(x); + const rounding: u32 = 0x7fff + ((u >> 16) & 1); + return @intCast((u +% rounding) >> 16); +} + +/// MLX affine-4, group 64, over bf16 rows: a group's scale is bf16((max - min) / 15) and bias bf16(min); each weight's +/// nibble is round((w - bias) / scale) clamped to [0, 15], eight to a word, the first in the low bits. +fn quantize(gpa: std.mem.Allocator, w: []align(1) const u16, n: usize, k: usize) !struct { words: []u32, scales: []u16, biases: []u16 } { + const groups = k / 64; + const words = try gpa.alloc(u32, n * k / 8); + errdefer gpa.free(words); + const scales = try gpa.alloc(u16, n * groups); + errdefer gpa.free(scales); + const biases = try gpa.alloc(u16, n * groups); + errdefer gpa.free(biases); + for (0..n) |r| { + for (0..groups) |g| { + const row = w[r * k + g * 64 ..][0..64]; + var lo: f32 = std.math.inf(f32); + var hi: f32 = -std.math.inf(f32); + for (row) |h| { + const x = bf16(h); + lo = @min(lo, x); + hi = @max(hi, x); + } + const scale = toBf16((hi - lo) / 15.0); + const bias = toBf16(lo); + scales[r * groups + g] = scale; + biases[r * groups + g] = bias; + const s = bf16(scale); + const b = bf16(bias); + const inv = if (s == 0) 0.0 else 1.0 / s; + for (0..8) |j| { + var word: u32 = 0; + for (0..8) |i| { + const q = std.math.clamp(@round((bf16(row[j * 8 + i]) - b) * inv), 0.0, 15.0); + word |= @as(u32, @intFromFloat(q)) << @intCast(4 * i); + } + words[(r * k + g * 64) / 8 + j] = word; + } + } + } + return .{ .words = words, .scales = scales, .biases = biases }; +} + +/// A bf16 (n, k) tensor quantized and packed through the target's staging buffer into the drafter's arena. +fn linear(gpa: std.mem.Allocator, e: *Engine, w: *Weights, ck: *core.Checkpoint, name: []const u8, n: usize, k: usize) !Affine4 { + const t = try ck.expect(name, .bf16, &.{ n, k }); + const q = try quantize(gpa, std.mem.bytesAsSlice(u16, t.bytes), n, k); + defer { + gpa.free(q.words); + gpa.free(q.scales); + gpa.free(q.biases); + } + const parts = [_][]const u8{ std.mem.sliceAsBytes(q.words), std.mem.sliceAsBytes(q.scales), std.mem.sliceAsBytes(q.biases) }; + var raw: [3]u64 = undefined; + var at: usize = 0; + for (parts, &raw) |p, *r| { + if (at + p.len > e.stage.len) return error.StageTooSmall; + try e.stage.upload(at, p); + r.* = e.stage.ptr + at; + at += p.len; + } + const l = Affine4.layout(n, k); + const out: Affine4 = .{ .w = try w.alloc(l.words), .s = try w.alloc(l.scales), .b = try w.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + try e.affine.pack(e.stream, raw[0], raw[1], raw[2], out); + try e.stream.synchronize(); + return out; +} + +fn vector(w: *Weights, ck: *core.Checkpoint, name: []const u8, shape: []const usize) !u64 { + const t = try ck.expect(name, .bf16, shape); + const p = try w.alloc(t.bytes.len); + try w.arena.upload(p - w.arena.ptr, t.bytes); + return p; +} + +fn hostBf16(gpa: std.mem.Allocator, ck: *core.Checkpoint, name: []const u8, shape: []const usize) ![]u16 { + const t = try ck.expect(name, .bf16, shape); + const out = try gpa.alloc(u16, t.bytes.len / 2); + @memcpy(out, std.mem.bytesAsSlice(u16, t.bytes)); + return out; +} + +pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *Engine, dir: []const u8, s: Shape) !Weights { + var ck = try core.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + var w: Weights = .{ .fc = undefined, .hidden_norm = 0, .final_norm = 0, .selector = undefined, .predecessor = &.{}, .successor = &.{}, .layers = &.{}, .arena = try DeviceBuffer.alloc(e.d, total(s)) }; + errdefer w.arena.free(); + try w.arena.fill8(0, null); + w.layers = try gpa.alloc(Layer, s.layers); + errdefer gpa.free(w.layers); + w.fc = try linear(gpa, e, &w, &ck, "fc.weight", s.hidden, s.taps * s.hidden); + w.hidden_norm = try vector(&w, &ck, "hidden_norm.weight", &.{s.hidden}); + w.final_norm = try vector(&w, &ck, "norm.weight", &.{s.hidden}); + w.selector = try linear(gpa, e, &w, &ck, "candidate_selector.hidden_projection.weight", s.rank, s.hidden); + w.predecessor = try hostBf16(gpa, &ck, "candidate_selector.predecessor_codebook", &.{ s.vocab, s.rank }); + errdefer gpa.free(w.predecessor); + w.successor = try hostBf16(gpa, &ck, "candidate_selector.successor_codebook", &.{ s.vocab, s.rank }); + errdefer gpa.free(w.successor); + for (w.layers, 0..) |*l, i| { + var buf: [96]u8 = undefined; + const Name = struct { + buf: *[96]u8, + i: usize, + fn of(self: @This(), suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(self.buf, "layers.{d}.{s}", .{ self.i, suffix }); + } + }; + const nm: Name = .{ .buf = &buf, .i = i }; + l.* = .{ + .input_norm = try vector(&w, &ck, try nm.of("input_layernorm.weight"), &.{s.hidden}), + .post_norm = try vector(&w, &ck, try nm.of("post_attention_layernorm.weight"), &.{s.hidden}), + .q = try linear(gpa, e, &w, &ck, try nm.of("self_attn.q_proj.weight"), s.q_width, s.hidden), + .k = try linear(gpa, e, &w, &ck, try nm.of("self_attn.k_proj.weight"), s.kv_width, s.hidden), + .v = try linear(gpa, e, &w, &ck, try nm.of("self_attn.v_proj.weight"), s.kv_width, s.hidden), + .o = try linear(gpa, e, &w, &ck, try nm.of("self_attn.o_proj.weight"), s.hidden, s.q_width), + .q_norm = try vector(&w, &ck, try nm.of("self_attn.q_norm.weight"), &.{s.head_dim}), + .k_norm = try vector(&w, &ck, try nm.of("self_attn.k_norm.weight"), &.{s.head_dim}), + .gate = try linear(gpa, e, &w, &ck, try nm.of("mlp.gate_proj.weight"), s.intermediate, s.hidden), + .up = try linear(gpa, e, &w, &ck, try nm.of("mlp.up_proj.weight"), s.intermediate, s.hidden), + .down = try linear(gpa, e, &w, &ck, try nm.of("mlp.down_proj.weight"), s.hidden, s.intermediate), + .attn_conv = try linear(gpa, e, &w, &ck, try nm.of("attention_conv.kernel_projection.weight"), s.dynamic, s.hidden), + .attn_base = try vector(&w, &ck, try nm.of("attention_conv.base_kernel"), &.{ 2, 2, s.hidden }), + .mlp_conv = try linear(gpa, e, &w, &ck, try nm.of("mlp_conv.kernel_projection.weight"), s.dynamic, s.hidden), + .mlp_base = try vector(&w, &ck, try nm.of("mlp_conv.base_kernel"), &.{ 2, 2, s.hidden }), + }; + } + return w; +} + +test "bf16 rounding of a quantized group's scale and bias" { + try std.testing.expectEqual(@as(u16, 0x3f80), toBf16(1.0)); + try std.testing.expectEqual(@as(f32, 1.0), bf16(0x3f80)); + const w: [64]u16 = @splat(0x3f80); // a flat group: scale 0, bias 1, every nibble 0 + const q = try quantize(std.testing.allocator, &w, 1, 64); + defer { + std.testing.allocator.free(q.words); + std.testing.allocator.free(q.scales); + std.testing.allocator.free(q.biases); + } + try std.testing.expectEqual(@as(u16, 0), q.scales[0]); + try std.testing.expectEqual(@as(u16, 0x3f80), q.biases[0]); + try std.testing.expectEqual(@as(u32, 0), q.words[0]); +} diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 05e69afa11..4ce0cfceb5 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -33,12 +33,17 @@ pub fn aligned(n: usize) usize { /// qwen35.cu's Geo and Seg. const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int }; -const Seg = extern struct { rec: u64, conv: u64, kv: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int }; +const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; + +/// Rows a drafter's tap ring holds a stream: its window and a prompt chunk more. +pub const tap_rows = 2048 + max_rows; /// One stream's state on the device and the logits of its newest prompt chunk on the host. pub const Cache = struct { state: DeviceBuffer, kv: DeviceBuffer, + /// A drafter's tapped rows (`Engine.taps`), a ring of tap_rows: the row at position p in slot p % tap_rows. + tap: ?DeviceBuffer = null, cap: usize, len: usize = 0, logits: []u16, @@ -52,7 +57,13 @@ pub const Cache = struct { var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); errdefer kv.free(); try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) - return .{ .state = state, .kv = kv, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + var tap: ?DeviceBuffer = null; + if (e.n_taps > 0) { + tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); + try tap.?.fill8(0, null); + } + errdefer if (tap) |*t| t.free(); + return .{ .state = state, .kv = kv, .tap = tap, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; } /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). @@ -65,11 +76,13 @@ pub const Cache = struct { pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { self.state.free(); self.kv.free(); + if (self.tap) |*t| t.free(); gpa.free(self.logits); } - fn seg(self: *const Cache, g: Geometry, row0: usize, rows: usize) Seg { - return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows) }; + fn seg(self: *const Cache, e: *const Engine, row0: usize, rows: usize) Seg { + const g = e.g; + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; } }; @@ -99,6 +112,7 @@ const Fns = struct { attn_gate: cuda.Function, mlp_act: cuda.Function, argmax: cuda.Function, + tap: cuda.Function, }; /// A window of one stream in a forward. @@ -175,6 +189,9 @@ pub const Engine = struct { tokens: [most_verify]u32 = undefined, /// The captured forwards by shape (`Shape.key`): the launches from the embedding to the head. graphs: std.AutoHashMapUnmanaged(u32, cuda.graph.Exec) = .empty, + /// A drafter's tapped layers (`setTaps`): each row's residual after them lands in its stream's tap ring. + taps: [8]u32 = undefined, + n_taps: usize = 0, const ids_at = 0; const segs_at = max_rows * 4; @@ -182,8 +199,8 @@ pub const Engine = struct { const rules_at = rowseg_at + max_rows * 4; const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; - const params_at = tokens_at + most_verify * 4; // the attention chunks and their width - const tables_len = params_at + 2 * 4; + const params_at = tokens_at + most_verify * 4; // the attention chunks, their width, the rows, the windows + const tables_len = params_at + 4 * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -200,6 +217,7 @@ pub const Engine = struct { e.window_rows = @min(most_window, e.verify_rows); e.arena_at = 0; e.graphs = .empty; + e.n_taps = 0; e.stream = try cuda.Stream.init(d, true); errdefer e.stream.deinit(); e.arena = try DeviceBuffer.alloc(d, weights.total(g) + try e.scratch(.measure)); @@ -260,9 +278,22 @@ pub const Engine = struct { e.gpa.destroy(e); } - /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values. + /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values, taps. pub fn streamBytes(e: *const Engine, cap: usize) usize { - return e.g.stateBytes() + e.g.kvBytes(cap); + return e.g.stateBytes() + e.g.kvBytes(cap) + e.tapBytes(); + } + + /// A stream's tap ring: tap_rows rows of the tapped layers' residuals. + pub fn tapBytes(e: *const Engine) usize { + return tap_rows * e.n_taps * e.g.hidden * 2; + } + + /// The layers a drafter taps, told before the streams' caches exist; each forward then fills their rings. + pub fn setTaps(e: *Engine, taps: []const u32) !void { + if (taps.len > e.taps.len) return error.TooManyTaps; + for (taps) |t| if (t >= e.g.layers) return error.BadTap; + @memcpy(e.taps[0..taps.len], taps); + e.n_taps = taps.len; } /// The next `bytes` of the arena, 256-byte aligned. @@ -309,7 +340,7 @@ pub const Engine = struct { return total; } - fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { + pub fn run(e: *Engine, f: cuda.Function, grid: [3]usize, block: [3]usize, args: anytype) !void { var a: cuda.Args = .{}; inline for (args) |v| a.add(v); const cfg: cuda.Config = .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = @intCast(block[0]), .y = @intCast(block[1]), .z = @intCast(block[2]) } }; @@ -318,9 +349,9 @@ pub const Engine = struct { /// out (rows, n) = x (rows, k) through `l`: decode tiles 16 rows at a time (each row's bits whatever the batch), /// or the prompt GEMM for a prompt chunk wider than one tile. - fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { + pub fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { if (prompt and rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); - try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)) }); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)), e.tables.ptr + params_at }); var at: usize = 0; while (at < rows) : (at += 16) { const n = @min(16, rows - at); @@ -329,7 +360,7 @@ pub const Engine = struct { } fn norm(e: *Engine, residual: bool, w: u64, rows: usize) !void { - try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)) }); + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); } /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. @@ -358,7 +389,7 @@ pub const Engine = struct { const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. total * 4]); var row0: usize = 0; for (parts, 0..) |p, i| { - segs[i] = p.cache.seg(g, row0, p.rows); + segs[i] = p.cache.seg(e, row0, p.rows); for (row0..row0 + p.rows) |r| rowseg[r] = @intCast(i); p.cache.last = if (record) .{ .start = p.cache.len, .row0 = row0, .rows = p.rows } else null; row0 += p.rows; @@ -373,11 +404,18 @@ pub const Engine = struct { } // attention chunks: at most max_chunks a row, each a multiple of 128 keys const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); - const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..8]); + const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..16]); params[0] = @intCast((widest + ck - 1) / ck); params[1] = @intCast(ck); + params[2] = @intCast(total); + params[3] = @intCast(parts.len); try e.tables.uploadAsync(0, st, e.stream.handle); - const shape: Shape = .{ .total = @intCast(total), .nseg = @intCast(parts.len), .record = record, .prompt = prompt }; + // a captured forward's launches cover a bucket of rows and every window; its kernels read the true counts + // (a lone stream's one-row step, the common case without a drafter, keeps its own exact graph) + const lone = total == 1 and parts.len == 1; + const bucket: usize = if (lone) 1 else if (total <= 16) 16 else most_verify; + const windows: usize = if (lone) 1 else max_segs; + const shape: Shape = .{ .total = @intCast(if (prompt) total else bucket), .nseg = @intCast(if (prompt) parts.len else windows), .record = record, .prompt = prompt, .taps = e.n_taps > 0 }; if (prompt) { // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length would be a graph of its own, // paying the driver's stall on a new graph's first run (about a second on a wide model) @@ -428,9 +466,10 @@ pub const Engine = struct { nseg: u8, record: bool, prompt: bool, + taps: bool, fn key(s: Shape) u32 { - return @as(u32, s.total) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17; + return @as(u32, s.total) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17 | @as(u32, @intFromBool(s.taps)) << 18; } }; @@ -449,7 +488,7 @@ pub const Engine = struct { const geo = e.geo; const group = g.q_heads / g.kv_heads; - try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)) }); + try e.run(e.fns.embed, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ t_ids, e.embed_w, e.embed_s, e.embed_b, e.h, @as(c_int, @intCast(g.hidden)), t_params }); var di: usize = 0; var ai: c_int = 0; for (e.layers, 0..) |l, i| { @@ -463,10 +502,10 @@ pub const Engine = struct { try e.proj(m.z, e.x, e.z, total, prompt); try e.proj(m.a, e.x, e.a, total, prompt); try e.proj(m.b, e.x, e.b, total, prompt); - try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta }); - try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, layer, conv_snap }); - try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, layer, e.y, rec_snap }); - try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix }); + try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta, t_params }); + try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, layer, conv_snap, t_params }); + try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, layer, e.y, rec_snap, t_params }); + try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix, t_params }); try e.proj(m.out, e.mix, e.r, total, prompt); di += 1; }, @@ -475,12 +514,12 @@ pub const Engine = struct { try e.proj(m.q, e.x, e.qraw, total, prompt); try e.proj(m.k, e.x, e.kraw, total, prompt); try e.proj(m.v, e.x, e.vraw, total, prompt); - try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps }); - try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps }); - try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai }); + try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps, t_params }); + try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai, t_params }); try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_params, e.y }); - try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)) }); + try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); try e.proj(m.o, e.mix, e.r, total, prompt); ai += 1; }, @@ -489,8 +528,11 @@ pub const Engine = struct { try e.proj(l.gate, e.x, e.gate, total, prompt); try e.proj(l.up, e.x, e.up, total, prompt); const n: c_longlong = @intCast(total * g.intermediate); - try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n }); + try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n, t_params, @as(c_int, @intCast(g.intermediate)) }); try e.proj(l.down, e.act, e.r, total, prompt); + if (shape.taps) for (e.taps[0..e.n_taps], 0..) |t, ti| if (t == i) { + try e.run(e.fns.tap, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, e.r, t_segs, t_rows, @as(c_int, @intCast(ti)), @as(c_int, @intCast(g.hidden)), t_params }); + }; } try e.norm(true, e.final_norm, total); // the head over the kept rows diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 7fdea31ae7..66b25c02c8 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -7,6 +7,7 @@ const lanes = @import("lanes"); const c = @import("config.zig"); const E = @import("cuda_engine.zig"); const draws = @import("draw.zig"); +const Draft = @import("cuda_draft.zig").Draft; const be = lanes.backend; pub const model_type = "qwen3_5"; @@ -14,6 +15,8 @@ pub const formats: []const []const u8 = &.{"mlx-q4g64"}; 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. +pub const takes_drafter = true; pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; @@ -38,7 +41,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer e.deinit(); if (e.g.vocab != c.vocab) return error.UnsupportedQwenVocab; // draw.zig's masks are sized by the 2B's vocabulary const l = try gpa.create(Lanes); - l.* = .{ .gpa = gpa, .e = e, .capacity = o.context }; + l.* = .{ .gpa = gpa, .io = io, .e = e, .capacity = o.context }; return .{ .backend = l.backend(), .facts = .{ .exact_width = @intCast(e.window_rows), .hidden_rows = true, .batch_rows = @intCast(e.verify_rows), .max_streams = E.max_segs, .streams_exact = true }, @@ -74,6 +77,30 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); } + // with a drafter: a lone stream's windows of the pending token and 1 to 7 drafts + if (l.e.n_taps > 0) for (2..@min(8, l.e.verify_rows) + 1) |rows| { + for (picks[0..rows]) |*pk| pk.* = .{ .rule = null, .position = 1 }; + try l.e.forward(&.{.{ .cache = l.pool.items[0], .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); + try l.pool.items[0].reset(l.e); + }; +} + +/// The DFlash2 drafter in `dir` over the loaded engine `p`, its weights in 4 bits (the only width served). +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 } { + if (bits != 4) return error.DrafterBitsUnsupported; + const l: *Lanes = @ptrCast(@alignCast(p)); + const d = try Draft.init(gpa, io, l.e, dir); + return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter }; +} + +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(); } /// A request this engine refuses, in words; null: none of its own. @@ -99,6 +126,7 @@ fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { const Lanes = struct { gpa: std.mem.Allocator, + io: std.Io, e: *E.Engine, capacity: usize, caches: std.AutoHashMapUnmanaged(*lanes.Stream, *E.Cache) = .empty, @@ -137,7 +165,7 @@ const Lanes = struct { } fn backend(self: *Lanes) be.Backend { - return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; + return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .prepare_features = prepareFeaturesFn, .features = featuresFn, .release = releaseFn } }; } fn cast(ptr: *anyopaque) *Lanes { @@ -199,6 +227,8 @@ const Lanes = struct { /// Each window (a stream's pending token and copies) checked, then as many as fit in each forward. fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { const self = cast(ptr); + const t0 = std.Io.Clock.awake.now(self.io); + defer @import("cuda_draft.zig").timing("verify", self.io, t0); if (windows.len != out.len) return error.InvalidQwenWindows; for (windows, out) |w, o| { const rows = w.rows(); @@ -243,6 +273,8 @@ const Lanes = struct { /// The accepted prefix of each window: its state as it was after that row. fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { const self = cast(ptr); + const t0 = std.Io.Clock.awake.now(self.io); + defer @import("cuda_draft.zig").timing("keep", self.io, t0); for (windows, paths) |w, path| { for (path, 0..) |r, i| if (r != i) return error.UnsupportedQwenTree; try self.e.rewind(try self.cacheOf(w.stream), path.len); @@ -253,6 +285,24 @@ const Lanes = struct { for (requests) |r| if (r.depth != 0) return error.QwenCheckpointHasNoDraftHead; } + /// A drafter's tapped layers, before any stream's caches exist: each forward fills their rings after. + fn prepareFeaturesFn(ptr: *anyopaque, taps: []const u32) !void { + const self = cast(ptr); + if (self.caches.count() != 0 or self.pool.items.len != 0) return error.TapsAfterStreams; + try self.e.setTaps(taps); + } + + /// Rows [start, start + count) of a stream's tap ring: `offset` carries the first row's position, the drafter + /// (cuda_draft.zig) reads slot position % tap_rows, so a request may span the ring's end. + fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { + const self = cast(ptr); + if (taps.len != self.e.n_taps) return error.TapsMismatch; + const cache = try self.cacheOf(stream); + const ring = cache.tap orelse return error.NoTaps; + if (start + count > cache.len or count > E.tap_rows) return error.RowsNotHeld; + return .{ .buffer = ring.ptr, .offset = start, .rows = count, .row_bytes = @intCast(self.e.n_taps * self.e.g.hidden * 2), .space = .device, .dtype = .bf16 }; + } + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { const self = cast(ptr); if (self.caches.fetchRemove(stream)) |entry| self.give(entry.value); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index cf64c4f8b9..27ca52012f 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -10,6 +10,7 @@ const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); const nemotron_slide = @import("nemotron_slide.zig"); +const drafted = @import("cuda_drafted.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. const registry = .{ nemotron.native, qwen35.native }; @@ -154,6 +155,7 @@ const Host = struct { slide: ?*anyopaque = null, // that learner's adapter, freed by slide_free before the family slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into + served: drafted.Served = undefined, // the family's backend, wrapped with --drafter's draft model when given fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); @@ -167,6 +169,7 @@ const Host = struct { } if (h.slide) |s| h.slide_free.?(h.gpa, s); h.gpa.free(h.dir); + h.served.close(h.gpa); h.release(h.family); if (h.gpu) |g| { g.ctx.deinit(); @@ -335,6 +338,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return null; }; defer if (!opened) loaded.deinit(loaded.ctx); + var served = try drafted.wrap(F, a, gpa, io, loaded, o, problem) orelse return null; + defer if (!opened) served.close(gpa); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; // what a family's Loaded may leave out: no learner, no lone driver, no prompt-cache copies @@ -342,24 +347,30 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const learn_bytes: usize = if (@hasField(L, "learn_bytes")) loaded.learn_bytes else 0; const room = after.room(model) -| learn_bytes; const free = loaded.free_streams; - const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { + const streams = budget.admit(room, served.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { - error.NoStream => try std.fmt.allocPrint(a, "the CUDA memory budget fits no stream: one at a {d}-token window takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; lower --context, or free device memory", .{ window, toGib(loaded.stream_bytes), toGib(room), toGib(model), toGib(after.reserve) }), - error.TooMany => try std.fmt.allocPrint(a, "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, and the CUDA memory budget leaves {d:.2} GiB: serve --parallel {d}, or lower --context", .{ o.lanes, toGib(loaded.stream_bytes * (o.lanes -| free)), window, toGib(room), free + room / loaded.stream_bytes }), + error.NoStream => try std.fmt.allocPrint(a, "the CUDA memory budget fits no stream: one at a {d}-token window takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; lower --context, or free device memory", .{ window, toGib(served.stream_bytes), toGib(room), toGib(model), toGib(after.reserve) }), + error.TooMany => try std.fmt.allocPrint(a, "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, and the CUDA memory budget leaves {d:.2} GiB: serve --parallel {d}, or lower --context", .{ o.lanes, toGib(served.stream_bytes * (o.lanes -| free)), window, toGib(room), free + room / served.stream_bytes }), }; return null; }; // a family that pre-allocates its streams' buffers does so now that their count is known - if (@hasField(L, "reserve")) loaded.reserve(loaded.ctx, streams) catch |e| { + if (@hasField(L, "reserve")) { + loaded.reserve(loaded.ctx, streams) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot reserve {d} streams' caches at a {d}-token window ({s}): serve fewer with --parallel, or lower --context", .{ streams, window, @errorName(e) }); + return null; + }; + } + served.reserve(streams) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot reserve {d} streams' caches at a {d}-token window ({s}): serve fewer with --parallel, or lower --context", .{ streams, window, @errorName(e) }); return null; }; const h = try gpa.create(Host); errdefer gpa.destroy(h); - h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; - const cache = if (!@hasField(L, "target")) 0 else budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); + h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = served.backend, .served = served, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = if (@hasField(L, "lone")) loaded.lone else null }; + const cache = if (!@hasField(L, "target")) 0 else budget.cacheBytes(o.prompt_cache_gib, room, served.stream_bytes * (streams -| free)); h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(served.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); @@ -394,7 +405,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, served.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = served.backend.vtable.first_row != null, .prompt_logprobs = served.backend.vtable.prompt_rows, .structures = served.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } diff --git a/zig/src/native/cuda_drafted.zig b/zig/src/native/cuda_drafted.zig new file mode 100644 index 0000000000..d259af0993 --- /dev/null +++ b/zig/src/native/cuda_drafted.zig @@ -0,0 +1,57 @@ +//! --drafter on the CUDA host: a family's draft model behind the lane core's Drafted wrapper over its backend. +const std = @import("std"); +const api = @import("engine_api"); +const lanes = @import("lanes"); +const Allocator = std.mem.Allocator; + +/// What the host serves: the family's backend as is, or wrapped with a drafter in front. +pub const Served = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + stream_bytes: usize, + held: ?*Held = null, + + /// The drafter's own per-stream buffers for `streams`, once the server knows the count. + pub fn reserve(s: *Served, streams: u32) !void { + if (s.held) |h| try h.reserve(h.ctx, streams); + } + + pub fn close(s: *Served, gpa: Allocator) void { + if (s.held) |h| { + h.drafted.deinit(); + h.close(h.ctx); + gpa.destroy(h); + } + s.held = null; + } +}; + +const Held = struct { drafted: lanes.drafted.Drafted, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void }; + +/// `loaded`'s backend, wrapped with the drafter `o.drafter` names when it does; null with `problem` set on refusal. +pub fn wrap(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, loaded: anytype, o: api.Open, problem: *[]const u8) !?Served { + var s: Served = .{ .backend = loaded.backend, .facts = loaded.facts, .stream_bytes = loaded.stream_bytes }; + const dir = o.drafter orelse return s; + if (!o.drafts) return s; + if (!@hasDecl(F, "openDrafter")) { + problem.* = try std.fmt.allocPrint(a, "--drafter: {s} checkpoints load no draft model on CUDA", .{F.model_type}); + return null; + } + const opened = F.openDrafter(gpa, io, loaded.ctx, dir, o.drafter_bits) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load the drafter {s} ({s})", .{ dir, @errorName(e) }); + return null; + }; + const h = try gpa.create(Held); + errdefer gpa.destroy(h); + h.* = .{ .drafted = undefined, .ctx = opened.ctx, .close = opened.close, .reserve = opened.reserve }; + h.drafted = lanes.drafted.Drafted.init(gpa, loaded.backend, opened.drafter) catch |e| { + opened.close(opened.ctx); + problem.* = try std.fmt.allocPrint(a, "the drafter {s} cannot attach to this model ({s})", .{ dir, @errorName(e) }); + return null; + }; + s.backend = h.drafted.backend(); + s.facts = h.drafted.facts(loaded.facts); + s.stream_bytes = h.drafted.streamBytes(loaded.stream_bytes); + s.held = h; + return s; +} From 288ba03d6a64292cf2e44bd26197c9051cebe70d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 06:05:12 +0000 Subject: [PATCH 10/35] qwen3_5: tree drafts on CUDA, the kept path replayed from a log A verify window's rows may form a tree (each row's parent given): the convolution taps, the delta rule and attention follow each row's own branch, and the DFlash2 drafter fills the lane core's planned lanes (each lane its parent's rank-r choice by the codebook selector), so a lone stream verifies up to 12 lanes a round and the planner learns which shapes land. A verify no longer snapshots every row's state (1.3 GB shared, which capped the 27B at 8 rows): each stream logs its rows' decay, write and key, and keep replays the kept path onto the recurrence in one pass (and moves a tree's kept keys down), so any window width fits. One-row steps and prompts update in place as before. The lane host times the target's windows at load, which the depth rule had lacked (it had been capping chains at 3 drafts). The runner's new tree mode verifies tree windows with junk branches against step logits (bit-exact on both models). --- zig/build/cuda.zig | 1 + zig/kernels/cuda/qwen35.cu | 164 +++++++++------- zig/kernels/cuda/qwen35.cuh | 81 ++++++++ zig/kernels/cuda/qwen35_draft.cu | 15 ++ zig/kernels/cuda/qwen35_keep.cu | 63 ++++++ zig/src/core/lanes/drafted.zig | 16 +- zig/src/core/lanes/drafter.zig | 11 +- zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/cuda_cache.zig | 104 ++++++++++ zig/src/families/qwen3_5/cuda_draft.zig | 149 ++++++++++---- zig/src/families/qwen3_5/cuda_engine.zig | 216 +++++++++------------ zig/src/families/qwen3_5/cuda_geometry.zig | 11 ++ zig/src/families/qwen3_5/cuda_native.zig | 61 ++++-- zig/src/native/cuda.zig | 4 + zig/tests/qwen35_cuda_main.zig | 47 ++++- 15 files changed, 684 insertions(+), 260 deletions(-) create mode 100644 zig/kernels/cuda/qwen35.cuh create mode 100644 zig/kernels/cuda/qwen35_keep.cu create mode 100644 zig/src/families/qwen3_5/cuda_cache.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 78fdc3e649..1f3e047eb4 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -34,6 +34,7 @@ const kernels = [_]Kernel{ .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: the Qwen3.5 hybrids' glue, the Metal kernels' rounding points .{ .name = "qwen35_draft", .flags = &.{"-O3"} }, // ours: the DFlash2 drafter's glue (draft_ops.metal's rounding) + .{ .name = "qwen35_keep", .flags = &.{"-O3"} }, // ours: a verified window's kept path applied to its stream's state .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index de830da26f..2186037011 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -2,50 +2,18 @@ // with the same bf16 rounding points, their widths from the model's Geo. Rows of one forward belong to segments (one // stream's window each); a segment's state lives in its stream's blocks. -#include -#include - -namespace tf_q35 { - -constexpr int DK = 128; // linear-attention key and value head width -constexpr int TAPS = 4; // causal convolution taps -constexpr int HD = 256; // attention head width -constexpr int MAXQ = 32; // attention query heads at most - -// The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads. -struct Geo { - int hid, nk, nv, qh, kvh; - __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees - __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence - __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } -}; - -// One stream's window in this forward: its state blocks, cache length before the window, and rows. -struct Seg { - unsigned long long rec; // float recurrences, one rec() block a linear layer - unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer - unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time - unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each - int cap, len, row0, rows, tap_rows, taps; -}; - -__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } -__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } - -__device__ __forceinline__ float warp_sum(float x) { -#pragma unroll - for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); - return x; -} - -} // namespace tf_q35 +#include "qwen35.cuh" using namespace tf_q35; // The per-forward counts a captured forward reads from the device, so one graph serves every shape: P[0] attention -// chunks, P[1] their width, P[2] rows, P[3] windows. A null P means the launch's grid is exact. +// chunks, P[1] their width, P[2] rows, P[3] windows, P[4] whether the rows are recorded for keep (a verify window: +// its state stays as it was until the kept path is applied). A null P means the launch's grid is exact. +// A window's rows may form a tree: rowpar names each forward row's parent forward row (-1: the window's first row, +// after the stream's cache), rowanc each row's ancestors and itself as bits over its window's rows. __device__ __forceinline__ bool past_rows(const int* P, int m) { return P != nullptr && m >= P[2]; } __device__ __forceinline__ bool past_windows(const int* P, int s) { return P != nullptr && s >= P[3]; } +__device__ __forceinline__ bool recorded(const int* P) { return P != nullptr && P[4] != 0; } // MLX affine 4-bit token rows: out = bf16(s * q + b), D wide. extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __restrict__ ids, const uint8_t* __restrict__ W, @@ -115,9 +83,11 @@ extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, floa } // The causal convolution, silu, q/k normalization and the decay and write strength of each row's linear heads. -// Grid (1, 2 * nk + nv, rows), a warp a (row, head): heads [0, nk) are q, [nk, 2nk) k, then nv of v. +// Grid (1, 2 * nk + nv, rows), a warp a (row, head): heads [0, nk) are q, [nk, 2nk) k, then nv of v. A row's taps +// are its parents (then the stream's tail), so a tree's rows each see their own branch. extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, const __nv_bfloat16* __restrict__ CW, + const int* __restrict__ rowseg, const int* __restrict__ rowpar, int layer, + const __nv_bfloat16* __restrict__ CW, const __nv_bfloat16* __restrict__ Ain, const __nv_bfloat16* __restrict__ Bin, const float* __restrict__ ALOG, const __nv_bfloat16* __restrict__ DT, __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ K, @@ -129,19 +99,22 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ if (past_rows(P, w)) return; const int NK = geo.nk, NV = geo.nv, C = geo.conv(); const Seg sg = segs[rowseg[w]]; - const int r = w - sg.row0; const __nv_bfloat16* CS = reinterpret_cast(sg.conv) + (size_t)layer * geo.conv_rows(); const bool isq = head < NK, isk = !isq && head < 2 * NK; const int c0 = head * DK; // q, k and v channels are consecutive heads of DK + int src[TAPS]; // tap t reads the row TAPS - 1 - t tokens back: a forward row, or -1 - (tail row) +#pragma unroll + for (int tap = 0; tap < TAPS; tap++) src[tap] = back(rowpar, w, TAPS - 1 - tap); constexpr int PER = DK / 32; float vals[PER]; #pragma unroll for (int j = 0; j < PER; j++) { const int c = c0 + lane * PER + j; float acc = 0.0f; +#pragma unroll for (int tap = 0; tap < TAPS; tap++) { - const int at = r + tap; // in [state rows (TAPS - 1)] ++ [this segment's rows] - const float xv = at < TAPS - 1 ? bf(CS[(size_t)at * C + c]) : bf(QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]); + const int at = src[tap]; + const float xv = at < 0 ? bf(CS[(size_t)(-1 - at) * C + c]) : bf(QKV[(size_t)at * C + c]); acc += bf(CW[c * TAPS + tap]) * xv; } const float conv = rbf(acc); @@ -171,22 +144,26 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ } } -// Each segment's convolution tail after its rows: the last TAPS - 1 of [tail ++ rows]. A thread a (segment, channel). -// SNAP, if given, also takes the tail after each row (forward row m at m * conv_rows()), read before the update. -extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, int layer, - __nv_bfloat16* __restrict__ SNAP, const int* __restrict__ P) { +// Each segment's convolution tail: a recorded window logs the tail after each of its rows (the last TAPS - 1 tokens of +// the row's branch) for keep, and leaves the stream's own; a prompt chunk or one-row window updates the stream's in +// place. A thread a (segment, channel). +extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, + const int* __restrict__ rowpar, int layer, const int* __restrict__ P) { const int C = geo.conv(); const int c = blockIdx.x * blockDim.x + threadIdx.x; if (c >= C || past_windows(P, blockIdx.y)) return; const Seg sg = segs[blockIdx.y]; __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * geo.conv_rows(); - if (SNAP != nullptr) - for (int j = 0; j < sg.rows; j++) + if (recorded(P)) { + for (int j = 0; j < sg.rows; j++) { + __nv_bfloat16* out = log_row(geo, sg, layer, j).conv; for (int r = 0; r < TAPS - 1; r++) { - const int at = r + j + 1; - SNAP[((size_t)(sg.row0 + j) * (TAPS - 1) + r) * C + c] = - at < TAPS - 1 ? CS[(size_t)at * C + c] : QKV[(size_t)(sg.row0 + at - (TAPS - 1)) * C + c]; + const int at = back(rowpar, sg.row0 + j, TAPS - 2 - r); + out[(size_t)r * C + c] = at < 0 ? CS[(size_t)(-1 - at) * C + c] : QKV[(size_t)at * C + c]; } + } + if (sg.rows != 1) return; + } // ascending: slot r takes [r + rows], which is either a later state slot (not yet written) or a row for (int r = 0; r < TAPS - 1; r++) { const int at = r + sg.rows; @@ -194,13 +171,17 @@ extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __re } } -// The gated delta rule over each segment's rows, its recurrence updated in place; SNAP, if given, takes the state after -// each row (forward row m at m * rec()). Value head hv reads key head hv / (nv / nk). Grid (1, DK / 4, nv * segments), -// block (32, 4): a warp a (value head, value row), its 32 lanes over the key width. +// The gated delta rule over each segment's rows, each row's state its parent's updated: a chain carries the state in +// registers, a branch rebuilds its parent's from the stream's recurrence and the logged updates of the ancestors. +// A recorded window logs each row's (decay, write, key) and leaves the recurrence as it was (keep applies the kept +// path, qwen35_keep.cu); a prompt chunk or one-row window updates it in place. Value head hv reads key head +// hv / (nv / nk). Grid (1, DK / 4, nv * segments), block (32, 4): a warp a (value head, value row), its 32 lanes +// over the key width. extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, - const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, int layer, - __nv_bfloat16* __restrict__ y, float* __restrict__ SNAP, const int* __restrict__ P) { + const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, + const int* __restrict__ rowpar, int layer, __nv_bfloat16* __restrict__ y, + const int* __restrict__ P) { const int NK = geo.nk, NV = geo.nv; if (past_windows(P, blockIdx.z / NV)) return; const int hv = blockIdx.z % NV; @@ -211,18 +192,40 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, cons constexpr int PT = DK / 32; const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; + const bool record = recorded(P); float state[PT]; #pragma unroll for (int i = 0; i < PT; i++) state[i] = st[i]; + int prev = -1; // the row the registers' state follows for (int r = 0; r < sg.rows; r++) { - const size_t row = sg.row0 + r; + const int row = sg.row0 + r; + const int parent = rowpar[row]; + if (parent != prev) { + // a branch: its parent's state from the stream's and the ancestors' updates, oldest first (their keys + // and decays from this forward's rows, their writes from the log, which this thread wrote itself) + int anc[MAXW]; + int n = 0; + for (int a = parent; a >= 0; a = rowpar[a]) anc[n++] = a; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int j = n - 1; j >= 0; j--) { + const int a = anc[j]; + const float ga = g[a * NV + hv], da = log_row(geo, sg, layer, a - sg.row0).delta[hv * DK + dv]; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * ga; + state[i] = fmaf(bf(k[((size_t)a * NK + hk) * DK + PT * dk + i]), da, state[i]); + } + } + } + prev = row; float kc[PT], qc[PT]; #pragma unroll for (int i = 0; i < PT; i++) { - kc[i] = bf(k[(row * NK + hk) * DK + PT * dk + i]); - qc[i] = bf(q[(row * NK + hk) * DK + PT * dk + i]); + kc[i] = bf(k[((size_t)row * NK + hk) * DK + PT * dk + i]); + qc[i] = bf(q[((size_t)row * NK + hk) * DK + PT * dk + i]); } - const float vc = bf(v[(row * NV + hv) * DK + dv]); + const float vc = bf(v[((size_t)row * NV + hv) * DK + dv]); const float gc = g[row * NV + hv]; const float bc = bf(beta[row * NV + hv]); float kv_mem = 0.0f; @@ -236,17 +239,23 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, cons float out = 0.0f; #pragma unroll for (int i = 0; i < PT; i++) { - state[i] = state[i] + kc[i] * delta; + state[i] = fmaf(kc[i], delta, state[i]); out += state[i] * qc[i]; } out = warp_sum(out); - if (dk == 0) y[(row * NV + hv) * DK + dv] = __float2bfloat16_rn(out); - if (SNAP != nullptr) { - float* sn = SNAP + row * geo.rec() + at; + if (dk == 0) y[((size_t)row * NV + hv) * DK + dv] = __float2bfloat16_rn(out); + if (record) { + const LogRow l = log_row(geo, sg, layer, r); + if (dk == 0) { + l.delta[hv * DK + dv] = delta; + if (dv == 0) l.g[hv] = gc; + } + if (dv == 0 && hv % (NV / NK) == 0) #pragma unroll - for (int i = 0; i < PT; i++) sn[i] = state[i]; + for (int i = 0; i < PT; i++) l.k[hk * DK + PT * dk + i] = k[((size_t)row * NK + hk) * DK + PT * dk + i]; } } + if (record && sg.rows != 1) return; #pragma unroll for (int i = 0; i < PT; i++) st[i] = state[i]; } @@ -316,17 +325,19 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); } -// Queries in place at their positions; keys rotated and values copied into their segment's cache. -// Grid (rows), HD threads. +// Queries in place at their positions (the cache length plus the row's depth in its window); keys rotated and values +// copied into their segment's cache at the row's own slot (the cache length plus its row). Grid (rows), HD threads. extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, const int* __restrict__ P) { + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, + const int* __restrict__ P) { const int QH = geo.qh, KVH = geo.kvh; const int d = threadIdx.x; if (past_rows(P, blockIdx.x)) return; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; - const int pos = sg.len + (int)m - sg.row0; + const int pos = sg.len + __popc(rowanc[m]) - 1; + const int slot = sg.len + (int)m - sg.row0; __shared__ __nv_bfloat16 qs[MAXQ * HD]; for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; __syncthreads(); @@ -334,7 +345,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + (size_t)layer * 2 * KVH * sg.cap * HD; __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; for (int h = 0; h < KVH; h++) { - const size_t dst = ((size_t)h * sg.cap + pos) * HD + d; + const size_t dst = ((size_t)h * sg.cap + slot) * HD + d; keys[dst] = rotate(Kn + (m * KVH + h) * HD, d, pos); values[dst] = Vr[(m * KVH + h) * HD + d]; } @@ -342,9 +353,11 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo // Causal attention over a segment's cache, P[1] keys a block over P[0] chunks (read from the device, so a captured // forward replays as caches grow): grid (chunks at most, kvh, rows), block qh / kvh warps, a warp a query head of -// the kv head's group. Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +// the kv head's group. A row sees the cache before its window and its window's ancestors and itself (rowanc). +// Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, - const int* __restrict__ rowseg, int layer, const int* __restrict__ P, float scale, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, + const int* __restrict__ P, float scale, float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { constexpr int DPL = HD / 32; const int nch = P[0], ck = P[1]; @@ -355,8 +368,8 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r const size_t m = blockIdx.z; const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; - const int last = sg.len + (int)m - sg.row0; // the row's own position, the last key it sees - const int k0 = c * ck, k1 = min(k0 + ck, last + 1); + const unsigned anc = rowanc[m]; + const int k0 = c * ck, k1 = min(k0 + ck, sg.len + sg.rows); const size_t slot = (m * QH + qh) * nch + c; if (k0 >= k1) { if (lane == 0) { @@ -375,6 +388,7 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r } float mx = -INFINITY, l = 0.0f; for (int pos = k0; pos < k1; pos++) { + if (pos >= sg.len && ((anc >> (pos - sg.len)) & 1u) == 0) continue; const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); const __nv_bfloat16* kb = reinterpret_cast(&kw); float dot = 0.0f; diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh new file mode 100644 index 0000000000..f3baa2be6c --- /dev/null +++ b/zig/kernels/cuda/qwen35.cuh @@ -0,0 +1,81 @@ +// What qwen35.cu and qwen35_keep.cu share: the model's widths (Geo), a stream's window in a forward (Seg), the per-row +// record a verify window leaves in its stream's log, and the walk from a row back through its parents. + +#pragma once +#include +#include + +namespace tf_q35 { + +constexpr int DK = 128; // linear-attention key and value head width +constexpr int TAPS = 4; // causal convolution taps +constexpr int HD = 256; // attention head width +constexpr int MAXQ = 32; // attention query heads at most +constexpr int MAXW = 16; // a window's rows at most (the log's rows a stream) + +// The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads, linear layers. +struct Geo { + int hid, nk, nv, qh, kvh, nl; + __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees + __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence + __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } + // one row's record of one linear layer: decay and write strength a value head, its write a value (head, dv), its + // key a key head, and the convolution tail after it + __device__ size_t log_bytes() const { return (size_t)nv * 4 + (size_t)nv * DK * 4 + (size_t)nk * DK * 2 + conv_rows() * 2; } +}; + +// One stream's window in this forward: its state blocks, cache length before the window, and rows. +struct Seg { + unsigned long long rec; // float recurrences, one rec() block a linear layer + unsigned long long conv; // bf16 convolution tails, conv_rows() a linear layer + unsigned long long kv; // bf16 keys then values, kvh * cap * HD each, an attention layer at a time + unsigned long long tap; // a drafter's ring of tapped rows (0: none): slot pos % tap_rows, taps * hid bf16 each + unsigned long long log; // the window's rows' records (MAXW rows of nl layers of log_bytes), for keep + int cap, len, row0, rows, tap_rows, taps; +}; + +// A window row's record of one linear layer (Seg.log). +struct LogRow { + float* g; // [nv] decay + float* delta; // [nv][DK] the write's strength along dv + __nv_bfloat16* k; // [nk][DK] the key + __nv_bfloat16* conv; // [TAPS - 1][conv()] the tail after the row +}; + +__device__ __forceinline__ LogRow log_row(const Geo& geo, const Seg& sg, int layer, int row) { + char* base = reinterpret_cast(sg.log) + ((size_t)row * geo.nl + layer) * geo.log_bytes(); + LogRow l; + l.g = reinterpret_cast(base); + l.delta = l.g + geo.nv; + l.k = reinterpret_cast<__nv_bfloat16*>(l.delta + (size_t)geo.nv * DK); + l.conv = l.k + (size_t)geo.nk * DK; + return l; +} + +// The rows of a window to keep, in order (row 0 first), passed by value. +struct Path { + int n; + int rows[MAXW]; +}; + +// The row `steps` tokens before forward row `w` through its parents (rowpar: each forward row's parent forward row, -1 +// for a window's first row): a forward row, or -1 - t for the stream's convolution tail row t (TAPS - 2 the newest). +__device__ __forceinline__ int back(const int* __restrict__ rowpar, int w, int steps) { + int a = w; + for (int s = 1; s <= steps; s++) { + a = rowpar[a]; + if (a < 0) return -1 - (TAPS - 2 - (steps - s)); + } + return a; +} + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q35 diff --git a/zig/kernels/cuda/qwen35_draft.cu b/zig/kernels/cuda/qwen35_draft.cu index 5d6286ea2b..a4726d0711 100644 --- a/zig/kernels/cuda/qwen35_draft.cu +++ b/zig/kernels/cuda/qwen35_draft.cu @@ -31,6 +31,21 @@ extern "C" __global__ void tf_q35d_gather_taps(const __nv_bfloat16* __restrict__ for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; } +// Rows of a window to gather, by value (a tree's kept path). +struct TapRows { + int n; + int rows[16]; +}; + +// Tap rows (start + rows.rows[i]) of a stream's ring into OUT rows 0..n: a kept path's rows in order. Grid (n). +extern "C" __global__ void tf_q35d_gather_rows(const __nv_bfloat16* __restrict__ RING, int R, long long start, TapRows rows, int width, + __nv_bfloat16* __restrict__ OUT) { + const long long row = start + rows.rows[blockIdx.x]; + const __nv_bfloat16* src = RING + (size_t)(row % R) * width; + __nv_bfloat16* dst = OUT + (size_t)blockIdx.x * width; + for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; +} + // y[t, c] = sum over lags l of (BASE[branch, l, c] + DYN[t, (branch * 2 + l) * groups + c / group]) * x[t - l, c], // each product and sum rounded to bf16 in turn; x[-1] is zero (the block's first row has no past); RES, if given, // is added last. D wide rows, `group` channels a dynamic coefficient. Grid (rows), 256 threads. diff --git a/zig/kernels/cuda/qwen35_keep.cu b/zig/kernels/cuda/qwen35_keep.cu new file mode 100644 index 0000000000..b2dae0b2b5 --- /dev/null +++ b/zig/kernels/cuda/qwen35_keep.cu @@ -0,0 +1,63 @@ +// Keeping a verified window's path on CUDA (zig/src/families/qwen3_5/cuda_cache.zig): the stream's recurrences take +// the path's logged updates in order, its convolution tails the path's last row's, and a tree's kept keys move down +// to the slots after the cache. The rows of a window were recorded by qwen35.cu's kernels (Seg.log). + +#include "qwen35.cuh" + +using namespace tf_q35; + +// Each linear layer's recurrence through the path's rows' (decay, key, write), as the forward's chain would have +// left it. Grid (1, DK / 4, nv * nl), block (32, 4): a warp a (layer, value head, value row), its lanes over the key. +extern "C" __global__ void __launch_bounds__(128) tf_q35k_apply(Geo geo, Seg sg, Path path) { + const int NK = geo.nk, NV = geo.nv; + const int layer = blockIdx.z / NV, hv = blockIdx.z % NV; + const int hk = hv / (NV / NK); + const int dk = threadIdx.x; + const int dv = blockIdx.y * 4 + threadIdx.y; + constexpr int PT = DK / 32; + const size_t at = ((size_t)hv * DK + dv) * DK + PT * dk; + float* st = reinterpret_cast(sg.rec) + (size_t)layer * geo.rec() + at; + float state[PT]; +#pragma unroll + for (int i = 0; i < PT; i++) state[i] = st[i]; + for (int j = 0; j < path.n; j++) { + const LogRow l = log_row(geo, sg, layer, path.rows[j]); + const float g = l.g[hv], d = l.delta[hv * DK + dv]; +#pragma unroll + for (int i = 0; i < PT; i++) { + state[i] = state[i] * g; + state[i] = fmaf(bf(l.k[hk * DK + PT * dk + i]), d, state[i]); + } + } +#pragma unroll + for (int i = 0; i < PT; i++) st[i] = state[i]; +} + +// Each linear layer's convolution tail as logged after window row `last`. Grid (nl, tail values / 256), 256 threads. +extern "C" __global__ void tf_q35k_conv(Geo geo, Seg sg, int last) { + const int layer = blockIdx.x; + const size_t n = geo.conv_rows(); + const size_t i = (size_t)blockIdx.y * blockDim.x + threadIdx.x; + if (i >= n) return; + __nv_bfloat16* CS = reinterpret_cast<__nv_bfloat16*>(sg.conv) + (size_t)layer * n; + CS[i] = log_row(geo, sg, layer, last).conv[i]; +} + +// The path's rows' keys and values (slot cache length + row) down to the slots cache length + 0, 1, ... in path +// order: row j of the path is at least j, so a move never lands on a slot a later move still reads. One block a +// (attention layer, kv head): grid (attention layers, kvh), HD threads, a barrier between moves. +extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact(Geo geo, Seg sg, Path path) { + const int KVH = geo.kvh; + const int layer = blockIdx.x, h = blockIdx.y; + __nv_bfloat16* keys = reinterpret_cast<__nv_bfloat16*>(sg.kv) + ((size_t)layer * 2 * KVH + h) * sg.cap * HD; + __nv_bfloat16* values = keys + (size_t)KVH * sg.cap * HD; + const int d = threadIdx.x; + for (int j = 0; j < path.n; j++) { + const int from = sg.len + path.rows[j], to = sg.len + j; + if (from != to) { + keys[(size_t)to * HD + d] = keys[(size_t)from * HD + d]; + values[(size_t)to * HD + d] = values[(size_t)from * HD + d]; + } + __syncthreads(); + } +} diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 3bdffe7c7a..e1a79fd2de 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -60,7 +60,9 @@ pub const Drafted = struct { m.plain_guard = d.plain_guard; m.draft_streams = d.batched; m.draft_probabilities = false; - m.head_trees = false; + m.head_trees = d.trees; + if (d.window_costs.len > 0) m.window_costs = d.window_costs; + if (d.lane_costs.len > 0) m.lane_costs = d.lane_costs; return m; } @@ -157,15 +159,17 @@ pub const Drafted = struct { x.holds.clearRetainingCapacity(); try x.firsts.resize(x.gpa, requests.len); for (requests, x.firsts.items) |r, *tok| { - if (r.lanes != null) return error.TreesNotBuilt; + if (r.lanes != null and !x.drafter.facts().trees) return error.TreesNotBuilt; const lane = x.lanes.getPtr(r.stream) orelse return error.NotDrafting; lane.filler = 0; // a new request supersedes fillers a copy round left unused var pending: u32 = undefined; if (r.rows) |rows| { - for (rows, 0..) |row, i| if (row != i) return error.TreesNotBuilt; - const n: u32 = @intCast(rows.len); + // a tree's kept path names the window's rows out of order: the features of all up to its last + var prefix = true; + for (rows, 0..) |row, i| prefix = prefix and row == i; + const n: u32 = @intCast(if (prefix) rows.len else rows[rows.len - 1] + 1); const f = try x.target.features(r.stream, x.drafter.taps(), r.start, n); - try x.absorbs.append(x.gpa, .{ .stream = r.stream, .features = f, .start = r.start, .follow = r.follow[0..rows.len] }); + try x.absorbs.append(x.gpa, .{ .stream = r.stream, .features = f, .start = r.start, .follow = r.follow[0..rows.len], .rows = if (prefix) null else rows }); pending = r.follow[rows.len - 1]; lane.rows = true; } else { @@ -187,7 +191,7 @@ pub const Drafted = struct { lane.filler = r.depth; continue; } - try x.holds.append(x.gpa, .{ .stream = r.stream, .pending = pending, .position = r.position, .depth = r.depth }); + try x.holds.append(x.gpa, .{ .stream = r.stream, .pending = pending, .position = r.position, .depth = r.depth, .lanes = r.lanes }); } if (x.absorbs.items.len > 0) try x.drafter.absorb(x.absorbs.items); if (x.holds.items.len > 0) try x.drafter.hold(x.holds.items); diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 0f8eab242a..038ee35782 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -1,6 +1,8 @@ //! A drafter outside the target family reading its tapped states; each call takes a round's streams together. const Stream = @import("stream.zig").Stream; const Features = @import("backend.zig").Features; +const Cost = @import("config.zig").Cost; +const Shape = @import("shape.zig").Shape; /// What the round loop's depth rule needs from the drafter itself (not from the target, whose own head may not exist). pub const Facts = struct { @@ -10,13 +12,18 @@ pub const Facts = struct { plain_guard: bool = true, // plain rounds compete with drafted depths (the rule may choose no drafts) batched: bool = false, // hold() runs a shared round's streams as one batch stream_bytes: usize = 0, // device bytes each opened stream holds (its caches), for memory admission + trees: bool = false, // hold() takes a tree of lanes (Hold.lanes), each lane its parent's rank-r token + window_costs: []const Cost = &.{}, // the target's window ms by rows (timed at load), for the depth rule + lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) }; /// Target rows [start, start + features.rows) and the token after each; rows at or past `start` replace what was held. -pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32 }; +/// `rows`, when given, picks the features' rows to absorb (a tree's kept path), in order; null: all of them. +pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32, rows: ?[]const u32 = null }; /// Draft `depth` tokens after `pending`, the first at `position`; asked only once row `position - 2` was absorbed. -pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32 }; +/// With `lanes`, a tree instead: lane i's token is the rank-r choice after its parent (depth unused). +pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32, lanes: ?*const Shape = null }; pub const Drafter = struct { ptr: *anyopaque, diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index ebc3722cf6..756e535a76 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -32,6 +32,7 @@ pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; pub const qwen35_draft: []const u8 = if (available) &Blob("fatbin_qwen35_draft").bytes else &.{}; +pub const qwen35_keep: []const u8 = if (available) &Blob("fatbin_qwen35_keep").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig new file mode 100644 index 0000000000..1d2a31d531 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -0,0 +1,104 @@ +//! One stream's state on the CUDA engine: its recurrences and convolution tails, keys and values, a drafter's tap +//! ring, and the log of its newest verify window's rows, from which `keep` lands any path of that window (a prefix +//! of a chain, or a branch of a tree) without the forward having touched the stream's state. +const std = @import("std"); +const cuda = @import("cuda"); +const geometry = @import("cuda_geometry.zig"); +const engine = @import("cuda_engine.zig"); + +const Engine = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows a drafter's tap ring holds a stream: its window and a prompt chunk more. +pub const tap_rows = 2048 + engine.max_rows; + +/// qwen35.cuh's Seg and Path. +pub const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, log: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; +pub const Path = extern struct { n: c_int, rows: [engine.most_window]c_int }; + +pub const Cache = struct { + state: DeviceBuffer, + kv: DeviceBuffer, + /// A drafter's tapped rows (`Engine.taps`), a ring of tap_rows: the row at position p in slot p % tap_rows. + tap: ?DeviceBuffer = null, + /// The newest verify window's rows' records (qwen35.cuh's LogRow), most_window rows a linear layer. + log: DeviceBuffer, + cap: usize, + len: usize = 0, + /// Positions the newest forward wrote up to: a tap ring holds them until the next forward, whatever `keep` lands. + written: usize = 0, + logits: []u16, + /// The newest verify window of more than one row, until its path is kept: where it started, its rows. + last: ?struct { start: usize, rows: usize } = null, + + pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { + var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); + errdefer state.free(); + try state.fill8(0, null); + var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); + errdefer kv.free(); + try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) + var log = try DeviceBuffer.alloc(e.d, e.g.logBytes(engine.most_window)); + errdefer log.free(); + try log.fill8(0, null); + var tap: ?DeviceBuffer = null; + if (e.n_taps > 0) { + tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); + try tap.?.fill8(0, null); + } + errdefer if (tap) |*t| t.free(); + return .{ .state = state, .kv = kv, .tap = tap, .log = log, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; + } + + /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). + pub fn reset(self: *Cache, e: *const Engine) !void { + self.len = 0; + self.written = 0; + self.last = null; + try self.state.fill8(0, e.stream.handle); + } + + pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { + self.state.free(); + self.kv.free(); + self.log.free(); + if (self.tap) |*t| t.free(); + gpa.free(self.logits); + } + + /// The stream's window of `rows` rows from forward row `row0`, after `len` cached positions. + pub fn seg(self: *const Cache, e: *const Engine, len: usize, row0: usize, rows: usize) Seg { + const g = e.g; + return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; + } + + /// Keeps `path` (the window's rows, the first 0, ascending) of the newest window: the recurrences and tails as a + /// chain of those rows would have left them, a tree's kept keys moved down to the slots after the cache, the + /// length at the path's end. Queued on the engine's stream; nothing waits. + pub fn keep(self: *Cache, e: *Engine, path: []const u32) !void { + const last = self.last orelse return error.NothingToKeep; + if (path.len == 0 or path.len > last.rows or path[0] != 0) return error.NothingToKeep; + var p: Path = .{ .n = @intCast(path.len), .rows = @splat(0) }; + var prefix = true; + for (path, 0..) |r, i| { + if (r >= last.rows or (i > 0 and r <= path[i - 1])) return error.NothingToKeep; + p.rows[i] = @intCast(r); + prefix = prefix and r == i; + } + self.last = null; + self.len = last.start + path.len; + const g = e.g; + const sg = self.seg(e, last.start, 0, last.rows); + try e.run(e.keeps.apply, .{ 1, geometry.dk / 4, g.v_heads * g.deltas() }, .{ 32, 4, 1 }, .{ e.geo, sg, p }); + try e.run(e.keeps.conv, .{ g.deltas(), (g.convBytes() / 2 + 255) / 256, 1 }, .{ 256, 1, 1 }, .{ e.geo, sg, @as(c_int, @intCast(path[path.len - 1])) }); + if (!prefix and g.attentions() > 0) try e.run(e.keeps.compact, .{ g.attentions(), g.kv_heads, 1 }, .{ geometry.head_dim, 1, 1 }, .{ e.geo, sg, p }); + } + + /// The newest window kept whole (the round loop keeps nothing when every row landed), before the next forward. + pub fn settle(self: *Cache, e: *Engine) !void { + const last = self.last orelse return; + var all: [engine.most_window]u32 = undefined; + for (all[0..last.rows], 0..) |*r, i| r.* = @intCast(i); + try self.keep(e, all[0..last.rows]); + } +}; diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index 90cce5c4cd..9b5a9c1c2c 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -1,6 +1,7 @@ //! The DFlash2 drafter on CUDA behind the lane core's Drafter contract (core/lanes/drafter.zig): it absorbs the //! target's tapped rows into a context ring a stream, then drafts a block of 7 tokens after the pending one, the -//! target's own embedding and head at its ends and a codebook selector tracing one path through the top 16 a row. +//! target's own embedding and head at its ends and a codebook selector scoring the top 16 a row under each parent: +//! one path (a chain), or the lanes of the round loop's tree (each lane its parent's rank-r choice). const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -113,14 +114,18 @@ fn num(o: std.json.ObjectMap, key: []const u8) !f64 { }; } -const Fns = struct { gather_taps: cuda.Function, dyn_conv: cuda.Function, norm_rope: cuda.Function, ring_put: cuda.Function, attn: cuda.Function }; +const Fns = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, dyn_conv: cuda.Function, norm_rope: cuda.Function, ring_put: cuda.Function, attn: cuda.Function }; -/// A stream's context: its layers' key and value rings over positions [begin, end), and the drafts it holds. +/// qwen35_draft.cu's TapRows. +const TapRows = extern struct { n: c_int, rows: [engine.most_window]c_int }; + +/// A stream's context: its layers' key and value rings over positions [begin, end), and the drafts it holds (a +/// chain, or a tree's lanes in the shape's order). const Ctx = struct { ring: DeviceBuffer, begin: u64 = 0, end: u64 = 0, - held: [8]u32 = undefined, + held: [engine.most_window]u32 = undefined, held_n: usize = 0, }; @@ -167,7 +172,11 @@ pub const Draft = struct { /// Contexts between streams: allocated and run through once (`reserve`), since the driver stalls a buffer's /// first kernel for seconds. pool: std.ArrayList(*Ctx) = .empty, - step_ms: f64 = 4.0, + /// One level of a hold (a hold drafts every level at once; timed at load, over the block's levels). + step_ms: f64 = 1.0, + /// The target's window ms by rows (cuda_native.zig times them), for the round loop's prices. + costs: [engine.most_window]lanes.config.Cost = undefined, + n_costs: usize = 0, const words_at = 64; const hp_at = words_at + 16 * lp_words * 4; @@ -191,9 +200,37 @@ pub const Draft = struct { d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); errdefer d.staged.free(); try d.picksZero(); + // the first context, and a hold timed on it (the quickest of several: the first ones pay the driver's stalls) + const ctx = try d.create(); + errdefer d.drop(ctx); + try d.absorb(ctx, d.taps_in, 0, 1, null); + var quickest: f64 = std.math.inf(f64); + for (0..8) |_| { + ctx.end = 1; + const t0 = std.Io.Clock.awake.now(io); + try d.hold(ctx, 0, 1, c.block - 1, null); + quickest = @min(quickest, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6); + } + d.step_ms = quickest / @as(f64, @floatFromInt(c.block - 1)); + ctx.* = .{ .ring = ctx.ring }; + try d.pool.append(gpa, ctx); return d; } + /// The drafter's taps (the target's layers it reads). + pub fn taps(d: *const Draft) []const u32 { + return d.c.taps[0..d.c.n_taps]; + } + + /// The target's window of `rows` rows takes `ms` (told by cuda_native.zig before the facts are read). + pub fn setCost(d: *Draft, rows: u32, ms: f64) void { + if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a hold level {d:.2})\n", .{ rows, ms, d.step_ms }); + if (d.n_costs < d.costs.len) { + d.costs[d.n_costs] = .{ .width = rows, .ms = ms }; + d.n_costs += 1; + } + } + pub fn deinit(d: *Draft) void { var it = d.streams.valueIterator(); while (it.next()) |ctx| d.drop(ctx.*); @@ -253,8 +290,8 @@ pub const Draft = struct { ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.c.ringBytes()) }; errdefer ctx.ring.free(); try ctx.ring.fill8(0, null); - try d.absorb(ctx, d.taps_in, 0, 1); // any tap rows do: the drafter's own zeroed staging - try d.hold(ctx, 0, 1, d.c.block - 1); + try d.absorb(ctx, d.taps_in, 0, 1, null); // any tap rows do: the drafter's own zeroed staging + try d.hold(ctx, 0, 1, d.c.block - 1, null); ctx.* = .{ .ring = ctx.ring }; return ctx; } @@ -295,11 +332,13 @@ pub const Draft = struct { } /// Target rows at positions [start, start + n) from the stream's tap ring into its context (replacing what was - /// held at or past `start`); rows older than the window are skipped. - fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize) !void { + /// held at or past `start`); rows older than the window are skipped. With `rows`, the ring's rows start + rows[i] + /// (a tree's kept path, which the target keeps at the positions start + i) instead, n of them. + fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32) !void { const c = d.c; var start = start_; var n = n_; + if (rows_) |rs| if (rs.len != n or n > engine.most_window) return error.BadDraftRows; if (start < ctx.end) ctx.end = start; if (start > ctx.end) ctx.begin = start; // a gap (a prompt restored from a cache): the context starts here if (n > c.window) { @@ -311,7 +350,11 @@ pub const Draft = struct { const e = d.e; while (n > 0) { const rows = @min(n, max_rows); - try e.run(d.fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); + if (rows_) |rs| { + var tr: TapRows = .{ .n = @intCast(rows), .rows = @splat(0) }; + for (rs, 0..) |r, i| tr.rows[i] = @intCast(r); + try e.run(d.fns.gather_rows, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), tr, @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); + } else try e.run(d.fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(c.n_taps * c.hidden)), d.taps_in }); try d.proj(d.w.fc, d.taps_in, d.th, rows); try d.norm(d.th, d.w.hidden_norm, d.thn, rows); for (d.w.layers, 0..) |l, i| { @@ -327,8 +370,9 @@ pub const Draft = struct { if (ctx.end > c.window and ctx.begin < ctx.end - c.window) ctx.begin = ctx.end - c.window; } - /// The block [pending, mask x 7] at positions pos0.. over the stream's context: `depth` drafts held. - fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize) !void { + /// The block [pending, mask x 7] at positions pos0.. over the stream's context: `depth` drafts held as a chain, + /// or the `lanes` tree's tokens. + fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { const c = d.c; const e = d.e; const B = c.block; @@ -369,34 +413,62 @@ pub const Draft = struct { try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + words_at, d.words, rows * lp_words * 4, e.stream.handle), "cuMemcpyDtoHAsync"); try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + hp_at, d.hp, rows * c.rank * 2, e.stream.handle), "cuMemcpyDtoHAsync"); try e.stream.synchronize(); - d.select(ctx, pending, rows); - ctx.held_n = @min(depth, rows); + if (lanes_) |l| { + try d.selectLanes(ctx, pending, rows, l); + if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ l.parents.len, l.levels() }); + } else { + d.select(ctx, pending, rows); + ctx.held_n = @min(depth, rows); + if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] chain {d}\n", .{ctx.held_n}); + } } - /// The path through the rows' candidates: each row's best of unary logit + predecessor . projected . successor. - fn select(d: *Draft, ctx: *Ctx, pending: u32, rows: usize) void { + /// Row `r`'s candidate of `rank` after the token `pred`: by unary logit + predecessor . projected . successor, + /// ties to the earlier candidate (the row's top-k order). + fn choice(d: *const Draft, r: usize, pred: u32, rank: usize) u32 { const c = d.c; - const words = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][0 .. rows * lp_words * 4]); - const hp = std.mem.bytesAsSlice(u16, d.staged.bytes[hp_at..][0 .. rows * c.rank * 2]); + const w = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][r * lp_words * 4 ..][0 .. lp_words * 4]); + const hp = std.mem.bytesAsSlice(u16, d.staged.bytes[hp_at..][r * c.rank * 2 ..][0 .. c.rank * 2]); + var scores: [lp_top]f32 = undefined; + for (0..c.topk) |j| { + const cand = w[2 + j]; + const unary: f32 = @bitCast(w[2 + lp_top + j]); + var edge: f32 = 0; + for (0..c.rank) |k| edge += bf16(d.w.predecessor[pred * c.rank + k]) * bf16(hp[k]) * bf16(d.w.successor[cand * c.rank + k]); + scores[j] = unary + edge; + } + var taken: [lp_top]bool = @splat(false); + var pick: usize = 0; + for (0..rank + 1) |_| { + var best: f32 = -std.math.inf(f32); + pick = 0; + for (0..c.topk) |j| if (!taken[j] and scores[j] > best) { + best = scores[j]; + pick = j; + }; + taken[pick] = true; + } + return w[2 + pick]; + } + + /// The path through the rows' candidates: each row's best after the one before. + fn select(d: *Draft, ctx: *Ctx, pending: u32, rows: usize) void { var pred = pending; for (0..rows) |r| { - const w = words[r * lp_words ..][0..lp_words]; - var best: f32 = -std.math.inf(f32); - var pick: u32 = w[2]; - for (0..c.topk) |j| { - const cand = w[2 + j]; - const unary: f32 = @bitCast(w[2 + lp_top + j]); - var edge: f32 = 0; - for (0..c.rank) |k| edge += bf16(d.w.predecessor[pred * c.rank + k]) * bf16(hp[r * c.rank + k]) * bf16(d.w.successor[cand * c.rank + k]); - const score = unary + edge; - if (score > best) { - best = score; - pick = cand; - } - } - ctx.held[r] = pick; - pred = pick; + ctx.held[r] = d.choice(r, pred, 0); + pred = ctx.held[r]; + } + } + + /// The tree's lanes in order: lane i the rank-r candidate at its depth's row after its parent lane's token. + fn selectLanes(d: *Draft, ctx: *Ctx, pending: u32, rows: usize, l: *const lanes.shape.Shape) !void { + if (l.parents.len > ctx.held.len) return error.TooManyLanes; + for (l.parents, l.depths, l.ranks, 0..) |par, depth, rank, i| { + if (depth >= rows or rank >= d.c.topk or par >= @as(i32, @intCast(i))) return error.BadLanes; + const pred = if (par < 0) pending else ctx.held[@intCast(par)]; + ctx.held[i] = d.choice(depth, pred, rank); } + ctx.held_n = l.parents.len; } // the Drafter contract @@ -416,7 +488,8 @@ pub const Draft = struct { fn factsFn(p: *anyopaque) dr.Facts { const d = cast(p); - return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes() }; + // TF_NO_TREES=1: chains only (for comparing) + return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes(), .trees = std.c.getenv("TF_NO_TREES") == null, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; } fn openFn(p: *anyopaque, s: *Stream) anyerror!void { @@ -436,7 +509,9 @@ pub const Draft = struct { const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) - try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows); + if (it.rows) |rs| { + try d.absorb(ctx, it.features.buffer, it.features.offset, rs.len, rs); + } else try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows, null); } } @@ -446,7 +521,7 @@ pub const Draft = struct { defer timing("hold", d.io, t0); for (items) |it| { const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; - try d.hold(ctx, it.pending, it.position - 1, it.depth); + try d.hold(ctx, it.pending, it.position - 1, it.depth, it.lanes); } } diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 4ce0cfceb5..be59726fc6 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -1,14 +1,20 @@ //! Qwen3.5's hybrid text models on CUDA (the 2B, Qwen3.8-27B): the affine-4 projections through cuda.qlinear, the //! glue in qwen35.cu, one forward for any mix of streams' windows. Each stream owns its recurrences, convolution tails -//! and KV cache (Cache); the widths come from config.json (cuda_geometry.zig), the weights from cuda_weights.zig. -//! A forward's launches are captured into a CUDA graph the first time its shape is seen and replayed after. +//! and KV cache (cuda_cache.zig); the widths come from config.json (cuda_geometry.zig), the weights from +//! cuda_weights.zig. A forward's launches are captured into a CUDA graph the first time its shape is seen and +//! replayed after. A verify window's rows may form a tree (each row's parent given): the forward computes every +//! branch, logs each row, and the stream's state stays as it was until a path is kept. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); const geometry = @import("cuda_geometry.zig"); const weights = @import("cuda_weights.zig"); +const caches = @import("cuda_cache.zig"); pub const Geometry = geometry.Geometry; +pub const Cache = caches.Cache; +pub const tap_rows = caches.tap_rows; +const Seg = caches.Seg; const Affine4 = cuda.qlinear.Affine4; const DeviceBuffer = cuda.DeviceBuffer; const head_dim = geometry.head_dim; @@ -17,12 +23,10 @@ const head_dim = geometry.head_dim; pub const max_rows = 128; /// Streams' windows in one forward at most. pub const max_segs = 32; -/// Rows of a verify forward at most: each row's recurrences are recorded for a partial acceptance. -const most_verify = 32; -/// A window's rows at most (exact: one-row steps' bits). -const most_window = 16; -/// Device bytes the recorded rows may take: a wide model records fewer rows a forward. -const snapshot_budget = 1280 << 20; +/// Rows of a verify forward at most: each row's logits are kept and drawn. +pub const most_verify = 32; +/// A window's rows at most (exact: one-row steps' bits; qwen35.cuh's MAXW, the rows a stream's log holds). +pub const most_window = 16; /// Partial attention chunks a row at most: the chunk width grows with the window instead. const max_chunks = 64; @@ -31,60 +35,8 @@ pub fn aligned(n: usize) usize { return std.mem.alignForward(usize, n, 256); } -/// qwen35.cu's Geo and Seg. -const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int }; -const Seg = extern struct { rec: u64, conv: u64, kv: u64, tap: u64, cap: c_int, len: c_int, row0: c_int, rows: c_int, tap_rows: c_int, taps: c_int }; - -/// Rows a drafter's tap ring holds a stream: its window and a prompt chunk more. -pub const tap_rows = 2048 + max_rows; - -/// One stream's state on the device and the logits of its newest prompt chunk on the host. -pub const Cache = struct { - state: DeviceBuffer, - kv: DeviceBuffer, - /// A drafter's tapped rows (`Engine.taps`), a ring of tap_rows: the row at position p in slot p % tap_rows. - tap: ?DeviceBuffer = null, - cap: usize, - len: usize = 0, - logits: []u16, - /// The newest verify window: where it started, its first row in that forward, its rows. - last: ?struct { start: usize, row0: usize, rows: usize } = null, - - pub fn init(gpa: std.mem.Allocator, e: *const Engine, cap: usize) !Cache { - var state = try DeviceBuffer.alloc(e.d, e.g.stateBytes()); - errdefer state.free(); - try state.fill8(0, null); - var kv = try DeviceBuffer.alloc(e.d, e.g.kvBytes(cap)); - errdefer kv.free(); - try kv.fill8(0, null); // touched now, not as the first request's keys land (the driver commits on first touch) - var tap: ?DeviceBuffer = null; - if (e.n_taps > 0) { - tap = try DeviceBuffer.alloc(e.d, e.tapBytes()); - try tap.?.fill8(0, null); - } - errdefer if (tap) |*t| t.free(); - return .{ .state = state, .kv = kv, .tap = tap, .cap = cap, .logits = try gpa.alloc(u16, e.g.vocab) }; - } - - /// Ready for another stream: no rows, the recurrences and tails zero (on the engine's stream, before its forwards). - pub fn reset(self: *Cache, e: *const Engine) !void { - self.len = 0; - self.last = null; - try self.state.fill8(0, e.stream.handle); - } - - pub fn deinit(self: *Cache, gpa: std.mem.Allocator) void { - self.state.free(); - self.kv.free(); - if (self.tap) |*t| t.free(); - gpa.free(self.logits); - } - - fn seg(self: *const Cache, e: *const Engine, row0: usize, rows: usize) Seg { - const g = e.g; - return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .cap = @intCast(self.cap), .len = @intCast(self.len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; - } -}; +/// qwen35.cuh's Geo. +pub const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh: c_int, nl: c_int }; pub const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; pub const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; @@ -115,12 +67,15 @@ const Fns = struct { tap: cuda.Function, }; -/// A window of one stream in a forward. -pub const Part = struct { cache: *Cache, rows: usize }; +/// qwen35_keep.cu's kernels (Cache.keep). +pub const Keeps = struct { apply: cuda.Function, conv: cuda.Function, compact: cuda.Function }; + +/// A window of one stream in a forward: `parents`, for a tree, names each row's parent row (-1: the first row). +pub const Part = struct { cache: *Cache, rows: usize, parents: ?[]const i32 = null }; /// What a forward keeps. A prompt chunk (`last`, one window) puts its last row's logits in its cache's `logits`; /// verify windows keep every row, drawn on the device into `tokens` (`draws`, one a row) or their logits for -/// `rowLogits` (`logits`), and record each row's state for `rewind`. +/// `rowLogits` (`logits`), and log each row for `Cache.keep`. pub const Keep = union(enum) { last, draws: []const Draw, logits }; /// A kept row's draw on the device: sample.cu's keyed rule at `position`, or greedy without one. @@ -132,16 +87,17 @@ pub const Engine = struct { g: Geometry, geo: Geo, stream: cuda.Stream, - mods: [6]cuda.Module, + mods: [7]cuda.Module, draws: cuda.sample.Draws, affine: cuda.qlinear.Affine4Kernels, fns: Fns, + keeps: Keeps, /// Every weight and scratch buffer in one allocation: the WDDM driver's launch cost grows with a process's /// allocation count, and a 64-layer model would have thousands. arena: DeviceBuffer, arena_at: usize = 0, stage: DeviceBuffer, // the load's staging of each projection as stored, packed from there - /// Rows a verify forward records (snapshot_budget over a row's states), and a window's rows at most. + /// Rows a verify forward takes at most, and a window's rows at most. verify_rows: usize, window_rows: usize, // weights @@ -179,9 +135,8 @@ pub const Engine = struct { pl: u64, po: u64, logits: u64, - rec_snap: u64, // recorded recurrences: verify_rows rows a linear layer - conv_snap: u64, // recorded convolution tails, the same - // the per-forward tables: ids, segments, each row's segment and the draws' rules, staged in pinned memory + // the per-forward tables: ids, segments, each row's segment, parent and ancestors, the draws' rules, staged in + // pinned memory tables: DeviceBuffer, staged: cuda.HostBuffer, host_logits: cuda.HostBuffer, @@ -196,11 +151,13 @@ pub const Engine = struct { const ids_at = 0; const segs_at = max_rows * 4; const rowseg_at = segs_at + max_segs * @sizeOf(Seg); - const rules_at = rowseg_at + max_rows * 4; + const rowpar_at = rowseg_at + max_rows * 4; + const rowanc_at = rowpar_at + max_rows * 4; + const rules_at = rowanc_at + max_rows * 4; const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; - const params_at = tokens_at + most_verify * 4; // the attention chunks, their width, the rows, the windows - const tables_len = params_at + 4 * 4; + const params_at = tokens_at + most_verify * 4; // the attention chunks, their width, the rows, the windows, record + const tables_len = params_at + 8 * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -212,9 +169,9 @@ pub const Engine = struct { e.gpa = gpa; e.d = d; e.g = g; - e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads) }; - e.verify_rows = std.math.clamp(snapshot_budget / (g.deltas() * (g.recBytes() + g.convBytes())), 1, most_verify); - e.window_rows = @min(most_window, e.verify_rows); + e.geo = .{ .hid = @intCast(g.hidden), .nk = @intCast(g.k_heads), .nv = @intCast(g.v_heads), .qh = @intCast(g.q_heads), .kvh = @intCast(g.kv_heads), .nl = @intCast(g.deltas()) }; + e.verify_rows = most_verify; + e.window_rows = most_window; e.arena_at = 0; e.graphs = .empty; e.n_taps = 0; @@ -226,7 +183,7 @@ pub const Engine = struct { e.stage = try DeviceBuffer.alloc(d, weights.stageBytes(g)); errdefer e.stage.free(); const kk = cuda.kernels; - const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample }; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample, kk.qwen35_keep }; var loaded: usize = 0; errdefer for (e.mods[0..loaded]) |*m| m.unload(); for (images, 0..) |img, i| { @@ -239,6 +196,7 @@ pub const Engine = struct { errdefer e.affine.deinit(); e.draws = try cuda.sample.Draws.resolve(e.mods[5]); inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(e.fns, name) = try e.mods[4].function("tf_q35_" ++ name); + inline for (@typeInfo(Keeps).@"struct".field_names) |name| @field(e.keeps, name) = try e.mods[6].function("tf_q35k_" ++ name); e.layers = try gpa.alloc(Layer, g.layers); errdefer gpa.free(e.layers); var ck = try core.Checkpoint.openModel(gpa, io, dir); @@ -251,17 +209,26 @@ pub const Engine = struct { errdefer e.staged.free(); e.host_logits = try cuda.HostBuffer.alloc(d, most_verify * g.vocab * 2); errdefer e.host_logits.free(); - // every path a request takes, once now: the WDDM driver stalls a process's first forwards for seconds - // (a 16 GiB model: a 20 s launch, then a 3 s one), which no request should pay - var warm = try Cache.init(gpa, e, 64); - defer warm.deinit(gpa); - const ids: [32]u32 = @splat(0); - try e.forward(&.{.{ .cache = &warm, .rows = ids.len }}, &ids, .last); // a prompt chunk: the prompt GEMMs - try e.forward(&.{.{ .cache = &warm, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = null, .position = 33 }} }); - try e.forward(&.{.{ .cache = &warm, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = cuda.sample.rule(1, 0.7, 0.95, -std.math.inf(f64), 20), .position = 34 }} }); + try e.warm(gpa); return e; } + /// Every path a request takes, once now: the WDDM driver stalls a process's first forwards for seconds (a 16 GiB + /// model: a 20 s launch, then a 3 s one), which no request should pay. + fn warm(e: *Engine, gpa: std.mem.Allocator) !void { + var warm_cache = try Cache.init(gpa, e, 64); + defer warm_cache.deinit(gpa); + const ids: [32]u32 = @splat(0); + const greedy: Draw = .{ .rule = null, .position = 33 }; + try e.forward(&.{.{ .cache = &warm_cache, .rows = ids.len }}, &ids, .last); // a prompt chunk: the prompt GEMMs + try e.forward(&.{.{ .cache = &warm_cache, .rows = 1 }}, ids[0..1], .{ .draws = &.{greedy} }); + try e.forward(&.{.{ .cache = &warm_cache, .rows = 1 }}, ids[0..1], .{ .draws = &.{.{ .rule = .{ .seed = 1, .inv_t = 1.0 / 0.7, .top_p = 0.95, .near = 20.0, .min_log = -std.math.inf(f32), .top_k = 20 }, .position = 34 }} }); + // a tree window and a branch of it kept: the keep kernels, and the forward's branch path + try e.forward(&.{.{ .cache = &warm_cache, .rows = 3, .parents = &.{ -1, 0, 0 } }}, ids[0..3], .{ .draws = &.{ greedy, greedy, greedy } }); + try warm_cache.keep(e, &.{ 0, 2 }); + try e.stream.synchronize(); + } + pub fn deinit(e: *Engine) void { var it = e.graphs.valueIterator(); while (it.next()) |x| x.deinit(); @@ -278,9 +245,9 @@ pub const Engine = struct { e.gpa.destroy(e); } - /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values, taps. + /// A stream's device bytes at a `cap`-token window: recurrences, convolution tails, keys and values, taps, log. pub fn streamBytes(e: *const Engine, cap: usize) usize { - return e.g.stateBytes() + e.g.kvBytes(cap) + e.tapBytes(); + return e.g.stateBytes() + e.g.kvBytes(cap) + e.tapBytes() + e.g.logBytes(most_window); } /// A stream's tap ring: tap_rows rows of the tapped layers' residuals. @@ -335,8 +302,6 @@ pub const Engine = struct { e.pl = try t.alloc(R * g.q_heads * max_chunks * 4); e.po = try t.alloc(R * g.q_heads * max_chunks * head_dim * 4); e.logits = try t.alloc(most_verify * g.vocab * 2); - e.rec_snap = try t.alloc(g.deltas() * e.verify_rows * g.recBytes()); - e.conv_snap = try t.alloc(g.deltas() * e.verify_rows * g.convBytes()); return total; } @@ -363,12 +328,34 @@ pub const Engine = struct { try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); } - /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. + /// The row tables of `parts`: each forward row's segment, parent forward row (-1: its window's first) and the + /// ancestors-and-self bits over its window's rows (a chain without `parents`). + fn rowTables(st: []u8, parts: []const Part) !void { + const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. max_rows * 4]); + const rowpar = std.mem.bytesAsSlice(c_int, st[rowpar_at..][0 .. max_rows * 4]); + const rowanc = std.mem.bytesAsSlice(u32, st[rowanc_at..][0 .. max_rows * 4]); + var row0: usize = 0; + for (parts, 0..) |p, i| { + if (p.parents) |ps| if (ps.len != p.rows or p.rows > most_window or ps[0] != -1) return error.InvalidQwenParents; + for (0..p.rows) |r| { + const parent: i32 = if (p.parents) |ps| ps[r] else if (r == 0) -1 else @intCast(r - 1); + if (parent >= @as(i32, @intCast(r)) or parent < -1) return error.InvalidQwenParents; + rowseg[row0 + r] = @intCast(i); + rowpar[row0 + r] = if (parent < 0) -1 else @intCast(row0 + @as(usize, @intCast(parent))); + rowanc[row0 + r] = (@as(u32, 1) << @intCast(r)) | (if (parent < 0) 0 else rowanc[row0 + @as(usize, @intCast(parent))]); + } + row0 += p.rows; + } + } + + /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. A + /// window not yet kept from the stream's last verify is kept whole first. pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { const g = e.g; var total: usize = 0; var widest: usize = 0; for (parts) |p| { + try p.cache.settle(e); if (p.rows == 0 or p.cache.len + p.rows > p.cache.cap) return error.QwenContextFull; total += p.rows; widest = @max(widest, p.cache.len + p.rows); @@ -379,6 +366,7 @@ pub const Engine = struct { if (!record and parts.len != 1) return error.InvalidQwenSegments; if (keep == .draws and keep.draws.len != total) return error.InvalidQwenSegments; for (ids) |id| if (id >= g.vocab) return error.InvalidQwenToken; + for (parts) |p| if (p.parents != null and !record) return error.InvalidQwenParents; const heads: usize = if (record) total else 1; // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps const prompt = !record; @@ -386,12 +374,10 @@ pub const Engine = struct { const st = e.staged.bytes; @memcpy(std.mem.bytesAsSlice(u32, st[ids_at..][0 .. total * 4]), ids); const segs = std.mem.bytesAsSlice(Seg, st[segs_at..][0 .. parts.len * @sizeOf(Seg)]); - const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. total * 4]); + try rowTables(st, parts); var row0: usize = 0; for (parts, 0..) |p, i| { - segs[i] = p.cache.seg(e, row0, p.rows); - for (row0..row0 + p.rows) |r| rowseg[r] = @intCast(i); - p.cache.last = if (record) .{ .start = p.cache.len, .row0 = row0, .rows = p.rows } else null; + segs[i] = p.cache.seg(e, p.cache.len, row0, p.rows); row0 += p.rows; } if (keep == .draws) { @@ -404,11 +390,12 @@ pub const Engine = struct { } // attention chunks: at most max_chunks a row, each a multiple of 128 keys const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); - const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..16]); + const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..32]); params[0] = @intCast((widest + ck - 1) / ck); params[1] = @intCast(ck); params[2] = @intCast(total); params[3] = @intCast(parts.len); + params[4] = @intFromBool(record); try e.tables.uploadAsync(0, st, e.stream.handle); // a captured forward's launches cover a bucket of rows and every window; its kernels read the true counts // (a lone stream's one-row step, the common case without a drafter, keeps its own exact graph) @@ -437,7 +424,11 @@ pub const Engine = struct { } try gop.value_ptr.launchOn(e.stream); } - for (parts) |p| p.cache.len += p.rows; + for (parts) |p| { + p.cache.last = if (record and p.rows > 1) .{ .start = p.cache.len, .rows = p.rows } else null; + p.cache.len += p.rows; + p.cache.written = p.cache.len; + } switch (keep) { .draws => |pk| { const toks = e.tables.ptr + tokens_at; @@ -484,6 +475,8 @@ pub const Engine = struct { const t_ids = e.tables.ptr + ids_at; const t_segs = e.tables.ptr + segs_at; const t_rows = e.tables.ptr + rowseg_at; + const t_par = e.tables.ptr + rowpar_at; + const t_anc = e.tables.ptr + rowanc_at; const t_params = e.tables.ptr + params_at; const geo = e.geo; const group = g.q_heads / g.kv_heads; @@ -496,15 +489,13 @@ pub const Engine = struct { switch (l.mixer) { .delta => |m| { const layer: c_int = @intCast(di); - const rec_snap: u64 = if (record) e.rec_snap + di * e.verify_rows * g.recBytes() else 0; - const conv_snap: u64 = if (record) e.conv_snap + di * e.verify_rows * g.convBytes() else 0; try e.proj(m.qkv, e.x, e.qkv, total, prompt); try e.proj(m.z, e.x, e.z, total, prompt); try e.proj(m.a, e.x, e.a, total, prompt); try e.proj(m.b, e.x, e.b, total, prompt); - try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta, t_params }); - try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, layer, conv_snap, t_params }); - try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, layer, e.y, rec_snap, t_params }); + try e.run(e.fns.gdn_pre, .{ 1, 2 * g.k_heads + g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.qkv, t_segs, t_rows, t_par, layer, m.conv, e.a, e.b, m.a_log, m.dt_bias, e.q, e.k, e.v, e.g_, e.beta, t_params }); + try e.run(e.fns.conv_commit, .{ (g.convWidth() + 255) / 256, nseg, 1 }, .{ 256, 1, 1 }, .{ geo, e.qkv, t_segs, t_par, layer, t_params }); + try e.run(e.fns.gdn_chain, .{ 1, geometry.dk / 4, g.v_heads * nseg }, .{ 32, 4, 1 }, .{ geo, e.q, e.k, e.v, e.g_, e.beta, t_segs, t_par, layer, e.y, t_params }); try e.run(e.fns.gdn_post, .{ 1, g.v_heads, total }, .{ 32, 1, 1 }, .{ geo, e.y, e.z, m.norm, g.eps, e.mix, t_params }); try e.proj(m.out, e.mix, e.r, total, prompt); di += 1; @@ -516,8 +507,8 @@ pub const Engine = struct { try e.proj(m.v, e.x, e.vraw, total, prompt); try e.run(e.fns.head_norm, .{ total * g.q_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.qraw, m.q_norm, e.qn, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), g.eps, t_params }); try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); - try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, ai, t_params }); - try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); + try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, t_anc, ai, t_params }); + try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, t_anc, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_params, e.y }); try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); try e.proj(m.o, e.mix, e.r, total, prompt); @@ -543,23 +534,4 @@ pub const Engine = struct { pub fn rowLogits(e: *const Engine, i: usize) []const u16 { return e.host_logits.slice(u16)[i * e.g.vocab ..][0..e.g.vocab]; } - - /// Keeps the first `rows` of the cache's newest window: its recurrences and convolution tails as they were after - /// that row (the forward recorded them), its length at that row. The window's later keys stay as unread slots. - pub fn rewind(e: *Engine, cache: *Cache, rows: usize) !void { - const g = e.g; - const last = cache.last orelse return error.NothingToKeep; - if (rows == 0 or rows > last.rows) return error.NothingToKeep; - cache.last = null; - cache.len = last.start + rows; - if (rows == last.rows) return; - const row = last.row0 + rows - 1; - const rb = g.recBytes(); - const cb = g.convBytes(); - for (0..g.deltas()) |di| { - try cache.state.copyFrom(di * rb, e.rec_snap + (di * e.verify_rows + row) * rb, rb, e.stream.handle); - try cache.state.copyFrom(g.deltas() * rb + di * cb, e.conv_snap + (di * e.verify_rows + row) * cb, cb, e.stream.handle); - } - try e.stream.synchronize(); - } }; diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig index 34207bacb5..7b327fd1a9 100644 --- a/zig/src/families/qwen3_5/cuda_geometry.zig +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -70,6 +70,17 @@ pub const Geometry = struct { return g.deltas() * (g.recBytes() + g.convBytes()); } + /// One verified row's record of one linear layer (qwen35.cuh's log_bytes): decay and write by value head, + /// key by key head, the tail after it. + pub fn logRowBytes(g: Geometry) usize { + return g.v_heads * 4 + g.v_heads * dk * 4 + g.k_heads * dk * 2 + g.convBytes(); + } + + /// A stream's log of a window's `rows` rows. + pub fn logBytes(g: Geometry, rows: usize) usize { + return rows * g.deltas() * g.logRowBytes(); + } + /// A stream's keys and values at a `cap`-token window. pub fn kvBytes(g: Geometry, cap: usize) usize { return g.attentions() * 2 * g.kv_heads * cap * head_dim * 2; diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 66b25c02c8..d4fa634cb1 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -77,19 +77,24 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { try l.e.forward(parts[0 .. n + 1], ids[0 .. n + 1], .{ .draws = picks[0 .. n + 1] }); for (l.pool.items[0 .. n + 1]) |cache| try cache.reset(l.e); } - // with a drafter: a lone stream's windows of the pending token and 1 to 7 drafts - if (l.e.n_taps > 0) for (2..@min(8, l.e.verify_rows) + 1) |rows| { + // with a drafter: a lone stream's windows of the pending token and its drafts, chains and trees + if (l.e.n_taps > 0) for (2..l.e.window_rows + 1) |rows| { for (picks[0..rows]) |*pk| pk.* = .{ .rule = null, .position = 1 }; try l.e.forward(&.{.{ .cache = l.pool.items[0], .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); try l.pool.items[0].reset(l.e); }; } -/// The DFlash2 drafter in `dir` over the loaded engine `p`, its weights in 4 bits (the only width served). +/// The DFlash2 drafter in `dir` over the loaded engine `p`, its weights in 4 bits (the only width served). The +/// target's windows of 1 to window_rows rows are timed here (with the drafter's taps on) for the round loop's prices. 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 } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); + 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()); + try l.timeWindows(d); return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter }; } @@ -121,7 +126,7 @@ fn release(p: *anyopaque) void { fn pick(sampling: ?lanes.Sampling, position: u64) E.Draw { const s = sampling orelse return .{ .rule = null, .position = position }; if (s.temperature <= 0) return .{ .rule = null, .position = position }; - return .{ .rule = cuda.sample.rule(s.seed, s.temperature, s.top_p, s.minLog(), s.top_k), .position = position }; + return .{ .rule = cuda.sample.Rule.from(lanes.gpu_rule.rule(s)), .position = position }; } const Lanes = struct { @@ -164,6 +169,26 @@ const Lanes = struct { self.pool.append(self.gpa, cache) catch self.drop(cache); } + /// Each window width's forward, timed on a cache of its own (its third run: the shape's graph captured and the + /// driver's stall behind it, the previous window's keep in front as in a round), into the drafter's facts. + fn timeWindows(self: *Lanes, d: *Draft) !void { + var cache = try E.Cache.init(self.gpa, self.e, 64); + defer cache.deinit(self.gpa); + const ids: [E.most_window]u32 = @splat(0); + var picks: [E.most_window]E.Draw = undefined; + for (&picks) |*pk| pk.* = .{ .rule = null, .position = 1 }; + for (1..self.e.window_rows + 1) |rows| { + var ms: f64 = 0; + for (0..3) |run| { + try cache.reset(self.e); + const t0 = std.Io.Clock.awake.now(self.io); + try self.e.forward(&.{.{ .cache = &cache, .rows = rows }}, ids[0..rows], .{ .draws = picks[0..rows] }); + if (run == 2) ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(self.io)).toNanoseconds())) / 1e6; + } + d.setCost(@intCast(rows), ms); + } + } + fn backend(self: *Lanes) be.Backend { return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .prepare_features = prepareFeaturesFn, .features = featuresFn, .release = releaseFn } }; } @@ -224,7 +249,8 @@ const Lanes = struct { return @intCast(handle); } - /// Each window (a stream's pending token and copies) checked, then as many as fit in each forward. + /// Each window (a stream's pending token and drafts, a chain or a tree) checked, then as many as fit in each + /// forward. A tree's rows are keyed at the cache length plus their depth, as the round loop positions them. fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { const self = cast(ptr); const t0 = std.Io.Clock.awake.now(self.io); @@ -232,11 +258,18 @@ const Lanes = struct { if (windows.len != out.len) return error.InvalidQwenWindows; for (windows, out) |w, o| { const rows = w.rows(); - if (w.parents != null or w.held != 0 or rows > self.e.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; + if (w.held != 0 or rows > self.e.window_rows or w.positions.len != rows) return error.UnsupportedQwenWindow; if (w.masks.len != 0 and w.masks.len != rows * draws.mask_words) return error.InvalidQwenMasks; if (o.sampled.len != rows or o.drafts.len != w.tokens.len) return error.InvalidQwenOutputs; const cache = try self.cacheOf(w.stream); - for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidQwenPosition; + var depth: [E.most_window]u64 = undefined; + for (w.positions, 0..) |p, r| { + const parent: i32 = if (w.parents) |ps| ps[r] else if (r == 0) -1 else @intCast(r - 1); + if (r == 0 and parent != -1) return error.UnsupportedQwenWindow; + if (r > 0 and (parent < 0 or parent >= @as(i32, @intCast(r)))) return error.UnsupportedQwenWindow; + depth[r] = if (parent < 0) 0 else depth[@intCast(parent)] + 1; + if (p != cache.len + depth[r] + 1) return error.InvalidQwenPosition; + } } var at: usize = 0; while (at < windows.len) { @@ -248,7 +281,7 @@ const Lanes = struct { var masked = false; while (at + n < windows.len and n < E.max_segs and total + windows[at + n].rows() <= self.e.verify_rows) : (n += 1) { const w = windows[at + n]; - parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows() }; + parts[n] = .{ .cache = self.caches.get(w.stream).?, .rows = w.rows(), .parents = w.parents }; ids[total] = w.pending; @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); for (w.positions, 0..) |p, r| picks[total + r] = pick(w.stream.sampling, p); @@ -270,14 +303,15 @@ const Lanes = struct { } } - /// The accepted prefix of each window: its state as it was after that row. + /// The accepted path of each window (a chain's prefix or a tree's branch): the state as that chain leaves it. fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { const self = cast(ptr); const t0 = std.Io.Clock.awake.now(self.io); defer @import("cuda_draft.zig").timing("keep", self.io, t0); for (windows, paths) |w, path| { - for (path, 0..) |r, i| if (r != i) return error.UnsupportedQwenTree; - try self.e.rewind(try self.cacheOf(w.stream), path.len); + const cache = try self.cacheOf(w.stream); + if (cache.last == null) continue; // a one-row window (or one kept already): nothing to roll back + try cache.keep(self.e, path); } } @@ -293,13 +327,14 @@ const Lanes = struct { } /// Rows [start, start + count) of a stream's tap ring: `offset` carries the first row's position, the drafter - /// (cuda_draft.zig) reads slot position % tap_rows, so a request may span the ring's end. + /// (cuda_draft.zig) reads slot position % tap_rows, so a request may span the ring's end. The newest window's + /// rows stay readable after a keep (a tree's kept path reads them by row) until the next forward. fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { const self = cast(ptr); if (taps.len != self.e.n_taps) return error.TapsMismatch; const cache = try self.cacheOf(stream); const ring = cache.tap orelse return error.NoTaps; - if (start + count > cache.len or count > E.tap_rows) return error.RowsNotHeld; + if (start + count > cache.written or count > E.tap_rows) return error.RowsNotHeld; return .{ .buffer = ring.ptr, .offset = start, .rows = count, .row_bytes = @intCast(self.e.n_taps * self.e.g.hidden * 2), .space = .device, .dtype = .bf16 }; } diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 27ca52012f..ae7e82d1e5 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -193,6 +193,10 @@ const Host = struct { fn serve(h: *Host, io: std.Io, facts: lanes.Model, rows: u32, info: api.Info, explain: ?api.Explain) !void { h.cfg = try lanes.Config.init(h.gpa, facts, rows, rows - 1); errdefer h.cfg.deinit(h.gpa); + if (facts.head_trees) { // a drafter's trees: the planner picks a lone stream's lanes, as deep as it drafts + h.cfg.head_lanes = rows - 1; + h.cfg.head_depth = @max(1, facts.drafts); + } h.clock = .{ .io = io }; h.core = lanes.Engine.init(h.gpa, &h.cfg, h.backend(), h.clock.clock()); errdefer h.core.deinit(); diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 78414fdb32..858dc34c07 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,7 +1,8 @@ -//! `tf-qwen35-cuda `: a Qwen3.5 hybrid's CUDA forward over comma-separated -//! token ids: one token a forward (step: every position's logits), windows of growing widths (window: the same rows, -//! which must equal step's) or the ids as one prompt (prompt: the last position's), as little-endian bf16 rows into -//! for a reference to compare. Step mode also prints the decode speed. +//! `tf-qwen35-cuda `: a Qwen3.5 hybrid's CUDA forward over +//! comma-separated token ids: one token a forward (step: every position's logits), windows of growing widths +//! (window: the same rows, which must equal step's), tree windows with junk branches kept to the true one (tree: +//! the same), or the ids as one prompt (prompt: the last position's), as little-endian bf16 rows into for a +//! reference to compare. Step mode also prints the decode speed. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -11,7 +12,7 @@ pub fn main(init: std.process.Init) !u8 { const a = init.arena.allocator(); const args = try init.minimal.args.toSlice(a); if (args.len != 5) { - std.debug.print("usage: tf-qwen35-cuda \n", .{}); + std.debug.print("usage: tf-qwen35-cuda \n", .{}); return 2; } var ids: std.ArrayList(u32) = .empty; @@ -49,6 +50,42 @@ pub fn main(init: std.process.Init) !u8 { for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); at += n; } + } else if (std.mem.eql(u8, args[2], "tree")) { + // tree windows of growing widths: the true continuation down one branch, a junk sibling before each of its + // rows and a junk grandchild, the true branch kept; its rows' logits must equal step's bit for bit + var at: usize = 0; + var w: usize = 2; + while (at < ids.items.len) : (w = w % e.window_rows + 1) { + var toks: [E.most_window]u32 = undefined; + var parents: [E.most_window]i32 = undefined; + var path: [E.most_window]u32 = undefined; + toks[0] = ids.items[at]; + parents[0] = -1; + path[0] = 0; + var n: usize = 1; + var kept: usize = 1; + var last: i32 = 0; + var depth: usize = 1; + while (n + 2 <= w and at + depth < ids.items.len) : (depth += 1) { + toks[n] = (ids.items[at + depth] + 7) % @as(u32, @intCast(e.g.vocab)); + parents[n] = last; + toks[n + 1] = ids.items[at + depth]; + parents[n + 1] = last; + path[kept] = @intCast(n + 1); + last = @intCast(n + 1); + n += 2; + kept += 1; + } + if (n < w and n >= 3) { + toks[n] = 1; + parents[n] = 1; + n += 1; + } + try e.forward(&.{.{ .cache = &cache, .rows = n, .parents = parents[0..n] }}, toks[0..n], .logits); + for (path[0..kept]) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); + if (n > 1) try cache.keep(e, path[0..kept]); + at += kept; + } } else { var at: usize = 0; while (at < ids.items.len) { From 61fc2bf9412eca769848af16eef7eed322c569dd Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 06:06:09 +0000 Subject: [PATCH 11/35] docs: the Qwen hybrids on CUDA in the Windows page and the README What the port serves beyond Nemotron: Qwen3.5-2B and Qwen3.8-27B with its DFlash2 drafter, the speeds measured on an RTX 4090, the load time and why, the timing and chain-only switches, and what is not there. --- README.md | 5 +++-- docs/WINDOWS.md | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index a6a018278f..c838faa144 100644 --- a/README.md +++ b/README.md @@ -73,8 +73,8 @@ The platform column names the hardware tested for each model. | Qwen3.8 Flash Next | MLX affine 6-bit, group 32 | Metal on M5 Ultra | | GLM-5.3-Flash | MLX affine 4-bit, group 64 | Metal on two M5 Ultras | | GLM-5.3 | EXLR8 `home` profile (trellis K2/K3/K4 experts) | Metal on two M5 Ultras | -| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max | -| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra | +| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max; CUDA on RTX 4090 (unofficial Windows port) | +| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (unofficial Windows port) | Nemotron's named checkpoint is `TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`. NVIDIA's NVFP4 checkpoint, `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4`, serves on CUDA as stored; its bf16 MTP head is quantized to 4-bit at load for drafts. @@ -90,6 +90,7 @@ Its prompt processing trails mlx_lm for now, and we are fixing it. An M3 Ultra r Bonsai, Gemma 4, Qwen3.6 and DeepSeek-V4 are still under qualification for 1.0.x. The Python 0.6.6 engine remains available at the `v0.6.6` tag for those models and other backends. On CUDA, 1.0.11 serves Nemotron on a GB10, greedy and sampled, with concurrent requests sharing each round. +The unofficial Windows port ([docs/WINDOWS.md](docs/WINDOWS.md)) also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees. The Linux x86_64 archive adds NVIDIA Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), tested on an RTX 3090, and Blackwell cards with compute capability 12.0 (the RTX 50 series and RTX PRO 6000 Blackwell), tested on an RTX PRO 6000; the A100 (8.0) is not supported yet. ## Exact decoding diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index fdb9b7f88f..b11d72f4fd 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -34,6 +34,32 @@ tensorfold-native serve --name local-model --parallel 1 --context 8 The server listens on `http://127.0.0.1:8888/v1`. Ctrl+C (or closing the console) stops it cleanly. +## Models on CUDA + +Nemotron 3.5 Lightning serves as on Linux. The Qwen3.5 hybrids have a CUDA engine of their own in this port +(`zig/src/families/qwen3_5/cuda*.zig`, kernels `zig/kernels/cuda/qwen35*.cu`), the Metal kernels' arithmetic with +the same rounding points, so a window of up to 16 rows reproduces one-token steps bit for bit: + +```bat +tensorfold-native serve --name qwen2b --context 32768 +tensorfold-native serve --name qwen27b --context 8192 --parallel 4 ^ + --drafter +``` + +- Qwen3.5-2B: about 340 tokens/s for one request on an RTX 4090, 1,500 across eight. +- Qwen3.8-27B (16 GiB): 40 tokens/s plain, at the card's memory bandwidth; 240 across eight requests. +- `--drafter` (4-bit only) loads z-lab's DFlash2 draft model: the target writes five layers' residuals into a ring + for it, and it drafts a tree of up to 12 lanes a round that the round loop's planner shapes from what the target + took. A lone greedy request decodes at 80 to 160 tokens/s by how predictable the text is (prose at the low end, + code at the high), sampled ones at 50 to 70, and four at once reach about 160 aggregate; replies stay identical to + plain decoding. Loading takes about 80 s: the drafter's weights are quantized to 4 bits on the host and every + shape is run once, as the driver stalls a process's first use of each kernel and buffer. +- A verify window's rows are not snapshotted: each stream logs its rows' recurrence updates and a keep replays the + accepted path, so the 27B records 32 rows a forward where 8 fit before. +- `TF_TIMING=1` prints each verify, keep, absorb and hold's milliseconds and the drafted shape (chain or lanes); + `TF_NO_TREES=1` drafts chains only. +- No prompt cache (`--prompt-cache-gib`), `--learn` or `--slide` for these two models. + ## How the port works - `zig/src/platform/sys.zig` (`@import("sys")`, wired into every module by `zig/build/sys.zig`) is the operating From 6aaf241affa27864e08898a4648bc4e32a491015 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:37 +0000 Subject: [PATCH 12/35] qwen3_5, platform: comments on one line The lean rule (tools/zig/lean_check.py) counts a comment that runs on to a second line; the Windows port's new files had a hundred such runs. Each is one line of at most 120 columns now, and the checker reports the same 539 problems as main. --- zig/build/sys.zig | 3 +- zig/kernels/cuda/qwen35.cu | 45 +++++-------------- zig/kernels/cuda/qwen35.cuh | 9 ++-- zig/kernels/cuda/qwen35_draft.cu | 15 ++----- zig/kernels/cuda/qwen35_keep.cu | 11 ++--- zig/src/cli/pull_file.zig | 3 +- zig/src/core/lanes/drafter.zig | 6 +-- zig/src/families/qwen3_5/cuda.zig | 3 +- zig/src/families/qwen3_5/cuda_cache.zig | 8 +--- zig/src/families/qwen3_5/cuda_draft.zig | 24 +++------- .../families/qwen3_5/cuda_draft_weights.zig | 7 +-- zig/src/families/qwen3_5/cuda_engine.zig | 35 +++++---------- zig/src/families/qwen3_5/cuda_geometry.zig | 6 +-- zig/src/families/qwen3_5/cuda_native.zig | 23 +++------- zig/src/families/qwen3_5/cuda_weights.zig | 3 +- zig/src/platform/sys.zig | 3 +- zig/src/platform/win_c.zig | 7 +-- zig/src/platform/win_clock.zig | 3 +- zig/src/platform/win_posix.zig | 7 +-- zig/tests/qwen35_cuda_main.zig | 9 +--- 20 files changed, 65 insertions(+), 165 deletions(-) diff --git a/zig/build/sys.zig b/zig/build/sys.zig index b34be717ed..1c1309af78 100644 --- a/zig/build/sys.zig +++ b/zig/build/sys.zig @@ -1,5 +1,4 @@ -//! `sys` (zig/src/platform/sys.zig) for every module the build compiles: one walk of the step graph after it is -//! declared, so no module's declaration needs to name it. +//! `sys` (zig/src/platform/sys.zig) for every module the build compiles, by one walk of the step graph once declared. const std = @import("std"); diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 2186037011..cef7c7b5a2 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -1,16 +1,10 @@ -// Qwen3.5's hybrid text models on CUDA (the 2B and Qwen3.8-27B): the Metal engine's kernels (zig/kernels/metal/qwen3_5) -// with the same bf16 rounding points, their widths from the model's Geo. Rows of one forward belong to segments (one -// stream's window each); a segment's state lives in its stream's blocks. +// Qwen3.5's hybrids (2B, 27B) on CUDA: the Metal kernels' arithmetic and rounding; a segment is one stream's window. #include "qwen35.cuh" using namespace tf_q35; -// The per-forward counts a captured forward reads from the device, so one graph serves every shape: P[0] attention -// chunks, P[1] their width, P[2] rows, P[3] windows, P[4] whether the rows are recorded for keep (a verify window: -// its state stays as it was until the kept path is applied). A null P means the launch's grid is exact. -// A window's rows may form a tree: rowpar names each forward row's parent forward row (-1: the window's first row, -// after the stream's cache), rowanc each row's ancestors and itself as bits over its window's rows. +// P (device counts, one graph for every shape): chunks, width, rows, windows, recorded; rowpar/rowanc give the tree. __device__ __forceinline__ bool past_rows(const int* P, int m) { return P != nullptr && m >= P[2]; } __device__ __forceinline__ bool past_windows(const int* P, int s) { return P != nullptr && s >= P[3]; } __device__ __forceinline__ bool recorded(const int* P) { return P != nullptr && P[4] != 0; } @@ -30,8 +24,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_embed(const int* __rest } } -// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)), D wide (a multiple of 256, at most 8192). One 256-thread -// block a row; each thread's elements in index order, warps summed in order. +// h = bf16(h + r) (r may be null), x = bf16(w * (h * inv)); D a multiple of 256 up to 8192; a 256-thread block a row. extern "C" __global__ void __launch_bounds__(256) tf_q35_norm(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ X, float eps, int D, const int* __restrict__ P) { @@ -82,9 +75,7 @@ extern "C" __global__ void tf_q35_sums(const __nv_bfloat16* __restrict__ X, floa } } -// The causal convolution, silu, q/k normalization and the decay and write strength of each row's linear heads. -// Grid (1, 2 * nk + nv, rows), a warp a (row, head): heads [0, nk) are q, [nk, 2nk) k, then nv of v. A row's taps -// are its parents (then the stream's tail), so a tree's rows each see their own branch. +// Conv, silu, q/k norms, the linear heads' decay and write: a warp a (row, head) over q, k, v; taps from the parents. extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const int* __restrict__ rowpar, int layer, const __nv_bfloat16* __restrict__ CW, @@ -144,9 +135,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_pre(Geo geo, const _ } } -// Each segment's convolution tail: a recorded window logs the tail after each of its rows (the last TAPS - 1 tokens of -// the row's branch) for keep, and leaves the stream's own; a prompt chunk or one-row window updates the stream's in -// place. A thread a (segment, channel). +// A segment's conv tail: a recorded window logs each row's for keep and leaves the stream's; others update it in place. extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __restrict__ QKV, const Seg* __restrict__ segs, const int* __restrict__ rowpar, int layer, const int* __restrict__ P) { const int C = geo.conv(); @@ -171,12 +160,7 @@ extern "C" __global__ void tf_q35_conv_commit(Geo geo, const __nv_bfloat16* __re } } -// The gated delta rule over each segment's rows, each row's state its parent's updated: a chain carries the state in -// registers, a branch rebuilds its parent's from the stream's recurrence and the logged updates of the ancestors. -// A recorded window logs each row's (decay, write, key) and leaves the recurrence as it was (keep applies the kept -// path, qwen35_keep.cu); a prompt chunk or one-row window updates it in place. Value head hv reads key head -// hv / (nv / nk). Grid (1, DK / 4, nv * segments), block (32, 4): a warp a (value head, value row), its 32 lanes -// over the key width. +// Gated delta rule, each row's state its parent's: a chain keeps it in registers, a branch rebuilds it from the log. extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const float* __restrict__ g, const __nv_bfloat16* __restrict__ beta, const Seg* __restrict__ segs, @@ -201,8 +185,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_q35_gdn_chain(Geo geo, cons const int row = sg.row0 + r; const int parent = rowpar[row]; if (parent != prev) { - // a branch: its parent's state from the stream's and the ancestors' updates, oldest first (their keys - // and decays from this forward's rows, their writes from the log, which this thread wrote itself) + // a branch: its parent's state from the stream's and the ancestors' logged updates, oldest first int anc[MAXW]; int n = 0; for (int a = parent; a >= 0; a = rowpar[a]) anc[n++] = a; @@ -286,8 +269,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q35_gdn_post(Geo geo, const } } -// Per-head RMSNorm of attention rows: X's head h of row m at m * stride + h * hstride, out packed (rows, heads, HD). -// Grid (rows * heads), 64 threads of 4 values. +// Per-head RMSNorm of attention rows, out packed (rows, heads, HD); grid (rows * heads), 64 threads of 4 values. extern "C" __global__ void __launch_bounds__(64) tf_q35_head_norm(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ OUT, int heads, int stride, int hstride, float eps, const int* __restrict__ P) { @@ -325,8 +307,7 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); } -// Queries in place at their positions (the cache length plus the row's depth in its window); keys rotated and values -// copied into their segment's cache at the row's own slot (the cache length plus its row). Grid (rows), HD threads. +// Queries rotated in place, keys rotated and values copied into the row's cache slot; grid (rows), HD threads. extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, @@ -351,10 +332,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo } } -// Causal attention over a segment's cache, P[1] keys a block over P[0] chunks (read from the device, so a captured -// forward replays as caches grow): grid (chunks at most, kvh, rows), block qh / kvh warps, a warp a query head of -// the kv head's group. A row sees the cache before its window and its window's ancestors and itself (rowanc). -// Partials (max, sum, out) at slot (row, qh, chunk); empty chunks leave sum 0. +// Causal attention over the cache in P[0] chunks of P[1] keys, a warp a query head; partials by (row, head, chunk). extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, const int* __restrict__ P, float scale, @@ -454,8 +432,7 @@ extern "C" __global__ void tf_q35_attn_gate(Geo geo, const __nv_bfloat16* __rest OUT[i] = __float2bfloat16_rn(bf(Y[i]) * s); } -// A tapped layer's residual stream, bf16(h + r) as the next norm rounds it, into each row's segment's tap ring at -// the row's position: plane `tap` of the row's taps. Grid (rows), 256 threads. +// A tapped layer's residual bf16(h + r) into the segment's tap ring at the row's position; a 256-thread block a row. extern "C" __global__ void __launch_bounds__(256) tf_q35_tap(const __nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ R, const Seg* __restrict__ segs, const int* __restrict__ rowseg, int tap, int D, const int* __restrict__ P) { diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh index f3baa2be6c..b4f789814f 100644 --- a/zig/kernels/cuda/qwen35.cuh +++ b/zig/kernels/cuda/qwen35.cuh @@ -1,5 +1,4 @@ -// What qwen35.cu and qwen35_keep.cu share: the model's widths (Geo), a stream's window in a forward (Seg), the per-row -// record a verify window leaves in its stream's log, and the walk from a row back through its parents. +// Shared by qwen35.cu and qwen35_keep.cu: the widths (Geo), a stream's window (Seg), a row's log record, parent walks. #pragma once #include @@ -19,8 +18,7 @@ struct Geo { __device__ int conv() const { return (2 * nk + nv) * DK; } // q, k and v channels the convolution sees __device__ size_t rec() const { return (size_t)nv * DK * DK; } // floats of one layer's recurrence __device__ size_t conv_rows() const { return (size_t)(TAPS - 1) * conv(); } - // one row's record of one linear layer: decay and write strength a value head, its write a value (head, dv), its - // key a key head, and the convolution tail after it + // a row's record of a linear layer: decay, strength and write a value head, key a key head, the tail after it __device__ size_t log_bytes() const { return (size_t)nv * 4 + (size_t)nv * DK * 4 + (size_t)nk * DK * 2 + conv_rows() * 2; } }; @@ -58,8 +56,7 @@ struct Path { int rows[MAXW]; }; -// The row `steps` tokens before forward row `w` through its parents (rowpar: each forward row's parent forward row, -1 -// for a window's first row): a forward row, or -1 - t for the stream's convolution tail row t (TAPS - 2 the newest). +// The row `steps` tokens before forward row `w` through rowpar: a forward row, or -1 - t for the stream's tail row t. __device__ __forceinline__ int back(const int* __restrict__ rowpar, int w, int steps) { int a = w; for (int s = 1; s <= steps; s++) { diff --git a/zig/kernels/cuda/qwen35_draft.cu b/zig/kernels/cuda/qwen35_draft.cu index a4726d0711..af2533bc51 100644 --- a/zig/kernels/cuda/qwen35_draft.cu +++ b/zig/kernels/cuda/qwen35_draft.cu @@ -1,6 +1,4 @@ -// The DFlash2 drafter's glue on CUDA (zig/src/families/qwen3_5/cuda_draft.zig): a block of rows attends over the -// target's tapped context, with two-tap dynamic convolutions around attention and MLP. Rounding points follow the -// Metal drafter's draft_ops.metal; the projections run through the target engine's affine-4 kernels. +// DFlash2 drafter glue on CUDA (cuda_draft.zig): a block of rows over the target's taps, draft_ops.metal's rounding. #include #include @@ -46,9 +44,7 @@ extern "C" __global__ void tf_q35d_gather_rows(const __nv_bfloat16* __restrict__ for (int i = threadIdx.x; i < width; i += blockDim.x) dst[i] = src[i]; } -// y[t, c] = sum over lags l of (BASE[branch, l, c] + DYN[t, (branch * 2 + l) * groups + c / group]) * x[t - l, c], -// each product and sum rounded to bf16 in turn; x[-1] is zero (the block's first row has no past); RES, if given, -// is added last. D wide rows, `group` channels a dynamic coefficient. Grid (rows), 256 threads. +// y[t, c] = sum over lags l of (BASE[l, c] + DYN[t, l, c / group]) * x[t - l, c] in bf16 steps; x[-1] = 0; RES last. extern "C" __global__ void tf_q35d_dyn_conv(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ DYN, const __nv_bfloat16* __restrict__ BASE, const __nv_bfloat16* __restrict__ RES, __nv_bfloat16* __restrict__ OUT, int D, int group, int branch) { @@ -69,8 +65,7 @@ extern "C" __global__ void tf_q35d_dyn_conv(const __nv_bfloat16* __restrict__ X, } } -// Per-head RMSNorm over HD then rotate-half RoPE over all HD at position pos0 + row: X's head h of row m at -// m * stride + h * HD, OUT packed (rows, heads, HD). Grid (rows * heads), 32 threads of 4 values. +// Per-head RMSNorm then rotate-half RoPE at position pos0 + row; OUT packed (rows, heads, HD), grid (rows * heads). extern "C" __global__ void __launch_bounds__(32) tf_q35d_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ OUT, int heads, int stride, long long pos0, float eps, float theta) { @@ -115,9 +110,7 @@ extern "C" __global__ void tf_q35d_ring_put(const __nv_bfloat16* __restrict__ K, } } -// The block's attention: query row m, head h sees the context ring's positions [max(begin, p - window + 1), pos0) -// where p = pos0 + m, and every block row (the drafter is not causal within its block). GQA over kv_heads; scale -// 1 / sqrt(HD). Grid (rows, heads), a warp each; online softmax in fp32, bf16 out (rows, heads * HD). +// Block attention: a row sees ring positions [max(begin, p - window + 1), pos0) and every block row; a warp a row. extern "C" __global__ void __launch_bounds__(32) tf_q35d_attn(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, const __nv_bfloat16* __restrict__ RK, const __nv_bfloat16* __restrict__ RV, int rows, int heads, int kv_heads, int W, diff --git a/zig/kernels/cuda/qwen35_keep.cu b/zig/kernels/cuda/qwen35_keep.cu index b2dae0b2b5..20d00fc212 100644 --- a/zig/kernels/cuda/qwen35_keep.cu +++ b/zig/kernels/cuda/qwen35_keep.cu @@ -1,13 +1,10 @@ -// Keeping a verified window's path on CUDA (zig/src/families/qwen3_5/cuda_cache.zig): the stream's recurrences take -// the path's logged updates in order, its convolution tails the path's last row's, and a tree's kept keys move down -// to the slots after the cache. The rows of a window were recorded by qwen35.cu's kernels (Seg.log). +// Keeping a verified window's path (cuda_cache.zig): the recurrences take its logged updates and its keys move down. #include "qwen35.cuh" using namespace tf_q35; -// Each linear layer's recurrence through the path's rows' (decay, key, write), as the forward's chain would have -// left it. Grid (1, DK / 4, nv * nl), block (32, 4): a warp a (layer, value head, value row), its lanes over the key. +// Each linear layer's recurrence through the path's (decay, key, write), as the forward's chain would have left it. extern "C" __global__ void __launch_bounds__(128) tf_q35k_apply(Geo geo, Seg sg, Path path) { const int NK = geo.nk, NV = geo.nv; const int layer = blockIdx.z / NV, hv = blockIdx.z % NV; @@ -43,9 +40,7 @@ extern "C" __global__ void tf_q35k_conv(Geo geo, Seg sg, int last) { CS[i] = log_row(geo, sg, layer, last).conv[i]; } -// The path's rows' keys and values (slot cache length + row) down to the slots cache length + 0, 1, ... in path -// order: row j of the path is at least j, so a move never lands on a slot a later move still reads. One block a -// (attention layer, kv head): grid (attention layers, kvh), HD threads, a barrier between moves. +// The path's keys and values move down to the slots after the cache; row j lands at or below j, never on one unread. extern "C" __global__ void __launch_bounds__(256) tf_q35k_compact(Geo geo, Seg sg, Path path) { const int KVH = geo.kvh; const int layer = blockIdx.x, h = blockIdx.y; diff --git a/zig/src/cli/pull_file.zig b/zig/src/cli/pull_file.zig index 1926f553ff..a5750e3e79 100644 --- a/zig/src/cli/pull_file.zig +++ b/zig/src/cli/pull_file.zig @@ -145,8 +145,7 @@ pub fn linkIntoSnapshot(a: Allocator, io: std.Io, snapshot_dir: []const u8, path std.Io.Dir.cwd().deleteFile(io, link_path) catch {}; const target = try std.fs.path.join(a, &.{ "..", "..", "blobs", blob_name }); std.Io.Dir.cwd().symLink(io, target, link_path, .{}) catch |err| { - // Windows grants symlinks only with Developer Mode or elevation (and refuses some): a hard link to the blob, - // which costs no space (huggingface_hub copies instead); learning refuses linked shards either way + // Windows grants symlinks only with Developer Mode or elevation: a hard link to the blob costs no space if (@import("builtin").os.tag != .windows) { std.log.err("cannot link {s} to its blob: {t}", .{ link_path, err }); return error.SymLinkFailed; diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 038ee35782..c956522f50 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -17,12 +17,10 @@ pub const Facts = struct { lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) }; -/// Target rows [start, start + features.rows) and the token after each; rows at or past `start` replace what was held. -/// `rows`, when given, picks the features' rows to absorb (a tree's kept path), in order; null: all of them. +/// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb. pub const Absorb = struct { stream: *Stream, features: Features, start: u64, follow: []const u32, rows: ?[]const u32 = null }; -/// Draft `depth` tokens after `pending`, the first at `position`; asked only once row `position - 2` was absorbed. -/// With `lanes`, a tree instead: lane i's token is the rank-r choice after its parent (depth unused). +/// `depth` drafts after `pending` at `position`; with `lanes` a tree, lane i the rank-r choice after its parent. pub const Hold = struct { stream: *Stream, pending: u32, position: u64, depth: u32, lanes: ?*const Shape = null }; pub const Drafter = struct { diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index f21778c7bb..435ebe162e 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -1,5 +1,4 @@ -//! The Qwen3.5 hybrids on CUDA (the 2B, Qwen3.8-27B): the engine and its native family entry (the Metal 2B engine is -//! qwen3_5.zig). +//! The Qwen3.5 hybrids on CUDA (2B, Qwen3.8-27B): the engine and its native family entry; Metal's 2B is qwen3_5.zig. pub const config = @import("config.zig"); pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); diff --git a/zig/src/families/qwen3_5/cuda_cache.zig b/zig/src/families/qwen3_5/cuda_cache.zig index 1d2a31d531..b4376c9490 100644 --- a/zig/src/families/qwen3_5/cuda_cache.zig +++ b/zig/src/families/qwen3_5/cuda_cache.zig @@ -1,6 +1,4 @@ -//! One stream's state on the CUDA engine: its recurrences and convolution tails, keys and values, a drafter's tap -//! ring, and the log of its newest verify window's rows, from which `keep` lands any path of that window (a prefix -//! of a chain, or a branch of a tree) without the forward having touched the stream's state. +//! A stream's state on CUDA: recurrences, conv tails, keys and values, the drafter's tap ring, the last window's log. const std = @import("std"); const cuda = @import("cuda"); const geometry = @import("cuda_geometry.zig"); @@ -72,9 +70,7 @@ pub const Cache = struct { return .{ .rec = self.state.ptr, .conv = self.state.ptr + g.deltas() * g.recBytes(), .kv = self.kv.ptr, .tap = if (self.tap) |t| t.ptr else 0, .log = self.log.ptr, .cap = @intCast(self.cap), .len = @intCast(len), .row0 = @intCast(row0), .rows = @intCast(rows), .tap_rows = tap_rows, .taps = @intCast(e.n_taps) }; } - /// Keeps `path` (the window's rows, the first 0, ascending) of the newest window: the recurrences and tails as a - /// chain of those rows would have left them, a tree's kept keys moved down to the slots after the cache, the - /// length at the path's end. Queued on the engine's stream; nothing waits. + /// Keeps `path` (window rows, ascending from 0): the state as that chain would leave it, a tree's keys moved down. pub fn keep(self: *Cache, e: *Engine, path: []const u32) !void { const last = self.last orelse return error.NothingToKeep; if (path.len == 0 or path.len > last.rows or path[0] != 0) return error.NothingToKeep; diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index 9b5a9c1c2c..5df707de3c 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -1,7 +1,4 @@ -//! The DFlash2 drafter on CUDA behind the lane core's Drafter contract (core/lanes/drafter.zig): it absorbs the -//! target's tapped rows into a context ring a stream, then drafts a block of 7 tokens after the pending one, the -//! target's own embedding and head at its ends and a codebook selector scoring the top 16 a row under each parent: -//! one path (a chain), or the lanes of the round loop's tree (each lane its parent's rank-r choice). +//! The DFlash2 drafter on CUDA behind core/lanes/drafter.zig: a tap ring a stream, 7 drafts as a chain or a tree. const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -119,8 +116,7 @@ const Fns = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, dyn /// qwen35_draft.cu's TapRows. const TapRows = extern struct { n: c_int, rows: [engine.most_window]c_int }; -/// A stream's context: its layers' key and value rings over positions [begin, end), and the drafts it holds (a -/// chain, or a tree's lanes in the shape's order). +/// A stream's context: its layers' key and value rings over [begin, end), and the drafts it holds (a chain or tree). const Ctx = struct { ring: DeviceBuffer, begin: u64 = 0, @@ -169,8 +165,7 @@ pub const Draft = struct { ids: u64, // u32 the block's ids staged: cuda.HostBuffer, // pinned: the ids up, the words and projected rows down streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, - /// Contexts between streams: allocated and run through once (`reserve`), since the driver stalls a buffer's - /// first kernel for seconds. + /// Contexts between streams, allocated and run once (`reserve`), as the driver stalls a buffer's first kernel. pool: std.ArrayList(*Ctx) = .empty, /// One level of a hold (a hold drafts every level at once; timed at load, over the block's levels). step_ms: f64 = 1.0, @@ -282,8 +277,7 @@ pub const Draft = struct { d.gpa.destroy(ctx); } - /// A new context, its ring touched by the kernels once: a row absorbed from the first tap ring at hand, then - /// a block drafted over it, so no request pays the driver's stall. + /// A new context, its ring touched once (an absorb, then a draft), so no request pays the driver's stall. fn create(d: *Draft) !*Ctx { const ctx = try d.gpa.create(Ctx); errdefer d.gpa.destroy(ctx); @@ -331,9 +325,7 @@ pub const Draft = struct { try d.e.run(d.fns.norm_rope, .{ rows * heads, 1, 1 }, .{ 32, 1, 1 }, .{ x, w, out, @as(c_int, @intCast(heads)), @as(c_int, @intCast(heads * d.c.head_dim)), @as(c_longlong, @intCast(pos0)), d.c.eps, d.c.theta }); } - /// Target rows at positions [start, start + n) from the stream's tap ring into its context (replacing what was - /// held at or past `start`); rows older than the window are skipped. With `rows`, the ring's rows start + rows[i] - /// (a tree's kept path, which the target keeps at the positions start + i) instead, n of them. + /// Target rows [start, start + n) from the tap ring into the context; with `rows`, ring rows start + rows[i]. fn absorb(d: *Draft, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32) !void { const c = d.c; var start = start_; @@ -370,8 +362,7 @@ pub const Draft = struct { if (ctx.end > c.window and ctx.begin < ctx.end - c.window) ctx.begin = ctx.end - c.window; } - /// The block [pending, mask x 7] at positions pos0.. over the stream's context: `depth` drafts held as a chain, - /// or the `lanes` tree's tokens. + /// The block [pending, mask x 7] at pos0.. over the context: `depth` drafts as a chain, or the `lanes` tree. fn hold(d: *Draft, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { const c = d.c; const e = d.e; @@ -423,8 +414,7 @@ pub const Draft = struct { } } - /// Row `r`'s candidate of `rank` after the token `pred`: by unary logit + predecessor . projected . successor, - /// ties to the earlier candidate (the row's top-k order). + /// Row `r`'s `rank`th candidate after `pred`: unary logit + pred . projected . successor, ties to the earlier one. fn choice(d: *const Draft, r: usize, pred: u32, rank: usize) u32 { const c = d.c; const w = std.mem.bytesAsSlice(u32, d.staged.bytes[words_at..][r * lp_words * 4 ..][0 .. lp_words * 4]); diff --git a/zig/src/families/qwen3_5/cuda_draft_weights.zig b/zig/src/families/qwen3_5/cuda_draft_weights.zig index 2b0cf9f350..6d614f7982 100644 --- a/zig/src/families/qwen3_5/cuda_draft_weights.zig +++ b/zig/src/families/qwen3_5/cuda_draft_weights.zig @@ -1,6 +1,4 @@ -//! The DFlash2 drafter's checkpoint onto the device: its bf16 projections quantized to affine-4 (group 64) on the -//! host and packed into qlinear's tiles through the target engine's staging buffer, its norms and base kernels as -//! they are, its codebooks kept on the host for the selector. +//! The DFlash2 drafter's weights onto the device: projections quantized to affine-4 on the host, codebooks kept there. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -82,8 +80,7 @@ fn toBf16(x: f32) u16 { return @intCast((u +% rounding) >> 16); } -/// MLX affine-4, group 64, over bf16 rows: a group's scale is bf16((max - min) / 15) and bias bf16(min); each weight's -/// nibble is round((w - bias) / scale) clamped to [0, 15], eight to a word, the first in the low bits. +/// MLX affine-4, group 64: scale bf16((max - min) / 15), bias bf16(min), nibbles round((w - bias) / scale), low first. fn quantize(gpa: std.mem.Allocator, w: []align(1) const u16, n: usize, k: usize) !struct { words: []u32, scales: []u16, biases: []u16 } { const groups = k / 64; const words = try gpa.alloc(u32, n * k / 8); diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index be59726fc6..c9cf79d096 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -1,9 +1,4 @@ -//! Qwen3.5's hybrid text models on CUDA (the 2B, Qwen3.8-27B): the affine-4 projections through cuda.qlinear, the -//! glue in qwen35.cu, one forward for any mix of streams' windows. Each stream owns its recurrences, convolution tails -//! and KV cache (cuda_cache.zig); the widths come from config.json (cuda_geometry.zig), the weights from -//! cuda_weights.zig. A forward's launches are captured into a CUDA graph the first time its shape is seen and -//! replayed after. A verify window's rows may form a tree (each row's parent given): the forward computes every -//! branch, logs each row, and the stream's state stays as it was until a path is kept. +//! Qwen3.5's hybrids on CUDA: qlinear projections, qwen35.cu glue, one captured forward over every stream's window. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -73,9 +68,7 @@ pub const Keeps = struct { apply: cuda.Function, conv: cuda.Function, compact: c /// A window of one stream in a forward: `parents`, for a tree, names each row's parent row (-1: the first row). pub const Part = struct { cache: *Cache, rows: usize, parents: ?[]const i32 = null }; -/// What a forward keeps. A prompt chunk (`last`, one window) puts its last row's logits in its cache's `logits`; -/// verify windows keep every row, drawn on the device into `tokens` (`draws`, one a row) or their logits for -/// `rowLogits` (`logits`), and log each row for `Cache.keep`. +/// What a forward keeps: a prompt chunk its last row's logits, a verify window every row's draw or logits and log. pub const Keep = union(enum) { last, draws: []const Draw, logits }; /// A kept row's draw on the device: sample.cu's keyed rule at `position`, or greedy without one. @@ -92,8 +85,7 @@ pub const Engine = struct { affine: cuda.qlinear.Affine4Kernels, fns: Fns, keeps: Keeps, - /// Every weight and scratch buffer in one allocation: the WDDM driver's launch cost grows with a process's - /// allocation count, and a 64-layer model would have thousands. + /// Every weight and scratch buffer in one allocation: WDDM's launch cost grows with the allocation count. arena: DeviceBuffer, arena_at: usize = 0, stage: DeviceBuffer, // the load's staging of each projection as stored, packed from there @@ -135,8 +127,7 @@ pub const Engine = struct { pl: u64, po: u64, logits: u64, - // the per-forward tables: ids, segments, each row's segment, parent and ancestors, the draws' rules, staged in - // pinned memory + // per-forward tables (ids, segments, rows' segment, parent and ancestors, the draws' rules) staged in pinned memory tables: DeviceBuffer, staged: cuda.HostBuffer, host_logits: cuda.HostBuffer, @@ -213,8 +204,7 @@ pub const Engine = struct { return e; } - /// Every path a request takes, once now: the WDDM driver stalls a process's first forwards for seconds (a 16 GiB - /// model: a 20 s launch, then a 3 s one), which no request should pay. + /// Every path a request takes, run once now: WDDM stalls a process's first forwards for seconds (a 20 s launch). fn warm(e: *Engine, gpa: std.mem.Allocator) !void { var warm_cache = try Cache.init(gpa, e, 64); defer warm_cache.deinit(gpa); @@ -312,8 +302,7 @@ pub const Engine = struct { try cuda.launch.launch(f, cfg, e.stream, &a); } - /// out (rows, n) = x (rows, k) through `l`: decode tiles 16 rows at a time (each row's bits whatever the batch), - /// or the prompt GEMM for a prompt chunk wider than one tile. + /// out (rows, n) = x (rows, k) through `l`: 16-row decode tiles (a row's bits batch-independent) or a prompt GEMM. pub fn proj(e: *Engine, l: Affine4, x: u64, out: u64, rows: usize, prompt: bool) !void { if (prompt and rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, e.xs, @as(c_int, @intCast(l.k)), e.tables.ptr + params_at }); @@ -328,8 +317,7 @@ pub const Engine = struct { try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); } - /// The row tables of `parts`: each forward row's segment, parent forward row (-1: its window's first) and the - /// ancestors-and-self bits over its window's rows (a chain without `parents`). + /// The row tables of `parts`: each row's segment, parent row (-1: a window's first) and ancestors-and-self bits. fn rowTables(st: []u8, parts: []const Part) !void { const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. max_rows * 4]); const rowpar = std.mem.bytesAsSlice(c_int, st[rowpar_at..][0 .. max_rows * 4]); @@ -348,8 +336,7 @@ pub const Engine = struct { } } - /// One forward over `parts`' windows, the rows of `ids` in order. `keep` names the rows it keeps and how. A - /// window not yet kept from the stream's last verify is kept whole first. + /// One forward over `parts`' windows, `ids` in order; `keep` names the kept rows; an unkept verify is kept first. pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { const g = e.g; var total: usize = 0; @@ -397,15 +384,13 @@ pub const Engine = struct { params[3] = @intCast(parts.len); params[4] = @intFromBool(record); try e.tables.uploadAsync(0, st, e.stream.handle); - // a captured forward's launches cover a bucket of rows and every window; its kernels read the true counts - // (a lone stream's one-row step, the common case without a drafter, keeps its own exact graph) + // a captured forward covers a row bucket and every window; kernels read the true counts; a lone step is exact const lone = total == 1 and parts.len == 1; const bucket: usize = if (lone) 1 else if (total <= 16) 16 else most_verify; const windows: usize = if (lone) 1 else max_segs; const shape: Shape = .{ .total = @intCast(if (prompt) total else bucket), .nseg = @intCast(if (prompt) parts.len else windows), .record = record, .prompt = prompt, .taps = e.n_taps > 0 }; if (prompt) { - // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length would be a graph of its own, - // paying the driver's stall on a new graph's first run (about a second on a wide model) + // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length's graph would pay the stall try e.body(shape); } else { const gop = try e.graphs.getOrPut(e.gpa, shape.key()); diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig index 7b327fd1a9..fdd6ecb0c2 100644 --- a/zig/src/families/qwen3_5/cuda_geometry.zig +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -1,5 +1,4 @@ -//! The widths of a Qwen3.5 hybrid text model (the 2B, Qwen3.8-27B) as the CUDA engine reads them from config.json: -//! what varies is read, what qwen35.cu fixes (head widths, taps, rotary layout, 4-bit groups of 64) is checked. +//! A Qwen3.5 hybrid's widths (2B, 27B) as the CUDA engine reads them from config.json; what qwen35.cu fixes is checked. const std = @import("std"); pub const max_layers = 128; @@ -70,8 +69,7 @@ pub const Geometry = struct { return g.deltas() * (g.recBytes() + g.convBytes()); } - /// One verified row's record of one linear layer (qwen35.cuh's log_bytes): decay and write by value head, - /// key by key head, the tail after it. + /// A row's record of a linear layer (qwen35.cuh log_bytes): decay and write a value head, key a key head, the tail. pub fn logRowBytes(g: Geometry) usize { return g.v_heads * 4 + g.v_heads * dk * 4 + g.k_heads * dk * 2 + g.convBytes(); } diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index d4fa634cb1..161b34a882 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -1,6 +1,4 @@ -//! The Qwen3.5 hybrids' CUDA lane backend (the 2B, Qwen3.8-27B) behind the native family interface (native/cuda.zig's -//! registry): windows of up to 16 rows (the lane core's copies), every stream's window in one forward, prompts in -//! chunks, draws on the device. +//! The Qwen3.5 hybrids' CUDA lane backend (native/cuda.zig's registry): 16-row windows, every stream in one forward. const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -30,8 +28,7 @@ pub const Loaded = struct { free_streams: u32 = 0, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, - /// Called with the admitted stream count: that many caches allocated and touched now, since the driver commits - /// memory on first touch (seconds for a wide model's keys), and a request should not pay it. + /// Called with the admitted stream count: that many caches allocated and touched now, as memory commits on touch. reserve: *const fn (*anyopaque, u32) anyerror!void, }; @@ -66,8 +63,7 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { try cache.reset(l.e); l.pool.appendAssumeCapacity(cache); } - // the shared rounds' shapes, a row of each of 1 to `streams` streams, captured now: a new graph's first run - // pays the driver's stall too (about a second on a wide model) + // the shared rounds' shapes, a row of 1 to `streams` streams, captured now: a new graph's first run stalls too var parts: [E.max_segs]E.Part = undefined; var ids: [E.max_segs]u32 = @splat(0); var picks: [E.max_segs]E.Draw = undefined; @@ -85,8 +81,7 @@ fn reserve(p: *anyopaque, streams: u32) anyerror!void { }; } -/// The DFlash2 drafter in `dir` over the loaded engine `p`, its weights in 4 bits (the only width served). The -/// target's windows of 1 to window_rows rows are timed here (with the drafter's taps on) for the round loop's prices. +/// 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 } { if (bits != 4) return error.DrafterBitsUnsupported; const l: *Lanes = @ptrCast(@alignCast(p)); @@ -169,8 +164,7 @@ const Lanes = struct { self.pool.append(self.gpa, cache) catch self.drop(cache); } - /// Each window width's forward, timed on a cache of its own (its third run: the shape's graph captured and the - /// driver's stall behind it, the previous window's keep in front as in a round), into the drafter's facts. + /// Each width's forward timed on its own cache, on its third run: graph captured, stall paid, a keep in front. fn timeWindows(self: *Lanes, d: *Draft) !void { var cache = try E.Cache.init(self.gpa, self.e, 64); defer cache.deinit(self.gpa); @@ -249,8 +243,7 @@ const Lanes = struct { return @intCast(handle); } - /// Each window (a stream's pending token and drafts, a chain or a tree) checked, then as many as fit in each - /// forward. A tree's rows are keyed at the cache length plus their depth, as the round loop positions them. + /// Each window (a pending token and drafts) checked, then as many as fit a forward; a tree's rows keyed by depth. fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { const self = cast(ptr); const t0 = std.Io.Clock.awake.now(self.io); @@ -326,9 +319,7 @@ const Lanes = struct { try self.e.setTaps(taps); } - /// Rows [start, start + count) of a stream's tap ring: `offset` carries the first row's position, the drafter - /// (cuda_draft.zig) reads slot position % tap_rows, so a request may span the ring's end. The newest window's - /// rows stay readable after a keep (a tree's kept path reads them by row) until the next forward. + /// Rows [start, start + count) of a stream's tap ring at slot position % tap_rows; a request may span its end. fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { const self = cast(ptr); if (taps.len != self.e.n_taps) return error.TapsMismatch; diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig index a105b876f4..2a2b75c361 100644 --- a/zig/src/families/qwen3_5/cuda_weights.zig +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -1,5 +1,4 @@ -//! The checkpoint's weights onto the device: each affine-4 projection staged as stored and packed into qlinear's -//! tiles, the vectors as they are, all from the engine's one arena (cuda_engine.zig). +//! The checkpoint's weights onto the device: affine-4 projections packed as qlinear tiles, out of the engine's arena. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); diff --git a/zig/src/platform/sys.zig b/zig/src/platform/sys.zig index 35bb853608..204eeb3d4a 100644 --- a/zig/src/platform/sys.zig +++ b/zig/src/platform/sys.zig @@ -1,5 +1,4 @@ -//! The operating system as TensorFold calls it: `c` is libc's surface (std.c on POSIX, win_c.zig on Windows) and -//! `DynLib` opens shared libraries (std.DynLib on POSIX, LoadLibraryW on Windows). Every module imports this as `sys`. +//! The OS as TensorFold calls it, imported as `sys`: `c` is std.c or win_c.zig, `DynLib` is std.DynLib or LoadLibraryW. const std = @import("std"); const builtin = @import("builtin"); diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig index eb4d0c746a..21e057d6f5 100644 --- a/zig/src/platform/win_c.zig +++ b/zig/src/platform/win_c.zig @@ -1,6 +1,4 @@ -//! The libc surface TensorFold calls through `sys.c`, on Windows: integer descriptors over the UCRT (files opened by -//! CreateFileW with every share mode, so unlink and rename behave as on POSIX), directory descriptors as remembered -//! paths, readdir over FindFirstFileW, and POSIX's names for the rest. Each call keeps libc's contract: -1 and errno. +//! libc's surface on Windows (`sys.c`): descriptors over the UCRT, CreateFileW with every share mode, -1 and errno. const std = @import("std"); @@ -165,8 +163,7 @@ fn lastError() E { const Wide = [32767:0]u16; -/// UTF-8 path to a NUL-terminated wide path in `buf`; null (errno set) when it does not fit or is not UTF-8. -/// POSIX's /tmp is the user's temporary folder and /proc/self/exe this executable, as Cygwin and MSYS map them. +/// UTF-8 path as a wide path in `buf`, null with errno if it does not fit; /tmp and /proc/self/exe map as in Cygwin. fn wide(buf: *Wide, path: []const u8) ?[*:0]const u16 { var at: usize = 0; var rest = path; diff --git a/zig/src/platform/win_clock.zig b/zig/src/platform/win_clock.zig index ec41325cfc..0c0c67ab4e 100644 --- a/zig/src/platform/win_clock.zig +++ b/zig/src/platform/win_clock.zig @@ -1,5 +1,4 @@ -//! The process, environment and clock part of the libc surface on Windows (win_c.zig re-exports these): the process -//! id, environment writes over the UCRT, and nanosleep and clock_gettime over the kernel's timers and counters. +//! Process, environment and clock calls of the Windows libc surface, re-exported by win_c.zig: UCRT and kernel timers. const std = @import("std"); const c = @import("win_c.zig"); diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index 04c82d4dba..c0ab151e25 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -1,6 +1,4 @@ -//! The std.posix surface the HTTP server calls through `sys.posix`, on Windows: Winsock sockets (SOCKET as an isize, -//! INVALID_SOCKET as -1), WSAPoll, console control events as SIGINT and SIGTERM, and no descriptor limits to raise. -//! Calls in `system` return -1 with errno set from WSAGetLastError, as libc's do. +//! std.posix as the server calls it, on Windows: Winsock sockets, WSAPoll, console events as SIGINT and SIGTERM. const std = @import("std"); const c = @import("win_c.zig"); @@ -29,8 +27,7 @@ pub const SOL = struct { pub const SOCKET: i32 = 0xffff; }; pub const SO = struct { - /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, which the listener - /// sets REUSEADDR to prevent; exclusive use is the POSIX meaning. + /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, the opposite of POSIX's. pub const REUSEADDR: u32 = @bitCast(@as(i32, -5)); }; pub const SHUT = struct { diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 858dc34c07..0d1f711252 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,8 +1,4 @@ -//! `tf-qwen35-cuda `: a Qwen3.5 hybrid's CUDA forward over -//! comma-separated token ids: one token a forward (step: every position's logits), windows of growing widths -//! (window: the same rows, which must equal step's), tree windows with junk branches kept to the true one (tree: -//! the same), or the ids as one prompt (prompt: the last position's), as little-endian bf16 rows into for a -//! reference to compare. Step mode also prints the decode speed. +//! `tf-qwen35-cuda `: the ids' logits as bf16 rows for a reference. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -51,8 +47,7 @@ pub fn main(init: std.process.Init) !u8 { at += n; } } else if (std.mem.eql(u8, args[2], "tree")) { - // tree windows of growing widths: the true continuation down one branch, a junk sibling before each of its - // rows and a junk grandchild, the true branch kept; its rows' logits must equal step's bit for bit + // tree windows of growing widths, junk branches beside the kept true one, whose logits must equal step's bits var at: usize = 0; var w: usize = 2; while (at < ids.items.len) : (w = w % e.window_rows + 1) { From 8107ee2b90c92f3d038c1db8e22caa0f20f0ca23 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:37 +0000 Subject: [PATCH 13/35] Windows: checkpoint reads skip the page cache direct_io opened files with O_DIRECT on Linux alone. On Windows the same open now asks for FILE_FLAG_NO_BUFFERING (which win_c maps O_DIRECT to), and the 4 KiB-aligned spans it already reads satisfy the sector alignment that needs; a file system that refuses it still falls back to buffered reads. The direct-read test runs on Windows too, checking the executable's bytes against buffered reads of it. --- docs/WINDOWS.md | 5 +++-- zig/src/core/direct_io.zig | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index b11d72f4fd..db0a6040e1 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -65,8 +65,9 @@ tensorfold-native serve --name qwen27b --context 8192 --p - `zig/src/platform/sys.zig` (`@import("sys")`, wired into every module by `zig/build/sys.zig`) is the operating system as the engine calls it. On Linux and macOS it is `std.c`, `std.posix` and `std.DynLib` unchanged; on Windows `win_c.zig` and `win_posix.zig` give the same names over the UCRT, Win32 and Winsock. -- Files are opened with every share mode, so rename-over and unlink-while-open behave as on POSIX. `O_DIRECT` maps to - unbuffered reads. `/tmp` is the user's temporary folder and `HOME` falls back to `USERPROFILE`. +- Files are opened with every share mode, so rename-over and unlink-while-open behave as on POSIX. Checkpoint reads + skip the page cache (`FILE_FLAG_NO_BUFFERING`, as `O_DIRECT` on Linux). `/tmp` is the user's temporary folder and + `HOME` falls back to `USERPROFILE`. - The Hugging Face cache links snapshot files to blobs. Windows allows symlinks only with Developer Mode or elevation, so `pull` falls back to hard links (no extra space). - Checkpoint reads use threads (io_uring is Linux's). diff --git a/zig/src/core/direct_io.zig b/zig/src/core/direct_io.zig index 2f4cc546f6..f5bba757e2 100644 --- a/zig/src/core/direct_io.zig +++ b/zig/src/core/direct_io.zig @@ -26,7 +26,7 @@ pub const File = struct { /// `path` read-only, direct when its file system allows it, else through the page cache. pub fn open(path: [:0]const u8) !File { - if (builtin.os.tag == .linux) { + if (builtin.os.tag == .linux or builtin.os.tag == .windows) { const fd = platform.c.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }); if (fd >= 0) return .{ .fd = fd, .direct = true }; } @@ -85,7 +85,7 @@ test "span and fits" { } test "direct reads return the same bytes as buffered reads" { - if (builtin.os.tag != .linux) return error.SkipZigTest; + if (builtin.os.tag != .linux and builtin.os.tag != .windows) return error.SkipZigTest; const gpa = std.testing.allocator; var f = try File.open("/proc/self/exe"); defer f.close(); From 693533e9d4d80ac3642bc5b89bebcbda2d70c9f5 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:37 +0000 Subject: [PATCH 14/35] Windows: the signal handler type allows address 1, for aarch64 SIG.IGN is a function pointer made from 1, as std.posix has it. aarch64 requires function pointers to be aligned unless the type says align(1), so the cross-build for aarch64-windows refused it. With align(1), as std's own handler type has, tensorfold-native cross-compiles for aarch64-windows with the GB10 (sm_121) fatbins; no hardware has run it. --- zig/src/platform/win_posix.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index c0ab151e25..c154eb7153 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -203,7 +203,7 @@ pub const SIG = enum(c_int) { pub const IGN: ?Handler = @ptrFromInt(1); pub const DFL: ?Handler = null; }; -pub const Handler = *const fn (SIG) callconv(.c) void; +pub const Handler = *align(1) const fn (SIG) callconv(.c) void; // align(1): IGN is address 1, as std.posix has it pub const sigset_t = u32; pub const Sigaction = struct { handler: extern union { handler: ?Handler, sigaction: ?*const anyopaque }, From f53ae20aba85a7043e90dd5dce23511c48d69f90 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:38 +0000 Subject: [PATCH 15/35] docs: Windows as a target in the README and the Windows page The Windows page opens with what the build is, what it is tested on and that it cross-compiles for aarch64-windows; the README's model table and the CUDA paragraph say Windows where they said unofficial port. --- README.md | 6 +++--- docs/WINDOWS.md | 5 +++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index c838faa144..48c7b8d27f 100644 --- a/README.md +++ b/README.md @@ -73,8 +73,8 @@ The platform column names the hardware tested for each model. | Qwen3.8 Flash Next | MLX affine 6-bit, group 32 | Metal on M5 Ultra | | GLM-5.3-Flash | MLX affine 4-bit, group 64 | Metal on two M5 Ultras | | GLM-5.3 | EXLR8 `home` profile (trellis K2/K3/K4 experts) | Metal on two M5 Ultras | -| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max; CUDA on RTX 4090 (unofficial Windows port) | -| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (unofficial Windows port) | +| Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max; CUDA on RTX 4090 (Windows) | +| Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (Windows) | Nemotron's named checkpoint is `TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`. NVIDIA's NVFP4 checkpoint, `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4`, serves on CUDA as stored; its bf16 MTP head is quantized to 4-bit at load for drafts. @@ -90,7 +90,7 @@ Its prompt processing trails mlx_lm for now, and we are fixing it. An M3 Ultra r Bonsai, Gemma 4, Qwen3.6 and DeepSeek-V4 are still under qualification for 1.0.x. The Python 0.6.6 engine remains available at the `v0.6.6` tag for those models and other backends. On CUDA, 1.0.11 serves Nemotron on a GB10, greedy and sampled, with concurrent requests sharing each round. -The unofficial Windows port ([docs/WINDOWS.md](docs/WINDOWS.md)) also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees. +On Windows ([docs/WINDOWS.md](docs/WINDOWS.md)), the native build also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees. The Linux x86_64 archive adds NVIDIA Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), tested on an RTX 3090, and Blackwell cards with compute capability 12.0 (the RTX 50 series and RTX PRO 6000 Blackwell), tested on an RTX PRO 6000; the A100 (8.0) is not supported yet. ## Exact decoding diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index db0a6040e1..52b17191c6 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -1,7 +1,8 @@ # TensorFold on Windows (native) -An unofficial port: `tensorfold-native.exe` runs on Windows x86_64 with an NVIDIA GPU through the same CUDA engine -the Linux build uses. Only the CUDA backend is ported; Metal is macOS's and HIP is not built. +`tensorfold-native.exe` runs on Windows with an NVIDIA GPU through the same CUDA engine the Linux build uses. Only the +CUDA backend is built on Windows; Metal is macOS's and HIP is not built. It is tested on an RTX 4090 (x86_64); the +host code also cross-compiles for `aarch64-windows` with the GB10 (sm_121) kernels, untested on hardware so far. ## What you need From e875ee2ac3e2fd3e742f67074cc8441de915afad Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:17:38 +0000 Subject: [PATCH 16/35] CI: the host unit tests on Windows The no-GPU workflow runs zig build test on windows-latest beside ubuntu-latest, so the platform layer is built and tested on every pull request. --- .github/workflows/no-gpu.yml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 6acdde321e..09dcde39d9 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -38,3 +38,13 @@ jobs: (cd ../lean-base && python3 tools/zig/lean_check.py > ../lean-base.txt || true) python3 tools/zig/lean_check.py > ../lean-head.txt || true python3 tools/zig/lean_new.py ../lean-base.txt ../lean-head.txt + zig-windows: + runs-on: windows-latest + steps: + - uses: actions/checkout@v4 + - uses: mlugg/setup-zig@v2 + with: + version: 0.17.0 + # The same host unit tests on Windows: the platform layer (zig/src/platform) over the UCRT, Win32 and Winsock. + - name: zig build test + run: zig build test From 38ecce53f7d3f678c29c313a8ee88100adf19708 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:21:56 +0000 Subject: [PATCH 17/35] qwen3_5: the drafter reads TF_TIMING through the platform layer Every other module calls libc through `sys`; the drafter's four getenv calls were the last direct std.c calls in a file the Windows build compiles. --- zig/src/families/qwen3_5/cuda_draft.zig | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index 5df707de3c..cfb8670cf5 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -1,5 +1,6 @@ //! The DFlash2 drafter on CUDA behind core/lanes/drafter.zig: a tap ring a stream, 7 drafts as a chain or a tree. const std = @import("std"); +const platform = @import("sys"); const cuda = @import("cuda"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); @@ -219,7 +220,7 @@ pub const Draft = struct { /// The target's window of `rows` rows takes `ms` (told by cuda_native.zig before the facts are read). pub fn setCost(d: *Draft, rows: u32, ms: f64) void { - if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a hold level {d:.2})\n", .{ rows, ms, d.step_ms }); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a hold level {d:.2})\n", .{ rows, ms, d.step_ms }); if (d.n_costs < d.costs.len) { d.costs[d.n_costs] = .{ .width = rows, .ms = ms }; d.n_costs += 1; @@ -406,11 +407,11 @@ pub const Draft = struct { try e.stream.synchronize(); if (lanes_) |l| { try d.selectLanes(ctx, pending, rows, l); - if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ l.parents.len, l.levels() }); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ l.parents.len, l.levels() }); } else { d.select(ctx, pending, rows); ctx.held_n = @min(depth, rows); - if (std.c.getenv("TF_TIMING") != null) std.debug.print("[timing] chain {d}\n", .{ctx.held_n}); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] chain {d}\n", .{ctx.held_n}); } } @@ -479,7 +480,7 @@ pub const Draft = struct { fn factsFn(p: *anyopaque) dr.Facts { const d = cast(p); // TF_NO_TREES=1: chains only (for comparing) - return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes(), .trees = std.c.getenv("TF_NO_TREES") == null, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; + return .{ .depth = @intCast(d.c.block - 1), .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.c.ringBytes(), .trees = platform.c.getenv("TF_NO_TREES") == null, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; } fn openFn(p: *anyopaque, s: *Stream) anyerror!void { @@ -539,7 +540,7 @@ fn bf16(h: u16) f32 { /// TF_TIMING=1: each call's wall time on stderr. pub fn timing(what: []const u8, io: std.Io, t0: std.Io.Timestamp) void { - if (std.c.getenv("TF_TIMING") == null) return; + if (platform.c.getenv("TF_TIMING") == null) return; const now = std.Io.Clock.awake.now(io); const ns: f64 = @floatFromInt(t0.durationTo(now).toNanoseconds()); const at: f64 = @floatFromInt(@mod(now.toNanoseconds(), 1_000_000_000_000)); From 181782342dbfbd4a508eec3ce77a39e2c36136e3 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:33:03 +0000 Subject: [PATCH 18/35] qwen3_5: the drafter takes a prompt longer than its tap ring After a prompt the lane core hands the drafter every prompt row's features; the CUDA backend refused any request wider than its tap ring (2,176 rows), so a 27B request with --drafter and a prompt past that ended with RowsNotHeld, where the same prompt without the drafter ran. The ring keeps the newest tap_rows rows, which hold the drafter's whole context window, and absorb already reads only that window of what it is handed; the backend now refuses only rows not yet written or all gone. On an RTX 4090 the 27B with its drafter takes a 6,041-token prompt in 7.4 s, and its reply equals the one with "draft": false. --- zig/src/families/qwen3_5/cuda_native.zig | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 161b34a882..8773dffc71 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -319,13 +319,15 @@ const Lanes = struct { try self.e.setTaps(taps); } - /// Rows [start, start + count) of a stream's tap ring at slot position % tap_rows; a request may span its end. + /// Rows [start, start + count) of a stream's tap ring (slot position % tap_rows), the newest tap_rows rows kept. fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { const self = cast(ptr); if (taps.len != self.e.n_taps) return error.TapsMismatch; const cache = try self.cacheOf(stream); const ring = cache.tap orelse return error.NoTaps; - if (start + count > cache.written or count > E.tap_rows) return error.RowsNotHeld; + // a long prompt asks for more rows than the ring keeps; the drafter reads only its window of the newest + const held_from = cache.written -| E.tap_rows; + if (start + count > cache.written or start + count <= held_from) return error.RowsNotHeld; return .{ .buffer = ring.ptr, .offset = start, .rows = count, .row_bytes = @intCast(self.e.n_taps * self.e.g.hidden * 2), .space = .device, .dtype = .bf16 }; } From b811a3b3f316c223fae6288b47418962b4b09f42 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 16:16:38 +0000 Subject: [PATCH 19/35] qwen3_5 cuda: prompt chunks past 32 rows see the right keys, and tree attention cuts chunks at logical positions - rowanc held a window's rows as bits of a u32, but prompt chunks run up to 128 rows: rows from 32 on rotated at the wrong positions and attended the wrong keys. A chain row (no parents) now carries its index in its window under the top bit; tree windows (at most 16 rows) keep their bits. - tf_q35_attn_partial cut its 128-key chunks at a key's cache row, so a kept tree path whose window crossed a chunk boundary summed in other chunks than one-row steps and changed bits. Chunks are now cut at the key's logical position (the stream's length plus the ancestors before it). Checked on an RTX 5090 (Linux, sm_120), tf-qwen35-cuda over 300 ids on Qwen3.5-2B and Qwen3.8-27B: step bits unchanged, window and tree equal step byte for byte (tree differed from row 126 on the 2B and 254 on the 27B before). Prompt mode's last row against step's at 16 to 300 ids: KL at most 0.003 and the same argmax in 20 of 20 (before: 27B at 40 ids KL 2.3 with a different argmax, at 128 ids KL 1.35). --- zig/kernels/cuda/qwen35.cu | 19 +++++++++++++++---- zig/src/families/qwen3_5/cuda_engine.zig | 5 ++++- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index cef7c7b5a2..558baa7908 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -307,6 +307,16 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i return __float2bfloat16_rn(d < 32 ? a * cs - b * sn : b * cs + a * sn); } +// A row's ancestors-and-self: its tree window's row bits, or with the top bit set a chain row's index in its window. +__device__ __forceinline__ int rows_seen(unsigned anc) { return (anc >> 31) ? (int)(anc & 0x7fffffffu) + 1 : __popc(anc); } + +// The window row holding a row's nth seen window key. +__device__ __forceinline__ int seen_row(unsigned anc, int n) { + if (anc >> 31) return n; + for (int j = 0; j < n; j++) anc &= anc - 1; + return __ffs(anc) - 1; +} + // Queries rotated in place, keys rotated and values copied into the row's cache slot; grid (rows), HD threads. extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bfloat16* __restrict__ Qn, const __nv_bfloat16* __restrict__ Kn, const __nv_bfloat16* __restrict__ Vr, const Seg* __restrict__ segs, @@ -317,7 +327,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo if (past_rows(P, blockIdx.x)) return; const size_t m = blockIdx.x; const Seg sg = segs[rowseg[m]]; - const int pos = sg.len + __popc(rowanc[m]) - 1; + const int pos = sg.len + rows_seen(rowanc[m]) - 1; const int slot = sg.len + (int)m - sg.row0; __shared__ __nv_bfloat16 qs[MAXQ * HD]; for (int h = 0; h < QH; h++) qs[h * HD + d] = Qn[(m * QH + h) * HD + d]; @@ -347,7 +357,8 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; const unsigned anc = rowanc[m]; - const int k0 = c * ck, k1 = min(k0 + ck, sg.len + sg.rows); + // chunks cut at logical positions (a window key's: the stream's length plus its ancestors), as a step's are + const int k0 = c * ck, k1 = min(k0 + ck, sg.len + rows_seen(anc)); const size_t slot = (m * QH + qh) * nch + c; if (k0 >= k1) { if (lane == 0) { @@ -365,8 +376,8 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (int pos = k0; pos < k1; pos++) { - if (pos >= sg.len && ((anc >> (pos - sg.len)) & 1u) == 0) continue; + for (int t = k0; t < k1; t++) { + const int pos = t < sg.len ? t : sg.len + seen_row(anc, t - sg.len); const uint4 kw = *reinterpret_cast(K + ((size_t)h * sg.cap + pos) * HD + lane * DPL); const __nv_bfloat16* kb = reinterpret_cast(&kw); float dot = 0.0f; diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index c9cf79d096..389b32f766 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -317,6 +317,9 @@ pub const Engine = struct { try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, if (residual) e.r else @as(u64, 0), w, e.x, e.g.eps, @as(c_int, @intCast(e.g.hidden)), e.tables.ptr + params_at }); } + /// A chain row's rowanc: its index in its window, as a prompt chunk's rows outnumber the bits. + const chain_row: u32 = 1 << 31; + /// The row tables of `parts`: each row's segment, parent row (-1: a window's first) and ancestors-and-self bits. fn rowTables(st: []u8, parts: []const Part) !void { const rowseg = std.mem.bytesAsSlice(c_int, st[rowseg_at..][0 .. max_rows * 4]); @@ -330,7 +333,7 @@ pub const Engine = struct { if (parent >= @as(i32, @intCast(r)) or parent < -1) return error.InvalidQwenParents; rowseg[row0 + r] = @intCast(i); rowpar[row0 + r] = if (parent < 0) -1 else @intCast(row0 + @as(usize, @intCast(parent))); - rowanc[row0 + r] = (@as(u32, 1) << @intCast(r)) | (if (parent < 0) 0 else rowanc[row0 + @as(usize, @intCast(parent))]); + rowanc[row0 + r] = if (p.parents == null) chain_row | @as(u32, @intCast(r)) else (@as(u32, 1) << @intCast(r)) | (if (parent < 0) 0 else rowanc[row0 + @as(usize, @intCast(parent))]); } row0 += p.rows; } From 0a0856ba97fd704b55ea39975821d76c0edcf8b2 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 17:47:32 +0000 Subject: [PATCH 20/35] Windows: the server goldens and the Cortheon fuzz step on Windows hosts zig build test-golden and zig build fuzz-cortheon exist on every non-macOS host since the golden step moved into zig/build/golden.zig, so the golden check compiles against the platform layer on Windows: - the platform's posix gains sockaddr.in, timeval, connect and kill for a child's handle (TerminateProcess; HUP has no equivalent), setsockopt turns a timeval for SO_RCVTIMEO or SO_SNDTIMEO into Winsock's milliseconds, and a socket read that times out is WouldBlock, as std.posix reports it. - the check carries SystemRoot into the spawned server's environment, which Winsock needs to start, and drains the server's stdout after the port line on every host: the fake server logs each request there, and a full pipe (4 KiB on Windows) blocked it after sixteen cases. - on Windows the check cannot send SIGTERM or SIGHUP, so the sigterm-exit-0 outcome and the cases after a key rewrite are skipped and printed as such. On Windows the goldens give 306 of 315 equal with the same nine known differences as Linux, and the fuzz step passes its eight tests; the no-GPU workflow's Windows job runs both. --- .github/workflows/no-gpu.yml | 5 ++++ zig/src/platform/win_posix.zig | 39 +++++++++++++++++++++++++++---- zig/tests/server/golden_check.zig | 33 +++++++++++++++++++++++++- 3 files changed, 72 insertions(+), 5 deletions(-) diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 09dcde39d9..e74bb9e833 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -48,3 +48,8 @@ jobs: # The same host unit tests on Windows: the platform layer (zig/src/platform) over the UCRT, Win32 and Winsock. - name: zig build test run: zig build test + # The server goldens and the Cortheon fuzz harness on Windows too; the goldens skip the signal-driven cases. + - name: zig build test-golden + run: zig build test-golden + - name: zig build fuzz-cortheon + run: zig build fuzz-cortheon diff --git a/zig/src/platform/win_posix.zig b/zig/src/platform/win_posix.zig index c154eb7153..4088f3214b 100644 --- a/zig/src/platform/win_posix.zig +++ b/zig/src/platform/win_posix.zig @@ -13,7 +13,13 @@ pub const S = c.S; pub const socket_t = isize; pub const socklen_t = c_int; -pub const sockaddr = extern struct { family: u16, data: [14]u8 }; +pub const sockaddr = extern struct { + family: u16, + data: [14]u8, + + pub const in = extern struct { family: u16, port: u16, addr: u32, zero: [8]u8 = @splat(0) }; +}; +pub const timeval = c.timeval; pub const AF = struct { pub const UNIX: u32 = 1; @@ -29,6 +35,8 @@ pub const SOL = struct { pub const SO = struct { /// SO_EXCLUSIVEADDRUSE: Windows' SO_REUSEADDR would let a second server take the port, the opposite of POSIX's. pub const REUSEADDR: u32 = @bitCast(@as(i32, -5)); + pub const RCVTIMEO: u32 = 0x1006; + pub const SNDTIMEO: u32 = 0x1005; }; pub const SHUT = struct { pub const RD: i32 = 0; @@ -124,6 +132,9 @@ pub const system = struct { pub fn bind(fd: socket_t, addr: *const anyopaque, len: anytype) c_int { return if (ws.bind(fd, addr, @intCast(len)) != 0) failed() else 0; } + pub fn connect(fd: socket_t, addr: *const anyopaque, len: anytype) c_int { + return if (ws.connect(fd, addr, @intCast(len)) != 0) failed() else 0; + } pub fn listen(fd: socket_t, backlog: u31) c_int { return if (ws.listen(fd, backlog) != 0) failed() else 0; } @@ -155,6 +166,13 @@ pub const system = struct { }; pub fn setsockopt(fd: socket_t, level: i32, opt: u32, value: []const u8) error{SetSockOptFailed}!void { + if (opt == SO.RCVTIMEO or opt == SO.SNDTIMEO) { + // POSIX gives a timeval; Winsock takes the milliseconds as a DWORD + const tv: *align(1) const timeval = @ptrCast(value.ptr); + const ms: u32 = @intCast(@as(i64, tv.sec) * 1000 + @divTrunc(@as(i64, tv.usec), 1000)); + if (ws.setsockopt(fd, level, @bitCast(opt), @ptrCast(&ms), @sizeOf(u32)) != 0) return error.SetSockOptFailed; + return; + } if (ws.setsockopt(fd, level, @bitCast(opt), value.ptr, @intCast(value.len)) != 0) return error.SetSockOptFailed; } @@ -184,13 +202,26 @@ pub fn poll(fds: []pollfd, timeout_ms: i32) error{PollFailed}!usize { extern "kernel32" fn Sleep(ms: u32) callconv(.winapi) void; -/// Receives from a socket; 0 is the peer's orderly close. -pub fn read(fd: socket_t, buf: []u8) error{ConnectionResetByPeer}!usize { +/// Receives from a socket; 0 is the peer's orderly close, and a receive timeout (SO.RCVTIMEO) is WouldBlock. +pub fn read(fd: socket_t, buf: []u8) error{ ConnectionResetByPeer, WouldBlock }!usize { const n = ws.recv(fd, buf.ptr, len32(buf.len), 0); - if (n < 0) return error.ConnectionResetByPeer; + if (n < 0) return switch (ws.WSAGetLastError()) { + 10035, 10060 => error.WouldBlock, + else => error.ConnectionResetByPeer, + }; return @intCast(n); } +// ---- processes ---- + +extern "kernel32" fn TerminateProcess(process: *anyopaque, code: c_uint) callconv(.winapi) c_int; + +/// kill for a spawned child's handle: TERM ends it with exit code 1; HUP and the rest have no Windows equivalent. +pub fn kill(process: std.os.windows.HANDLE, sig: SIG) error{ Unsupported, Unexpected }!void { + if (sig != .TERM) return error.Unsupported; + if (TerminateProcess(process, 1) == 0) return error.Unexpected; +} + // ---- signals: console control events ---- pub const SIG = enum(c_int) { diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index d7f916f4a6..4976658e6b 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -1,7 +1,8 @@ //! Frozen server goldens for `zig build test-golden`, with the ten differences that are still known. const std = @import("std"); +const builtin = @import("builtin"); const platform = @import("sys"); -const posix = std.posix; +const posix = platform.posix; const Io = std.Io; const wire = @import("wire.zig"); const Reply = wire.Reply; @@ -22,6 +23,18 @@ const known_diffs = [_]Known{ .{ .group = "anthropic", .name = "tool-use-stream" }, }; +/// Reads a spawned server's stdout to its end, so its log lines never fill the pipe (4 KiB on Windows). +fn drainStdout(io: Io, out: Io.File) void { + var buf: [4096]u8 = undefined; + var reader = out.reader(io, &buf); + while (true) { + _ = reader.interface.discardShort(buf.len) catch return; + } +} + +/// The parent's SystemRoot, which a Windows child needs for Winsock; null elsewhere. +var system_root: ?[]const u8 = null; + fn isKnown(group: []const u8, name: []const u8) bool { for (known_diffs) |k| { if (std.mem.eql(u8, k.group, group) and std.mem.eql(u8, k.name, name)) return true; @@ -355,6 +368,7 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con } else try argv.append(a, arg); } var map = std.process.Environ.Map.init(a); + if (system_root) |root| try map.put("SystemRoot", root); // Winsock refuses to start without it try map.put("TENSORFOLD_NO_LIVE", "1"); try map.put("TF_FAKE_CONTEXT", context_tokens); var it = env.iterator(); @@ -398,12 +412,16 @@ fn spawnServer(a: std.mem.Allocator, io: Io, binary: []const u8, fixtures: []con const line = reader.interface.takeDelimiterExclusive('\n') catch return error.NoPortLine; if (std.mem.startsWith(u8, line, "PORT ")) { const port = try std.fmt.parseInt(u16, std.mem.trim(u8, line[5..], " \r"), 10); + // the server's log lines keep coming on stdout: drained, or a full pipe blocks the server + const drain = try std.Thread.spawn(.{}, drainStdout, .{ io, child.stdout.? }); + drain.detach(); return .{ .child = child, .port = port, .log_path = log_path, .key_path = key_path }; } } } pub fn main(init: std.process.Init) !void { + if (builtin.os.tag == .windows) system_root = init.environ_map.get("SystemRoot"); const gpa = init.gpa; const io = init.io; const args = try init.minimal.args.toSlice(init.arena.allocator()); @@ -436,15 +454,26 @@ pub fn main(init: std.process.Init) !void { if (obj.get("cases")) |cases| { for (cases.array.items) |case| try runCase(init.arena.allocator(), io, golden_dir, server.?.port, case); } + var after_signal = false; // Windows: a step's cases after a signal it cannot send are skipped, and said so for (obj.get("then").?.array.items) |then| { const t = then.object; const kind = t.get("kind").?.string; const s = &server.?; if (std.mem.eql(u8, kind, "cases")) { + if (after_signal) { + std.debug.print("SKIPPED {d} cases after a signal Windows cannot send\n", .{t.get("cases").?.array.items.len}); + continue; + } for (t.get("cases").?.array.items) |case| try runCase(init.arena.allocator(), io, golden_dir, s.port, case); } else if (std.mem.eql(u8, kind, "sleep")) { std.Io.sleep(io, .fromMilliseconds(@intFromFloat(t.get("seconds").?.float * 1000.0)), .awake) catch {}; } else if (std.mem.eql(u8, kind, "sigterm")) { + if (builtin.os.tag == .windows) { + std.debug.print("SKIPPED lifecycle/sigterm-exit-0: Windows has no SIGTERM to send, the server is ended\n", .{}); + posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; + _ = s.child.wait(io) catch {}; + continue; + } posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; const term = s.child.wait(io) catch return error.WaitFailed; try outcomes.append(init.arena.allocator(), .{ @@ -471,8 +500,10 @@ pub fn main(init: std.process.Init) !void { try fw.interface.writeAll(t.get("content").?.string); try fw.interface.flush(); posix.kill(s.child.id.?, posix.SIG.HUP) catch {}; // the server rereads the key file, as the reference does + after_signal = builtin.os.tag == .windows; } else if (std.mem.eql(u8, kind, "sighup")) { posix.kill(s.child.id.?, posix.SIG.HUP) catch {}; + after_signal = builtin.os.tag == .windows; } else if (std.mem.eql(u8, kind, "terminate")) { posix.kill(s.child.id.?, posix.SIG.TERM) catch {}; _ = s.child.wait(io) catch {}; From ccd06b94acfe63c82e3497a4d6e762b245b6a3a8 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 17:56:21 +0000 Subject: [PATCH 21/35] cuda: cuBLASLt opens by its Windows DLL name The loader tried only libcublasLt.so.13 and libcublasLt.so, so the cublaslt suite of tf-cuda-test reported LibraryUnavailable on Windows; there it opens cublasLt64_13.dll from the toolkit's bin folder. --- zig/src/cuda/cublaslt.zig | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/zig/src/cuda/cublaslt.zig b/zig/src/cuda/cublaslt.zig index 9e1f154172..cf5b167fc3 100644 --- a/zig/src/cuda/cublaslt.zig +++ b/zig/src/cuda/cublaslt.zig @@ -58,9 +58,10 @@ pub const Library = struct { lib: platform.DynLib, api: Api, - /// The CUDA 13 soname first, then the unversioned link name. + /// The CUDA 13 soname first, then the unversioned link name; on Windows the toolkit's DLL. pub fn open() Error!Library { - for ([_][]const u8{ "libcublasLt.so.13", "libcublasLt.so" }) |path| { + const names: []const []const u8 = if (@import("builtin").os.tag == .windows) &.{"cublasLt64_13.dll"} else &.{ "libcublasLt.so.13", "libcublasLt.so" }; + for (names) |path| { return openPath(path) catch |e| switch (e) { error.LibraryUnavailable => continue, else => return e, From 045abbd023a31553b006c8fb12c090acfbc102a8 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:20:52 +0000 Subject: [PATCH 22/35] Windows: writes on an O_APPEND descriptor land at the end whole The platform handed O_APPEND to the C runtime as _O_APPEND, whose write seeks to the end and then writes, so two threads appending at once could overwrite each other's bytes: the request log lost lines under concurrent requests, which the server goldens' request-log case caught on a slow runner. A descriptor opened with O_APPEND now writes through WriteFile at the end-of-file offset, which the kernel places whole, as POSIX does; the file keeps its write access, so ftruncate on such a descriptor still works (prompt_imprint repairs a torn index that way). A request_log test appends 200 lines from each of 8 threads and expects every line whole; it lost lines on Windows before this change. --- zig/src/platform/win_c.zig | 13 +++++++++++ zig/src/server/request_log.zig | 42 ++++++++++++++++++++++++++++++++++ zig/src/server/root.zig | 1 + 3 files changed, 56 insertions(+) diff --git a/zig/src/platform/win_c.zig b/zig/src/platform/win_c.zig index 21e057d6f5..bfb41928c5 100644 --- a/zig/src/platform/win_c.zig +++ b/zig/src/platform/win_c.zig @@ -284,6 +284,7 @@ fn openFile(path: [*:0]const u8, flags: O) c_int { _ = CloseHandle(@ptrFromInt(@as(usize, @bitCast(h)))); return fail(.MFILE); } + if (fd < append_fds.len) @atomicStore(u8, &append_fds[@intCast(fd)], @intFromBool(flags.APPEND), .release); return fd; } @@ -293,7 +294,11 @@ pub fn openat(dir: fd_t, path: [*:0]const u8, flags: O, ...) callconv(.c) c_int return openFile(p.ptr, flags); } +/// Descriptors opened with O_APPEND, whose writes go to the end of the file whole, as POSIX appends do. +var append_fds: [8192]u8 = @splat(0); + pub fn close(fd: fd_t) c_int { + if (fd >= 0 and fd < append_fds.len) @atomicStore(u8, &append_fds[@intCast(fd)], 0, .release); if (fd >= dir_base and fd < dir_base + dir_paths.len) { lockDirs(); defer dirs_lock.unlock(); @@ -313,6 +318,14 @@ pub fn read(fd: fd_t, buf: [*]u8, len: usize) isize { } pub fn write(fd: fd_t, buf: [*]const u8, len: usize) isize { + if (fd >= 0 and fd < append_fds.len and @atomicLoad(u8, &append_fds[@intCast(fd)], .acquire) != 0) { + // WriteFile at the end-of-file offset lands whole; the CRT's seek-then-write lets writers interleave + const h = osHandle(fd) orelse return -1; + var ov: OVERLAPPED = .{ .offset = 0xFFFF_FFFF, .offset_high = 0xFFFF_FFFF }; + var put: u32 = 0; + if (WriteFile(h, buf, @intCast(@min(len, chunk)), &put, &ov) == 0) return fail(lastError()); + return put; + } return _write(fd, buf, @intCast(@min(len, chunk))); } diff --git a/zig/src/server/request_log.zig b/zig/src/server/request_log.zig index 02c26a0c0d..6d6513113e 100644 --- a/zig/src/server/request_log.zig +++ b/zig/src/server/request_log.zig @@ -40,3 +40,45 @@ pub fn append(a: Allocator, path: []const u8, body: Value) void { const framed = std.mem.concat(a, u8, &.{ line, "\n" }) catch return; _ = platform.c.write(fd, framed.ptr, framed.len); } + +test "appends from several threads keep every line whole" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const cwd = try std.process.currentPathAlloc(std.testing.io, a); + defer a.free(cwd); + const path = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}/log.jsonl", .{ cwd, tmp.sub_path }); + defer a.free(path); + const threads = 8; + const lines = 200; + const head = "{\"line\": \""; + const Worker = struct { + fn run(p: []const u8, k: u8) void { + var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator); + defer arena.deinit(); + const text = arena.allocator().alloc(u8, 300) catch return; + @memset(text, 'a' + k); + for (0..lines) |_| { + const o = json.newObject(arena.allocator()) catch return; + o.put(arena.allocator(), "line", .{ .string = text }) catch return; + append(arena.allocator(), p, .{ .object = o }); + } + } + }; + var workers: [threads]std.Thread = undefined; + for (&workers, 0..) |*w, k| w.* = try std.Thread.spawn(.{}, Worker.run, .{ path, @as(u8, @intCast(k)) }); + for (workers) |w| w.join(); + const got = try std.Io.Dir.cwd().readFileAlloc(std.testing.io, path, a, .limited(1 << 24)); + defer a.free(got); + var counts: [threads]usize = @splat(0); + var it = std.mem.splitScalar(u8, got, '\n'); + while (it.next()) |line| { + if (line.len == 0) continue; + try std.testing.expectEqual(head.len + 300 + 2, line.len); + const k = line[head.len] - 'a'; + try std.testing.expect(k < threads); + for (line[head.len..][0..300]) |c| try std.testing.expectEqual('a' + k, c); + counts[k] += 1; + } + for (counts) |n| try std.testing.expectEqual(@as(usize, lines), n); +} diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 7630f5217a..9cfa2dc1d4 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -44,6 +44,7 @@ test { _ = @import("logprobs_test.zig"); _ = @import("late_system_test.zig"); _ = @import("compact_test.zig"); + _ = @import("request_log.zig"); _ = @import("completion_prompt_test.zig"); _ = @import("server_init_test.zig"); _ = @import("decisions.zig"); From bdfeb8464af099b0be544c4b6552034aa937a830 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:30:32 +0000 Subject: [PATCH 23/35] golden check: a request-log mismatch names the first differing line The lifecycle/request-log outcome said only that the lines differ; it now gives both logs' line counts and the first line that differs, cut to 100 bytes each, so a mismatch seen on a CI runner can be read from its log. --- zig/tests/server/golden_check.zig | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index 4976658e6b..13dc2d2b13 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -23,6 +23,26 @@ const known_diffs = [_]Known{ .{ .group = "anthropic", .name = "tool-use-stream" }, }; +/// Why two request logs differ: their line counts and the first line that does, each cut to 100 bytes. +fn logDiff(a: std.mem.Allocator, want: []const u8, got: []const u8) ![]const []const u8 { + var out: std.ArrayList([]const u8) = .empty; + var wl = std.mem.splitScalar(u8, want, '\n'); + var gl = std.mem.splitScalar(u8, got, '\n'); + var n: usize = 0; + while (true) : (n += 1) { + const w = wl.next(); + const g = gl.next(); + if (w == null and g == null) break; + if (w != null and g != null and std.mem.eql(u8, w.?, g.?)) continue; + const ws = w orelse ""; + const gs = g orelse ""; + try out.append(a, try std.fmt.allocPrint(a, "line {d}: golden [{s}] zig [{s}]", .{ n + 1, ws[0..@min(ws.len, 100)], gs[0..@min(gs.len, 100)] })); + break; + } + try out.append(a, try std.fmt.allocPrint(a, "TENSORFOLD_REQUEST_LOG: golden {d} lines, zig {d} lines", .{ std.mem.count(u8, want, "\n"), std.mem.count(u8, got, "\n") })); + return out.items; +} + /// Reads a spawned server's stdout to its end, so its log lines never fill the pipe (4 KiB on Windows). fn drainStdout(io: Io, out: Io.File) void { var buf: [4096]u8 = undefined; @@ -490,7 +510,7 @@ pub fn main(init: std.process.Init) !void { .group = "lifecycle", .name = "request-log", .verdict = if (std.mem.eql(u8, want, got)) .equal else .differ, - .why = if (std.mem.eql(u8, want, got)) &.{} else &.{"TENSORFOLD_REQUEST_LOG lines differ from the frozen ones"}, + .why = if (std.mem.eql(u8, want, got)) &.{} else try logDiff(init.arena.allocator(), want, got), }); } else if (std.mem.eql(u8, kind, "rewrite_key")) { var file = try Io.Dir.cwd().createFile(io, s.key_path.?, .{}); From 153d6b438625a65a3c2cf64da767f5519cd74c5b Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:41:58 +0000 Subject: [PATCH 24/35] golden check: the request log compares with LF endings on either side A checkout with core.autocrlf, as GitHub's Windows runners make, gives the frozen request log CRLF line endings while the server writes LF, so the lifecycle/request-log case differed at its first line's end on windows-latest and nowhere else. Both logs compare with CRLF read as LF; the other goldens are parsed as JSON and never saw the endings. --- zig/tests/server/golden_check.zig | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/zig/tests/server/golden_check.zig b/zig/tests/server/golden_check.zig index 13dc2d2b13..73d607350f 100644 --- a/zig/tests/server/golden_check.zig +++ b/zig/tests/server/golden_check.zig @@ -504,8 +504,9 @@ pub fn main(init: std.process.Init) !void { }); } else if (std.mem.eql(u8, kind, "compare_log")) { const golden_log = try std.fmt.allocPrint(init.arena.allocator(), "{s}/{s}", .{ golden_dir, t.get("golden").?.string }); - const want = Io.Dir.cwd().readFileAlloc(io, golden_log, init.arena.allocator(), .limited(1 << 28)) catch ""; - const got = Io.Dir.cwd().readFileAlloc(io, s.log_path.?, init.arena.allocator(), .limited(1 << 28)) catch ""; + // a checkout with autocrlf gives the frozen log CRLF endings, the server writes LF: both compare as LF + const want = try std.mem.replaceOwned(u8, init.arena.allocator(), Io.Dir.cwd().readFileAlloc(io, golden_log, init.arena.allocator(), .limited(1 << 28)) catch "", "\r\n", "\n"); + const got = try std.mem.replaceOwned(u8, init.arena.allocator(), Io.Dir.cwd().readFileAlloc(io, s.log_path.?, init.arena.allocator(), .limited(1 << 28)) catch "", "\r\n", "\n"); try outcomes.append(init.arena.allocator(), .{ .group = "lifecycle", .name = "request-log", From 8570e0688ae6118e0e90eee0655940aeef7614bf Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:16:46 +0000 Subject: [PATCH 25/35] qwen3_5: attention chunks follow each row's own keys, not the forward's widest stream The chunk width came from the longest stream or window in the forward, so a 201-key row cut its keys into 128-key chunks alone and 256-key chunks beside a stream past 8192, and a window crossing 8192 recut its earlier rows; either changed the fp32 sums. Each row now takes chunk_keys(its logical keys), the width its one-row step uses, in both the partial and the merge kernels, and the partials sit at a fixed max_chunks stride. tf-qwen35-cuda takes TF_Q35_PREFILL=N (start after N cycled ids) and TF_Q35_BESIDE=N (step an N-token stream in each forward), which cover both cases: before this change, windows from 8151 keys kept 36 of 120 rows and trees 35, and a 151-key stream beside one from 8189 kept 4 to 11 of 120; after it all of them equal one-row steps on the 2B and 27B. --- zig/kernels/cuda/qwen35.cu | 25 +++++++----- zig/kernels/cuda/qwen35.cuh | 1 + zig/src/families/qwen3_5/cuda_engine.zig | 15 +++---- zig/tests/qwen35_cuda_main.zig | 50 ++++++++++++++++++++++-- 4 files changed, 68 insertions(+), 23 deletions(-) diff --git a/zig/kernels/cuda/qwen35.cu b/zig/kernels/cuda/qwen35.cu index 558baa7908..9217a19fa3 100644 --- a/zig/kernels/cuda/qwen35.cu +++ b/zig/kernels/cuda/qwen35.cu @@ -310,6 +310,9 @@ __device__ __forceinline__ __nv_bfloat16 rotate(const __nv_bfloat16* x, int d, i // A row's ancestors-and-self: its tree window's row bits, or with the top bit set a chain row's index in its window. __device__ __forceinline__ int rows_seen(unsigned anc) { return (anc >> 31) ? (int)(anc & 0x7fffffffu) + 1 : __popc(anc); } +// A row's chunk width from its own logical keys alone, as its one-row step's: at most MAXCH chunks, 128-key multiples. +__device__ __forceinline__ int chunk_keys(int keys) { return max(128, (keys + MAXCH * 128 - 1) / (MAXCH * 128) * 128); } + // The window row holding a row's nth seen window key. __device__ __forceinline__ int seen_row(unsigned anc, int n) { if (anc >> 31) return n; @@ -342,24 +345,24 @@ extern "C" __global__ void __launch_bounds__(256) tf_q35_rope(Geo geo, __nv_bflo } } -// Causal attention over the cache in P[0] chunks of P[1] keys, a warp a query head; partials by (row, head, chunk). +// Causal attention in chunks of chunk_keys(the row's keys), a warp a query head; partials by (row, head, chunk). extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __restrict__ Q, const Seg* __restrict__ segs, const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, int layer, const int* __restrict__ P, float scale, float* __restrict__ PM, float* __restrict__ PL, float* __restrict__ PO) { constexpr int DPL = HD / 32; - const int nch = P[0], ck = P[1]; const int QH = geo.qh, KVH = geo.kvh; const int lane = threadIdx.x & 31, g = threadIdx.x >> 5; const int c = blockIdx.x, h = blockIdx.y; - if (c >= nch || past_rows(P, blockIdx.z)) return; + if (c >= MAXCH || past_rows(P, blockIdx.z)) return; const size_t m = blockIdx.z; const int qh = h * (QH / KVH) + g; const Seg sg = segs[rowseg[m]]; const unsigned anc = rowanc[m]; // chunks cut at logical positions (a window key's: the stream's length plus its ancestors), as a step's are - const int k0 = c * ck, k1 = min(k0 + ck, sg.len + rows_seen(anc)); - const size_t slot = (m * QH + qh) * nch + c; + const int keys = sg.len + rows_seen(anc), ck = chunk_keys(keys); + const int k0 = c * ck, k1 = min(k0 + ck, keys); + const size_t slot = (m * QH + qh) * MAXCH + c; if (k0 >= k1) { if (lane == 0) { PM[slot] = -INFINITY; @@ -401,17 +404,19 @@ extern "C" __global__ void tf_q35_attn_partial(Geo geo, const __nv_bfloat16* __r for (int i = 0; i < DPL; i++) PO[slot * HD + lane * DPL + i] = o[i]; } -// The P[0] chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. +// A row's chunks' partials into each (row, query head)'s output, rows qh * HD wide. Grid (qh, rows), a warp each. extern "C" __global__ void __launch_bounds__(32) tf_q35_attn_merge(Geo geo, const float* __restrict__ PM, const float* __restrict__ PL, - const float* __restrict__ PO, const int* __restrict__ P, - __nv_bfloat16* __restrict__ OUT) { + const float* __restrict__ PO, const Seg* __restrict__ segs, + const int* __restrict__ rowseg, const unsigned* __restrict__ rowanc, + const int* __restrict__ P, __nv_bfloat16* __restrict__ OUT) { constexpr int DPL = HD / 32; - const int nch = P[0]; const int QH = geo.qh; const int lane = threadIdx.x, qh = blockIdx.x; if (past_rows(P, blockIdx.y)) return; const size_t m = blockIdx.y; - const size_t base = (m * QH + qh) * nch; + const int keys = segs[rowseg[m]].len + rows_seen(rowanc[m]); + const int ck = chunk_keys(keys), nch = (keys + ck - 1) / ck; + const size_t base = (m * QH + qh) * MAXCH; float mx = -INFINITY; for (int c = 0; c < nch; c++) if (PL[base + c] > 0.0f) mx = fmaxf(mx, PM[base + c]); diff --git a/zig/kernels/cuda/qwen35.cuh b/zig/kernels/cuda/qwen35.cuh index b4f789814f..5df37e308e 100644 --- a/zig/kernels/cuda/qwen35.cuh +++ b/zig/kernels/cuda/qwen35.cuh @@ -11,6 +11,7 @@ constexpr int TAPS = 4; // causal convolution taps constexpr int HD = 256; // attention head width constexpr int MAXQ = 32; // attention query heads at most constexpr int MAXW = 16; // a window's rows at most (the log's rows a stream) +constexpr int MAXCH = 64; // attention chunks a row at most (cuda_engine.zig's max_chunks) // The model's widths: hidden, linear key heads, linear value heads, attention query and key/value heads, linear layers. struct Geo { diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 389b32f766..15975b8d42 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -22,7 +22,7 @@ pub const max_segs = 32; pub const most_verify = 32; /// A window's rows at most (exact: one-row steps' bits; qwen35.cuh's MAXW, the rows a stream's log holds). pub const most_window = 16; -/// Partial attention chunks a row at most: the chunk width grows with the window instead. +/// Partial attention chunks a row at most (qwen35.cuh's MAXCH): each row's chunk width grows with its own keys. const max_chunks = 64; /// Device allocations are 256-byte aligned, as the driver's are. @@ -147,7 +147,7 @@ pub const Engine = struct { const rules_at = rowanc_at + max_rows * 4; const meta_at = rules_at + most_verify * @sizeOf(cuda.sample.Rule); const tokens_at = meta_at + most_verify * 4; - const params_at = tokens_at + most_verify * 4; // the attention chunks, their width, the rows, the windows, record + const params_at = tokens_at + most_verify * 4; // two unused words, the rows, the windows, record const tables_len = params_at + 8 * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8) !*Engine { @@ -343,12 +343,10 @@ pub const Engine = struct { pub fn forward(e: *Engine, parts: []const Part, ids: []const u32, keep: Keep) !void { const g = e.g; var total: usize = 0; - var widest: usize = 0; for (parts) |p| { try p.cache.settle(e); if (p.rows == 0 or p.cache.len + p.rows > p.cache.cap) return error.QwenContextFull; total += p.rows; - widest = @max(widest, p.cache.len + p.rows); } if (parts.len == 0 or parts.len > max_segs or total != ids.len or total > max_rows) return error.InvalidQwenSegments; const record = keep != .last; @@ -378,11 +376,10 @@ pub const Engine = struct { m.* = @intCast(p.position -% 1); // sample.cu keys row r at meta + r + 1 } } - // attention chunks: at most max_chunks a row, each a multiple of 128 keys - const ck = @max(128, (widest + max_chunks * 128 - 1) / (max_chunks * 128) * 128); + // words 0 and 1 are reserved (each row's attention chunks follow its own keys: qwen35.cu's chunk_keys) const params = std.mem.bytesAsSlice(c_int, st[params_at..][0..32]); - params[0] = @intCast((widest + ck - 1) / ck); - params[1] = @intCast(ck); + params[0] = 0; + params[1] = 0; params[2] = @intCast(total); params[3] = @intCast(parts.len); params[4] = @intFromBool(record); @@ -497,7 +494,7 @@ pub const Engine = struct { try e.run(e.fns.head_norm, .{ total * g.kv_heads, 1, 1 }, .{ 64, 1, 1 }, .{ e.kraw, m.k_norm, e.kn, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(g.kv_heads * head_dim)), @as(c_int, head_dim), g.eps, t_params }); try e.run(e.fns.rope, .{ total, 1, 1 }, .{ head_dim, 1, 1 }, .{ geo, e.qn, e.kn, e.vraw, t_segs, t_rows, t_anc, ai, t_params }); try e.run(e.fns.attn_partial, .{ max_chunks, g.kv_heads, total }, .{ 32 * group, 1, 1 }, .{ geo, e.qn, t_segs, t_rows, t_anc, ai, t_params, @as(f32, 1.0 / 16.0), e.pm, e.pl, e.po }); - try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_params, e.y }); + try e.run(e.fns.attn_merge, .{ g.q_heads, total, 1 }, .{ 32, 1, 1 }, .{ geo, e.pm, e.pl, e.po, t_segs, t_rows, t_anc, t_params, e.y }); try e.run(e.fns.attn_gate, .{ (total * aw + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ geo, e.y, e.qraw, e.mix, @as(c_int, @intCast(total)), t_params }); try e.proj(m.o, e.mix, e.r, total, prompt); ai += 1; diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 0d1f711252..e0bb3ac025 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -20,15 +20,23 @@ pub fn main(init: std.process.Init) !u8 { defer ctx.deinit(); const e = try E.Engine.init(init.gpa, init.io, &ctx, args[1]); defer e.deinit(); - var cache = try E.Cache.init(init.gpa, e, 4096); + // TF_Q35_PREFILL=N starts the stream after N cycled ids; TF_Q35_BESIDE=N steps an N-token stream beside it + const prefill = try envCount("TF_Q35_PREFILL"); + const beside = try envCount("TF_Q35_BESIDE"); + var cache = try E.Cache.init(init.gpa, e, 4096 + prefill); defer cache.deinit(init.gpa); + try fill(e, &cache, ids.items, prefill); + var other = try E.Cache.init(init.gpa, e, 4096 + beside); + defer other.deinit(init.gpa); + try fill(e, &other, ids.items, beside); + var f: Fwd = .{ .e = e, .other = if (beside > 0) &other else null, .ids = ids.items }; var out: std.ArrayList(u8) = .empty; if (std.mem.eql(u8, args[2], "step")) { var first: i64 = 0; var rest: i64 = 0; for (ids.items, 0..) |id, i| { const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); - try e.forward(&.{.{ .cache = &cache, .rows = 1 }}, &.{id}, .logits); + try f.run(.{ .cache = &cache, .rows = 1 }, &.{id}); const dt: i64 = @intCast(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0); if (i == 0) first = dt else rest += dt; @memcpy(cache.logits, e.rowLogits(0)); @@ -42,7 +50,7 @@ pub fn main(init: std.process.Init) !u8 { var w: usize = 1; while (at < ids.items.len) : (w = w % e.window_rows + 1) { const n = @min(w, ids.items.len - at); - try e.forward(&.{.{ .cache = &cache, .rows = n }}, ids.items[at..][0..n], .logits); + try f.run(.{ .cache = &cache, .rows = n }, ids.items[at..][0..n]); for (0..n) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); at += n; } @@ -76,7 +84,7 @@ pub fn main(init: std.process.Init) !u8 { parents[n] = 1; n += 1; } - try e.forward(&.{.{ .cache = &cache, .rows = n, .parents = parents[0..n] }}, toks[0..n], .logits); + try f.run(.{ .cache = &cache, .rows = n, .parents = parents[0..n] }, toks[0..n]); for (path[0..kept]) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r))); if (n > 1) try cache.keep(e, path[0..kept]); at += kept; @@ -93,3 +101,37 @@ pub fn main(init: std.process.Init) !u8 { try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = args[4], .data = out.items }); return 0; } + +/// One forward keeping logits: the stream's part first (its rows' logits come first), then the second stream's step. +const Fwd = struct { + e: *E.Engine, + other: ?*E.Cache, + ids: []const u32, + n: usize = 0, + + fn run(f: *Fwd, part: E.Part, toks: []const u32) !void { + const o = f.other orelse return f.e.forward(&.{part}, toks, .logits); + var all: [E.most_window + 1]u32 = undefined; + @memcpy(all[0..toks.len], toks); + all[toks.len] = f.ids[f.n % f.ids.len]; + f.n += 1; + try f.e.forward(&.{ part, .{ .cache = o, .rows = 1 } }, all[0 .. toks.len + 1], .logits); + } +}; + +fn envCount(name: [:0]const u8) !usize { + const v = std.c.getenv(name) orelse return 0; + return std.fmt.parseInt(usize, std.mem.span(v), 10); +} + +/// `n` positions of the ids cycled, through prompt forwards. +fn fill(e: *E.Engine, cache: *E.Cache, ids: []const u32, n: usize) !void { + var buf: [E.max_rows]u32 = undefined; + var at: usize = 0; + while (at < n) { + const k = @min(E.max_rows, n - at); + for (buf[0..k], at..) |*t, i| t.* = ids[i % ids.len]; + try e.forward(&.{.{ .cache = cache, .rows = k }}, buf[0..k], .last); + at += k; + } +} From 5e1c3feaa79e40f7659537d96b939129abe39df9 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:16:46 +0000 Subject: [PATCH 26/35] qwen3_5: the DFlash2 config's divisors must be positive, and staging fits its selector rank The integer reader accepted zero, so a conv_group_size or num_key_value_heads of 0 reached a remainder at admission; sizes and divisors are now positive (mask_token_id alone may be 0). The host staging area held 16 rows at rank 256, so a larger selector_rank overran it; it is now sized for the admitted rank. Config.parse takes the bytes, so a host test covers both. --- zig/src/families/qwen3_5/cuda.zig | 1 + zig/src/families/qwen3_5/cuda_draft.zig | 42 +++++++++++++++++++++++-- 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 435ebe162e..47158adda9 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -10,4 +10,5 @@ test { _ = geometry; _ = @import("draw.zig"); _ = @import("cuda_draft_weights.zig"); + _ = draft; } diff --git a/zig/src/families/qwen3_5/cuda_draft.zig b/zig/src/families/qwen3_5/cuda_draft.zig index cfb8670cf5..7f0e4e61be 100644 --- a/zig/src/families/qwen3_5/cuda_draft.zig +++ b/zig/src/families/qwen3_5/cuda_draft.zig @@ -58,6 +58,10 @@ pub const Config = struct { defer gpa.free(path); const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); defer gpa.free(bytes); + return parse(gpa, bytes, target); + } + + fn parse(gpa: std.mem.Allocator, bytes: []const u8, target: engine.Geometry) !Config { const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); defer doc.deinit(); if (doc.value != .object) return error.BadDraftConfig; @@ -73,7 +77,7 @@ pub const Config = struct { .kv_heads = try int(o, "num_key_value_heads"), .head_dim = try int(o, "head_dim"), .block = try int(d.object, "block_size"), - .mask = @intCast(try int(d.object, "mask_token_id")), + .mask = @intCast(try whole(d.object, "mask_token_id")), .window = try int(o, "sliding_window"), .group = try int(d.object, "conv_group_size"), .rank = try int(d.object, "selector_rank"), @@ -98,7 +102,15 @@ pub const Config = struct { } }; +/// 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); + if (v == 0) return error.BadDraftConfig; + return v; +} + +/// A token id: any whole number. +fn whole(o: std.json.ObjectMap, key: []const u8) !usize { const v = o.get(key) orelse return error.BadDraftConfig; if (v != .integer or v.integer < 0) return error.BadDraftConfig; return @intCast(v.integer); @@ -176,7 +188,10 @@ pub const Draft = struct { const words_at = 64; const hp_at = words_at + 16 * lp_words * 4; - const staged_len = hp_at + 16 * 256 * 2; + /// The host staging area: the words then 16 rows of `rank` selector projections (the block's mask rows fit). + fn stagedLen(rank: usize) usize { + return hp_at + 16 * rank * 2; + } pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8) !*Draft { const c = try Config.read(gpa, io, dir, e.g); @@ -193,7 +208,7 @@ pub const Draft = struct { try d.scratch.fill8(0, null); d.at = 0; _ = try d.layout(true); - d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); + d.staged = try cuda.HostBuffer.alloc(e.d, stagedLen(c.rank)); errdefer d.staged.free(); try d.picksZero(); // the first context, and a hold timed on it (the quickest of several: the first ones pay the driver's stalls) @@ -546,3 +561,24 @@ pub fn timing(what: []const u8, io: std.Io, t0: std.Io.Timestamp) void { const at: f64 = @floatFromInt(@mod(now.toNanoseconds(), 1_000_000_000_000)); std.debug.print("[timing] {s} {d:.2} ms at {d:.1}\n", .{ what, ns / 1e6, at / 1e6 }); } + +test "a drafter config with a zero divisor is refused, and any admitted rank's selector rows fit the staging area" { + const gpa = std.testing.allocator; + const target: engine.Geometry = .{ .hidden = 5120, .intermediate = 17408, .layers = 64, .vocab = 248320, .k_heads = 16, .v_heads = 48, .q_heads = 24, .kv_heads = 4, .eps = 1e-6, .tied = false, .context = 262144, .linear = @splat(false) }; + const text = + \\{"dflash_config":{"block_size":8,"conv_group_size":16,"conv_kernel_size":2,"mask_token_id":248070,"selector_rank":256, + \\ "selector_top_k":16,"target_layer_ids":[5,19,33,47,61]},"head_dim":128,"hidden_size":5120,"intermediate_size":17408, + \\ "num_attention_heads":32,"num_hidden_layers":5,"num_key_value_heads":8,"rms_norm_eps":1e-06, + \\ "rope_parameters":{"rope_theta":10000000},"sliding_window":2048,"vocab_size":248320} + ; + try std.testing.expectEqual(@as(usize, 256), (try Config.parse(gpa, text, target)).rank); + for ([_][2][]const u8{ .{ "\"conv_group_size\":16", "\"conv_group_size\":0" }, .{ "\"num_key_value_heads\":8", "\"num_key_value_heads\":0" } }) |r| { + const bad = try std.mem.replaceOwned(u8, gpa, text, r[0], r[1]); + defer gpa.free(bad); + try std.testing.expectError(error.BadDraftConfig, Config.parse(gpa, bad, target)); + } + const wide = try std.mem.replaceOwned(u8, gpa, text, "\"selector_rank\":256", "\"selector_rank\":1024"); + defer gpa.free(wide); + const c = try Config.parse(gpa, wide, target); + try std.testing.expect(Draft.stagedLen(c.rank) >= Draft.hp_at + (c.block - 1) * c.rank * 2); +} From 330ef45547e7143ff79001d87a7f42770075c1d1 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:13:51 +0000 Subject: [PATCH 27/35] cuda: the affine-4 experts' SwiGLU pair form, and softmax top-k routing with a shared expert The grouped affine-4 expert kernel ran only Nemotron's relu^2 up; its SwiGLU epilogue over two accumulators is now instantiated for decode and prompt rows (M = 2, gate and up interleaved per group), and the pack kernel takes a table stride and index so gate and up pack into one interleaved table. Layer.swiglu picks that form; the shared expert sits last in the same tables. cuda/moe.zig and moe.cu are the routing the frozen tensorfold/cuda/moe.py did: fp32 logits over the experts and the shared expert's gate row, each row's top-k by logit with ties to the lower id, weights exp(l - l_top) over the k rounded to bf16, the shared expert as a slot weighted by bf16(sigmoid(bf16(gate))), and the slots' weighted sum rounded once to bf16. tf-cuda-test moe checks each against a host reference: the router and the combine byte-equal, the picks equal with weights within a bf16 step (a tie, a flat row and a NaN included), and the pair experts against fp64 over random affine-4 tables. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/experts.cu | 4 + zig/kernels/cuda/experts_pack.cu | 7 +- zig/kernels/cuda/experts_prefill.cu | 3 + zig/kernels/cuda/moe.cu | 113 ++++++++++ zig/src/cuda/experts.zig | 49 +++-- zig/src/cuda/kernels.zig | 1 + zig/src/cuda/moe.zig | 65 ++++++ zig/src/cuda/root.zig | 1 + zig/tests/cuda/main.zig | 3 + zig/tests/cuda/moe_tests.zig | 309 ++++++++++++++++++++++++++++ 11 files changed, 539 insertions(+), 17 deletions(-) create mode 100644 zig/kernels/cuda/moe.cu create mode 100644 zig/src/cuda/moe.zig create mode 100644 zig/tests/cuda/moe_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 1f3e047eb4..3a994fd300 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -21,6 +21,7 @@ const kernels = [_]Kernel{ .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7 .{ .name = "experts_prefill", .flags = &.{"-O3"} }, .{ .name = "experts_pack", .flags = &.{"-O3"} }, + .{ .name = "moe", .flags = &.{"-O3"} }, // cuda/moe.py: softmax top-k routing with a sigmoid-gated shared expert .{ .name = "prefill_attention", .flags = &.{ "-O3", "--fmad=false" } }, // tensorfold_prefill_attention_v1 .{ .name = "scan_rows", .flags = &.{ "-O3", "--fmad=false" } }, // nemotron_h/cuda/mamba.py .{ .name = "nemotron_ops", .flags = &.{"-O3"} }, // ours: Nemotron's serial feed, routed plan and rests diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index d19f6ee90c..3969b15cee 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -280,3 +280,7 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); TF_EXPERT(1) TF_EXPERT(0) + +// Pair form: gate and up blocks interleaved per group (M = 2), SwiGLU in the epilogue (2) as fp8/nvfp4's up does. +#define TF_EXPERT_PAIR(EPI) template __global__ void tf_experts::expert_kernel<64, 2, EPI, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); +TF_EXPERT_PAIR(2) diff --git a/zig/kernels/cuda/experts_pack.cu b/zig/kernels/cuda/experts_pack.cu index b47ded555f..0ca3d77fda 100644 --- a/zig/kernels/cuda/experts_pack.cu +++ b/zig/kernels/cuda/experts_pack.cu @@ -16,10 +16,11 @@ __device__ __forceinline__ uint32_t shuffle(uint32_t w) { template // group size / 32 __global__ void __launch_bounds__(32 * 4 * H + 32) pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ scales, - const uint16_t* __restrict__ biases, uint32_t* __restrict__ out, int n, int k8, int kg, int nb) { + const uint16_t* __restrict__ biases, uint32_t* __restrict__ out, int n, int k8, int kg, int nb, int m, int mi) { constexpr int WORDS = 32 * 4 * H; // a block's weight words: NTW * H words a lane const int g = blockIdx.x, b = blockIdx.y, e = blockIdx.z, i = threadIdx.x; - uint32_t* dst = out + ((static_cast(e) * nb + b) * kg + g) * (WORDS + 32); + // block (e, b, g) of m interleaved tables, this one the mi-th: [E][nb][kg][m][WORDS + 32] + uint32_t* dst = out + ((((static_cast(e) * nb + b) * kg + g) * m) + mi) * (WORDS + 32); if (i < WORDS) { const int row = i / (4 * H), kk = i % (4 * H); // row = t * 8 + r; kk = q * H + j const int t = row / 8, r = row % 8, q = kk / H, j = kk % H; @@ -39,4 +40,4 @@ pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ sca // Groups of 64 inputs. template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t* __restrict__, const uint16_t* __restrict__, const uint16_t* __restrict__, - uint32_t* __restrict__, int, int, int, int); + uint32_t* __restrict__, int, int, int, int, int, int); diff --git a/zig/kernels/cuda/experts_prefill.cu b/zig/kernels/cuda/experts_prefill.cu index 3a1df1dfee..f56c4704a8 100644 --- a/zig/kernels/cuda/experts_prefill.cu +++ b/zig/kernels/cuda/experts_prefill.cu @@ -155,3 +155,6 @@ __global__ void __launch_bounds__(WM * WN * 32) const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); TF_PREFILL(1) TF_PREFILL(3) + +// Pair form for prompt rows: gate and up interleaved (M = 2) under SwiGLU (epilogue 2). +template __global__ void tf_experts_prefill::prefill_kernel<64, 2, 2, 2, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float); diff --git a/zig/kernels/cuda/moe.cu b/zig/kernels/cuda/moe.cu new file mode 100644 index 0000000000..2adb06a220 --- /dev/null +++ b/zig/kernels/cuda/moe.cu @@ -0,0 +1,113 @@ +// Softmax top-k routing with a sigmoid-gated shared expert on fp32 logits, as tensorfold/cuda/moe.py routes. + +#include +#include +#include + +namespace { + +__device__ __forceinline__ float bf(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float rbf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +constexpr int TILE = 16; // rows and experts a block covers +constexpr int KC = 256; // inputs a shared-memory chunk holds + +} // namespace + +// OUT[r][e] = fp32 x[r] . w[e] over D in index order (one fma chain a pair); grid (rows / 16, E / 16), block 256. +extern "C" __global__ void __launch_bounds__(256) tf_moe_router(const __nv_bfloat16* __restrict__ X, int x_stride, + const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, + int R, int D, int E) { + __shared__ __nv_bfloat16 xs[TILE][KC + 2]; + __shared__ __nv_bfloat16 ws[TILE][KC + 2]; + const int row0 = blockIdx.x * TILE, e0 = blockIdx.y * TILE; + const int rr = threadIdx.x / TILE, ee = threadIdx.x % TILE; + float acc = 0.0f; + for (int k0 = 0; k0 < D; k0 += KC) { + __syncthreads(); + for (int c = threadIdx.x; c < TILE * KC; c += 256) { + const int r = c / KC, k = c % KC; + const bool kin = k0 + k < D; + xs[r][k] = kin && row0 + r < R ? X[static_cast(row0 + r) * x_stride + k0 + k] : __float2bfloat16_rn(0.0f); + ws[r][k] = kin && e0 + r < E ? W[static_cast(e0 + r) * D + k0 + k] : __float2bfloat16_rn(0.0f); + } + __syncthreads(); + const int kn = min(KC, D - k0); + for (int k = 0; k < kn; ++k) acc = __fmaf_rn(bf(xs[rr][k]), bf(ws[ee][k]), acc); + } + if (row0 + rr < R && e0 + ee < E) OUT[static_cast(row0 + rr) * E + e0 + ee] = acc; +} + +namespace { + +constexpr int MAXC = 8; // experts a lane holds: up to 256 routed experts + +// The larger of two (value, index) picks: a higher value, or the lower index on a tie (values are never NaN). +__device__ __forceinline__ bool beats(float v, int i, float bv, int bi) { return v > bv || (v == bv && i < bi); } + +} // namespace + +// Row r by a warp: TOPK experts by logit (ties to the lower id), weights bf16(exp(l - top) / sum), the shared slot. +extern "C" __global__ void __launch_bounds__(32) tf_moe_topk(const float* __restrict__ L, int* __restrict__ PICK, + float* __restrict__ WT, int NE, int TOPK) { + const int64_t r = blockIdx.x; + const int lane = threadIdx.x; + const int NL = NE + 1, SLOTS = TOPK + 1; // L[r][NE], the shared gate's logit, weighs slot TOPK by bf16(sigmoid(bf16(it))) + float v[MAXC]; +#pragma unroll + for (int c = 0; c < MAXC; ++c) { + const int e = lane + 32 * c; + v[c] = e < NE ? L[r * NL + e] : -INFINITY; + } + float top = -INFINITY; +#pragma unroll + for (int c = 0; c < MAXC; ++c) top = fmaxf(top, v[c]); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) top = fmaxf(top, __shfl_xor_sync(0xffffffffu, top, off)); + float total = 0.0f; + float ex[MAXC]; + int id[MAXC]; + for (int k = 0; k < TOPK; ++k) { + float bv = -INFINITY; + int bi = 0x7fffffff; +#pragma unroll + for (int c = 0; c < MAXC; ++c) + if (lane + 32 * c < NE && beats(v[c], lane + 32 * c, bv, bi)) bv = v[c], bi = lane + 32 * c; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (beats(ov, oi, bv, bi)) bv = ov, bi = oi; + } + const float e = expf(bv - top); + total += e; + if (lane == k % 32) ex[k / 32] = e, id[k / 32] = bi; // slot k lives in lane k % 32 (TOPK <= 32) +#pragma unroll + for (int c = 0; c < MAXC; ++c) + if (lane + 32 * c == bi) v[c] = -INFINITY; + } + if (lane < TOPK) { + PICK[r * SLOTS + lane] = id[0]; + WT[r * SLOTS + lane] = rbf(ex[0] / total); + } + if (lane == 0) { + const float sg = rbf(L[r * NL + NE]); + PICK[r * SLOTS + TOPK] = NE; + WT[r * SLOTS + TOPK] = rbf(1.0f / (1.0f + expf(-sg))); + } +} + +// OUT[r][d] = bf16(sum over slots in order of fma(w[r][s], y[r][s][d])); y fp32 (decode) or bf16 (prompt rows). +extern "C" __global__ void __launch_bounds__(256) tf_moe_combine(const void* __restrict__ Y, int y_f32, const float* __restrict__ WT, + __nv_bfloat16* __restrict__ OUT, int S, int D) { + const int64_t r = blockIdx.x; + for (int d = blockIdx.y * 256 + threadIdx.x; d < D; d += gridDim.y * 256) { + float acc = 0.0f; + for (int s = 0; s < S; ++s) { + const int64_t at = (r * S + s) * D + d; + const float y = y_f32 ? static_cast(Y)[at] : bf(static_cast(Y)[at]); + acc = __fmaf_rn(WT[r * S + s], y, acc); + } + OUT[r * D + d] = __float2bfloat16_rn(acc); + } +} diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index c940202bea..86b1086501 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -27,6 +27,7 @@ pub const Layer = struct { experts: usize, limit: f32 = 0.0, // fp8g, nvfp4: the SwiGLU clamp, 0 for none relu2: bool = false, // nvfp4: up is one projection under relu^2 (Nemotron), not SwiGLU's gate and up + swiglu: bool = false, // affine4: up is gate and up interleaved per group ([E][NI/32][D/64][2][288]) under SwiGLU }; /// A call's epilogue: up with its format's activation (affine4 relu^2, else SwiGLU), down to fp32 or to bf16 sums. @@ -51,7 +52,7 @@ pub const Grouped = struct { const up = epi == .up; const nb = (if (up) l.width else l.dims) / 32; const w = Affine4Experts.view(l, up); - return a.run(s, up, in.x, in.stride, if (up) in.slots else 0, w.ptr, w.kg, nb, p, out, w.n, items * nb); + return a.run(s, up, up and l.swiglu, in.x, in.stride, if (up) in.slots else 0, w.ptr, w.kg, nb, p, out, w.n, items * nb); }, .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, false, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), @@ -68,7 +69,7 @@ pub const Grouped = struct { const up = epi == .up; const nb = (if (up) l.width else l.dims) / 32; const w = Affine4Experts.view(l, up); - return a.runPrompt(s, up, in.x, in.stride, if (up) in.slots else 0, w.ptr, w.kg, nb, p, out, w.n, items); + return a.runPrompt(s, up, up and l.swiglu, in.x, in.stride, if (up) in.slots else 0, w.ptr, w.kg, nb, p, out, w.n, items); }, .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, true, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), @@ -91,20 +92,25 @@ fn int(x: usize) c_int { pub const Affine4Experts = struct { up_fn: Function, down_fn: Function, + pair_fn: Function, pre_up_fn: Function, pre_down_fn: Function, + pre_pair_fn: Function, pack_fn: Function, - resident: [2]usize, // blocks the decode up and down kernels fill: per SM times SMs + resident: [3]usize, // blocks the decode up, down and pair kernels fill: per SM times SMs pub const prompt_smem: u32 = 38400; // Pre<64, 1, 2, 2, 4>: three stages of 800 uint4 + pub const pair_prompt_smem: u32 = 34816; // Pre<64, 2, 2, 2, 4>: two stages of 1088 uint4 /// The instances' names (cuobjdump -symbols of the experts, experts_prefill and experts_pack fatbins). pub const symbols = struct { pub const up = "_ZN10tf_experts13expert_kernelILi64ELi1ELi1ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; pub const down = "_ZN10tf_experts13expert_kernelILi64ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + pub const pair = "_ZN10tf_experts13expert_kernelILi64ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; pub const pre_up = "_ZN18tf_experts_prefill14prefill_kernelILi64ELi1ELi1ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; pub const pre_down = "_ZN18tf_experts_prefill14prefill_kernelILi64ELi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; - pub const pack = "_ZN15tf_experts_pack11pack_kernelILi2EEEvPKjPKtS4_Pjiiii"; + pub const pre_pair = "_ZN18tf_experts_prefill14prefill_kernelILi64ELi2ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + pub const pack = "_ZN15tf_experts_pack11pack_kernelILi2EEEvPKjPKtS4_Pjiiiiii"; }; /// From loaded experts, experts_prefill and experts_pack modules. @@ -112,12 +118,15 @@ pub const Affine4Experts = struct { var a: Affine4Experts = undefined; a.up_fn = try decode_mod.function(symbols.up); a.down_fn = try decode_mod.function(symbols.down); + a.pair_fn = try decode_mod.function(symbols.pair); a.pre_up_fn = try prompt_mod.function(symbols.pre_up); a.pre_down_fn = try prompt_mod.function(symbols.pre_down); + a.pre_pair_fn = try prompt_mod.function(symbols.pre_pair); a.pack_fn = try pack_mod.function(symbols.pack); try a.pre_up_fn.allowDynamicShared(prompt_smem); try a.pre_down_fn.allowDynamicShared(prompt_smem); - a.resident = .{ @max(1, try a.up_fn.occupancy(128, 0)) * sms, @max(1, try a.down_fn.occupancy(128, 0)) * sms }; + try a.pre_pair_fn.allowDynamicShared(pair_prompt_smem); + a.resident = .{ @max(1, try a.up_fn.occupancy(128, 0)) * sms, @max(1, try a.down_fn.occupancy(128, 0)) * sms, @max(1, try a.pair_fn.occupancy(128, 0)) * sms }; return a; } @@ -140,31 +149,38 @@ pub const Affine4Experts = struct { return a; } - /// experts.run (decode form): `up` takes token rows (relu^2, bf16 out), else pair rows (fp32 out). - pub fn run(e: *const Affine4Experts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_units: usize) !void { - const grid = @min((max_units + 3) / 4, e.resident[if (up) 0 else 1]); + /// experts.run (decode form): `up` takes token rows (relu^2, or SwiGLU over a `pair` table), else pair rows to f32. + pub fn run(e: *const Affine4Experts, s: Stream, up: bool, pair: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_units: usize) !void { + const which: usize = if (!up) 1 else if (pair) 2 else 0; + const grid = @min((max_units + 3) / 4, e.resident[which]); if (grid < 1) return; var a = args(x, x_stride, slots, w, kg, nb, p, out, n); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 }, .shared = 0 }; - try launch_.launch(if (up) e.up_fn else e.down_fn, cfg, s, &a); + try launch_.launch(if (!up) e.down_fn else if (pair) e.pair_fn else e.up_fn, cfg, s, &a); } - /// experts.prefill: 64 pairs x 128 columns a CTA; `up` relu^2, else bf16 sums (epilogue 3). - pub fn runPrompt(e: *const Affine4Experts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_items: usize) !void { + /// experts.prefill: 64 pairs x 128 columns a CTA; `up` relu^2 (or SwiGLU over a `pair` table), else bf16 sums. + pub fn runPrompt(e: *const Affine4Experts, s: Stream, up: bool, pair: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_items: usize) !void { const grid = max_items * ((nb + 3) / 4); if (grid < 1) return; var a = args(x, x_stride, slots, w, kg, nb, p, out, n); - const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 256 }, .shared = prompt_smem }; - try launch_.launch(if (up) e.pre_up_fn else e.pre_down_fn, cfg, s, &a); + const smem: u32 = if (up and pair) pair_prompt_smem else prompt_smem; + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 256 }, .shared = smem }; + try launch_.launch(if (!up) e.pre_down_fn else if (pair) e.pre_pair_fn else e.pre_up_fn, cfg, s, &a); } /// experts.make: e experts' MLX words (n, k/8), scales and biases (n, k/64) on the device into `out`'s blocks. pub fn pack(e: *const Affine4Experts, s: Stream, words: u64, scales: u64, biases: u64, out: u64, count: usize, n: usize, k: usize) !void { + return e.packInto(s, words, scales, biases, out, count, n, k, 1, 0); + } + + /// pack into the `mi`-th of `m` interleaved tables ([E][n/32][k/64][m][288]): gate then up for the pair form. + pub fn packInto(e: *const Affine4Experts, s: Stream, words: u64, scales: u64, biases: u64, out: u64, count: usize, n: usize, k: usize, m: usize, mi: usize) !void { const kg = k / 64; const nb = n / 32; var a: launch_.Args = .{}; for ([_]u64{ words, scales, biases, out }) |v| a.add(v); - for ([_]usize{ n, k / 8, kg, nb }) |v| a.add(int(v)); + for ([_]usize{ n, k / 8, kg, nb, m, mi }) |v| a.add(int(v)); try launch_.launch(e.pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(count) }, .block = .{ .x = 288 } }, s, &a); } @@ -172,6 +188,11 @@ pub const Affine4Experts = struct { pub fn bytes(count: usize, n: usize, k: usize) usize { return count * (n / 32) * (k / 64) * 288 * 4; } + + /// Bytes of e experts' gate-and-up pair tables ([E][n/32][k/64][2][288] int32). + pub fn pairBytes(count: usize, n: usize, k: usize) usize { + return 2 * bytes(count, n, k); + } }; /// Block-FP8 experts (fp8/experts.cu): the Python packing, kernels and grid choice, decode and prompt forms. diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 756e535a76..65ae27b3d7 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -19,6 +19,7 @@ pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").b pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack").bytes else &.{}; +pub const moe: []const u8 = if (available) &Blob("fatbin_moe").bytes else &.{}; pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_attention").bytes else &.{}; pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes else &.{}; pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; diff --git a/zig/src/cuda/moe.zig b/zig/src/cuda/moe.zig new file mode 100644 index 0000000000..d293ae770d --- /dev/null +++ b/zig/src/cuda/moe.zig @@ -0,0 +1,65 @@ +//! Softmax top-k routing with a sigmoid-gated shared expert (moe.cu), by the rules of tensorfold/cuda/moe.py. + +const std = @import("std"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// Routed experts a row may pick from (a lane holds eight), and the picks a row may hold (a lane a slot). +pub const max_experts = 256; +pub const max_top_k = 31; + +pub const Fns = struct { + router_fn: Function, + topk_fn: Function, + combine_fn: Function, + + /// From the loaded moe module (kernels.moe). + pub fn resolve(mod: Module) !Fns { + return .{ .router_fn = try mod.function("tf_moe_router"), .topk_fn = try mod.function("tf_moe_topk"), .combine_fn = try mod.function("tf_moe_combine") }; + } + + /// logits[r][e] = fp32 x[r] . w[e]: `rows` bf16 rows `x_stride` apart against `count` rows of `dims`, gate last. + pub fn route(f: Fns, s: Stream, x: u64, x_stride: usize, rows: usize, w: u64, logits: u64, dims: usize, count: usize) !void { + var a: launch_.Args = .{}; + a.add(x); + a.add(int(x_stride)); + a.add(w); + a.add(logits); + for ([_]usize{ rows, dims, count }) |v| a.add(int(v)); + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast((rows + 15) / 16), .y = @intCast((count + 15) / 16), .z = 1 }, .block = .{ .x = 256 } }; + try launch_.launch(f.router_fn, cfg, s, &a); + } + + /// Each row's top-k picks and bf16-rounded weights ([rows][top_k + 1]); slot top_k is the shared expert. + pub fn pick(f: Fns, s: Stream, logits: u64, picks: u64, wts: u64, rows: usize, experts: usize, top_k: usize) !void { + if (experts > max_experts or top_k > max_top_k) return error.RoutingTooWide; + var a: launch_.Args = .{}; + a.add(logits); + a.add(picks); + a.add(wts); + a.add(int(experts)); + a.add(int(top_k)); + try launch_.launch(f.topk_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 32 } }, s, &a); + } + + /// out[r] = bf16(sum of w[r][s] * y[r][s]) over the slots in order; y fp32 (decode rows) or bf16 (prompt rows). + pub fn combine(f: Fns, s: Stream, y: u64, y_f32: bool, wts: u64, out: u64, rows: usize, slots: usize, dims: usize) !void { + var a: launch_.Args = .{}; + a.add(y); + a.add(@as(c_int, @intFromBool(y_f32))); + a.add(wts); + a.add(out); + a.add(int(slots)); + a.add(int(dims)); + const blocks = @min((dims + 255) / 256, 16); + try launch_.launch(f.combine_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(blocks), .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } +}; + +fn int(x: usize) c_int { + return @intCast(x); +} diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..629a66bfe7 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -34,6 +34,7 @@ pub const qlinear = @import("qlinear.zig"); pub const norms = @import("norms.zig"); pub const embed_rows = @import("embed_rows.zig"); pub const experts = @import("experts.zig"); +pub const moe = @import("moe.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aec6b9ec1b..c83cdf5123 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -15,6 +15,7 @@ const resume_tests = @import("resume_tests.zig"); const prompt_cuts = @import("prompt_cuts.zig"); const chunk_costs = @import("chunk_costs.zig"); const grouped_tests = @import("grouped_tests.zig"); +const moe_tests = @import("moe_tests.zig"); const fp8_tests = @import("fp8_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); @@ -58,6 +59,7 @@ const usage = \\ prompt-cuts MODEL IDS_FILE a prompt keeping states inside chunks and at their ends (background passes too), and resumed from each, equals one uncut pass \\ recapture MODEL graphs captured again while another sequence is bound: both keep their sampling rules \\ grouped-plan the shared expert plan against experts.route's bytes (oracle/grouped_plan.py) + \\ moe softmax top-k routing and the SwiGLU pair experts against host references (synthetic) \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes @@ -141,6 +143,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "prompt-cuts")) return prompt_cuts.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "recapture")) return recapture_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "grouped-plan")) return grouped_tests.plan(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "moe")) return moe_tests.run(gpu); if (std.mem.eql(u8, cmd, "carveout")) { const mib = if (rest.len > 0) try std.fmt.parseInt(usize, rest[0], 10) else cuda.carveout.default_bytes >> 20; return carveout_tests.run(gpu, if (rest.len > 1) rest[1] else cuda.carveout.default_card, mib); diff --git a/zig/tests/cuda/moe_tests.zig b/zig/tests/cuda/moe_tests.zig new file mode 100644 index 0000000000..2bfd33a415 --- /dev/null +++ b/zig/tests/cuda/moe_tests.zig @@ -0,0 +1,309 @@ +//! Softmax top-k routing (cuda/moe.zig) and the SwiGLU pair form of the affine-4 experts against host references. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +fn bf16f(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +fn f32bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @intCast((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +fn rbf(x: f32) f32 { + return bf16f(f32bf16(x)); +} + +fn randBf16(r: std.Random, scale: f32) u16 { + return f32bf16((r.float(f32) * 2.0 - 1.0) * scale); +} + +fn upload(gpu: Gpu, bytes: []const u8) !cuda.DeviceBuffer { + return cuda.DeviceBuffer.fromHost(gpu.d, bytes); +} + +/// The router's fp32 logits: one fma chain a (row, expert) in index order, as tf_moe_router sums. +fn router(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { + const a = gpu.gpa; + const rows = 5; + const dims = 2048; + const count = 257; + const x = try a.alloc(u16, rows * dims); + defer a.free(x); + const w = try a.alloc(u16, count * dims); + defer a.free(w); + for (x) |*v| v.* = randBf16(r, 1.0); + for (w) |*v| v.* = randBf16(r, 0.1); + const want = try a.alloc(f32, rows * count); + defer a.free(want); + for (0..rows) |i| for (0..count) |e| { + var acc: f32 = 0; + for (0..dims) |k| acc = @mulAdd(f32, bf16f(x[i * dims + k]), bf16f(w[e * dims + k]), acc); + want[i * count + e] = acc; + }; + var dx = try upload(gpu, std.mem.sliceAsBytes(x)); + defer dx.free(); + var dw = try upload(gpu, std.mem.sliceAsBytes(w)); + defer dw.free(); + var dl = try cuda.DeviceBuffer.alloc(gpu.d, rows * count * 4); + defer dl.free(); + try f.route(s, dx.ptr, dims, rows, dw.ptr, dl.ptr, dims, count); + try s.synchronize(); + const got = try check.download(gpu, dl); + defer a.free(got); + try check.sameBytes("router logits", got, std.mem.sliceAsBytes(want)); + check.pass("moe router: {d} rows x {d} experts over {d} inputs equal the host's fma chains", .{ rows, count, dims }); +} + +/// Each row's picks and weights by cuda/moe.py's rule: top-k by logit, ties to the lower id, exp(l - top) / sum. +fn topk(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { + const a = gpu.gpa; + const rows = 7; + const experts = 256; + const top_k = 8; + const slots = top_k + 1; + const logits = try a.alloc(f32, rows * (experts + 1)); + defer a.free(logits); + for (logits) |*v| v.* = (r.float(f32) * 2.0 - 1.0) * 4.0; + // a planted tie, a row of equal logits and a NaN that never wins + logits[1 * (experts + 1) + 3] = logits[1 * (experts + 1) + 200]; + for (0..experts) |e| logits[2 * (experts + 1) + e] = 0.5; + logits[3 * (experts + 1) + 7] = std.math.nan(f32); + const want_ids = try a.alloc(i32, rows * slots); + defer a.free(want_ids); + const want_w = try a.alloc(f32, rows * slots); + defer a.free(want_w); + for (0..rows) |i| { + const row = logits[i * (experts + 1) ..][0..experts]; + var v: [experts]f32 = undefined; + for (&v, row) |*d, l| d.* = if (l != l) -std.math.inf(f32) else l; + var top = -std.math.inf(f32); + for (v) |l| top = @max(top, l); + var total: f32 = 0; + var ex: [top_k]f32 = undefined; + for (0..top_k) |k| { + var bi: usize = 0; + var bv = -std.math.inf(f32); + for (v, 0..) |l, e| if (l > bv) { + bv = l; + bi = e; + }; + ex[k] = @exp(bv - top); + total += ex[k]; + want_ids[i * slots + k] = @intCast(bi); + v[bi] = -std.math.inf(f32); + } + for (0..top_k) |k| want_w[i * slots + k] = rbf(ex[k] / total); + want_ids[i * slots + top_k] = experts; + const sg = rbf(logits[i * (experts + 1) + experts]); + want_w[i * slots + top_k] = rbf(1.0 / (1.0 + @exp(-sg))); + } + var dl = try upload(gpu, std.mem.sliceAsBytes(logits)); + defer dl.free(); + var dp = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer dp.free(); + var dw = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer dw.free(); + try f.pick(s, dl.ptr, dp.ptr, dw.ptr, rows, experts, top_k); + try s.synchronize(); + const got_ids = try check.download(gpu, dp); + defer a.free(got_ids); + try check.sameBytes("top-k ids", got_ids, std.mem.sliceAsBytes(want_ids)); + const got_w = try check.download(gpu, dw); + defer a.free(got_w); + var off: usize = 0; + for (std.mem.bytesAsSlice(f32, got_w), want_w) |g, wv| { + // exp on the device and the host may differ in the last fp32 place: the bf16 weight then moves one step at most + if (g != wv and @abs(g - wv) > @abs(wv) * (1.0 / 128.0)) off += 1; + } + try check.expect(off == 0, "top-k weights: {d} of {d} differ by more than a bf16 step", .{ off, want_w.len }); + check.pass("moe top-k: {d} rows' picks of {d} over {d} experts equal the host's (a tie, a flat row and a NaN included), weights within a bf16 step", .{ rows, top_k, experts }); +} + +/// The weighted sum of a row's slots in slot order, fma in fp32, rounded once to bf16. +fn combine(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { + const a = gpu.gpa; + const rows = 3; + const slots = 4; + const dims = 2048; + const y = try a.alloc(f32, rows * slots * dims); + defer a.free(y); + const w = try a.alloc(f32, rows * slots); + defer a.free(w); + for (y) |*v| v.* = (r.float(f32) * 2.0 - 1.0) * 3.0; + for (w) |*v| v.* = rbf(r.float(f32)); + const want = try a.alloc(u16, rows * dims); + defer a.free(want); + for (0..rows) |i| for (0..dims) |d| { + var acc: f32 = 0; + for (0..slots) |k| acc = @mulAdd(f32, w[i * slots + k], y[(i * slots + k) * dims + d], acc); + want[i * dims + d] = f32bf16(acc); + }; + var dy = try upload(gpu, std.mem.sliceAsBytes(y)); + defer dy.free(); + var dw = try upload(gpu, std.mem.sliceAsBytes(w)); + defer dw.free(); + var dout = try cuda.DeviceBuffer.alloc(gpu.d, rows * dims * 2); + defer dout.free(); + try f.combine(s, dy.ptr, true, dw.ptr, dout.ptr, rows, slots, dims); + try s.synchronize(); + const got = try check.download(gpu, dout); + defer a.free(got); + try check.sameBytes("combined rows", got, std.mem.sliceAsBytes(want)); + check.pass("moe combine: {d} rows of {d} slots equal the host's fma chains rounded to bf16", .{ rows, slots }); +} + +/// Random affine-4 experts (words, scales, biases) as the checkpoint stores them. +const Table = struct { words: []u32, scales: []u16, biases: []u16, count: usize, n: usize, k: usize }; + +fn table(a: std.mem.Allocator, r: std.Random, count: usize, n: usize, k: usize) !Table { + const t: Table = .{ .words = try a.alloc(u32, count * n * k / 8), .scales = try a.alloc(u16, count * n * k / 64), .biases = try a.alloc(u16, count * n * k / 64), .count = count, .n = n, .k = k }; + for (t.words) |*w| w.* = r.int(u32); + for (t.scales) |*v| v.* = f32bf16(0.02 + r.float(f32) * 0.03); + for (t.biases) |*v| v.* = randBf16(r, 0.2); + return t; +} + +/// Weight (e, j, i) dequantized: scale * code + bias. +fn weight(t: Table, e: usize, j: usize, i: usize) f64 { + const row = (e * t.n + j); + const code: f64 = @floatFromInt((t.words[row * (t.k / 8) + i / 8] >> @intCast(4 * (i % 8))) & 15); + return code * bf16f(t.scales[row * (t.k / 64) + i / 64]) + bf16f(t.biases[row * (t.k / 64) + i / 64]); +} + +fn silu(x: f32) f32 { + return x / (1.0 + @exp(-x)); +} + +/// The pair form: gate and up interleaved per group, SwiGLU with the kernel's bf16 rounding points, down to fp32 sums. +fn pairExperts(gpu: Gpu, s: cuda.Stream, r: std.Random) !void { + const a = gpu.gpa; + var mods: [3]cuda.Module = undefined; + mods[0] = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mods[0].unload(); + mods[1] = try cuda.Module.load(gpu.d, cuda.kernels.experts_prefill); + defer mods[1].unload(); + mods[2] = try cuda.Module.load(gpu.d, cuda.kernels.experts_pack); + defer mods[2].unload(); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + const ex = try cuda.experts.Affine4Experts.resolve(mods[0], mods[1], mods[2], sms); + const plan_fns = try cuda.grouped.Router.resolve(mods[0]); + const count = 6; + const ni = 64; + const dims = 128; + const rows = 3; + const slots = 4; + const pairs = rows * slots; + const gate = try table(a, r, count, ni, dims); + defer { + a.free(gate.words); + a.free(gate.scales); + a.free(gate.biases); + } + const up = try table(a, r, count, ni, dims); + defer { + a.free(up.words); + a.free(up.scales); + a.free(up.biases); + } + const down = try table(a, r, count, dims, ni); + defer { + a.free(down.words); + a.free(down.scales); + a.free(down.biases); + } + const x = try a.alloc(u16, rows * dims); + defer a.free(x); + for (x) |*v| v.* = randBf16(r, 1.0); + const picks = [pairs]i32{ 0, 2, 5, 5, 1, 2, 3, 5, 4, 4, 0, 5 }; + // the host reference + const want = try a.alloc(f64, pairs * dims); + defer a.free(want); + const act = try a.alloc(f32, ni); + defer a.free(act); + for (0..pairs) |p| { + const e: usize = @intCast(picks[p]); + const xr = x[(p / slots) * dims ..][0..dims]; + for (0..ni) |j| { + var g: f64 = 0; + var u: f64 = 0; + for (0..dims) |i| { + g += weight(gate, e, j, i) * bf16f(xr[i]); + u += weight(up, e, j, i) * bf16f(xr[i]); + } + const gv = rbf(@floatCast(g)); + const uv = rbf(@floatCast(u)); + act[j] = rbf(rbf(silu(gv)) * uv); + } + for (0..dims) |d| { + var y: f64 = 0; + for (0..ni) |j| y += weight(down, e, d, j) * act[j]; + want[p * dims + d] = y; + } + } + // the device: packed tables, the plan, up (pair) then down + var dpair = try cuda.DeviceBuffer.alloc(gpu.d, cuda.experts.Affine4Experts.pairBytes(count, ni, dims)); + defer dpair.free(); + var ddown = try cuda.DeviceBuffer.alloc(gpu.d, cuda.experts.Affine4Experts.bytes(count, dims, ni)); + defer ddown.free(); + for ([_]struct { Table, u64, usize, usize }{ .{ gate, dpair.ptr, 2, 0 }, .{ up, dpair.ptr, 2, 1 }, .{ down, ddown.ptr, 1, 0 } }) |item| { + const t = item[0]; + var w = try upload(gpu, std.mem.sliceAsBytes(t.words)); + defer w.free(); + var sc = try upload(gpu, std.mem.sliceAsBytes(t.scales)); + defer sc.free(); + var bi = try upload(gpu, std.mem.sliceAsBytes(t.biases)); + defer bi.free(); + try ex.packInto(s, w.ptr, sc.ptr, bi.ptr, item[1], t.count, t.n, t.k, item[2], item[3]); + try s.synchronize(); + } + var dx = try upload(gpu, std.mem.sliceAsBytes(x)); + defer dx.free(); + var dpicks = try upload(gpu, std.mem.sliceAsBytes(&picks)); + defer dpicks.free(); + var members = try cuda.DeviceBuffer.alloc(gpu.d, pairs * 4); + defer members.free(); + var items = try cuda.DeviceBuffer.alloc(gpu.d, cuda.grouped.maxItems(pairs, count, 16) * 12); + defer items.free(); + var counts = try cuda.DeviceBuffer.alloc(gpu.d, 8); + defer counts.free(); + var dact = try cuda.DeviceBuffer.alloc(gpu.d, pairs * ni * 2); + defer dact.free(); + var dy = try cuda.DeviceBuffer.alloc(gpu.d, pairs * dims * 4); + defer dy.free(); + const plan: cuda.grouped.Plan = .{ .members = members.ptr, .items = items.ptr, .counts = counts.ptr, .rank = 0, .hist = 0 }; + try plan_fns.route(s, dpicks.ptr, pairs, count, 16, plan); + const layer: cuda.experts.Layer = .{ .format = .affine4, .up = dpair.ptr, .down = ddown.ptr, .width = ni, .dims = dims, .experts = count, .swiglu = true }; + const grouped: cuda.experts.Grouped = .{ .affine4 = &ex }; + const items_n = cuda.grouped.maxItems(pairs, count, 16); + try grouped.decode(s, .up, .{ .x = dx.ptr, .stride = dims, .slots = slots }, layer, plan, dact.ptr, items_n, -1); + try grouped.decode(s, .down_f32, .{ .x = dact.ptr, .stride = ni }, layer, plan, dy.ptr, items_n, -1); + try s.synchronize(); + const got = try check.download(gpu, dy); + defer a.free(got); + var worst: f64 = 0; + var scale: f64 = 0; + for (std.mem.bytesAsSlice(f32, got), want) |g, wv| { + worst = @max(worst, @abs(@as(f64, g) - wv)); + scale = @max(scale, @abs(wv)); + } + try check.expect(worst <= scale * 2e-3, "pair experts: worst error {e} against a largest sum of {e}", .{ worst, scale }); + check.pass("moe experts: {d} pairs over {d} SwiGLU experts (gate and up interleaved) within {e} of the fp64 reference (largest sum {e})", .{ pairs, count, worst, scale }); +} + +pub fn run(gpu: Gpu) !void { + var prng = std.Random.DefaultPrng.init(0x5eed); + const r = prng.random(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.moe); + defer mod.unload(); + const f = try cuda.moe.Fns.resolve(mod); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + try router(gpu, f, s, r); + try topk(gpu, f, s, r); + try combine(gpu, f, s, r); + try pairExperts(gpu, s, r); +} From 73ff8925d9bad870d7808c55a0850d260d092984 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:13:52 +0000 Subject: [PATCH 28/35] qwen3_5: Qwen3.6-35B-A3B's routed experts on the CUDA engine model_type qwen3_5_moe serves on the same engine as the 2B and the 27B: the layers are theirs, with routed experts in place of the dense MLP. The geometry reads num_experts, num_experts_per_tok and moe_intermediate_size (the shared expert's width must equal it, as it runs in the routed experts' tables). The loader dequantizes the 8-bit router and shared-expert gate rows to bf16 on the host, and packs each layer's 256 experts plus the shared one as a SwiGLU pair table and a down table. A layer's MLP is a union of the dense projections and the routed set; the forward routes, plans, runs the pair and down kernels and combines into the residual, inside the captured graphs as before. Prompt chunks past a window's rows take the prefill kernels, shorter ones the decode form, so a short chunk keeps one-row steps' bits. A drafter is refused (takes_drafter false): the checkpoint drafts with its own MTP layer, which is not on CUDA yet. tf-qwen35-cuda gains a probe mode that prints one layer's picks, weights and logits for a reference to compare. On an RTX 4090, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP: one-row steps decode at 141 tok/s; windows and trees of up to 16 rows equal the steps byte for byte over 300 ids; the 2B's and 27B's bits are unchanged. Against an fp32 PyTorch reference the steps agree within the model's own rounding sensitivity (mean KL 0.015 over 40 steps, where two references differing only by bf16 row rounding differ by 0.016; the dense 2B agrees at 0.0005). --- zig/src/cli/hub.zig | 1 + zig/src/families/qwen3_5/cuda.zig | 1 + zig/src/families/qwen3_5/cuda_engine.zig | 83 +++++++++++++++--- zig/src/families/qwen3_5/cuda_geometry.zig | 53 +++++++++++- zig/src/families/qwen3_5/cuda_native_moe.zig | 13 +++ zig/src/families/qwen3_5/cuda_weights.zig | 88 +++++++++++++++++++- zig/src/native/cuda.zig | 2 +- zig/tests/qwen35_cuda_main.zig | 42 +++++++++- 8 files changed, 264 insertions(+), 19 deletions(-) create mode 100644 zig/src/families/qwen3_5/cuda_native_moe.zig diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 9b5edacb0e..7474d2f077 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -13,6 +13,7 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "qwen4_exp", .title = "Qwen3.8 Flash Next" }, .{ .model_type = "glm5_next", .title = "GLM-5.3-Flash" }, .{ .model_type = "qwen3_5", .title = "Qwen3.5-2B or Qwen3.8-27B" }, + .{ .model_type = "qwen3_5_moe", .title = "Qwen3.6-35B-A3B" }, }; /// The family that serves `model_type` from the engine registry, or null when Zig cannot serve it (the 0.6 line may). diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 47158adda9..0f2d62fcce 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -3,6 +3,7 @@ pub const config = @import("config.zig"); pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); +pub const native_moe = @import("cuda_native_moe.zig"); pub const draft = @import("cuda_draft.zig"); test { diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index 15975b8d42..c3efc4525d 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -35,12 +35,12 @@ pub const Geo = extern struct { hid: c_int, nk: c_int, nv: c_int, qh: c_int, kvh pub const Delta = struct { qkv: Affine4, z: Affine4, a: Affine4, b: Affine4, out: Affine4, conv: u64, norm: u64, a_log: u64, dt_bias: u64 }; pub const Attention = struct { q: Affine4, k: Affine4, v: Affine4, o: Affine4, q_norm: u64, k_norm: u64 }; +/// A dense SwiGLU MLP, or routed experts (the shared one last in their tables) behind a router of bf16 rows. +pub const Mlp = union(enum) { dense: struct { gate: Affine4, up: Affine4, down: Affine4 }, moe: struct { router: u64, experts: cuda.experts.Layer } }; pub const Layer = struct { input_norm: u64, post_norm: u64, - gate: Affine4, - up: Affine4, - down: Affine4, + mlp: Mlp, mixer: union(enum) { delta: Delta, attention: Attention }, }; @@ -80,9 +80,12 @@ pub const Engine = struct { g: Geometry, geo: Geo, stream: cuda.Stream, - mods: [7]cuda.Module, + mods: [11]cuda.Module, draws: cuda.sample.Draws, affine: cuda.qlinear.Affine4Kernels, + experts: cuda.experts.Affine4Experts, + plan: cuda.grouped.Router, + moe: cuda.moe.Fns, fns: Fns, keeps: Keeps, /// Every weight and scratch buffer in one allocation: WDDM's launch cost grows with the allocation count. @@ -127,6 +130,17 @@ pub const Engine = struct { pl: u64, po: u64, logits: u64, + // routed experts' scratch (qwen3_5_moe): logits, picks and weights a row, the plan, the pairs' activations and sums + rlogits: u64 = 0, + picks: u64 = 0, + wts: u64 = 0, + members: u64 = 0, + items: u64 = 0, + counts: u64 = 0, + rank: u64 = 0, + hist: u64 = 0, + eact: u64 = 0, + ey: u64 = 0, // per-forward tables (ids, segments, rows' segment, parent and ancestors, the draws' rules) staged in pinned memory tables: DeviceBuffer, staged: cuda.HostBuffer, @@ -174,7 +188,7 @@ pub const Engine = struct { e.stage = try DeviceBuffer.alloc(d, weights.stageBytes(g)); errdefer e.stage.free(); const kk = cuda.kernels; - const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample, kk.qwen35_keep }; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.qwen35, kk.sample, kk.qwen35_keep, kk.experts, kk.experts_prefill, kk.experts_pack, kk.moe }; var loaded: usize = 0; errdefer for (e.mods[0..loaded]) |*m| m.unload(); for (images, 0..) |img, i| { @@ -186,6 +200,9 @@ pub const Engine = struct { e.affine = try cuda.qlinear.Affine4Kernels.resolve(d, e.mods[0], e.mods[1], e.mods[2], e.mods[3], sms, gb10); errdefer e.affine.deinit(); e.draws = try cuda.sample.Draws.resolve(e.mods[5]); + e.experts = try cuda.experts.Affine4Experts.resolve(e.mods[7], e.mods[8], e.mods[9], sms); + e.plan = try cuda.grouped.Router.resolve(e.mods[7]); + e.moe = try cuda.moe.Fns.resolve(e.mods[10]); inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(e.fns, name) = try e.mods[4].function("tf_q35_" ++ name); inline for (@typeInfo(Keeps).@"struct".field_names) |name| @field(e.keeps, name) = try e.mods[6].function("tf_q35k_" ++ name); e.layers = try gpa.alloc(Layer, g.layers); @@ -292,6 +309,20 @@ pub const Engine = struct { e.pl = try t.alloc(R * g.q_heads * max_chunks * 4); e.po = try t.alloc(R * g.q_heads * max_chunks * head_dim * 4); e.logits = try t.alloc(most_verify * g.vocab * 2); + if (g.moe()) { + const pairs = R * g.slots(); + const count = g.experts + 1; + e.rlogits = try t.alloc(R * count * 4); + e.picks = try t.alloc(pairs * 4); + e.wts = try t.alloc(pairs * 4); + e.members = try t.alloc(pairs * 4); + e.items = try t.alloc(3 * cuda.grouped.maxItems(pairs, count, 16) * 4); + e.counts = try t.alloc(2 * 4); + e.rank = try t.alloc(pairs * 4); + e.hist = try t.alloc(((pairs + 1023) / 1024) * count * 4); + e.eact = try t.alloc(pairs * g.intermediate * 2); + e.ey = try t.alloc(pairs * g.hidden * 4); + } return total; } @@ -501,11 +532,16 @@ pub const Engine = struct { }, } try e.norm(true, l.post_norm, total); - try e.proj(l.gate, e.x, e.gate, total, prompt); - try e.proj(l.up, e.x, e.up, total, prompt); - const n: c_longlong = @intCast(total * g.intermediate); - try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n, t_params, @as(c_int, @intCast(g.intermediate)) }); - try e.proj(l.down, e.act, e.r, total, prompt); + switch (l.mlp) { + .dense => |m| { + try e.proj(m.gate, e.x, e.gate, total, prompt); + try e.proj(m.up, e.x, e.up, total, prompt); + const n: c_longlong = @intCast(total * g.intermediate); + try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n, t_params, @as(c_int, @intCast(g.intermediate)) }); + try e.proj(m.down, e.act, e.r, total, prompt); + }, + .moe => |m| try e.routed(m.router, m.experts, total, prompt), + } if (shape.taps) for (e.taps[0..e.n_taps], 0..) |t, ti| if (t == i) { try e.run(e.fns.tap, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, e.r, t_segs, t_rows, @as(c_int, @intCast(ti)), @as(c_int, @intCast(g.hidden)), t_params }); }; @@ -515,6 +551,33 @@ pub const Engine = struct { try e.proj(e.head, e.x + (total - heads) * g.hidden * 2, e.logits, heads, false); } + /// The routed experts over the normed rows in x: the picks, their plan, SwiGLU up, down, the weighted sum into r. + fn routed(e: *Engine, router: u64, layer: cuda.experts.Layer, total: usize, prompt: bool) !void { + const g = e.g; + const slots = g.slots(); + const count = g.experts + 1; + const pairs = total * slots; + // a short prompt chunk takes the decode form, as the dense projections do: its rows keep one-row steps' bits + const wide = prompt and total > most_window; + const tile: usize = if (wide) 64 else 16; + const items = cuda.grouped.maxItems(pairs, count, tile); + const plan: cuda.grouped.Plan = .{ .members = e.members, .items = e.items, .counts = e.counts, .rank = e.rank, .hist = e.hist }; + const grouped: cuda.experts.Grouped = .{ .affine4 = &e.experts }; + try e.moe.route(e.stream, e.x, g.hidden, total, router, e.rlogits, g.hidden, count); + try e.moe.pick(e.stream, e.rlogits, e.picks, e.wts, total, g.experts, g.top_k); + try e.plan.route(e.stream, e.picks, pairs, count, tile, plan); + const rows_in: cuda.experts.Rows = .{ .x = e.x, .stride = g.hidden, .slots = slots }; + const act_in: cuda.experts.Rows = .{ .x = e.eact, .stride = g.intermediate }; + if (wide) { + try grouped.prompt(e.stream, .up, rows_in, layer, plan, e.eact, items, -1); + try grouped.prompt(e.stream, .down_bf16, act_in, layer, plan, e.ey, items, -1); + } else { + try grouped.decode(e.stream, .up, rows_in, layer, plan, e.eact, items, -1); + try grouped.decode(e.stream, .down_f32, act_in, layer, plan, e.ey, items, -1); + } + try e.moe.combine(e.stream, e.ey, !wide, e.wts, e.r, total, slots, g.hidden); + } + /// Row `i`'s logits from the newest forward that kept logits. pub fn rowLogits(e: *const Engine, i: usize) []const u16 { return e.host_logits.slice(u16)[i * e.g.vocab ..][0..e.g.vocab]; diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig index fdd6ecb0c2..fecc4f4c01 100644 --- a/zig/src/families/qwen3_5/cuda_geometry.zig +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -21,6 +21,9 @@ pub const Geometry = struct { tied: bool, context: usize, linear: [max_layers]bool, + /// Routed experts (qwen3_5_moe): their count and the picks a row makes; 0: a dense MLP of `intermediate`. + experts: usize = 0, + top_k: usize = 0, pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Geometry { const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); @@ -30,6 +33,14 @@ pub const Geometry = struct { return parse(gpa, bytes); } + /// Whether the MLPs are routed experts, whose width is `intermediate` (the shared expert's too). + pub fn moe(g: Geometry) bool { + return g.experts != 0; + } + /// A row's expert slots: its picks, then the shared expert. + pub fn slots(g: Geometry) usize { + return g.top_k + 1; + } /// Value-head width of the gated delta rule's rows (and of z and the out projection's input). pub fn valueWidth(g: Geometry) usize { return g.v_heads * dk; @@ -118,9 +129,15 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); defer doc.deinit(); const root = try object(doc.value); - try is(root, "model_type", []const u8, "qwen3_5"); + const moe = blk: { + const v = root.get("model_type") orelse return error.UnsupportedQwenConfig; + if (v != .string) return error.UnsupportedQwenConfig; + if (std.mem.eql(u8, v.string, "qwen3_5")) break :blk false; + if (std.mem.eql(u8, v.string, "qwen3_5_moe")) break :blk true; + return error.UnsupportedQwenConfig; + }; const t = try object(root.get("text_config") orelse return error.BadQwenConfig); - try is(t, "model_type", []const u8, "qwen3_5_text"); + try is(t, "model_type", []const u8, if (moe) "qwen3_5_moe_text" else "qwen3_5_text"); try is(t, "attention_bias", bool, false); try is(t, "attn_output_gate", bool, true); try is(t, "hidden_act", []const u8, "silu"); @@ -136,7 +153,7 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { if (q.get("mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "affine")) return error.UnsupportedQwenQuantization; var g: Geometry = .{ .hidden = try integer(t, "hidden_size"), - .intermediate = try integer(t, "intermediate_size"), + .intermediate = try integer(t, if (moe) "moe_intermediate_size" else "intermediate_size"), .layers = try integer(t, "num_hidden_layers"), .vocab = try integer(t, "vocab_size"), .k_heads = try integer(t, "linear_num_key_heads"), @@ -153,6 +170,13 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { .linear = @splat(false), }; if (g.layers > max_layers or g.hidden % 256 != 0 or g.hidden > 8192 or g.intermediate % 64 != 0) return error.UnsupportedQwenGeometry; + if (moe) { + g.experts = try integer(t, "num_experts"); + g.top_k = try integer(t, "num_experts_per_tok"); + // the shared expert runs in the routed experts' table, so it needs their width + if (try integer(t, "shared_expert_intermediate_size") != g.intermediate) return error.UnsupportedQwenGeometry; + if (g.experts > 256 or g.top_k > 31 or g.top_k >= g.experts) return error.UnsupportedQwenGeometry; + } if (g.v_heads % g.k_heads != 0 or g.q_heads % g.kv_heads != 0 or g.q_heads > max_query_heads) return error.UnsupportedQwenGeometry; const types = t.get("layer_types") orelse return error.UnsupportedQwenGeometry; if (types != .array or types.array.items.len != g.layers) return error.UnsupportedQwenGeometry; @@ -163,6 +187,29 @@ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Geometry { return g; } +test "the 35B-A3B's routed experts: 256 of width 512, 8 a row, the shared one the same width" { + const text = + \\{"model_type":"qwen3_5_moe","quantization":{"group_size":64,"bits":4,"mode":"affine"},"tie_word_embeddings":false, + \\ "text_config":{"model_type":"qwen3_5_moe_text","attention_bias":false,"attn_output_gate":true,"hidden_act":"silu", + \\ "head_dim":256,"linear_key_head_dim":128,"linear_value_head_dim":128,"linear_conv_kernel_dim":4, + \\ "hidden_size":2048,"moe_intermediate_size":512,"shared_expert_intermediate_size":512,"num_experts":256, + \\ "num_experts_per_tok":8,"num_hidden_layers":4,"vocab_size":248320,"linear_num_key_heads":16, + \\ "linear_num_value_heads":32,"num_attention_heads":16,"num_key_value_heads":2,"rms_norm_eps":1e-06, + \\ "max_position_embeddings":262144, + \\ "layer_types":["linear_attention","linear_attention","linear_attention","full_attention"], + \\ "rope_parameters":{"rope_type":"default","rope_theta":10000000,"partial_rotary_factor":0.25}}} + ; + const g = try parse(std.testing.allocator, text); + try std.testing.expect(g.moe()); + try std.testing.expectEqual(@as(usize, 256), g.experts); + try std.testing.expectEqual(@as(usize, 9), g.slots()); + try std.testing.expectEqual(@as(usize, 512), g.intermediate); + try std.testing.expectEqual(@as(usize, 8192), g.convWidth()); + const odd = try std.mem.replaceOwned(u8, std.testing.allocator, text, "\"shared_expert_intermediate_size\":512", "\"shared_expert_intermediate_size\":1024"); + defer std.testing.allocator.free(odd); + try std.testing.expectError(error.UnsupportedQwenGeometry, parse(std.testing.allocator, odd)); +} + test "the 27B's widths: 48 value heads over 16 key heads, 24 query heads over 4" { const text = \\{"model_type":"qwen3_5","quantization":{"group_size":64,"bits":4,"mode":"affine"},"tie_word_embeddings":false, diff --git a/zig/src/families/qwen3_5/cuda_native_moe.zig b/zig/src/families/qwen3_5/cuda_native_moe.zig new file mode 100644 index 0000000000..c8cda02d6f --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_native_moe.zig @@ -0,0 +1,13 @@ +//! The qwen3_5_moe model type (Qwen3.6-35B-A3B) on the same CUDA engine: routed experts in place of the dense MLP. +const native = @import("cuda_native.zig"); + +pub const model_type = "qwen3_5_moe"; +pub const formats = native.formats; +pub const default_context = native.default_context; +pub const max_segments = native.max_segments; +pub const prompt_rows = native.prompt_rows; +/// Drafts come from the checkpoint's MTP layer, not a separate draft model. +pub const takes_drafter = false; +pub const open = native.open; +pub const openDrafter = native.openDrafter; +pub const explain = native.explain; diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig index 2a2b75c361..ea596a964b 100644 --- a/zig/src/families/qwen3_5/cuda_weights.zig +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -30,7 +30,11 @@ pub fn total(g: Geometry) usize { t += packedBytes(g.vocab, g.hidden); // the head, tied or not t += aligned(g.hidden * 2); for (g.linear[0..g.layers]) |lin| { - t += 2 * aligned(g.hidden * 2) + 2 * packedBytes(g.intermediate, g.hidden) + packedBytes(g.hidden, g.intermediate); + t += 2 * aligned(g.hidden * 2); + if (g.moe()) { + const ex = cuda.experts.Affine4Experts; + t += aligned(ex.pairBytes(g.experts + 1, g.intermediate, g.hidden)) + aligned(ex.bytes(g.experts + 1, g.hidden, g.intermediate)) + aligned((g.experts + 1) * g.hidden * 2); + } else t += 2 * packedBytes(g.intermediate, g.hidden) + packedBytes(g.hidden, g.intermediate); if (lin) { t += packedBytes(g.convWidth(), g.hidden) + packedBytes(g.valueWidth(), g.hidden) + 2 * packedBytes(g.v_heads, g.hidden) + packedBytes(g.hidden, g.valueWidth()); t += aligned(g.convWidth() * geometry.taps * 2) + aligned(geometry.dk * 2) + aligned(g.v_heads * 4) + aligned(g.v_heads * 2); @@ -46,6 +50,7 @@ pub fn total(g: Geometry) usize { pub fn stageBytes(g: Geometry) usize { var most: usize = 0; for ([_][2]usize{ .{ g.vocab, g.hidden }, .{ g.intermediate, g.hidden }, .{ g.hidden, g.intermediate }, .{ g.convWidth(), g.hidden }, .{ g.valueWidth(), g.hidden }, .{ 2 * g.attnWidth(), g.hidden } }) |nk| most = @max(most, rawBytes(nk[0], nk[1])); + if (g.moe()) most = @max(most, @max(rawBytes(g.experts * g.intermediate, g.hidden), rawBytes(g.experts * g.hidden, g.intermediate))); return most; } @@ -95,6 +100,79 @@ fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8 return q; } +fn bf16f(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// f32 to bf16, round to nearest even (finite values). +fn f32bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @intCast((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +/// 8-bit affine rows (four codes a word, a scale and bias a group of 64) dequantized to bf16: the router rows. +fn dequant8(dst: []u16, words: []align(1) const u32, scales: []align(1) const u16, biases: []align(1) const u16, k: usize) void { + const groups = k / 64; + for (0..dst.len / k) |r| { + for (0..k) |i| { + const code: f32 = @floatFromInt((words[r * (k / 4) + i / 4] >> @intCast(8 * (i % 4))) & 0xff); + dst[r * k + i] = f32bf16(code * bf16f(scales[r * groups + i / 64]) + bf16f(biases[r * groups + i / 64])); + } + } +} + +/// The router's rows as bf16 [experts + 1][hidden]: the routed experts' gate, then the shared expert's gate row. +fn routerRows(e: *Engine, ck: *core.Checkpoint, prefix: []const u8) !u64 { + const g = e.g; + const rows = try e.gpa.alloc(u16, (g.experts + 1) * g.hidden); + defer e.gpa.free(rows); + var buf: [160]u8 = undefined; + const names = [_]struct { []const u8, usize, usize }{ .{ "mlp.gate", 0, g.experts }, .{ "mlp.shared_expert_gate", g.experts, 1 } }; + for (names) |n| { + const w = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, n[0] }), .u32, &.{ n[2], g.hidden / 4 }); + const s = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.scales", .{ prefix, n[0] }), .bf16, &.{ n[2], g.hidden / 64 }); + const b = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.biases", .{ prefix, n[0] }), .bf16, &.{ n[2], g.hidden / 64 }); + dequant8(rows[n[1] * g.hidden ..][0 .. n[2] * g.hidden], std.mem.bytesAsSlice(u32, w.bytes), std.mem.bytesAsSlice(u16, s.bytes), std.mem.bytesAsSlice(u16, b.bytes), g.hidden); + } + return upload(e, std.mem.sliceAsBytes(rows)); +} + +/// `count` stacked experts' projection (one at count 0) staged as stored and packed into table `out` as the mi-th of m. +fn expertTable(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, count: usize, n: usize, k: usize, out: u64, m: usize, mi: usize) !void { + var buf: [160]u8 = undefined; + var raw: [3]u64 = undefined; + var at: usize = 0; + const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (parts, &raw) |p, *r| { + const stacked: []const usize = &.{ count, n, p[2] }; + const one: []const usize = &.{ n, p[2] }; + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], if (count == 0) one else stacked); + try e.stage.upload(at, t.bytes); + r.* = e.stage.ptr + at; + at += t.bytes.len; + } + try e.experts.packInto(e.stream, raw[0], raw[1], raw[2], out, @max(count, 1), n, k, m, mi); + try e.stream.synchronize(); // the next table's staging overwrites this one's +} + +/// A layer's routed experts: gate and up interleaved as one pair table, down apart, the shared expert last in each. +fn routed(e: *Engine, ck: *core.Checkpoint, p: []const u8) !engine.Mlp { + const g = e.g; + const ex = cuda.experts.Affine4Experts; + const ni = g.intermediate; + const pair = try e.alloc(ex.pairBytes(g.experts + 1, ni, g.hidden)); + const down = try e.alloc(ex.bytes(g.experts + 1, g.hidden, ni)); + try expertTable(e, ck, p, "mlp.switch_mlp.gate_proj", g.experts, ni, g.hidden, pair, 2, 0); + try expertTable(e, ck, p, "mlp.switch_mlp.up_proj", g.experts, ni, g.hidden, pair, 2, 1); + try expertTable(e, ck, p, "mlp.switch_mlp.down_proj", g.experts, g.hidden, ni, down, 1, 0); + const shared_pair = pair + ex.pairBytes(g.experts, ni, g.hidden); + try expertTable(e, ck, p, "mlp.shared_expert.gate_proj", 0, ni, g.hidden, shared_pair, 2, 0); + try expertTable(e, ck, p, "mlp.shared_expert.up_proj", 0, ni, g.hidden, shared_pair, 2, 1); + try expertTable(e, ck, p, "mlp.shared_expert.down_proj", 0, g.hidden, ni, down + ex.bytes(g.experts, g.hidden, ni), 1, 0); + const layer: cuda.experts.Layer = .{ .format = .affine4, .up = pair, .down = down, .width = ni, .dims = g.hidden, .experts = g.experts + 1, .swiglu = true }; + return .{ .moe = .{ .router = try routerRows(e, ck, p), .experts = layer } }; +} + pub fn load(e: *Engine, ck: *core.Checkpoint) !void { const g = e.g; const m = "language_model.model."; @@ -111,9 +189,11 @@ pub fn load(e: *Engine, ck: *core.Checkpoint) !void { l.* = .{ .input_norm = try tensor(e, ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{g.hidden}), .post_norm = try tensor(e, ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{g.hidden}), - .gate = try linear(e, ck, p, "mlp.gate_proj", g.intermediate, g.hidden), - .up = try linear(e, ck, p, "mlp.up_proj", g.intermediate, g.hidden), - .down = try linear(e, ck, p, "mlp.down_proj", g.hidden, g.intermediate), + .mlp = if (g.moe()) try routed(e, ck, p) else .{ .dense = .{ + .gate = try linear(e, ck, p, "mlp.gate_proj", g.intermediate, g.hidden), + .up = try linear(e, ck, p, "mlp.up_proj", g.intermediate, g.hidden), + .down = try linear(e, ck, p, "mlp.down_proj", g.hidden, g.intermediate), + } }, .mixer = undefined, }; if (g.linear[i]) { diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index ae7e82d1e5..8ce4e0bebb 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -13,7 +13,7 @@ const nemotron_slide = @import("nemotron_slide.zig"); const drafted = @import("cuda_drafted.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, qwen35.native }; +const registry = .{ nemotron.native, qwen35.native, qwen35.native_moe }; /// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index e0bb3ac025..50db48a217 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -1,4 +1,4 @@ -//! `tf-qwen35-cuda `: the ids' logits as bf16 rows for a reference. +//! `tf-qwen35-cuda `: bf16 logit rows for a reference, or picks. const std = @import("std"); const cuda = @import("cuda"); const q = @import("qwen35"); @@ -89,6 +89,46 @@ pub fn main(init: std.process.Init) !u8 { if (n > 1) try cache.keep(e, path[0..kept]); at += kept; } + } else if (std.mem.eql(u8, args[2], "probe")) { + // the first id through layers 0..N (N = the out argument): that layer's routing and rows, for a reference + const layer = try std.fmt.parseInt(usize, args[4], 10); + const all = e.layers; + e.layers = all[0 .. layer + 1]; + e.graphs.clearRetainingCapacity(); // the warm-up captured every layer; this forward captures the cut + try e.forward(&.{.{ .cache = &cache, .rows = 1 }}, ids.items[0..1], .logits); + e.layers = all; + const g = e.g; + if (!g.moe()) return error.NotRouted; + const count = g.experts + 1; + const logits_f = try a.alloc(f32, count); + const picks = try a.alloc(i32, g.slots()); + const wts = try a.alloc(f32, g.slots()); + var xrow: [8]u16 = undefined; + var rrow: [8]u16 = undefined; + try e.arena.download(e.rlogits - e.arena.ptr, std.mem.sliceAsBytes(logits_f)); + try e.arena.download(e.picks - e.arena.ptr, std.mem.sliceAsBytes(picks)); + try e.arena.download(e.wts - e.arena.ptr, std.mem.sliceAsBytes(wts)); + try e.arena.download(e.x - e.arena.ptr, std.mem.sliceAsBytes(&xrow)); + try e.arena.download(e.r - e.arena.ptr, std.mem.sliceAsBytes(&rrow)); + std.debug.print("probe layer {d} x[:8]", .{layer}); + for (xrow) |h| std.debug.print(" {d:.5}", .{@as(f32, @bitCast(@as(u32, h) << 16))}); + std.debug.print("\nprobe top logits", .{}); + for (0..10) |_| { + var best: usize = 0; + for (logits_f[0 .. count - 1], 0..) |v, i| if (v > logits_f[best]) { + best = i; + }; + std.debug.print(" ({d}, {d:.5})", .{ best, logits_f[best] }); + logits_f[best] = -std.math.inf(f32); + } + std.debug.print(" shared gate {d:.5}\nprobe picks", .{logits_f[count - 1]}); + for (picks) |pk| std.debug.print(" {d}", .{pk}); + std.debug.print(" weights", .{}); + for (wts) |w| std.debug.print(" {d:.5}", .{w}); + std.debug.print("\nprobe r[:8]", .{}); + for (rrow) |h| std.debug.print(" {d:.5}", .{@as(f32, @bitCast(@as(u32, h) << 16))}); + std.debug.print("\n", .{}); + return 0; } else { var at: usize = 0; while (at < ids.items.len) { From 6cec3f42e27781c4fee016462f6ce966bbc471bc Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:34:41 +0000 Subject: [PATCH 29/35] qwen3_5: the MTP layer drafts chains for Qwen3.6-35B-A3B on CUDA --drafter loads mtp-4bit.safetensors beside the weights: one attention-plus-experts layer over [norm(embedding of the next token) | norm(the row's final normed state)], the target's head for its draft, as the frozen engine's cuda/mtp.py ran it. The drafter taps the target's last layer and applies the final norm itself; each stream keeps the layer's keys and values over its whole window in a ring, absorbs the kept rows with the token after each (the newest row's output is the first draft), and holds a chain of up to TF_MTP_DRAFTS (default 3) more, each row the previous output with the draft before it. Drafts only propose: the target's windows verify them, so replies equal "draft": false ones. qwen36_mtp.cu holds the attention at the target's head width (256) with its partial rotary and output gate over the ring and the causal block; the ring kernels are the DFlash2 drafter's. The expert tables and router rows pack through the same loader as the target's, now writing into any arena. On an RTX 4090, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP with --parallel 2 --context 8192: one stream decodes 258 / 229 tok/s greedy (code / chat) and 252 / 181 sampled, against 101 / 98 plain; two streams 423 tok/s aggregate on code. bench_concurrent --alone --serial: every drafted reply equals its "draft": false reply and its lone run, 0 unequal. A chained draft costs 0.5 ms; a hold 1.2 ms, an absorb 0.7 ms. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/qwen36_mtp.cu | 115 +++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/qwen3_5/cuda.zig | 1 + zig/src/families/qwen3_5/cuda_engine.zig | 14 +- zig/src/families/qwen3_5/cuda_geometry.zig | 2 + zig/src/families/qwen3_5/cuda_mtp.zig | 417 ++++++++++++++++++ zig/src/families/qwen3_5/cuda_mtp_weights.zig | 80 ++++ zig/src/families/qwen3_5/cuda_native.zig | 25 +- zig/src/families/qwen3_5/cuda_native_moe.zig | 6 +- zig/src/families/qwen3_5/cuda_weights.zig | 127 ++++-- 11 files changed, 729 insertions(+), 60 deletions(-) create mode 100644 zig/kernels/cuda/qwen36_mtp.cu create mode 100644 zig/src/families/qwen3_5/cuda_mtp.zig create mode 100644 zig/src/families/qwen3_5/cuda_mtp_weights.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 3a994fd300..210b37a8aa 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -36,6 +36,7 @@ const kernels = [_]Kernel{ .{ .name = "qwen35", .flags = &.{"-O3"} }, // ours: the Qwen3.5 hybrids' glue, the Metal kernels' rounding points .{ .name = "qwen35_draft", .flags = &.{"-O3"} }, // ours: the DFlash2 drafter's glue (draft_ops.metal's rounding) .{ .name = "qwen35_keep", .flags = &.{"-O3"} }, // ours: a verified window's kept path applied to its stream's state + .{ .name = "qwen36_mtp", .flags = &.{"-O3"} }, // ours: the Qwen3.6 MTP drafter's attention at the target's head width .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, diff --git a/zig/kernels/cuda/qwen36_mtp.cu b/zig/kernels/cuda/qwen36_mtp.cu new file mode 100644 index 0000000000..c996c1fe72 --- /dev/null +++ b/zig/kernels/cuda/qwen36_mtp.cu @@ -0,0 +1,115 @@ +// The Qwen3.6 MTP drafter's glue on CUDA (cuda_mtp.zig): its attention at the target's head width over a ring. + +#include +#include + +namespace tf_q36m { + +constexpr int HD = 256; // the head width (q, k, v) +constexpr int ROT = 64; // rotary dims: pairs (i, i + 32), the target's partial rotary + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { +#pragma unroll + for (int m = 16; m; m >>= 1) x += __shfl_xor_sync(0xffffffffu, x, m); + return x; +} + +} // namespace tf_q36m + +using namespace tf_q36m; + +// A warp a (row, head): RMS norm over HD, rope on [0, ROT) as pairs (i, i + 32) at pos0 + row; OUT (rows, heads, HD). +extern "C" __global__ void __launch_bounds__(32) tf_q36m_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, + __nv_bfloat16* __restrict__ OUT, int heads, int stride, int in_head, + long long pos0, float eps, float theta) { + constexpr int DPL = HD / 32; + const int head = blockIdx.x % heads; + const size_t m = blockIdx.x / heads; + const int t = threadIdx.x; + const __nv_bfloat16* x = X + m * stride + (size_t)head * in_head; // in_head 512 for q beside its gate, 256 for k + float v[DPL], acc = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) { + v[i] = bf(x[DPL * t + i]); + acc = fmaf(v[i], v[i], acc); + } + acc = warp_sum(acc); + const float inv = rsqrtf(acc / (float)HD + eps); + float n[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) n[i] = rbf(bf(W[DPL * t + i]) * (v[i] * inv)); + // a rotary pair's partner sits 32 elements away: DPL = 8 a lane, so 4 lanes away at the same i + const float pos = (float)(pos0 + (long long)m); + __nv_bfloat16* out = OUT + (m * heads + head) * HD; +#pragma unroll + for (int i = 0; i < DPL; i++) { + const int d = DPL * t + i; + const float partner = __shfl_xor_sync(0xffffffffu, n[i], 4); + float r = n[i]; + if (d < ROT) { + const int j = d < ROT / 2 ? d : d - ROT / 2; + const float angle = pos * powf(theta, -(float)j / (float)(ROT / 2)); + float sn, cs; + sincosf(angle, &sn, &cs); + r = d < ROT / 2 ? n[i] * cs - partner * sn : n[i] * cs + partner * sn; + } + out[d] = __float2bfloat16_rn(r); + } +} + +// A warp a (row, head): ring slots [begin, pos0), then block rows 0..m; OUT = softmax-weighted values x sigmoid(gate). +extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ G, int gstride, + const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, + const __nv_bfloat16* __restrict__ RK, const __nv_bfloat16* __restrict__ RV, int heads, + int kv_heads, int W, long long begin, long long pos0, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + const float scale = 1.0f / 16.0f; + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + for (long long pos = begin; pos < pos0; pos++) { + const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + for (int r = 0; r <= m; r++) { + const __nv_bfloat16* k = KB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = VB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; // the q projection's gate +#pragma unroll + for (int i = 0; i < DPL; i++) { + const float gate = 1.0f / (1.0f + __expf(-bf(g[i]))); + OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(rbf(o[i] / l) * gate); + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 65ae27b3d7..c3014d0ed4 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -34,6 +34,7 @@ pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &. pub const qwen35: []const u8 = if (available) &Blob("fatbin_qwen35").bytes else &.{}; pub const qwen35_draft: []const u8 = if (available) &Blob("fatbin_qwen35_draft").bytes else &.{}; pub const qwen35_keep: []const u8 = if (available) &Blob("fatbin_qwen35_keep").bytes else &.{}; +pub const qwen36_mtp: []const u8 = if (available) &Blob("fatbin_qwen36_mtp").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 0f2d62fcce..e79beb5a2b 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -4,6 +4,7 @@ pub const geometry = @import("cuda_geometry.zig"); pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); pub const native_moe = @import("cuda_native_moe.zig"); +pub const mtp = @import("cuda_mtp.zig"); pub const draft = @import("cuda_draft.zig"); test { diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index c3efc4525d..8cca782f8a 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -540,7 +540,7 @@ pub const Engine = struct { try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n, t_params, @as(c_int, @intCast(g.intermediate)) }); try e.proj(m.down, e.act, e.r, total, prompt); }, - .moe => |m| try e.routed(m.router, m.experts, total, prompt), + .moe => |m| try e.routedRows(m.router, m.experts, e.x, e.r, total, prompt and total > most_window), } if (shape.taps) for (e.taps[0..e.n_taps], 0..) |t, ti| if (t == i) { try e.run(e.fns.tap, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, e.r, t_segs, t_rows, @as(c_int, @intCast(ti)), @as(c_int, @intCast(g.hidden)), t_params }); @@ -551,22 +551,20 @@ pub const Engine = struct { try e.proj(e.head, e.x + (total - heads) * g.hidden * 2, e.logits, heads, false); } - /// The routed experts over the normed rows in x: the picks, their plan, SwiGLU up, down, the weighted sum into r. - fn routed(e: *Engine, router: u64, layer: cuda.experts.Layer, total: usize, prompt: bool) !void { + /// Routed experts over `total` normed rows at x into out; `wide` rows take the prefill kernels, else decode's. + pub fn routedRows(e: *Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u64, total: usize, wide: bool) !void { const g = e.g; const slots = g.slots(); const count = g.experts + 1; const pairs = total * slots; - // a short prompt chunk takes the decode form, as the dense projections do: its rows keep one-row steps' bits - const wide = prompt and total > most_window; const tile: usize = if (wide) 64 else 16; const items = cuda.grouped.maxItems(pairs, count, tile); const plan: cuda.grouped.Plan = .{ .members = e.members, .items = e.items, .counts = e.counts, .rank = e.rank, .hist = e.hist }; const grouped: cuda.experts.Grouped = .{ .affine4 = &e.experts }; - try e.moe.route(e.stream, e.x, g.hidden, total, router, e.rlogits, g.hidden, count); + try e.moe.route(e.stream, x, g.hidden, total, router, e.rlogits, g.hidden, count); try e.moe.pick(e.stream, e.rlogits, e.picks, e.wts, total, g.experts, g.top_k); try e.plan.route(e.stream, e.picks, pairs, count, tile, plan); - const rows_in: cuda.experts.Rows = .{ .x = e.x, .stride = g.hidden, .slots = slots }; + const rows_in: cuda.experts.Rows = .{ .x = x, .stride = g.hidden, .slots = slots }; const act_in: cuda.experts.Rows = .{ .x = e.eact, .stride = g.intermediate }; if (wide) { try grouped.prompt(e.stream, .up, rows_in, layer, plan, e.eact, items, -1); @@ -575,7 +573,7 @@ pub const Engine = struct { try grouped.decode(e.stream, .up, rows_in, layer, plan, e.eact, items, -1); try grouped.decode(e.stream, .down_f32, act_in, layer, plan, e.ey, items, -1); } - try e.moe.combine(e.stream, e.ey, !wide, e.wts, e.r, total, slots, g.hidden); + try e.moe.combine(e.stream, e.ey, !wide, e.wts, out, total, slots, g.hidden); } /// Row `i`'s logits from the newest forward that kept logits. diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig index fecc4f4c01..3640559862 100644 --- a/zig/src/families/qwen3_5/cuda_geometry.zig +++ b/zig/src/families/qwen3_5/cuda_geometry.zig @@ -7,6 +7,8 @@ pub const dk = 128; pub const taps = 4; pub const head_dim = 256; pub const max_query_heads = 32; +/// The rotary base the kernels fix (parse checks the config says so). +pub const theta: f32 = 10000000; pub const Geometry = struct { hidden: usize, diff --git a/zig/src/families/qwen3_5/cuda_mtp.zig b/zig/src/families/qwen3_5/cuda_mtp.zig new file mode 100644 index 0000000000..3c581efc6f --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_mtp.zig @@ -0,0 +1,417 @@ +//! The Qwen3.6 MTP layer as the drafter on CUDA: chains from a kept row's final normed state and the token after it. +const std = @import("std"); +const platform = @import("sys"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const geometry = @import("cuda_geometry.zig"); +const weights = @import("cuda_mtp_weights.zig"); +const timing = @import("cuda_draft.zig").timing; + +const E = engine.Engine; +const DeviceBuffer = cuda.DeviceBuffer; +const dr = lanes.drafter; +const Stream = lanes.Stream; +const head_dim = geometry.head_dim; +const max_rows = engine.max_rows; +const most_window = engine.most_window; + +/// Drafts a round at most (the frozen engine's DEPTH); TF_MTP_DRAFTS=N sets it, up to a window's rows less one. +const default_depth = 3; + +const Fns = struct { norm_rope: cuda.Function, attn: cuda.Function }; +/// qwen35_draft.cu's ring kernels, shared with the DFlash2 drafter. +const Ring = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, ring_put: cuda.Function }; +const TapRows = extern struct { n: c_int, rows: [most_window]c_int }; + +/// A stream's context: its ring of keys and values over [begin, end), the newest row's state, and its held chain. +const Ctx = struct { + ring: DeviceBuffer, + last: DeviceBuffer, // the newest absorbed row's normed output (bf16 hidden): the first chain row's state + begin: u64 = 0, + end: u64 = 0, + first: u32 = 0, // the newest absorbed row's draft, its logits' argmax + held: [most_window]u32 = undefined, + held_n: usize = 0, +}; + +pub const Mtp = struct { + gpa: std.mem.Allocator, + io: std.Io, + e: *E, + w: weights.Weights, + mod: cuda.Module, + dmod: cuda.Module, + fns: Fns, + ring_fns: Ring, + scratch: DeviceBuffer, + at: usize = 0, + capacity: usize, // ring rows: every position of the stream's window + depth: u32, + tap: [1]u32, + // rows of max_rows, bf16 unless said + taps_in: u64, + hin: u64, + hn: u64, + emb: u64, + en: u64, + h: u64, + r: u64, + x: u64, + qraw: u64, + q: u64, + kraw: u64, + kb: u64, + vb: u64, + att: u64, + xs: u64, // f32 group sums for the projections + logits: u64, // one row + ids: u64, // u32 the rows' tokens + tok: u64, // u32 the drawn token + staged: cuda.HostBuffer, // pinned: the ids up, the token down + streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, + /// Contexts between streams, allocated and run once (`reserve`), as the driver stalls a buffer's first kernel. + pool: std.ArrayList(*Ctx) = .empty, + /// One chained draft's cost (timed at load). + step_ms: f64 = 1.0, + /// The target's window ms by rows (cuda_native.zig times them), for the round loop's prices. + costs: [most_window]lanes.config.Cost = undefined, + n_costs: usize = 0, + + const ids_at = 0; + const tok_at = max_rows * 4; + const staged_len = tok_at + 64; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8, capacity: usize) !*Mtp { + if (!e.g.moe()) return error.NotRouted; + const d = try gpa.create(Mtp); + errdefer gpa.destroy(d); + d.* = .{ .gpa = gpa, .io = io, .e = e, .w = undefined, .mod = undefined, .dmod = undefined, .fns = undefined, .ring_fns = undefined, .scratch = undefined, .capacity = capacity, .depth = default_depth, .tap = .{@intCast(e.g.layers - 1)}, .taps_in = 0, .hin = 0, .hn = 0, .emb = 0, .en = 0, .h = 0, .r = 0, .x = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .xs = 0, .logits = 0, .ids = 0, .tok = 0, .staged = undefined }; + if (platform.c.getenv("TF_MTP_DRAFTS")) |v| d.depth = std.fmt.parseInt(u32, std.mem.span(v), 10) catch default_depth; + d.depth = @min(@max(d.depth, 1), most_window - 1); + d.mod = try cuda.Module.load(e.d, cuda.kernels.qwen36_mtp); + errdefer d.mod.unload(); + d.dmod = try cuda.Module.load(e.d, cuda.kernels.qwen35_draft); + errdefer d.dmod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(d.fns, name) = try d.mod.function("tf_q36m_" ++ name); + inline for (@typeInfo(Ring).@"struct".field_names) |name| @field(d.ring_fns, name) = try d.dmod.function("tf_q35d_" ++ name); + d.w = try weights.load(gpa, io, e, dir); + errdefer d.w.deinit(); + d.scratch = try DeviceBuffer.alloc(e.d, d.scratchBytes()); + errdefer d.scratch.free(); + try d.scratch.fill8(0, null); + d.at = 0; + _ = try d.layout(true); + d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); + errdefer d.staged.free(); + // the first context, and a chain timed on it (the quickest of several: the first ones pay the driver's stalls) + const ctx = try d.create(); + errdefer d.drop(ctx); + var quickest: f64 = std.math.inf(f64); + for (0..6) |_| { + try d.absorb(ctx, d.taps_in, 0, 1, null, &.{0}); + const t0 = std.Io.Clock.awake.now(io); + try d.hold(ctx, 0, 1, d.depth); + quickest = @min(quickest, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6); + } + d.step_ms = quickest / @as(f64, @floatFromInt(d.depth)); + d.resetCtx(ctx); + try d.pool.append(gpa, ctx); + return d; + } + + /// The drafter's tap: the target's last layer, whose residual the final norm then reads. + pub fn taps(d: *const Mtp) []const u32 { + return &d.tap; + } + + pub fn setCost(d: *Mtp, rows: u32, ms: f64) void { + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a chained draft {d:.2})\n", .{ rows, ms, d.step_ms }); + if (d.n_costs < d.costs.len) { + d.costs[d.n_costs] = .{ .width = rows, .ms = ms }; + d.n_costs += 1; + } + } + + pub fn deinit(d: *Mtp) void { + var it = d.streams.valueIterator(); + while (it.next()) |ctx| d.drop(ctx.*); + d.streams.deinit(d.gpa); + for (d.pool.items) |ctx| d.drop(ctx); + d.pool.deinit(d.gpa); + d.staged.free(); + d.scratch.free(); + d.w.deinit(); + d.dmod.unload(); + d.mod.unload(); + d.gpa.destroy(d); + } + + fn kvWidth(d: *const Mtp) usize { + return d.e.g.kv_heads * head_dim; + } + + fn ringBytes(d: *const Mtp) usize { + return 2 * d.capacity * d.kvWidth() * 2; + } + + fn scratchBytes(d: *Mtp) usize { + d.at = 0; + return d.layout(false) catch unreachable; + } + + /// The scratch rows, measured (nothing assigned) or assigned from the scratch buffer. + fn layout(d: *Mtp, assign: bool) !usize { + const g = d.e.g; + const R = max_rows; + const aw = g.attnWidth(); + const kw = d.kvWidth(); + const sizes = .{ + .{ "taps_in", R * g.hidden * 2 }, .{ "hin", R * g.hidden * 2 }, .{ "hn", R * g.hidden * 2 }, .{ "emb", R * g.hidden * 2 }, + .{ "en", R * g.hidden * 2 }, .{ "h", R * g.hidden * 2 }, .{ "r", R * g.hidden * 2 }, .{ "x", R * g.hidden * 2 }, + .{ "qraw", R * 2 * aw * 2 }, .{ "q", R * aw * 2 }, .{ "kraw", R * kw * 2 }, .{ "kb", R * kw * 2 }, + .{ "vb", R * kw * 2 }, .{ "att", R * aw * 2 }, .{ "xs", R * (@max(g.hidden, aw) / 64) * 4 }, .{ "logits", g.vocab * 2 }, + .{ "ids", R * 4 }, .{ "tok", 64 }, + }; + var total: usize = 0; + inline for (sizes) |s| { + const at = engine.aligned(total); + total = at + s[1]; + if (assign) @field(d, s[0]) = d.scratch.ptr + at; + } + return total; + } + + fn drop(d: *Mtp, ctx: *Ctx) void { + ctx.last.free(); + ctx.ring.free(); + d.gpa.destroy(ctx); + } + + fn resetCtx(d: *Mtp, ctx: *Ctx) void { + _ = d; + ctx.* = .{ .ring = ctx.ring, .last = ctx.last }; + } + + /// A new context, touched once (an absorb, then a chain), so no request pays the driver's stall. + fn create(d: *Mtp) !*Ctx { + const ctx = try d.gpa.create(Ctx); + errdefer d.gpa.destroy(ctx); + ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.ringBytes()), .last = undefined }; + errdefer ctx.ring.free(); + ctx.last = try DeviceBuffer.alloc(d.e.d, d.e.g.hidden * 2); + errdefer ctx.last.free(); + try ctx.ring.fill8(0, null); + try ctx.last.fill8(0, null); + try d.absorb(ctx, d.taps_in, 0, 1, null, &.{0}); // any tap rows do: the drafter's own zeroed staging + try d.hold(ctx, 0, 1, d.depth); + d.resetCtx(ctx); + return ctx; + } + + /// `n` contexts ready in the pool, for the streams the server admits. + pub fn reserve(d: *Mtp, n: usize) !void { + try d.pool.ensureTotalCapacity(d.gpa, n); + while (d.pool.items.len < n) d.pool.appendAssumeCapacity(try d.create()); + } + + fn ringK(d: *Mtp, ctx: *Ctx) u64 { + _ = d; + return ctx.ring.ptr; + } + + fn ringV(d: *Mtp, ctx: *Ctx) u64 { + return ctx.ring.ptr + d.capacity * d.kvWidth() * 2; + } + + /// out (rows, n) = x (rows, k) through `l`, as the target's decode tiles or prompt GEMM, with the drafter's sums. + fn proj(d: *Mtp, l: cuda.qlinear.Affine4, x: u64, out: u64, rows: usize) !void { + const e = d.e; + if (rows > 16) return e.affine.prompt(e.stream, x, l, out, rows); + try e.run(e.fns.sums, .{ rows, 1, 1 }, .{ 32, 1, 1 }, .{ x, d.xs, @as(c_int, @intCast(l.k)), @as(u64, 0) }); + try e.affine.decode(e.stream, x, d.xs, l, out, rows); + } + + /// out = norm(h + r) with w; h takes the sum (r 0: no residual). + fn norm(d: *Mtp, h: u64, r: u64, w: u64, out: u64, rows: usize) !void { + const e = d.e; + try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ h, r, w, out, e.g.eps, @as(c_int, @intCast(e.g.hidden)), @as(u64, 0) }); + } + + /// The layer over `rows` rows at slots pos0..: d.ids their tokens, d.hin their states; x gets their normed outputs. + fn layer(d: *Mtp, ctx: *Ctx, rows: usize, pos0: u64) !void { + const e = d.e; + const g = e.g; + const w = &d.w; + const aw = g.attnWidth(); + const kw = d.kvWidth(); + const p0: c_longlong = @intCast(pos0); + try e.run(e.fns.embed, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.ids, e.embed_w, e.embed_s, e.embed_b, d.emb, @as(c_int, @intCast(g.hidden)), @as(u64, 0) }); + try d.norm(d.emb, 0, w.norm_e, d.en, rows); + try d.norm(d.hin, 0, w.norm_h, d.hn, rows); + try d.proj(w.fc_e, d.en, d.h, rows); + try d.proj(w.fc_h, d.hn, d.r, rows); + try d.norm(d.h, d.r, w.input_norm, d.x, rows); + try d.proj(w.q, d.x, d.qraw, rows); + try d.proj(w.k, d.x, d.kraw, rows); + try d.proj(w.v, d.x, d.vb, rows); + try e.run(d.fns.norm_rope, .{ rows * g.q_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.qraw, w.q_norm, d.q, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), p0, g.eps, geometry.theta }); + try e.run(d.fns.norm_rope, .{ rows * g.kv_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.kraw, w.k_norm, d.kb, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(kw)), @as(c_int, head_dim), p0, g.eps, geometry.theta }); + try e.run(d.fns.attn, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(d.capacity)), @as(c_longlong, @intCast(ctx.begin)), p0, d.att }); + try d.proj(w.o, d.att, d.r, rows); + try d.norm(d.h, d.r, w.post_norm, d.x, rows); + try e.routedRows(w.mlp.moe.router, w.mlp.moe.experts, d.x, d.r, rows, rows > most_window); + try d.norm(d.h, d.r, w.norm, d.x, rows); + try e.run(d.ring_fns.ring_put, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), @as(c_int, @intCast(kw)), p0, @as(c_int, @intCast(d.capacity)) }); + } + + /// Row `row` of x through the target's head: its argmax into `tok` (the host sees it after a synchronize). + fn draw(d: *Mtp, row: usize) !void { + const e = d.e; + try d.proj(e.head, d.x + row * e.g.hidden * 2, d.logits, 1); + try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ d.logits, @as(c_int, @intCast(e.g.vocab)), d.tok }); + } + + fn token(d: *Mtp) !u32 { + const e = d.e; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(d.staged.bytes.ptr + tok_at, d.tok, 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + return std.mem.bytesToValue(u32, d.staged.bytes[tok_at..][0..4]); + } + + fn putIds(d: *Mtp, ids: []const u32) !void { + const e = d.e; + @memcpy(std.mem.bytesAsSlice(u32, d.staged.bytes[ids_at..][0 .. ids.len * 4]), ids); + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(d.ids, d.staged.bytes.ptr + ids_at, ids.len * 4, e.stream.handle), "cuMemcpyHtoDAsync"); + } + + /// Target rows [start, start + n) from the tap ring with the token after each; `rows`, when given, picks them. + fn absorb(d: *Mtp, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32, follow: []const u32) !void { + const e = d.e; + const g = e.g; + var start = start_; + var n = n_; + var fi: usize = 0; + if (follow.len < n) return error.BadDraftRows; + if (rows_) |rs| if (rs.len != n or n > most_window) return error.BadDraftRows; + if (start < ctx.end) ctx.end = start; + if (start > ctx.end) ctx.begin = start; // a gap (a prompt restored from a cache): the context starts here + const held = engine.tap_rows - max_rows; // the target's ring holds this many rows behind its newest + if (n > held) { + fi = n - held; + start += n - held; + n = held; + ctx.begin = start; + } + ctx.end = start; + while (n > 0) { + const rows = @min(n, max_rows); + if (rows_) |rs| { + var tr: TapRows = .{ .n = @intCast(rows), .rows = @splat(0) }; + for (rs, 0..) |r, i| tr.rows[i] = @intCast(r); + try e.run(d.ring_fns.gather_rows, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), tr, @as(c_int, @intCast(g.hidden)), d.taps_in }); + } else try e.run(d.ring_fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(g.hidden)), d.taps_in }); + try d.norm(d.taps_in, 0, e.final_norm, d.hin, rows); // the target's final norm: the rows the head read + try d.putIds(follow[fi..][0..rows]); + try d.layer(ctx, rows, start); + start += rows; + n -= rows; + fi += rows; + ctx.end = start; + // the newest row: its state kept for the chain, its draft drawn + try ctx.last.copyFrom(0, d.x + (rows - 1) * g.hidden * 2, g.hidden * 2, e.stream.handle); + try d.draw(rows - 1); + } + ctx.first = try d.token(); + } + + /// `depth` drafts after `pending` at pos0 (the pending token's slot): the first came with the newest absorbed row. + fn hold(d: *Mtp, ctx: *Ctx, pending: u32, pos0: u64, depth: usize) !void { + const e = d.e; + const g = e.g; + if (ctx.end != pos0) return error.DraftContextBehind; + ctx.held[0] = ctx.first; + ctx.held_n = 1; + _ = pending; // the newest absorbed row's follow token, whose draft ctx.first is + var k: usize = 1; + while (k < depth and k < ctx.held.len) : (k += 1) { + // row k: (the previous row's normed output, the draft before) at slot pos0 + k - 1 + const src = if (k == 1) ctx.last.ptr else d.x; + try d.scratch.copyFrom(d.hin - d.scratch.ptr, src, g.hidden * 2, e.stream.handle); + try d.putIds(&.{ctx.held[k - 1]}); + try d.layer(ctx, 1, pos0 + k - 1); + try d.draw(0); + ctx.held[k] = try d.token(); + ctx.held_n = k + 1; + } + } + + // the Drafter contract + + pub fn drafter(d: *Mtp) dr.Drafter { + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(p: *anyopaque) *Mtp { + return @ptrCast(@alignCast(p)); + } + + fn tapsFn(p: *anyopaque) []const u32 { + return cast(p).taps(); + } + + fn factsFn(p: *anyopaque) dr.Facts { + const d = cast(p); + return .{ .depth = d.depth, .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.ringBytes() + d.e.g.hidden * 2, .trees = false, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; + } + + fn openFn(p: *anyopaque, s: *Stream) anyerror!void { + const d = cast(p); + if (d.streams.contains(s)) return error.StreamOpen; + const ctx = d.pool.pop() orelse try d.create(); + errdefer d.pool.append(d.gpa, ctx) catch d.drop(ctx); + d.resetCtx(ctx); + try d.streams.put(d.gpa, s, ctx); + } + + fn absorbFn(p: *anyopaque, items: []const dr.Absorb) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("mtp absorb", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; + // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) + if (it.rows) |rs| { + try d.absorb(ctx, it.features.buffer, it.features.offset, rs.len, rs, it.follow); + } else try d.absorb(ctx, it.features.buffer, it.features.offset, it.features.rows, null, it.follow); + } + } + + fn holdFn(p: *anyopaque, items: []const dr.Hold) anyerror!void { + const d = cast(p); + const t0 = std.Io.Clock.awake.now(d.io); + defer timing("mtp hold", d.io, t0); + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + try d.hold(ctx, it.pending, it.position - 1, it.depth); + } + } + + fn heldFn(p: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const d = cast(p); + for (streams, out) |s, o| { + const ctx = d.streams.get(s) orelse return error.UnknownStream; + if (o.len > ctx.held_n) return error.TooFewDrafts; + @memcpy(o, ctx.held[0..o.len]); + } + } + + fn releaseFn(p: *anyopaque, s: *Stream) void { + const d = cast(p); + if (d.streams.fetchRemove(s)) |kv| d.pool.append(d.gpa, kv.value) catch d.drop(kv.value); + } +}; + +/// Acceptance by depth until a stream has its own (the frozen engine's chains kept about three of four). +const prior = [_]f64{ 0.7, 0.55, 0.45, 0.4, 0.35, 0.3, 0.3 }; diff --git a/zig/src/families/qwen3_5/cuda_mtp_weights.zig b/zig/src/families/qwen3_5/cuda_mtp_weights.zig new file mode 100644 index 0000000000..09b9770951 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_mtp_weights.zig @@ -0,0 +1,80 @@ +//! The Qwen3.6 MTP layer's weights (mtp-4bit.safetensors beside the checkpoint) onto the device, their own arena. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const engine = @import("cuda_engine.zig"); +const geometry = @import("cuda_geometry.zig"); +const weights = @import("cuda_weights.zig"); + +const Affine4 = cuda.qlinear.Affine4; +const DeviceBuffer = cuda.DeviceBuffer; +const Engine = engine.Engine; +const aligned = engine.aligned; +const head_dim = geometry.head_dim; + +pub const file = "mtp-4bit.safetensors"; + +pub const Weights = struct { + arena: DeviceBuffer, + arena_at: usize = 0, + fc_e: Affine4, // fc's halves: the embedding's inputs, then the hidden state's + fc_h: Affine4, + norm_e: u64, + norm_h: u64, + input_norm: u64, + post_norm: u64, + q_norm: u64, + k_norm: u64, + norm: u64, + q: Affine4, + k: Affine4, + v: Affine4, + o: Affine4, + mlp: engine.Mlp, + + pub fn alloc(w: *Weights, bytes: usize) !u64 { + const at = aligned(w.arena_at); + if (at + bytes > w.arena.len) return error.ArenaExhausted; + w.arena_at = at + bytes; + return w.arena.ptr + at; + } + + pub fn deinit(w: *Weights) void { + w.arena.free(); + } +}; + +/// Device bytes of everything `load` packs. +pub fn total(g: geometry.Geometry) usize { + var t = 2 * weights.packedBytes(g.hidden, g.hidden) + 7 * aligned(@max(g.hidden, head_dim) * 2); + t += weights.packedBytes(2 * g.attnWidth(), g.hidden) + 2 * weights.packedBytes(g.kv_heads * head_dim, g.hidden) + weights.packedBytes(g.hidden, g.attnWidth()); + t += weights.routedBytes(g.experts, g.intermediate, g.hidden); + return t; +} + +pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *Engine, dir: []const u8) !Weights { + const g = e.g; + var ck: core.Checkpoint = .{ .gpa = gpa, .io = io }; + defer ck.close(); + try ck.add(dir, file); + var w: Weights = .{ .arena = try DeviceBuffer.alloc(e.d, total(g)), .fc_e = undefined, .fc_h = undefined, .norm_e = 0, .norm_h = 0, .input_norm = 0, .post_norm = 0, .q_norm = 0, .k_norm = 0, .norm = 0, .q = undefined, .k = undefined, .v = undefined, .o = undefined, .mlp = undefined }; + errdefer w.arena.free(); + try w.arena.fill8(0, null); + const m = "language_model.mtp."; + const p = "language_model.mtp.layers.0."; + w.fc_e = try weights.linearSlice(e, &w, &ck, m, "fc", g.hidden, 2 * g.hidden, 0, g.hidden); + w.fc_h = try weights.linearSlice(e, &w, &ck, m, "fc", g.hidden, 2 * g.hidden, g.hidden, g.hidden); + w.norm_e = try weights.tensor(&w, &ck, "{s}pre_fc_norm_embedding.weight", .{m}, .bf16, &.{g.hidden}); + w.norm_h = try weights.tensor(&w, &ck, "{s}pre_fc_norm_hidden.weight", .{m}, .bf16, &.{g.hidden}); + w.norm = try weights.tensor(&w, &ck, "{s}norm.weight", .{m}, .bf16, &.{g.hidden}); + w.input_norm = try weights.tensor(&w, &ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{g.hidden}); + w.post_norm = try weights.tensor(&w, &ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{g.hidden}); + w.q_norm = try weights.tensor(&w, &ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{head_dim}); + w.k_norm = try weights.tensor(&w, &ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{head_dim}); + w.q = try weights.linear(e, &w, &ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden); + w.k = try weights.linear(e, &w, &ck, p, "self_attn.k_proj", g.kv_heads * head_dim, g.hidden); + w.v = try weights.linear(e, &w, &ck, p, "self_attn.v_proj", g.kv_heads * head_dim, g.hidden); + w.o = try weights.linear(e, &w, &ck, p, "self_attn.o_proj", g.hidden, g.attnWidth()); + w.mlp = try weights.routed(e, &w, &ck, p); + return w; +} diff --git a/zig/src/families/qwen3_5/cuda_native.zig b/zig/src/families/qwen3_5/cuda_native.zig index 8773dffc71..3f463e6c76 100644 --- a/zig/src/families/qwen3_5/cuda_native.zig +++ b/zig/src/families/qwen3_5/cuda_native.zig @@ -6,6 +6,7 @@ const c = @import("config.zig"); const E = @import("cuda_engine.zig"); const draws = @import("draw.zig"); const Draft = @import("cuda_draft.zig").Draft; +const Mtp = @import("cuda_mtp.zig").Mtp; const be = lanes.backend; pub const model_type = "qwen3_5"; @@ -93,6 +94,28 @@ pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []con return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveDrafter }; } +/// 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 } { + if (bits != 4) return error.DrafterBitsUnsupported; + const l: *Lanes = @ptrCast(@alignCast(p)); + 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()); + try l.timeWindows(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 reserveDrafter(p: *anyopaque, streams: u32) anyerror!void { const d: *Draft = @ptrCast(@alignCast(p)); try d.reserve(streams); @@ -165,7 +188,7 @@ const Lanes = struct { } /// Each width's forward timed on its own cache, on its third run: graph captured, stall paid, a keep in front. - fn timeWindows(self: *Lanes, d: *Draft) !void { + fn timeWindows(self: *Lanes, d: anytype) !void { var cache = try E.Cache.init(self.gpa, self.e, 64); defer cache.deinit(self.gpa); const ids: [E.most_window]u32 = @splat(0); diff --git a/zig/src/families/qwen3_5/cuda_native_moe.zig b/zig/src/families/qwen3_5/cuda_native_moe.zig index c8cda02d6f..54901476c9 100644 --- a/zig/src/families/qwen3_5/cuda_native_moe.zig +++ b/zig/src/families/qwen3_5/cuda_native_moe.zig @@ -6,8 +6,8 @@ pub const formats = native.formats; pub const default_context = native.default_context; pub const max_segments = native.max_segments; pub const prompt_rows = native.prompt_rows; -/// Drafts come from the checkpoint's MTP layer, not a separate draft model. -pub const takes_drafter = false; +/// --drafter names the checkpoint's own directory: its MTP layer (mtp-4bit.safetensors) drafts chains. +pub const takes_drafter = true; pub const open = native.open; -pub const openDrafter = native.openDrafter; +pub const openDrafter = native.openMtp; pub const explain = native.explain; diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig index ea596a964b..35af7ef55e 100644 --- a/zig/src/families/qwen3_5/cuda_weights.zig +++ b/zig/src/families/qwen3_5/cuda_weights.zig @@ -13,7 +13,7 @@ const aligned = engine.aligned; const head_dim = geometry.head_dim; /// The packed tiles' bytes of an n-by-k affine-4 projection (qlinear's words, scales and biases). -fn packedBytes(n: usize, k: usize) usize { +pub fn packedBytes(n: usize, k: usize) usize { const l = Affine4.layout(n, k); return aligned(l.words) + 2 * aligned(l.scales); } @@ -23,6 +23,12 @@ fn rawBytes(n: usize, k: usize) usize { return n * k / 8 * 4 + 2 * (n * k / 64 * 2); } +/// A layer's routed experts' device bytes: the pair table, the down table and the router rows. +pub fn routedBytes(experts: usize, intermediate: usize, hidden: usize) usize { + const ex = cuda.experts.Affine4Experts; + return aligned(ex.pairBytes(experts + 1, intermediate, hidden)) + aligned(ex.bytes(experts + 1, hidden, intermediate)) + aligned((experts + 1) * hidden * 2); +} + /// What `load` takes from the arena: every packed projection and vector, each allocation aligned. pub fn total(g: Geometry) usize { var t: usize = 0; @@ -31,10 +37,7 @@ pub fn total(g: Geometry) usize { t += aligned(g.hidden * 2); for (g.linear[0..g.layers]) |lin| { t += 2 * aligned(g.hidden * 2); - if (g.moe()) { - const ex = cuda.experts.Affine4Experts; - t += aligned(ex.pairBytes(g.experts + 1, g.intermediate, g.hidden)) + aligned(ex.bytes(g.experts + 1, g.hidden, g.intermediate)) + aligned((g.experts + 1) * g.hidden * 2); - } else t += 2 * packedBytes(g.intermediate, g.hidden) + packedBytes(g.hidden, g.intermediate); + if (g.moe()) t += routedBytes(g.experts, g.intermediate, g.hidden) else t += 2 * packedBytes(g.intermediate, g.hidden) + packedBytes(g.hidden, g.intermediate); if (lin) { t += packedBytes(g.convWidth(), g.hidden) + packedBytes(g.valueWidth(), g.hidden) + 2 * packedBytes(g.v_heads, g.hidden) + packedBytes(g.hidden, g.valueWidth()); t += aligned(g.convWidth() * geometry.taps * 2) + aligned(geometry.dk * 2) + aligned(g.v_heads * 4) + aligned(g.v_heads * 2); @@ -54,14 +57,15 @@ pub fn stageBytes(g: Geometry) usize { return most; } -fn upload(e: *Engine, bytes: []const u8) !u64 { - const p = try e.alloc(bytes.len); - try e.arena.upload(p - e.arena.ptr, bytes); +/// `into` is the engine or a drafter's weights: anything with `alloc(bytes)` out of its `arena`. +pub fn upload(into: anytype, bytes: []const u8) !u64 { + const p = try into.alloc(bytes.len); + try into.arena.upload(p - into.arena.ptr, bytes); return p; } -fn put(e: *Engine, t: core.checkpoint.Tensor) !u64 { - return upload(e, t.bytes); +fn put(into: anytype, t: core.checkpoint.Tensor) !u64 { + return upload(into, t.bytes); } /// A vector the kernels read as f32, stored as f32 (the 2B) or bf16 (the 27B's A_log). @@ -76,30 +80,60 @@ fn f32vec(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anyt return upload(e, std.mem.sliceAsBytes(words)); } -fn tensor(e: *Engine, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { +pub fn tensor(into: anytype, ck: *core.Checkpoint, comptime fmt: []const u8, args: anytype, dtype: core.checkpoint.DType, shape: []const usize) !u64 { var buf: [160]u8 = undefined; - return put(e, try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); + return put(into, try ck.expect(try std.fmt.bufPrint(&buf, fmt, args), dtype, shape)); } -/// An MLX affine-4 projection packed into qlinear's tiles from the staging buffer, which holds it as stored. -fn linear(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { - var buf: [160]u8 = undefined; +/// Raw words, scales and biases staged in order; each part's device pointer. +fn stage(e: *Engine, parts: [3][]const u8) ![3]u64 { var raw: [3]u64 = undefined; var at: usize = 0; - const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; for (parts, &raw) |p, *r| { - const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); - try e.stage.upload(at, t.bytes); + if (at + p.len > e.stage.len) return error.StageTooSmall; + try e.stage.upload(at, p); r.* = e.stage.ptr + at; - at += t.bytes.len; + at += p.len; } + return raw; +} + +fn packStaged(e: *Engine, into: anytype, raw: [3]u64, n: usize, k: usize) !Affine4 { const l = Affine4.layout(n, k); - const q: Affine4 = .{ .w = try e.alloc(l.words), .s = try e.alloc(l.scales), .b = try e.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; + const q: Affine4 = .{ .w = try into.alloc(l.words), .s = try into.alloc(l.scales), .b = try into.alloc(l.scales), .n = n, .k = k, .npad = l.npad }; try e.affine.pack(e.stream, raw[0], raw[1], raw[2], q); try e.stream.synchronize(); // the next projection's staging overwrites this one's return q; } +/// An MLX affine-4 projection packed into qlinear's tiles from the staging buffer, which holds it as stored. +pub fn linear(e: *Engine, into: anytype, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize) !Affine4 { + var buf: [160]u8 = undefined; + var parts: [3][]const u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (specs, &parts) |p, *part| part.* = (try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] })).bytes; + return packStaged(e, into, try stage(e, parts), n, k); +} + +/// Inputs [from, from + len) of an n-by-k projection as their own packed projection (the MTP's fc halves). +pub fn linearSlice(e: *Engine, into: anytype, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, n: usize, k: usize, from: usize, len: usize) !Affine4 { + var buf: [160]u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize, usize }{ .{ "weight", .u32, k / 8, 4 }, .{ "scales", .bf16, k / 64, 2 }, .{ "biases", .bf16, k / 64, 2 } }; + var parts: [3][]u8 = undefined; + var taken: usize = 0; + errdefer for (parts[0..taken]) |p| e.gpa.free(p); + for (specs, &parts) |p, *part| { + const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], &.{ n, p[2] }); + const per = k / p[2]; // inputs a column holds + const cols = len / per; + part.* = try e.gpa.alloc(u8, n * cols * p[3]); + taken += 1; + for (0..n) |r| @memcpy(part.*[r * cols * p[3] ..][0 .. cols * p[3]], t.bytes[(r * p[2] + from / per) * p[3] ..][0 .. cols * p[3]]); + } + defer for (parts) |p| e.gpa.free(p); + return packStaged(e, into, try stage(e, .{ parts[0], parts[1], parts[2] }), n, len); +} + fn bf16f(h: u16) f32 { return @bitCast(@as(u32, h) << 16); } @@ -122,7 +156,7 @@ fn dequant8(dst: []u16, words: []align(1) const u32, scales: []align(1) const u1 } /// The router's rows as bf16 [experts + 1][hidden]: the routed experts' gate, then the shared expert's gate row. -fn routerRows(e: *Engine, ck: *core.Checkpoint, prefix: []const u8) !u64 { +fn routerRows(e: *Engine, into: anytype, ck: *core.Checkpoint, prefix: []const u8) !u64 { const g = e.g; const rows = try e.gpa.alloc(u16, (g.experts + 1) * g.hidden); defer e.gpa.free(rows); @@ -134,34 +168,31 @@ fn routerRows(e: *Engine, ck: *core.Checkpoint, prefix: []const u8) !u64 { const b = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.biases", .{ prefix, n[0] }), .bf16, &.{ n[2], g.hidden / 64 }); dequant8(rows[n[1] * g.hidden ..][0 .. n[2] * g.hidden], std.mem.bytesAsSlice(u32, w.bytes), std.mem.bytesAsSlice(u16, s.bytes), std.mem.bytesAsSlice(u16, b.bytes), g.hidden); } - return upload(e, std.mem.sliceAsBytes(rows)); + return upload(into, std.mem.sliceAsBytes(rows)); } /// `count` stacked experts' projection (one at count 0) staged as stored and packed into table `out` as the mi-th of m. fn expertTable(e: *Engine, ck: *core.Checkpoint, prefix: []const u8, name: []const u8, count: usize, n: usize, k: usize, out: u64, m: usize, mi: usize) !void { var buf: [160]u8 = undefined; - var raw: [3]u64 = undefined; - var at: usize = 0; - const parts = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; - for (parts, &raw) |p, *r| { + var parts: [3][]const u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (specs, &parts) |p, *part| { const stacked: []const usize = &.{ count, n, p[2] }; const one: []const usize = &.{ n, p[2] }; - const t = try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], if (count == 0) one else stacked); - try e.stage.upload(at, t.bytes); - r.* = e.stage.ptr + at; - at += t.bytes.len; + part.* = (try ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, p[0] }), p[1], if (count == 0) one else stacked)).bytes; } + const raw = try stage(e, parts); try e.experts.packInto(e.stream, raw[0], raw[1], raw[2], out, @max(count, 1), n, k, m, mi); try e.stream.synchronize(); // the next table's staging overwrites this one's } /// A layer's routed experts: gate and up interleaved as one pair table, down apart, the shared expert last in each. -fn routed(e: *Engine, ck: *core.Checkpoint, p: []const u8) !engine.Mlp { +pub fn routed(e: *Engine, into: anytype, ck: *core.Checkpoint, p: []const u8) !engine.Mlp { const g = e.g; const ex = cuda.experts.Affine4Experts; const ni = g.intermediate; - const pair = try e.alloc(ex.pairBytes(g.experts + 1, ni, g.hidden)); - const down = try e.alloc(ex.bytes(g.experts + 1, g.hidden, ni)); + const pair = try into.alloc(ex.pairBytes(g.experts + 1, ni, g.hidden)); + const down = try into.alloc(ex.bytes(g.experts + 1, g.hidden, ni)); try expertTable(e, ck, p, "mlp.switch_mlp.gate_proj", g.experts, ni, g.hidden, pair, 2, 0); try expertTable(e, ck, p, "mlp.switch_mlp.up_proj", g.experts, ni, g.hidden, pair, 2, 1); try expertTable(e, ck, p, "mlp.switch_mlp.down_proj", g.experts, g.hidden, ni, down, 1, 0); @@ -170,7 +201,7 @@ fn routed(e: *Engine, ck: *core.Checkpoint, p: []const u8) !engine.Mlp { try expertTable(e, ck, p, "mlp.shared_expert.up_proj", 0, ni, g.hidden, shared_pair, 2, 1); try expertTable(e, ck, p, "mlp.shared_expert.down_proj", 0, g.hidden, ni, down + ex.bytes(g.experts, g.hidden, ni), 1, 0); const layer: cuda.experts.Layer = .{ .format = .affine4, .up = pair, .down = down, .width = ni, .dims = g.hidden, .experts = g.experts + 1, .swiglu = true }; - return .{ .moe = .{ .router = try routerRows(e, ck, p), .experts = layer } }; + return .{ .moe = .{ .router = try routerRows(e, into, ck, p), .experts = layer } }; } pub fn load(e: *Engine, ck: *core.Checkpoint) !void { @@ -179,7 +210,7 @@ pub fn load(e: *Engine, ck: *core.Checkpoint) !void { e.embed_w = try tensor(e, ck, "{s}embed_tokens.weight", .{m}, .u32, &.{ g.vocab, g.hidden / 8 }); e.embed_s = try tensor(e, ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 }); e.embed_b = try tensor(e, ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 }); - e.head = if (g.tied) try linear(e, ck, m, "embed_tokens", g.vocab, g.hidden) else try linear(e, ck, "language_model.", "lm_head", g.vocab, g.hidden); + e.head = if (g.tied) try linear(e, e, ck, m, "embed_tokens", g.vocab, g.hidden) else try linear(e, e, ck, "language_model.", "lm_head", g.vocab, g.hidden); e.final_norm = try tensor(e, ck, "{s}norm.weight", .{m}, .bf16, &.{g.hidden}); const conv = g.convWidth(); const vw = g.valueWidth(); @@ -189,20 +220,20 @@ pub fn load(e: *Engine, ck: *core.Checkpoint) !void { l.* = .{ .input_norm = try tensor(e, ck, "{s}input_layernorm.weight", .{p}, .bf16, &.{g.hidden}), .post_norm = try tensor(e, ck, "{s}post_attention_layernorm.weight", .{p}, .bf16, &.{g.hidden}), - .mlp = if (g.moe()) try routed(e, ck, p) else .{ .dense = .{ - .gate = try linear(e, ck, p, "mlp.gate_proj", g.intermediate, g.hidden), - .up = try linear(e, ck, p, "mlp.up_proj", g.intermediate, g.hidden), - .down = try linear(e, ck, p, "mlp.down_proj", g.hidden, g.intermediate), + .mlp = if (g.moe()) try routed(e, e, ck, p) else .{ .dense = .{ + .gate = try linear(e, e, ck, p, "mlp.gate_proj", g.intermediate, g.hidden), + .up = try linear(e, e, ck, p, "mlp.up_proj", g.intermediate, g.hidden), + .down = try linear(e, e, ck, p, "mlp.down_proj", g.hidden, g.intermediate), } }, .mixer = undefined, }; if (g.linear[i]) { l.mixer = .{ .delta = .{ - .qkv = try linear(e, ck, p, "linear_attn.in_proj_qkv", conv, g.hidden), - .z = try linear(e, ck, p, "linear_attn.in_proj_z", vw, g.hidden), - .a = try linear(e, ck, p, "linear_attn.in_proj_a", g.v_heads, g.hidden), - .b = try linear(e, ck, p, "linear_attn.in_proj_b", g.v_heads, g.hidden), - .out = try linear(e, ck, p, "linear_attn.out_proj", g.hidden, vw), + .qkv = try linear(e, e, ck, p, "linear_attn.in_proj_qkv", conv, g.hidden), + .z = try linear(e, e, ck, p, "linear_attn.in_proj_z", vw, g.hidden), + .a = try linear(e, e, ck, p, "linear_attn.in_proj_a", g.v_heads, g.hidden), + .b = try linear(e, e, ck, p, "linear_attn.in_proj_b", g.v_heads, g.hidden), + .out = try linear(e, e, ck, p, "linear_attn.out_proj", g.hidden, vw), .conv = try tensor(e, ck, "{s}linear_attn.conv1d.weight", .{p}, .bf16, &.{ conv, geometry.taps, 1 }), .norm = try tensor(e, ck, "{s}linear_attn.norm.weight", .{p}, .bf16, &.{geometry.dk}), .a_log = try f32vec(e, ck, "{s}linear_attn.A_log", .{p}, g.v_heads), @@ -210,10 +241,10 @@ pub fn load(e: *Engine, ck: *core.Checkpoint) !void { } }; } else { l.mixer = .{ .attention = .{ - .q = try linear(e, ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden), - .k = try linear(e, ck, p, "self_attn.k_proj", g.kv_heads * head_dim, g.hidden), - .v = try linear(e, ck, p, "self_attn.v_proj", g.kv_heads * head_dim, g.hidden), - .o = try linear(e, ck, p, "self_attn.o_proj", g.hidden, g.attnWidth()), + .q = try linear(e, e, ck, p, "self_attn.q_proj", 2 * g.attnWidth(), g.hidden), + .k = try linear(e, e, ck, p, "self_attn.k_proj", g.kv_heads * head_dim, g.hidden), + .v = try linear(e, e, ck, p, "self_attn.v_proj", g.kv_heads * head_dim, g.hidden), + .o = try linear(e, e, ck, p, "self_attn.o_proj", g.hidden, g.attnWidth()), .q_norm = try tensor(e, ck, "{s}self_attn.q_norm.weight", .{p}, .bf16, &.{head_dim}), .k_norm = try tensor(e, ck, "{s}self_attn.k_norm.weight", .{p}, .bf16, &.{head_dim}), } }; From d289747d4051fdf2d1e8eece95163a4aa54713bd Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:35:30 +0000 Subject: [PATCH 30/35] docs: Qwen3.6-35B-A3B on CUDA in the Windows page, the README and its recipe --- README.md | 3 ++- docs/WINDOWS.md | 8 ++++++++ docs/recipes/qwen3.6-moe.md | 3 ++- 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 48c7b8d27f..78b99d0745 100644 --- a/README.md +++ b/README.md @@ -75,6 +75,7 @@ The platform column names the hardware tested for each model. | GLM-5.3 | EXLR8 `home` profile (trellis K2/K3/K4 experts) | Metal on two M5 Ultras | | Qwen3.5-2B | Pinned MLX affine 4-bit, group 64, tied embeddings | Metal on M5 Max; CUDA on RTX 4090 (Windows) | | Qwen3.8-27B | MLX affine 4-bit, group 64, own output head; DFlash2 drafter | Metal on M5 Max and M3 Ultra; CUDA on RTX 4090 (Windows) | +| Qwen3.6-35B-A3B | MLX affine 4-bit, group 64, 8-bit routers; its MTP layer in `mtp-4bit.safetensors` | CUDA on RTX 4090 (Windows) | Nemotron's named checkpoint is `TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`. NVIDIA's NVFP4 checkpoint, `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4`, serves on CUDA as stored; its bf16 MTP head is quantized to 4-bit at load for drafts. @@ -90,7 +91,7 @@ Its prompt processing trails mlx_lm for now, and we are fixing it. An M3 Ultra r Bonsai, Gemma 4, Qwen3.6 and DeepSeek-V4 are still under qualification for 1.0.x. The Python 0.6.6 engine remains available at the `v0.6.6` tag for those models and other backends. On CUDA, 1.0.11 serves Nemotron on a GB10, greedy and sampled, with concurrent requests sharing each round. -On Windows ([docs/WINDOWS.md](docs/WINDOWS.md)), the native build also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees. +On Windows ([docs/WINDOWS.md](docs/WINDOWS.md)), the native build also serves the 2B and the 27B on CUDA, the 27B with its DFlash2 drafter drafting trees, and Qwen3.6-35B-A3B with its MTP layer drafting chains. The Linux x86_64 archive adds NVIDIA Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), tested on an RTX 3090, and Blackwell cards with compute capability 12.0 (the RTX 50 series and RTX PRO 6000 Blackwell), tested on an RTX PRO 6000; the A100 (8.0) is not supported yet. ## Exact decoding diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index 52b17191c6..165a8a9033 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -45,6 +45,8 @@ the same rounding points, so a window of up to 16 rows reproduces one-token step tensorfold-native serve --name qwen2b --context 32768 tensorfold-native serve --name qwen27b --context 8192 --parallel 4 ^ --drafter +tensorfold-native serve --name qwen36 --context 8192 --parallel 2 ^ + --drafter ``` - Qwen3.5-2B: about 340 tokens/s for one request on an RTX 4090, 1,500 across eight. @@ -55,6 +57,12 @@ tensorfold-native serve --name qwen27b --context 8192 --p code at the high), sampled ones at 50 to 70, and four at once reach about 160 aggregate; replies stay identical to plain decoding. Loading takes about 80 s: the drafter's weights are quantized to 4 bits on the host and every shape is run once, as the driver stalls a process's first use of each kernel and buffer. +- Qwen3.6-35B-A3B (`qwen3_5_moe`, 18.6 GiB): the same layers with 256 routed experts, 8 a token, in place of the + dense MLP. Routing follows the frozen engine's `cuda/moe.py` rule, and the experts run through the grouped affine-4 + kernel's SwiGLU form, the shared expert last in the same tables. About 100 tokens/s plain for one request. With + `--drafter` naming the checkpoint's own directory, its MTP layer (`mtp-4bit.safetensors`) drafts chains of three + from each kept row's final normed state: 230 to 260 tokens/s greedy, 180 to 250 sampled, 420 aggregate across two + requests; replies stay identical to plain decoding. `TF_MTP_DRAFTS=N` sets the chain's length. - A verify window's rows are not snapshotted: each stream logs its rows' recurrence updates and a keep replays the accepted path, so the 27B records 32 rows a forward where 8 fit before. - `TF_TIMING=1` prints each verify, keep, absorb and hold's milliseconds and the drafted shape (chain or lanes); diff --git a/docs/recipes/qwen3.6-moe.md b/docs/recipes/qwen3.6-moe.md index 15b756169b..28d7696b96 100644 --- a/docs/recipes/qwen3.6-moe.md +++ b/docs/recipes/qwen3.6-moe.md @@ -6,7 +6,8 @@ engine. @cshintov brought the model to Macs in [#138](https://github.com/ashhart @philip-pentatonic ([#78](https://github.com/ashhart/TensorFold/pull/78), [#84](https://github.com/ashhart/TensorFold/pull/84)) and @BHCC2025 ([#270](https://github.com/ashhart/TensorFold/pull/270)) worked on its CUDA path; see [credits](../credits.md). -The native binary has no `qwen3_5_moe` family yet, so `tensorfold pull` and `tensorfold info` refuse the checkpoint. +The native CUDA engine serves the checkpoint with its MTP layer ([docs/WINDOWS.md](../WINDOWS.md)); the Metal native +binary has no `qwen3_5_moe` family yet, so `tensorfold pull` and `tensorfold info` refuse the checkpoint there. The `qwen3_5_moe` family serves Qwen3.6-35B-A3B on Apple Silicon and on one NVIDIA GPU. Its layers are the 27B's (Gated DeltaNet and gated full attention, every fourth layer attention) with routed experts in place of the dense From 637f339058809f0f74213733e2a2e2fc3a10eb54 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:15:19 +0000 Subject: [PATCH 31/35] qwen3_5: the MTP drafter can draft the planner's trees, opt-in TF_MTP_TREES=1 lets the lane planner shape a tree of lanes (up to 12, the ranks of a row's best tokens) for the MTP drafter: each depth's lanes run as one pass of the layer, every row attending the ring and its own ancestor lanes' keys (tf_q36m_attn_tree over a per-row ancestor table), their outputs and keys kept by lane for their children, and each lane's token its parent's rank-r candidate from the logprob rows. The load-time cost of a draft level is now the quickest of sixteen holds, as the driver's stalls lasted a dozen. Measured on an RTX 4090, Qwen3.6-35B-A3B, --parallel 2 --context 8192, bench_openai median tok/s (code T=1 / chat T=1 / code greedy / chat greedy): chains 248 / 197 / 271 / 235, trees 237 / 195 / 250 / 236; two streams 387 to 450 aggregate with chains, 381 to 439 with trees. Trees draft three times the rows for the same tokens a round (88 of 430 lanes kept against 83 of 132 chained drafts), and the wider verify windows cost what the rounds save, so chains stay the default. Replies equal "draft": false either way. --- docs/WINDOWS.md | 3 +- zig/kernels/cuda/qwen36_mtp.cu | 61 +++++++- zig/src/families/qwen3_5/cuda_mtp.zig | 198 ++++++++++++++++++++------ 3 files changed, 215 insertions(+), 47 deletions(-) diff --git a/docs/WINDOWS.md b/docs/WINDOWS.md index 165a8a9033..8bd8bf5845 100644 --- a/docs/WINDOWS.md +++ b/docs/WINDOWS.md @@ -62,7 +62,8 @@ tensorfold-native serve --name qwen36 --context 8 kernel's SwiGLU form, the shared expert last in the same tables. About 100 tokens/s plain for one request. With `--drafter` naming the checkpoint's own directory, its MTP layer (`mtp-4bit.safetensors`) drafts chains of three from each kept row's final normed state: 230 to 260 tokens/s greedy, 180 to 250 sampled, 420 aggregate across two - requests; replies stay identical to plain decoding. `TF_MTP_DRAFTS=N` sets the chain's length. + requests; replies stay identical to plain decoding. `TF_MTP_DRAFTS=N` sets the chain's length, and `TF_MTP_TREES=1` + drafts the planner's trees instead, which measured no faster than chains on this head. - A verify window's rows are not snapshotted: each stream logs its rows' recurrence updates and a keep replays the accepted path, so the 27B records 32 rows a forward where 8 fit before. - `TF_TIMING=1` prints each verify, keep, absorb and hold's milliseconds and the drafted shape (chain or lanes); diff --git a/zig/kernels/cuda/qwen36_mtp.cu b/zig/kernels/cuda/qwen36_mtp.cu index c996c1fe72..5388000be4 100644 --- a/zig/kernels/cuda/qwen36_mtp.cu +++ b/zig/kernels/cuda/qwen36_mtp.cu @@ -106,7 +106,66 @@ extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn(const __nv_bfloat1 for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); mx = mn; } - const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; // the q projection's gate + const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; // the gate half +#pragma unroll + for (int i = 0; i < DPL; i++) { + const float gate = 1.0f / (1.0f + __expf(-bf(g[i]))); + OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(rbf(o[i] / l) * gate); + } +} + +// Tree rows at one depth, a warp a (row, head): ring [begin, pos0), the ancestor lanes ANC[m] names, then its own key. +extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn_tree(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ G, int gstride, + const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, + const __nv_bfloat16* __restrict__ LK, const __nv_bfloat16* __restrict__ LV, + const int* __restrict__ ANC, int anc_stride, const __nv_bfloat16* __restrict__ RK, + const __nv_bfloat16* __restrict__ RV, int heads, int kv_heads, int W, long long begin, + long long pos0, __nv_bfloat16* __restrict__ OUT) { + constexpr int DPL = HD / 32; + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + const float scale = 1.0f / 16.0f; + float q[DPL], o[DPL]; +#pragma unroll + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + for (long long pos = begin; pos < pos0; pos++) { + const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + const int* anc = ANC + (size_t)m * anc_stride; + const int n = anc[0]; + for (int j = 0; j <= n; j++) { + const bool self = j == n; + const __nv_bfloat16* k = (self ? KB + (size_t)m * width : LK + (size_t)anc[1 + j] * width) + kvh * HD + lane * DPL; + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + const float s = warp_sum(dot) * scale; + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; + const __nv_bfloat16* v = (self ? VB + (size_t)m * width : LV + (size_t)anc[1 + j] * width) + kvh * HD + lane * DPL; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; + } + const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; #pragma unroll for (int i = 0; i < DPL; i++) { const float gate = 1.0f / (1.0f + __expf(-bf(g[i]))); diff --git a/zig/src/families/qwen3_5/cuda_mtp.zig b/zig/src/families/qwen3_5/cuda_mtp.zig index 3c581efc6f..0f87e52acf 100644 --- a/zig/src/families/qwen3_5/cuda_mtp.zig +++ b/zig/src/families/qwen3_5/cuda_mtp.zig @@ -1,4 +1,4 @@ -//! The Qwen3.6 MTP layer as the drafter on CUDA: chains from a kept row's final normed state and the token after it. +//! The Qwen3.6 MTP layer as the drafter on CUDA: chains or trees from a kept row's final normed state and next token. const std = @import("std"); const platform = @import("sys"); const cuda = @import("cuda"); @@ -18,23 +18,31 @@ const most_window = engine.most_window; /// Drafts a round at most (the frozen engine's DEPTH); TF_MTP_DRAFTS=N sets it, up to a window's rows less one. const default_depth = 3; +/// A tree's lanes at most (the lane planner's widest shape within a window), and the ranks a row's candidates cover. +const max_lanes = most_window; +const ranks = lanes.shape.ranks; +/// The logprob row's words (sample.cu's LP_WORDS): the ids of a row's best tokens sit at words 2 onward. +const lp_words = 3 + 2 * 20; -const Fns = struct { norm_rope: cuda.Function, attn: cuda.Function }; +const Fns = struct { norm_rope: cuda.Function, attn: cuda.Function, attn_tree: cuda.Function }; /// qwen35_draft.cu's ring kernels, shared with the DFlash2 drafter. const Ring = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, ring_put: cuda.Function }; const TapRows = extern struct { n: c_int, rows: [most_window]c_int }; -/// A stream's context: its ring of keys and values over [begin, end), the newest row's state, and its held chain. +/// A stream's context: its ring of keys and values over [begin, end), the newest row's state, and its held drafts. const Ctx = struct { ring: DeviceBuffer, - last: DeviceBuffer, // the newest absorbed row's normed output (bf16 hidden): the first chain row's state + last: DeviceBuffer, // the newest absorbed row's normed output (bf16 hidden): the first draft row's state begin: u64 = 0, end: u64 = 0, - first: u32 = 0, // the newest absorbed row's draft, its logits' argmax - held: [most_window]u32 = undefined, + cands: [ranks]u32 = @splat(0), // the newest absorbed row's best tokens by rank (rank 0: its chain draft) + held: [max_lanes]u32 = undefined, held_n: usize = 0, }; +/// How a layer pass attends: absorbed or chained rows over the ring and the block's earlier rows, into the ring too. +const Pass = union(enum) { ring, tree: u64 }; + pub const Mtp = struct { gpa: std.mem.Allocator, io: std.Io, @@ -48,6 +56,7 @@ pub const Mtp = struct { at: usize = 0, capacity: usize, // ring rows: every position of the stream's window depth: u32, + trees: bool, // TF_MTP_TREES=1: the planner's trees instead of chains (measured a wash on this head) tap: [1]u32, // rows of max_rows, bf16 unless said taps_in: u64, @@ -65,28 +74,37 @@ pub const Mtp = struct { vb: u64, att: u64, xs: u64, // f32 group sums for the projections - logits: u64, // one row + logits: u64, // max_lanes rows + words: u64, // u32 logprob words a lane row + picks: u64, // u32 zeros: the logprob rows' pick (unused) + lane_state: u64, // each tree lane's normed output + lane_k: u64, // each tree lane's key and value rows + lane_v: u64, + anc: u64, // int32 ancestor table a row: a count, then lane indices ids: u64, // u32 the rows' tokens tok: u64, // u32 the drawn token - staged: cuda.HostBuffer, // pinned: the ids up, the token down + staged: cuda.HostBuffer, // pinned: the ids and ancestors up, the token and words down streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, /// Contexts between streams, allocated and run once (`reserve`), as the driver stalls a buffer's first kernel. pool: std.ArrayList(*Ctx) = .empty, - /// One chained draft's cost (timed at load). + /// One draft level's cost (timed at load): a chained row, or a tree's rows at one depth. step_ms: f64 = 1.0, /// The target's window ms by rows (cuda_native.zig times them), for the round loop's prices. costs: [most_window]lanes.config.Cost = undefined, n_costs: usize = 0, + const anc_stride = 1 + max_lanes; const ids_at = 0; const tok_at = max_rows * 4; - const staged_len = tok_at + 64; + const words_at = tok_at + 64; + const anc_at = words_at + max_lanes * lp_words * 4; + const staged_len = anc_at + max_lanes * anc_stride * 4; pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8, capacity: usize) !*Mtp { if (!e.g.moe()) return error.NotRouted; const d = try gpa.create(Mtp); errdefer gpa.destroy(d); - d.* = .{ .gpa = gpa, .io = io, .e = e, .w = undefined, .mod = undefined, .dmod = undefined, .fns = undefined, .ring_fns = undefined, .scratch = undefined, .capacity = capacity, .depth = default_depth, .tap = .{@intCast(e.g.layers - 1)}, .taps_in = 0, .hin = 0, .hn = 0, .emb = 0, .en = 0, .h = 0, .r = 0, .x = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .xs = 0, .logits = 0, .ids = 0, .tok = 0, .staged = undefined }; + d.* = .{ .gpa = gpa, .io = io, .e = e, .w = undefined, .mod = undefined, .dmod = undefined, .fns = undefined, .ring_fns = undefined, .scratch = undefined, .capacity = capacity, .depth = default_depth, .trees = platform.c.getenv("TF_MTP_TREES") != null, .tap = .{@intCast(e.g.layers - 1)}, .taps_in = 0, .hin = 0, .hn = 0, .emb = 0, .en = 0, .h = 0, .r = 0, .x = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .xs = 0, .logits = 0, .words = 0, .picks = 0, .lane_state = 0, .lane_k = 0, .lane_v = 0, .anc = 0, .ids = 0, .tok = 0, .staged = undefined }; if (platform.c.getenv("TF_MTP_DRAFTS")) |v| d.depth = std.fmt.parseInt(u32, std.mem.span(v), 10) catch default_depth; d.depth = @min(@max(d.depth, 1), most_window - 1); d.mod = try cuda.Module.load(e.d, cuda.kernels.qwen36_mtp); @@ -104,15 +122,16 @@ pub const Mtp = struct { _ = try d.layout(true); d.staged = try cuda.HostBuffer.alloc(e.d, staged_len); errdefer d.staged.free(); - // the first context, and a chain timed on it (the quickest of several: the first ones pay the driver's stalls) + // the first context, and a chain timed on it: the quickest of sixteen, as the driver's stalls last a dozen const ctx = try d.create(); errdefer d.drop(ctx); var quickest: f64 = std.math.inf(f64); - for (0..6) |_| { + for (0..16) |_| { try d.absorb(ctx, d.taps_in, 0, 1, null, &.{0}); const t0 = std.Io.Clock.awake.now(io); - try d.hold(ctx, 0, 1, d.depth); - quickest = @min(quickest, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6); + try d.hold(ctx, 0, 1, d.depth, null); + const ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6; + quickest = @min(quickest, ms); } d.step_ms = quickest / @as(f64, @floatFromInt(d.depth)); d.resetCtx(ctx); @@ -126,7 +145,7 @@ pub const Mtp = struct { } pub fn setCost(d: *Mtp, rows: u32, ms: f64) void { - if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a chained draft {d:.2})\n", .{ rows, ms, d.step_ms }); + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] window {d} rows {d:.2} ms (a draft level {d:.2})\n", .{ rows, ms, d.step_ms }); if (d.n_costs < d.costs.len) { d.costs[d.n_costs] = .{ .width = rows, .ms = ms }; d.n_costs += 1; @@ -164,14 +183,16 @@ pub const Mtp = struct { fn layout(d: *Mtp, assign: bool) !usize { const g = d.e.g; const R = max_rows; + const L = max_lanes; const aw = g.attnWidth(); const kw = d.kvWidth(); const sizes = .{ - .{ "taps_in", R * g.hidden * 2 }, .{ "hin", R * g.hidden * 2 }, .{ "hn", R * g.hidden * 2 }, .{ "emb", R * g.hidden * 2 }, - .{ "en", R * g.hidden * 2 }, .{ "h", R * g.hidden * 2 }, .{ "r", R * g.hidden * 2 }, .{ "x", R * g.hidden * 2 }, - .{ "qraw", R * 2 * aw * 2 }, .{ "q", R * aw * 2 }, .{ "kraw", R * kw * 2 }, .{ "kb", R * kw * 2 }, - .{ "vb", R * kw * 2 }, .{ "att", R * aw * 2 }, .{ "xs", R * (@max(g.hidden, aw) / 64) * 4 }, .{ "logits", g.vocab * 2 }, - .{ "ids", R * 4 }, .{ "tok", 64 }, + .{ "taps_in", R * g.hidden * 2 }, .{ "hin", R * g.hidden * 2 }, .{ "hn", R * g.hidden * 2 }, .{ "emb", R * g.hidden * 2 }, + .{ "en", R * g.hidden * 2 }, .{ "h", R * g.hidden * 2 }, .{ "r", R * g.hidden * 2 }, .{ "x", R * g.hidden * 2 }, + .{ "qraw", R * 2 * aw * 2 }, .{ "q", R * aw * 2 }, .{ "kraw", R * kw * 2 }, .{ "kb", R * kw * 2 }, + .{ "vb", R * kw * 2 }, .{ "att", R * aw * 2 }, .{ "xs", R * (@max(g.hidden, aw) / 64) * 4 }, .{ "logits", L * g.vocab * 2 }, + .{ "words", L * lp_words * 4 }, .{ "picks", L * 4 }, .{ "lane_state", L * g.hidden * 2 }, .{ "lane_k", L * kw * 2 }, + .{ "lane_v", L * kw * 2 }, .{ "anc", L * anc_stride * 4 }, .{ "ids", R * 4 }, .{ "tok", 64 }, }; var total: usize = 0; inline for (sizes) |s| { @@ -204,7 +225,7 @@ pub const Mtp = struct { try ctx.ring.fill8(0, null); try ctx.last.fill8(0, null); try d.absorb(ctx, d.taps_in, 0, 1, null, &.{0}); // any tap rows do: the drafter's own zeroed staging - try d.hold(ctx, 0, 1, d.depth); + try d.hold(ctx, 0, 1, d.depth, null); d.resetCtx(ctx); return ctx; } @@ -239,7 +260,7 @@ pub const Mtp = struct { } /// The layer over `rows` rows at slots pos0..: d.ids their tokens, d.hin their states; x gets their normed outputs. - fn layer(d: *Mtp, ctx: *Ctx, rows: usize, pos0: u64) !void { + fn layer(d: *Mtp, ctx: *Ctx, rows: usize, pos0: u64, pass: Pass) !void { const e = d.e; const g = e.g; const w = &d.w; @@ -257,26 +278,41 @@ pub const Mtp = struct { try d.proj(w.v, d.x, d.vb, rows); try e.run(d.fns.norm_rope, .{ rows * g.q_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.qraw, w.q_norm, d.q, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), p0, g.eps, geometry.theta }); try e.run(d.fns.norm_rope, .{ rows * g.kv_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.kraw, w.k_norm, d.kb, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(kw)), @as(c_int, head_dim), p0, g.eps, geometry.theta }); - try e.run(d.fns.attn, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(d.capacity)), @as(c_longlong, @intCast(ctx.begin)), p0, d.att }); + const heads: c_int = @intCast(g.q_heads); + const kvh: c_int = @intCast(g.kv_heads); + const cap: c_int = @intCast(d.capacity); + const begin: c_longlong = @intCast(ctx.begin); + switch (pass) { + .ring => try e.run(d.fns.attn, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), heads, kvh, cap, begin, p0, d.att }), + .tree => |anc| try e.run(d.fns.attn_tree, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.lane_k, d.lane_v, anc, @as(c_int, anc_stride), d.ringK(ctx), d.ringV(ctx), heads, kvh, cap, begin, p0, d.att }), + } try d.proj(w.o, d.att, d.r, rows); try d.norm(d.h, d.r, w.post_norm, d.x, rows); try e.routedRows(w.mlp.moe.router, w.mlp.moe.experts, d.x, d.r, rows, rows > most_window); try d.norm(d.h, d.r, w.norm, d.x, rows); - try e.run(d.ring_fns.ring_put, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), @as(c_int, @intCast(kw)), p0, @as(c_int, @intCast(d.capacity)) }); + if (pass == .ring) try e.run(d.ring_fns.ring_put, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), @as(c_int, @intCast(kw)), p0, cap }); } - /// Row `row` of x through the target's head: its argmax into `tok` (the host sees it after a synchronize). - fn draw(d: *Mtp, row: usize) !void { + /// Rows [first, first + n) of x through the target's head: their logits and best tokens' words, row 0's pick. + fn draw(d: *Mtp, first: usize, n: usize, cands: bool) !void { const e = d.e; - try d.proj(e.head, d.x + row * e.g.hidden * 2, d.logits, 1); + try d.proj(e.head, d.x + first * e.g.hidden * 2, d.logits, n); try e.run(e.fns.argmax, .{ 1, 1, 1 }, .{ 1024, 1, 1 }, .{ d.logits, @as(c_int, @intCast(e.g.vocab)), d.tok }); + if (cands and d.trees) try e.draws.logprobRows(e.stream, d.logits, e.g.vocab, d.picks, ranks, false, d.words, n); } - fn token(d: *Mtp) !u32 { + /// The drawn token, and the first `n` rows' candidates by rank, after a draw. + fn fetch(d: *Mtp, n: usize, cands: [][ranks]u32) !u32 { const e = d.e; - try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(d.staged.bytes.ptr + tok_at, d.tok, 4, e.stream.handle), "cuMemcpyDtoHAsync"); + const st = d.staged.bytes; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + tok_at, d.tok, 4, e.stream.handle), "cuMemcpyDtoHAsync"); + if (d.trees and n > 0) try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + words_at, d.words, n * lp_words * 4, e.stream.handle), "cuMemcpyDtoHAsync"); try e.stream.synchronize(); - return std.mem.bytesToValue(u32, d.staged.bytes[tok_at..][0..4]); + for (cands[0..n], 0..) |*c, r| { + const w = std.mem.bytesAsSlice(u32, st[words_at + r * lp_words * 4 ..][0 .. lp_words * 4]); + for (c, 0..) |*t, j| t.* = if (d.trees) w[2 + j] else std.mem.bytesToValue(u32, st[tok_at..][0..4]); + } + return std.mem.bytesToValue(u32, st[tok_at..][0..4]); } fn putIds(d: *Mtp, ids: []const u32) !void { @@ -313,39 +349,111 @@ pub const Mtp = struct { } else try e.run(d.ring_fns.gather_taps, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ tap_ring, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(g.hidden)), d.taps_in }); try d.norm(d.taps_in, 0, e.final_norm, d.hin, rows); // the target's final norm: the rows the head read try d.putIds(follow[fi..][0..rows]); - try d.layer(ctx, rows, start); + try d.layer(ctx, rows, start, .ring); start += rows; n -= rows; fi += rows; ctx.end = start; - // the newest row: its state kept for the chain, its draft drawn + // the newest row: its state kept for the drafts, its best tokens drawn try ctx.last.copyFrom(0, d.x + (rows - 1) * g.hidden * 2, g.hidden * 2, e.stream.handle); - try d.draw(rows - 1); + try d.draw(rows - 1, 1, true); } - ctx.first = try d.token(); + var cands: [1][ranks]u32 = undefined; + ctx.cands[0] = try d.fetch(1, &cands); + if (d.trees) ctx.cands = cands[0]; } - /// `depth` drafts after `pending` at pos0 (the pending token's slot): the first came with the newest absorbed row. - fn hold(d: *Mtp, ctx: *Ctx, pending: u32, pos0: u64, depth: usize) !void { + /// `depth` drafts after `pending` at pos0 (the pending token's slot): a chain, or the planner's tree of lanes. + fn hold(d: *Mtp, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { + _ = pending; // the newest absorbed row's follow token, whose candidates ctx.cands are + if (ctx.end != pos0) return error.DraftContextBehind; + if (lanes_) |l| return d.holdTree(ctx, pos0, l); const e = d.e; const g = e.g; - if (ctx.end != pos0) return error.DraftContextBehind; - ctx.held[0] = ctx.first; + ctx.held[0] = ctx.cands[0]; ctx.held_n = 1; - _ = pending; // the newest absorbed row's follow token, whose draft ctx.first is var k: usize = 1; while (k < depth and k < ctx.held.len) : (k += 1) { // row k: (the previous row's normed output, the draft before) at slot pos0 + k - 1 const src = if (k == 1) ctx.last.ptr else d.x; try d.scratch.copyFrom(d.hin - d.scratch.ptr, src, g.hidden * 2, e.stream.handle); try d.putIds(&.{ctx.held[k - 1]}); - try d.layer(ctx, 1, pos0 + k - 1); - try d.draw(0); - ctx.held[k] = try d.token(); + try d.layer(ctx, 1, pos0 + k - 1, .ring); + try d.draw(0, 1, false); + var none: [0][ranks]u32 = undefined; + ctx.held[k] = try d.fetch(0, &none); ctx.held_n = k + 1; } } + /// The tree's lanes, lane i the rank-r candidate of its parent's row; a depth's rows run as one pass. + fn holdTree(d: *Mtp, ctx: *Ctx, pos0: u64, l: *const lanes.shape.Shape) !void { + const e = d.e; + const g = e.g; + const n = l.parents.len; + if (n > max_lanes) return error.TooManyLanes; + var cands: [max_lanes][ranks]u32 = undefined; // each lane's row's candidates, once its row has run + var internal: [max_lanes]bool = @splat(false); + for (l.parents, l.depths, l.ranks, 0..) |par, dd, rank, i| { + if (rank >= ranks or par >= @as(i32, @intCast(i)) or dd >= most_window) return error.BadLanes; + if (par >= 0) internal[@intCast(par)] = true else ctx.held[i] = ctx.cands[rank]; + } + // internal lanes by depth: each pass's rows are the lanes at that depth whose rows later lanes read + var done: [max_lanes]bool = @splat(false); + var depth: u8 = 0; + while (depth < most_window) : (depth += 1) { + var batch: [max_lanes]usize = undefined; + var rows: usize = 0; + for (l.depths, 0..) |dd, i| if (dd == depth and internal[i]) { + batch[rows] = i; + rows += 1; + }; + if (rows == 0) { + var any = false; + for (l.depths, 0..) |dd, i| any = any or (dd > depth and internal[i]); + if (!any) break; + continue; + } + // each row's inputs: its parent's output state, its own token; its ancestors as lane indices + const st = d.staged.bytes; + const anc = std.mem.bytesAsSlice(i32, st[anc_at..][0 .. rows * anc_stride * 4]); + var toks: [max_lanes]u32 = undefined; + for (batch[0..rows], 0..) |i, m| { + const par = l.parents[i]; + const src = if (par < 0) ctx.last.ptr else d.lane_state + @as(u64, @intCast(par)) * g.hidden * 2; + try d.scratch.copyFrom(d.hin - d.scratch.ptr + m * g.hidden * 2, src, g.hidden * 2, e.stream.handle); + toks[m] = ctx.held[i]; + var count: usize = 0; + var a = par; + while (a >= 0) : (a = l.parents[@intCast(a)]) { + anc[m * anc_stride + 1 + count] = a; + count += 1; + } + anc[m * anc_stride] = @intCast(count); + } + try d.putIds(toks[0..rows]); + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(d.anc, st.ptr + anc_at, rows * anc_stride * 4, e.stream.handle), "cuMemcpyHtoDAsync"); + try d.layer(ctx, rows, pos0 + depth, .{ .tree = d.anc }); + const kw = d.kvWidth(); + for (batch[0..rows], 0..) |i, m| { + try d.scratch.copyFrom(d.lane_state - d.scratch.ptr + i * g.hidden * 2, d.x + m * g.hidden * 2, g.hidden * 2, e.stream.handle); + try d.scratch.copyFrom(d.lane_k - d.scratch.ptr + i * kw * 2, d.kb + m * kw * 2, kw * 2, e.stream.handle); + try d.scratch.copyFrom(d.lane_v - d.scratch.ptr + i * kw * 2, d.vb + m * kw * 2, kw * 2, e.stream.handle); + done[i] = true; + } + try d.draw(0, rows, true); + var got: [max_lanes][ranks]u32 = undefined; + _ = try d.fetch(rows, got[0..rows]); + for (batch[0..rows], 0..) |i, m| cands[i] = got[m]; + // the lanes whose parent's row just ran take their tokens now + for (l.parents, l.ranks, 0..) |par, rank, i| if (par >= 0 and done[@intCast(par)]) { + ctx.held[i] = cands[@intCast(par)][rank]; + }; + } + ctx.held_n = n; + if (platform.c.getenv("TF_TIMING") != null) std.debug.print("[timing] lanes {d} levels {d}\n", .{ n, l.levels() }); + } + // the Drafter contract pub fn drafter(d: *Mtp) dr.Drafter { @@ -362,7 +470,7 @@ pub const Mtp = struct { fn factsFn(p: *anyopaque) dr.Facts { const d = cast(p); - return .{ .depth = d.depth, .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.ringBytes() + d.e.g.hidden * 2, .trees = false, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; + return .{ .depth = d.depth, .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.ringBytes() + d.e.g.hidden * 2, .trees = d.trees, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; } fn openFn(p: *anyopaque, s: *Stream) anyerror!void { @@ -394,7 +502,7 @@ pub const Mtp = struct { defer timing("mtp hold", d.io, t0); for (items) |it| { const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; - try d.hold(ctx, it.pending, it.position - 1, it.depth); + try d.hold(ctx, it.pending, it.position - 1, it.depth, it.lanes); } } From f65cd734da2f17da9e98252f2ccc645356be9883 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 07:33:24 +0000 Subject: [PATCH 32/35] qwen3_5: MTP tree lanes at one depth share their rotary slot and read the ring to its committed end A tree pass packs a depth's sibling lanes as consecutive rows, and norm_rope rotated row m at pos0 + m, so a depth's second lane took the next slot's rotation, and so did the key copied out of it and the child drawn from it. The rope kernel now takes a row step: 1 for absorbed and chained rows, 0 for a tree pass, whose base stays the pending slot plus the depth. Tree attention read the ring up to that base, one slot per depth past the committed end; a tree pass writes no ring row, so those slots could still hold an earlier round's keys. It now reads the ring to the context's end. --- zig/kernels/cuda/qwen36_mtp.cu | 12 ++++++------ zig/src/families/qwen3_5/cuda_mtp.zig | 8 +++++--- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/zig/kernels/cuda/qwen36_mtp.cu b/zig/kernels/cuda/qwen36_mtp.cu index 5388000be4..5118a8aee7 100644 --- a/zig/kernels/cuda/qwen36_mtp.cu +++ b/zig/kernels/cuda/qwen36_mtp.cu @@ -21,10 +21,10 @@ __device__ __forceinline__ float warp_sum(float x) { using namespace tf_q36m; -// A warp a (row, head): RMS norm over HD, rope on [0, ROT) as pairs (i, i + 32) at pos0 + row; OUT (rows, heads, HD). +// A warp a (row, head): RMS norm, rope on [0, ROT) as pairs (i, i + 32) at pos0 + row * step; OUT (rows, heads, HD) extern "C" __global__ void __launch_bounds__(32) tf_q36m_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ OUT, int heads, int stride, int in_head, - long long pos0, float eps, float theta) { + long long pos0, int step, float eps, float theta) { constexpr int DPL = HD / 32; const int head = blockIdx.x % heads; const size_t m = blockIdx.x / heads; @@ -42,7 +42,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q36m_norm_rope(const __nv_bf #pragma unroll for (int i = 0; i < DPL; i++) n[i] = rbf(bf(W[DPL * t + i]) * (v[i] * inv)); // a rotary pair's partner sits 32 elements away: DPL = 8 a lane, so 4 lanes away at the same i - const float pos = (float)(pos0 + (long long)m); + const float pos = (float)(pos0 + (long long)m * step); __nv_bfloat16* out = OUT + (m * heads + head) * HD; #pragma unroll for (int i = 0; i < DPL; i++) { @@ -114,13 +114,13 @@ extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn(const __nv_bfloat1 } } -// Tree rows at one depth, a warp a (row, head): ring [begin, pos0), the ancestor lanes ANC[m] names, then its own key. +// Tree rows at one depth, a warp a (row, head): ring [begin, end), the ancestor lanes ANC[m] names, then its own key. extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn_tree(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ G, int gstride, const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, const __nv_bfloat16* __restrict__ LK, const __nv_bfloat16* __restrict__ LV, const int* __restrict__ ANC, int anc_stride, const __nv_bfloat16* __restrict__ RK, const __nv_bfloat16* __restrict__ RV, int heads, int kv_heads, int W, long long begin, - long long pos0, __nv_bfloat16* __restrict__ OUT) { + long long end, __nv_bfloat16* __restrict__ OUT) { constexpr int DPL = HD / 32; const int lane = threadIdx.x; const int m = blockIdx.x, h = blockIdx.y; @@ -134,7 +134,7 @@ extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn_tree(const __nv_bf o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (long long pos = begin; pos < pos0; pos++) { + for (long long pos = begin; pos < end; pos++) { const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; float dot = 0.0f; #pragma unroll diff --git a/zig/src/families/qwen3_5/cuda_mtp.zig b/zig/src/families/qwen3_5/cuda_mtp.zig index 0f87e52acf..1523cd51ba 100644 --- a/zig/src/families/qwen3_5/cuda_mtp.zig +++ b/zig/src/families/qwen3_5/cuda_mtp.zig @@ -267,6 +267,8 @@ pub const Mtp = struct { const aw = g.attnWidth(); const kw = d.kvWidth(); const p0: c_longlong = @intCast(pos0); + // a tree pass's rows are sibling lanes at one depth, so they share pos0; chained or absorbed rows step by one + const step: c_int = @intFromBool(pass == .ring); try e.run(e.fns.embed, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.ids, e.embed_w, e.embed_s, e.embed_b, d.emb, @as(c_int, @intCast(g.hidden)), @as(u64, 0) }); try d.norm(d.emb, 0, w.norm_e, d.en, rows); try d.norm(d.hin, 0, w.norm_h, d.hn, rows); @@ -276,15 +278,15 @@ pub const Mtp = struct { try d.proj(w.q, d.x, d.qraw, rows); try d.proj(w.k, d.x, d.kraw, rows); try d.proj(w.v, d.x, d.vb, rows); - try e.run(d.fns.norm_rope, .{ rows * g.q_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.qraw, w.q_norm, d.q, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), p0, g.eps, geometry.theta }); - try e.run(d.fns.norm_rope, .{ rows * g.kv_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.kraw, w.k_norm, d.kb, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(kw)), @as(c_int, head_dim), p0, g.eps, geometry.theta }); + try e.run(d.fns.norm_rope, .{ rows * g.q_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.qraw, w.q_norm, d.q, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), p0, step, g.eps, geometry.theta }); + try e.run(d.fns.norm_rope, .{ rows * g.kv_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.kraw, w.k_norm, d.kb, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(kw)), @as(c_int, head_dim), p0, step, g.eps, geometry.theta }); const heads: c_int = @intCast(g.q_heads); const kvh: c_int = @intCast(g.kv_heads); const cap: c_int = @intCast(d.capacity); const begin: c_longlong = @intCast(ctx.begin); switch (pass) { .ring => try e.run(d.fns.attn, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), heads, kvh, cap, begin, p0, d.att }), - .tree => |anc| try e.run(d.fns.attn_tree, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.lane_k, d.lane_v, anc, @as(c_int, anc_stride), d.ringK(ctx), d.ringV(ctx), heads, kvh, cap, begin, p0, d.att }), + .tree => |anc| try e.run(d.fns.attn_tree, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.lane_k, d.lane_v, anc, @as(c_int, anc_stride), d.ringK(ctx), d.ringV(ctx), heads, kvh, cap, begin, @as(c_longlong, @intCast(ctx.end)), d.att }), } try d.proj(w.o, d.att, d.r, rows); try d.norm(d.h, d.r, w.post_norm, d.x, rows); From 33b55a9e329204d130875b8085511fcc1eaf9305 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:26:19 +0000 Subject: [PATCH 33/35] qwen3_5: the MTP drafter's id uploads wait for the previous one putIds wrote the pinned staging area and queued its copy without waiting, and absorb calls it once a prompt chunk, so an absorb past 128 rows could overwrite a chunk's follow tokens before the GPU had read them. Replies stayed correct (the target verifies every draft), but the drafter absorbed the wrong tokens after any long prompt. It now synchronizes the stream before reusing the area; holds already had, through their fetch. --- zig/src/families/qwen3_5/cuda_mtp.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/src/families/qwen3_5/cuda_mtp.zig b/zig/src/families/qwen3_5/cuda_mtp.zig index 1523cd51ba..4beac9a97f 100644 --- a/zig/src/families/qwen3_5/cuda_mtp.zig +++ b/zig/src/families/qwen3_5/cuda_mtp.zig @@ -319,6 +319,7 @@ pub const Mtp = struct { fn putIds(d: *Mtp, ids: []const u32) !void { const e = d.e; + try e.stream.synchronize(); // the previous upload from this staging area has run: an absorb queues one a chunk @memcpy(std.mem.bytesAsSlice(u32, d.staged.bytes[ids_at..][0 .. ids.len * 4]), ids); try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(d.ids, d.staged.bytes.ptr + ids_at, ids.len * 4, e.stream.handle), "cuMemcpyHtoDAsync"); } From f5c1247c3b829d2f82ae63c7ac7e973d492f0999 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:26:28 +0000 Subject: [PATCH 34/35] qwen3_5: GPU regressions for the MTP drafter's own rows, and tree ancestors root first `tf-qwen35-cuda mtp ` absorbs the ids from the target's tap ring and compares the drafter's rows themselves, not the replies the target verifies: one absorb over several prompt chunks against the same rows a chunk a call (ring, newest state, candidates); a tree with two siblings at depth 0, six at depth 1 and a depth-2 lane with two ancestors, run again with the ring poisoned at and past its committed end; and every internal lane's state, key and value against its path chained one row a slot. All are byte-equal. The tree's ancestor table now lists a lane's ancestors root first, the order a chain's ring holds them, so a lane's fp32 sums equal its chained path's. Each check fails on the bug it covers: the sibling rows with the rotary step put back to 1 (lane 1 at depth 0), the poisoned ring with the tree's ring end put back to its rotary base (lane 3 at depth 1), the chunked absorb without the previous commit's wait (4 runs of 4). --- zig/src/families/qwen3_5/cuda.zig | 1 + zig/src/families/qwen3_5/cuda_mtp.zig | 22 ++-- zig/src/families/qwen3_5/cuda_mtp_check.zig | 133 ++++++++++++++++++++ zig/tests/qwen35_cuda_main.zig | 14 ++- 4 files changed, 159 insertions(+), 11 deletions(-) create mode 100644 zig/src/families/qwen3_5/cuda_mtp_check.zig diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index e79beb5a2b..653029762c 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -5,6 +5,7 @@ pub const engine = @import("cuda_engine.zig"); pub const native = @import("cuda_native.zig"); pub const native_moe = @import("cuda_native_moe.zig"); pub const mtp = @import("cuda_mtp.zig"); +pub const mtp_check = @import("cuda_mtp_check.zig"); pub const draft = @import("cuda_draft.zig"); test { diff --git a/zig/src/families/qwen3_5/cuda_mtp.zig b/zig/src/families/qwen3_5/cuda_mtp.zig index 4beac9a97f..fddb55f77a 100644 --- a/zig/src/families/qwen3_5/cuda_mtp.zig +++ b/zig/src/families/qwen3_5/cuda_mtp.zig @@ -30,7 +30,7 @@ const Ring = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, ri const TapRows = extern struct { n: c_int, rows: [most_window]c_int }; /// A stream's context: its ring of keys and values over [begin, end), the newest row's state, and its held drafts. -const Ctx = struct { +pub const Ctx = struct { ring: DeviceBuffer, last: DeviceBuffer, // the newest absorbed row's normed output (bf16 hidden): the first draft row's state begin: u64 = 0, @@ -166,7 +166,7 @@ pub const Mtp = struct { d.gpa.destroy(d); } - fn kvWidth(d: *const Mtp) usize { + pub fn kvWidth(d: *const Mtp) usize { return d.e.g.kv_heads * head_dim; } @@ -203,7 +203,7 @@ pub const Mtp = struct { return total; } - fn drop(d: *Mtp, ctx: *Ctx) void { + pub fn drop(d: *Mtp, ctx: *Ctx) void { ctx.last.free(); ctx.ring.free(); d.gpa.destroy(ctx); @@ -215,7 +215,7 @@ pub const Mtp = struct { } /// A new context, touched once (an absorb, then a chain), so no request pays the driver's stall. - fn create(d: *Mtp) !*Ctx { + pub fn create(d: *Mtp) !*Ctx { const ctx = try d.gpa.create(Ctx); errdefer d.gpa.destroy(ctx); ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.ringBytes()), .last = undefined }; @@ -236,12 +236,12 @@ pub const Mtp = struct { while (d.pool.items.len < n) d.pool.appendAssumeCapacity(try d.create()); } - fn ringK(d: *Mtp, ctx: *Ctx) u64 { + pub fn ringK(d: *Mtp, ctx: *Ctx) u64 { _ = d; return ctx.ring.ptr; } - fn ringV(d: *Mtp, ctx: *Ctx) u64 { + pub fn ringV(d: *Mtp, ctx: *Ctx) u64 { return ctx.ring.ptr + d.capacity * d.kvWidth() * 2; } @@ -260,7 +260,7 @@ pub const Mtp = struct { } /// The layer over `rows` rows at slots pos0..: d.ids their tokens, d.hin their states; x gets their normed outputs. - fn layer(d: *Mtp, ctx: *Ctx, rows: usize, pos0: u64, pass: Pass) !void { + pub fn layer(d: *Mtp, ctx: *Ctx, rows: usize, pos0: u64, pass: Pass) !void { const e = d.e; const g = e.g; const w = &d.w; @@ -317,7 +317,7 @@ pub const Mtp = struct { return std.mem.bytesToValue(u32, st[tok_at..][0..4]); } - fn putIds(d: *Mtp, ids: []const u32) !void { + pub fn putIds(d: *Mtp, ids: []const u32) !void { const e = d.e; try e.stream.synchronize(); // the previous upload from this staging area has run: an absorb queues one a chunk @memcpy(std.mem.bytesAsSlice(u32, d.staged.bytes[ids_at..][0 .. ids.len * 4]), ids); @@ -325,7 +325,7 @@ pub const Mtp = struct { } /// Target rows [start, start + n) from the tap ring with the token after each; `rows`, when given, picks them. - fn absorb(d: *Mtp, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32, follow: []const u32) !void { + pub fn absorb(d: *Mtp, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32, follow: []const u32) !void { const e = d.e; const g = e.g; var start = start_; @@ -390,7 +390,7 @@ pub const Mtp = struct { } /// The tree's lanes, lane i the rank-r candidate of its parent's row; a depth's rows run as one pass. - fn holdTree(d: *Mtp, ctx: *Ctx, pos0: u64, l: *const lanes.shape.Shape) !void { + pub fn holdTree(d: *Mtp, ctx: *Ctx, pos0: u64, l: *const lanes.shape.Shape) !void { const e = d.e; const g = e.g; const n = l.parents.len; @@ -432,6 +432,8 @@ pub const Mtp = struct { anc[m * anc_stride + 1 + count] = a; count += 1; } + // root first, as a chain's ring holds them: a lane's sums then equal its path's chained rows + std.mem.reverse(i32, anc[m * anc_stride + 1 ..][0..count]); anc[m * anc_stride] = @intCast(count); } try d.putIds(toks[0..rows]); diff --git a/zig/src/families/qwen3_5/cuda_mtp_check.zig b/zig/src/families/qwen3_5/cuda_mtp_check.zig new file mode 100644 index 0000000000..1eda62b862 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_mtp_check.zig @@ -0,0 +1,133 @@ +//! GPU regressions for the MTP drafter: absorbs in chunks, a tree's lanes against their paths chained, a poisoned ring. +const std = @import("std"); +const lanes = @import("lanes"); +const Mtp = @import("cuda_mtp.zig").Mtp; +const max_rows = @import("cuda_engine.zig").max_rows; + +/// Siblings 0 and 1 at depth 0, six siblings 3 to 8 at depth 1, lane 9 at depth 2 under 3 and 0; each has a child. +const parents = [_]i32{ -1, -1, -1, 0, 0, 0, 0, 1, 1, 3, 4, 5, 6, 7, 8, 9 }; +const ranks = [_]u8{ 0, 1, 2, 0, 1, 2, 3, 0, 1, 0, 0, 0, 0, 0, 0, 0 }; +const depths = [_]u8{ 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3 }; +const internal = [_]usize{ 0, 1, 3, 4, 5, 6, 7, 8, 9 }; + +/// A lane's rows from a pass: its normed output, its key and its value. +const Rows = struct { state: []u8, k: []u8, v: []u8 }; + +/// Absorbs `ids` from the target's tap ring (`tap`, filled by forwards over them), then checks the tree twice. +pub fn run(gpa: std.mem.Allocator, d: *Mtp, tap: u64, ids: []const u32) !bool { + const e = d.e; + const g = e.g; + d.trees = true; + const ctx = try d.create(); + defer d.drop(ctx); + try d.absorb(ctx, tap, 0, ids.len - 1, null, ids[1..]); + const pos0 = ctx.end; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + + // one absorb over several prompt chunks against the same rows absorbed a chunk a call + const apart = try d.create(); + defer d.drop(apart); + var from: usize = 0; + while (from < pos0) : (from += max_rows) { + const n = @min(max_rows, pos0 - from); + try d.absorb(apart, tap, from, n, null, ids[from + 1 ..][0..n]); + } + const ring = 2 * d.capacity * d.kvWidth() * 2; + const whole = try a.alloc(u8, ring + g.hidden * 2); + const parts = try a.alloc(u8, ring + g.hidden * 2); + try download(d, whole[0..ring], ctx.ring.ptr); + try download(d, whole[ring..], ctx.last.ptr); + try download(d, parts[0..ring], apart.ring.ptr); + try download(d, parts[ring..], apart.last.ptr); + var ok = std.mem.eql(u8, whole, parts) and std.mem.eql(u32, &ctx.cands, &apart.cands); + std.debug.print("{s} {d} rows absorbed at once vs a chunk of {d} a call: ring, state and candidates\n", .{ if (ok) "PASS" else "FAIL", pos0, max_rows }); + var p = parents; + var r = ranks; + var dp = depths; + const shape: lanes.shape.Shape = .{ .parents = &p, .ranks = &r, .depths = &dp, .expected = 0 }; + const hidden = g.hidden * 2; + const kw = d.kvWidth() * 2; + + try d.holdTree(ctx, pos0, &shape); + const tree = try lanesRows(a, d, hidden, kw); + const held = ctx.held; + + // the ring at and past the committed end poisoned: a tree pass reads none of it, so nothing may change + const junk = try a.alloc(u8, 16 * kw); + @memset(junk, 0x7f); + for (0..16) |i| { + const slot = (pos0 + i) % d.capacity; + try upload(d, d.ringK(ctx) + slot * kw, junk[0..kw]); + try upload(d, d.ringV(ctx) + slot * kw, junk[0..kw]); + } + try d.holdTree(ctx, pos0, &shape); + const poisoned = try lanesRows(a, d, hidden, kw); + ok = report("poisoned ring suffix, tree rows", poisoned, tree, &internal) and ok; + if (!std.mem.eql(u32, &ctx.held, &held)) { + std.debug.print("FAIL poisoned ring suffix: the lanes' tokens changed\n", .{}); + ok = false; + } + + // each internal lane's path chained from the newest absorbed row, one row a slot, against the tree's rows + var chain: Rows = .{ .state = try a.alloc(u8, tree.state.len), .k = try a.alloc(u8, tree.k.len), .v = try a.alloc(u8, tree.v.len) }; + for (internal) |lane| { + var path: [4]usize = undefined; + var n: usize = 0; + var at: i32 = @intCast(lane); + while (at >= 0) : (at = parents[@intCast(at)]) { + path[n] = @intCast(at); + n += 1; + } + std.mem.reverse(usize, path[0..n]); + for (path[0..n], 0..) |step, k| { + const src = if (k == 0) ctx.last.ptr else d.x; + try d.scratch.copyFrom(d.hin - d.scratch.ptr, src, hidden, e.stream.handle); + try d.putIds(&.{held[step]}); + try d.layer(ctx, 1, pos0 + k, .ring); + } + try download(d, chain.state[lane * hidden ..][0..hidden], d.x); + try download(d, chain.k[lane * kw ..][0..kw], d.kb); + try download(d, chain.v[lane * kw ..][0..kw], d.vb); + } + ok = report("six siblings a depth, tree rows vs each lane's path chained", tree, chain, &internal) and ok; + return ok; +} + +fn lanesRows(a: std.mem.Allocator, d: *Mtp, hidden: usize, kw: usize) !Rows { + const n = parents.len; + const out: Rows = .{ .state = try a.alloc(u8, n * hidden), .k = try a.alloc(u8, n * kw), .v = try a.alloc(u8, n * kw) }; + try download(d, out.state, d.lane_state); + try download(d, out.k, d.lane_k); + try download(d, out.v, d.lane_v); + return out; +} + +/// PASS when every internal lane's state, key and value bytes are equal; else the first lane that differs. +fn report(what: []const u8, got: Rows, want: Rows, lanes_: []const usize) bool { + const hidden = got.state.len / parents.len; + const kw = got.k.len / parents.len; + for (lanes_) |l| { + const same = std.mem.eql(u8, got.state[l * hidden ..][0..hidden], want.state[l * hidden ..][0..hidden]) and + std.mem.eql(u8, got.k[l * kw ..][0..kw], want.k[l * kw ..][0..kw]) and std.mem.eql(u8, got.v[l * kw ..][0..kw], want.v[l * kw ..][0..kw]); + if (!same) { + std.debug.print("FAIL {s}: lane {d} (depth {d}) differs\n", .{ what, l, depths[l] }); + return false; + } + } + std.debug.print("PASS {s}: {d} lanes byte-equal\n", .{ what, lanes_.len }); + return true; +} + +fn upload(d: *Mtp, dst: u64, bytes: []const u8) !void { + const e = d.e; + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, e.stream.handle), "cuMemcpyHtoDAsync"); + try e.stream.synchronize(); +} + +fn download(d: *Mtp, out: []u8, src: u64) !void { + const e = d.e; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(out.ptr, src, out.len, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); +} diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig index 50db48a217..847240302e 100644 --- a/zig/tests/qwen35_cuda_main.zig +++ b/zig/tests/qwen35_cuda_main.zig @@ -8,7 +8,7 @@ pub fn main(init: std.process.Init) !u8 { const a = init.arena.allocator(); const args = try init.minimal.args.toSlice(a); if (args.len != 5) { - std.debug.print("usage: tf-qwen35-cuda \n", .{}); + std.debug.print("usage: tf-qwen35-cuda \n", .{}); return 2; } var ids: std.ArrayList(u32) = .empty; @@ -20,6 +20,7 @@ pub fn main(init: std.process.Init) !u8 { defer ctx.deinit(); const e = try E.Engine.init(init.gpa, init.io, &ctx, args[1]); defer e.deinit(); + if (std.mem.eql(u8, args[2], "mtp")) return mtpCheck(init.gpa, init.io, e, args[1], ids.items); // TF_Q35_PREFILL=N starts the stream after N cycled ids; TF_Q35_BESIDE=N steps an N-token stream beside it const prefill = try envCount("TF_Q35_PREFILL"); const beside = try envCount("TF_Q35_BESIDE"); @@ -159,6 +160,17 @@ const Fwd = struct { } }; +/// The MTP drafter's tree regressions (cuda_mtp_check.zig) after a prompt over the ids fills the target's tap ring. +fn mtpCheck(gpa: std.mem.Allocator, io: std.Io, e: *E.Engine, dir: []const u8, ids: []const u32) !u8 { + const d = try q.mtp.Mtp.init(gpa, io, e, dir, 4096); + defer d.deinit(); + try e.setTaps(d.taps()); + var cache = try E.Cache.init(gpa, e, 4096); + defer cache.deinit(gpa); + try fill(e, &cache, ids, ids.len); + return if (try q.mtp_check.run(gpa, d, cache.tap.?.ptr, ids)) 0 else 1; +} + fn envCount(name: [:0]const u8) !usize { const v = std.c.getenv(name) orelse return 0; return std.fmt.parseInt(usize, std.mem.span(v), 10); From 3d6f96ee0ce824814d7385343514d4ce066021be Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:07:37 +0000 Subject: [PATCH 35/35] qwen3_5: the qwen3_5_moe entry exports Options #853's host builds each family's options from F.Options; with it, image input from #885 reaches the 35B-A3B unchanged. --- zig/src/families/qwen3_5/cuda_native_moe.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/src/families/qwen3_5/cuda_native_moe.zig b/zig/src/families/qwen3_5/cuda_native_moe.zig index 54901476c9..16d71b8d2b 100644 --- a/zig/src/families/qwen3_5/cuda_native_moe.zig +++ b/zig/src/families/qwen3_5/cuda_native_moe.zig @@ -8,6 +8,7 @@ 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. pub const takes_drafter = true; +pub const Options = native.Options; pub const open = native.open; pub const openDrafter = native.openMtp; pub const explain = native.explain;