From f0b3b9deea586cfdcd0fdd9899f59ef720afe83c Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 9 Nov 2026 13:58:59 +0000 Subject: [PATCH 01/12] feat(cuda): a shared engine for dense GQA decoders dots.mocr, Nemotron Parse and GLM-OCR need a dense decoder on CUDA, which had none, and every CUDA family carried its own copy of the lane plumbing. The cuda_dense module now holds that once: paged keys and values in core pages, prompt and decode attention picked from the head layout, the lane backend with copy drafts and cancel, device draws that refuse non-finite rows before any launch, host uploads kept alive to completion, failures drained before teardown, a recipe's own rotary table, and Native(R), which gives the CUDA host a family from a recipe of one layer's arithmetic. No family uses it yet. --- zig/build/cuda_dense.zig | 13 + zig/build/cuda_small_hosts.zig | 4 +- zig/src/cuda_dense/backend.zig | 350 ++++++++++++++++++++++++ zig/src/cuda_dense/device_fail_test.zig | 202 ++++++++++++++ zig/src/cuda_dense/draws.zig | 106 +++++++ zig/src/cuda_dense/kernels.zig | 145 ++++++++++ zig/src/cuda_dense/model.zig | 251 +++++++++++++++++ zig/src/cuda_dense/native.zig | 152 ++++++++++ zig/src/cuda_dense/native_test.zig | 105 +++++++ zig/src/cuda_dense/pages.zig | 66 +++++ zig/src/cuda_dense/pass.zig | 150 ++++++++++ zig/src/cuda_dense/put_test.zig | 72 +++++ zig/src/cuda_dense/recipe.zig | 136 +++++++++ zig/src/cuda_dense/root.zig | 30 ++ zig/src/cuda_dense/rope.zig | 31 +++ 15 files changed, 1812 insertions(+), 1 deletion(-) create mode 100644 zig/build/cuda_dense.zig create mode 100644 zig/src/cuda_dense/backend.zig create mode 100644 zig/src/cuda_dense/device_fail_test.zig create mode 100644 zig/src/cuda_dense/draws.zig create mode 100644 zig/src/cuda_dense/kernels.zig create mode 100644 zig/src/cuda_dense/model.zig create mode 100644 zig/src/cuda_dense/native.zig create mode 100644 zig/src/cuda_dense/native_test.zig create mode 100644 zig/src/cuda_dense/pages.zig create mode 100644 zig/src/cuda_dense/pass.zig create mode 100644 zig/src/cuda_dense/put_test.zig create mode 100644 zig/src/cuda_dense/recipe.zig create mode 100644 zig/src/cuda_dense/root.zig create mode 100644 zig/src/cuda_dense/rope.zig diff --git a/zig/build/cuda_dense.zig b/zig/build/cuda_dense.zig new file mode 100644 index 00000000000..4c8671f92ba --- /dev/null +++ b/zig/build/cuda_dense.zig @@ -0,0 +1,13 @@ +//! The shared CUDA dense-decoder engine (zig/src/cuda_dense) over the runtime, core and lane core; its host tests. +const std = @import("std"); + +/// The engine's module, made once a compilation and passed to each dense recipe's family module. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/cuda_dense/root.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes } } }); +} + +pub fn host(b: *std.Build, m: *std.Build.Module, all: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = m })); + b.step("test-cuda-dense-host", "The CUDA dense engine: attention shapes, pages, RoPE frequencies, Native decls, refusals without a GPU").dependOn(&run.step); + all.dependOn(&run.step); +} diff --git a/zig/build/cuda_small_hosts.zig b/zig/build/cuda_small_hosts.zig index d3b3f3d7196..0e6b160af17 100644 --- a/zig/build/cuda_small_hosts.zig +++ b/zig/build/cuda_small_hosts.zig @@ -1,7 +1,9 @@ -//! Host test steps of the small CUDA families (Spark 2.5 and LFM2 MoE), registered from cuda.zig's host step. +//! Host test steps of the small CUDA families (Spark 2.5 and LFM2 MoE) and the dense engine. const std = @import("std"); pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, all: *std.Build.Step) void { @import("spark25.zig").host(b, @import("spark25.zig").module(b, host, .debug, cuda, core, lanes), all); @import("lfm2_moe_cuda.zig").host(b, @import("lfm2_moe_cuda.zig").module(b, host, .debug, cuda, core, lanes), all); + const dense = @import("cuda_dense.zig").module(b, host, .debug, cuda, core, lanes); + @import("cuda_dense.zig").host(b, dense, all); } diff --git a/zig/src/cuda_dense/backend.zig b/zig/src/cuda_dense/backend.zig new file mode 100644 index 00000000000..7a1e5ba3be0 --- /dev/null +++ b/zig/src/cuda_dense/backend.zig @@ -0,0 +1,350 @@ +//! A dense decoder behind lanes.Backend: a slot and pool pages a stream, a chunk a prompt step, one forward a round. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const kv = @import("core").kv_pool; +const pass = @import("pass.zig"); +const mdl = @import("model.zig"); +const kern = @import("kernels.zig"); +const pages = @import("pages.zig"); +const draws_ = @import("draws.zig"); + +const Draws = draws_.Draws; +const ring = draws_.ring; +const be = lanes.backend; +const Stream = lanes.Stream; + +/// Prompt rows a step takes: the forward's chunk grid, from position 0. +pub const prompt_rows = mdl.prompt_rows; +/// Window rows a round takes over every stream. +pub const round_rows = mdl.round_rows; +/// The widest window a stream verifies: one-row steps and every wider window give a row the same bits. +pub const max_window = pages.window; + +const Lane = struct { + slot: usize, + len: u64 = 0, // the cache length: positions [0, len) hold this stream's keys + base: u64 = 0, // the cache length the last verify started from + rows: std.ArrayList(u32) = .empty, // the last verify's row tokens + last: u64, // the prompt's last row's bf16 logits [V] (the first draw reads them) + rule: u64, // sample.cu's Rule, when the stream samples + sampled: bool = false, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts + lease: kv.Lease = .{}, // the pages it holds and is promised +}; + +pub fn Dense(comptime R: type) type { + return struct { + const D = @This(); + gpa: std.mem.Allocator, + d: *const cuda.Driver, + e: pass.Engine(R), + m: *mdl.Model, + lanes: std.AutoHashMapUnmanaged(*const Stream, Lane) = .empty, + free: []bool, // slot i is free + draws: Draws, + slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot + pool: ?kv.Pool = null, // the pages every slot leases (attach) + + /// The device bytes `init` allocates: the draws and a round's slots' last rows and rules. + pub fn openBytes(vocab: usize) usize { + return Draws.openBytes() + round_rows * (vocab * 2 + @sizeOf(cuda.sample.Rule)); + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *mdl.Model, k: *const kern.Kernels, own: *const R.Kernels, w: *const R.Weights, s: cuda.Stream) !D { + const free = try gpa.alloc(bool, round_rows); + errdefer gpa.free(free); + @memset(free, true); + var draws = try Draws.init(gpa, d, k, s, m.g.vocab); + errdefer draws.deinit(gpa); + const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * round_rows); + errdefer gpa.free(bufs); + var made: usize = 0; + errdefer for (bufs[0..made]) |*b| b.free(); + for (bufs, 0..) |*b, i| { + b.* = try cuda.DeviceBuffer.alloc(d, if (i % 2 == 0) m.g.vocab * 2 else @sizeOf(cuda.sample.Rule)); + made += 1; + } + return .{ .gpa = gpa, .d = d, .e = .{ .m = m, .k = k, .own = own, .w = w, .s = s }, .m = m, .free = free, .draws = draws, .slot_bufs = bufs }; + } + + pub fn deinit(x: *D) void { + var it = x.lanes.valueIterator(); + while (it.next()) |l| { + if (x.pool) |*p| p.finish(&l.lease); + l.rows.deinit(x.gpa); + } + if (x.pool) |*p| p.deinit(); + x.lanes.deinit(x.gpa); + for (x.slot_bufs) |*b| b.free(); + x.gpa.free(x.slot_bufs); + x.gpa.free(x.free); + x.draws.deinit(x.gpa); + } + + /// The model's arena of `n` pages for `extra + 1` slots, and the pool that leases them. + pub fn attach(x: *D, extra: usize, n: usize) !void { + if (x.pool != null) return error.BadAttach; + try x.m.attach(x.e.s, extra, n); + x.pool = try kv.Pool.init(x.gpa, n, x.m.pageBytes()); + } + + /// The pool's free pages in a seeded random order, leased in that order (a check's hook). + pub fn shufflePages(x: *D, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + if (x.pool) |*p| prng.random().shuffle(u32, p.free[0..p.nfree]); + } + + pub fn backend(x: *D) be.Backend { + return .{ .ptr = x, .vtable = &vtable }; + } + + const vtable: be.Backend.VTable = .{ .prefill = prefill, .prefill_step = prefillStep, .cancel_prefill = cancelPrefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .reports_fill = true }; + + /// The round loop's facts: no draft head; host copies in windows up to 16 rows, streams sharing a forward. + pub fn facts(_: *const D) lanes.Model { + return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = round_rows }; + } + + fn self(ptr: *anyopaque) *D { + return @ptrCast(@alignCast(ptr)); + } + + fn forward(x: *D, chains: []const pass.Chain, prompt: bool, head0: usize, heads: usize) !void { + try pass.forward(R, x.e, chains, prompt, head0, heads); + } + + fn freeSlot(x: *const D) ?usize { + return std.mem.indexOfScalar(bool, x.free[0..x.m.slots], true); + } + + /// Whether a request could start now: a free slot and its pages; one that never fits passes (refused later). + pub fn admits(x: *const D, prompt: usize, max_new: usize) bool { + if (prompt + max_new > x.m.context) return true; + const p = x.pool orelse return false; + return x.freeSlot() != null and p.fits(pages.pagesFor(prompt, max_new)); + } + + const Map = struct { + x: *D, + slot: usize, + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const m: *const Map = @ptrCast(@alignCast(ctx)); + try m.x.m.mapPages(m.x.e.s, m.slot, ids[from..], from); + } + }; + + /// The lane's pages cover positions [0, end): pages taken from its promise as its context grows. + fn reach(x: *D, l: *Lane, end: usize) !void { + const p = if (x.pool) |*p| p else return error.NoPages; + const n = pages.pagesOf(end); + if (n <= l.lease.ids.items.len) return; + var map: Map = .{ .x = x, .slot = l.slot }; + try p.cover(&l.lease, n, .{ .ctx = &map, .write = Map.write }); + } + + /// A free slot for the request, its pages promised now (taken as its rows need them). + fn takeSlot(x: *D, lease: *kv.Lease, prompt: usize, max_new: usize) !usize { + const p = if (x.pool) |*p| p else return error.NoPages; + const slot = x.freeSlot() orelse return error.NoFreeSlot; + try p.promise(lease, pages.pagesFor(prompt, max_new)); + return slot; + } + + /// A prompt pass begins: a slot, its pages, the stream's rule; no prompt cache, so its first chunk starts at 0. + fn opening(x: *D, s: *Stream) !*Lane { + const p = s.prompt(); + if (p.len == 0) return error.EmptyPrompt; + if (s.images.len > 0) return error.DenseTextOnly; + if (p.len + @as(usize, s.max_new) > x.m.context) return error.PastContext; + const got = try x.lanes.getOrPut(x.gpa, s); + if (!got.found_existing) { + var lease: kv.Lease = .{}; + const slot = x.takeSlot(&lease, p.len, s.max_new) catch |e| { + x.lanes.removeByPtr(got.key_ptr); + return e; + }; + x.free[slot] = false; + got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr, .lease = lease }; + } + const l = got.value_ptr; + try x.reach(l, p.len); + l.sampled = false; + if (s.sampling) |sp| if (sp.temperature > 0) { + const r = cuda.sample.Rule.from(lanes.gpu_rule.rule(sp)); + try mdl.put(x.d, x.e.s, l.rule, std.mem.asBytes(&r)); + l.sampled = true; + }; + s.cached = 0; + if (s.reuse.saved != null) s.reuse_failed = true; // no prompt cache: every pass starts at 0 + l.fill = 0; + s.fill_at.store(0, .monotonic); + return l; + } + + /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row's logits are written. + fn chunk(x: *D, s: *Stream, l: *Lane) !bool { + const p = s.prompt(); + const at = l.fill.?; + if (s.isCancelled()) { + s.keepInterruptedPrefill(@intCast(at)); + return error.Cancelled; + } + var end = @min(p.len, (at / prompt_rows + 1) * prompt_rows); + for (s.reuse.marks) |mk| if (mk > at and mk < end) { + end = mk; + }; + const done = end == p.len; + try x.forward(&.{.{ .slot = l.slot, .p0 = at, .tokens = p[at..end] }}, true, end - at - 1, @intFromBool(done)); + if (done) try x.d.check(x.d.api.cuMemcpyDtoDAsync_v2(l.last, x.m.s.logits, x.m.g.vocab * 2, x.e.s.handle), "cuMemcpyDtoDAsync"); + s.fill_at.store(@intCast(end), .monotonic); + l.fill = if (done) null else end; + if (done) l.len = p.len; + if (s.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(end)) != null) hk.at(hk.ptr, s, @intCast(end)); + return done; + } + + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const l = try x.opening(s); + errdefer x.dropLane(s); + while (!try x.chunk(s, l)) {} + } + + fn prefillStep(ptr: *anyopaque, s: *Stream, _: u32) anyerror!bool { + const x = self(ptr); + const going = if (x.lanes.getPtr(s)) |l| l.fill != null else false; + const l = if (going) x.lanes.getPtr(s).? else try x.opening(s); + errdefer x.dropLane(s); + return x.chunk(s, l); + } + + fn cancelPrefill(ptr: *anyopaque, s: *Stream) !void { + const l = self(ptr).lanes.getPtr(s) orelse return; + if (l.fill) |at| s.keepInterruptedPrefill(@intCast(at)); + } + + fn first(ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if (position != l.len or l.fill != null) return error.PositionMismatch; + const h = x.draws.handle(); + try x.draws.rows(l.sampled, l.rule, l.last, 0, 1, position, h.at, 0); + return h.h; + } + + fn value(x: *D, feed: be.Feed) !u32 { + return switch (feed) { + .value => |v| v, + .handle => |h| x.draws.read(h), + }; + } + + fn queue(ptr: *anyopaque, s: *Stream, feed: be.Feed, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + const t = try x.value(feed); + if (position != l.len + 1) return error.PositionMismatch; + try x.reach(l, l.len + 1); + try x.forward(&.{.{ .slot = l.slot, .p0 = l.len, .tokens = &.{t} }}, false, 0, 1); + l.len += 1; + const h = x.draws.handle(); + try x.draws.rows(l.sampled, l.rule, x.m.s.logits, 0, 1, position, h.at, 0); + return h.h; + } + + fn read(ptr: *anyopaque, h: u64) anyerror!u32 { + return self(ptr).draws.read(h); + } + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const x = self(ptr); + if (windows.len > round_rows) return error.RoundTooWide; + const chains = try x.gpa.alloc(pass.Chain, windows.len); + defer x.gpa.free(chains); + var total: usize = 0; + for (windows, chains) |w, *c| { + if (w.held != 0) return error.NoHeldDrafts; // no draft head: every draft comes from the host + if (w.parents != null) return error.ChainsOnly; + if (w.masks.len != 0) return error.StructuresUnsupported; + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (w.rows() > max_window) return error.WindowTooWide; + l.base = l.len; + l.rows.clearRetainingCapacity(); + try l.rows.append(x.gpa, w.pending); + try l.rows.appendSlice(x.gpa, w.tokens); + for (w.positions, 0..) |pos, r| if (pos != l.base + r + 1) return error.PositionMismatch; + try x.reach(l, l.base + l.rows.items.len); + c.* = .{ .slot = l.slot, .p0 = l.base, .tokens = l.rows.items }; + total += l.rows.items.len; + } + if (total > round_rows) return error.RoundTooWide; + try x.forward(chains, false, 0, total); + var at: usize = 0; + var sampled = false; + for (windows) |w| { + const l = x.lanes.getPtr(w.stream).?; + try x.draws.pick(l.sampled, x.m.s.logits, at, l.rows.items.len, ring + at); + sampled = sampled or l.sampled; + at += l.rows.items.len; + } + if (sampled) try x.draws.finite(ring, total); // one wait a round, before any draw reads a row + at = 0; + for (windows, 0..) |w, i| { + const l = x.lanes.getPtr(w.stream).?; + if (l.sampled) try x.draws.draw(l.rule, x.m.s.logits, at, l.rows.items.len, w.positions[0], ring + at, i); + at += l.rows.items.len; + } + var got: [round_rows]u32 = undefined; + try x.draws.tokens(ring, got[0..total]); + at = 0; + for (windows, out) |w, o| { + const l = x.lanes.getPtr(w.stream).?; + const n = l.rows.items.len; + @memcpy(o.sampled[0..n], got[at..][0..n]); + @memcpy(o.drafts[0 .. n - 1], l.rows.items[1..]); + l.len = l.base + n; // a chain keeps every row unless rolled back + at += n; + } + } + + fn keep(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const x = self(ptr); + for (windows, paths) |w, path| { + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (path.len == 0 or path.len > l.rows.items.len) return error.EmptyPath; + for (path, 0..) |r, j| if (r != j) return error.ChainsOnly; // a chain keeps a prefix of its rows + l.len = l.base + path.len; // caches are positional: later rows overwrite what the rejected ones wrote + } + } + + /// No draft head: absorbing is a no-op, a request for drafts is refused by name (copies come from the host). + fn draft(_: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + for (requests) |r| if (r.depth != 0 or r.lanes != null) return error.NoDraftHead; + } + + fn dropLane(x: *D, s: *const Stream) void { + if (x.lanes.fetchRemove(s)) |entry| { + var l = entry.value; + if (x.pool) |*p| p.finish(&l.lease); // one CUDA stream: a later lease's writes queue behind these reads + x.m.unmapPages(l.slot); + x.free[l.slot] = true; + l.rows.deinit(x.gpa); + if (x.lanes.count() == 0) x.report(); + } + } + + /// The idle server's line: nothing promised, every page free. + fn report(x: *const D) void { + if (@import("builtin").is_test) return; + const p = x.pool orelse return; + const c = p.census(); + const ok = p.idle() and c.free == c.total; + std.log.info("kv pool idle{s}: promised {d}, free {d} == total {d}", .{ if (ok) "" else " (UNEVEN)", c.promised, c.free, c.total }); + } + + fn release(ptr: *anyopaque, s: *Stream) void { + self(ptr).dropLane(s); + } + }; +} diff --git a/zig/src/cuda_dense/device_fail_test.zig b/zig/src/cuda_dense/device_fail_test.zig new file mode 100644 index 00000000000..7fbc24e48b8 --- /dev/null +++ b/zig/src/cuda_dense/device_fail_test.zig @@ -0,0 +1,202 @@ +//! Failure paths on a recording fake device: no buffer freed before completion, no draw over non-finite logits. +const std = @import("std"); +const cuda = @import("cuda"); +const mdl = @import("model.zig"); +const kern = @import("kernels.zig"); +const draws = @import("draws.zig"); +const Geometry = @import("recipe.zig").Geometry; +const t = std.testing; + +const Fake = struct { + const queue: cuda.abi.Stream = @ptrFromInt(0x1800); + const base: u64 = 0x100000; + var mem: [4096]u8 = undefined; // device bytes from `base` on, for downloads + var next: u64 = 0x200000; + var live: usize = 0; + var pending = false; + var freed_pending: usize = 0; + var proofs: usize = 0; + var fail_wait = false; + var argmax: usize = 0; + var drawn: usize = 0; + fn reset() void { + live = 0; + pending = false; + freed_pending = 0; + proofs = 0; + fail_wait = false; + argmax = 0; + drawn = 0; + } + fn alloc(out: *u64, len: usize) callconv(.c) c_int { + out.* = next; + next += len + 256; + live += 1; + return 0; + } + fn free(_: u64) callconv(.c) c_int { + freed_pending += @intFromBool(pending); + live -= 1; + return 0; + } + fn fill(_: u64, _: u8, _: usize, _: cuda.abi.Stream) callconv(.c) c_int { + pending = true; + return 0; + } + fn host(out: *?*anyopaque, len: usize, _: c_uint) callconv(.c) c_int { + out.* = std.c.malloc(len) orelse return cuda.abi.error_not_ready; + return 0; + } + fn hostFree(p: ?*anyopaque) callconv(.c) c_int { + std.c.free(p); + return 0; + } + fn copy(_: u64, _: ?*const anyopaque, _: usize, _: cuda.abi.Stream) callconv(.c) c_int { + pending = true; + return 0; + } + fn wait(_: cuda.abi.Stream) callconv(.c) c_int { + if (fail_wait) return cuda.abi.error_not_ready; + pending = false; + return 0; + } + fn context() callconv(.c) c_int { + proofs += 1; + pending = false; + return 0; + } + fn event(out: *cuda.abi.Event, _: c_uint) callconv(.c) c_int { + out.* = @ptrFromInt(1); + return 0; + } + fn record(_: cuda.abi.Event, _: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + fn query(_: cuda.abi.Event) callconv(.c) c_int { + return if (pending) cuda.abi.error_not_ready else 0; + } + fn destroy(_: cuda.abi.Event) callconv(.c) c_int { + return 0; + } + fn launch(f: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { + argmax += @intFromBool(f == @as(cuda.abi.Function, @ptrFromInt(1))); + drawn += @intFromBool(f == @as(cuda.abi.Function, @ptrFromInt(2))); + return 0; + } + fn down(dst: ?*anyopaque, src: u64, len: usize) callconv(.c) c_int { + @memcpy(@as([*]u8, @ptrCast(dst.?))[0..len], mem[src - base ..][0..len]); + return 0; + } + fn device(out: *cuda.abi.Device) callconv(.c) c_int { + out.* = 0; + return 0; + } + fn attribute(out: *c_int, _: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) c_int { + out.* = 0; + return 0; + } + fn driver() cuda.Driver { + reset(); + var d: cuda.Driver = .{ .lib = undefined, .api = undefined }; + d.api.cuMemAlloc_v2 = alloc; + d.api.cuMemFree_v2 = free; + d.api.cuMemsetD8Async = fill; + d.api.cuMemHostAlloc = host; + d.api.cuMemFreeHost = hostFree; + d.api.cuMemcpyHtoDAsync_v2 = copy; + d.api.cuStreamSynchronize = wait; + d.api.cuCtxSynchronize = context; + d.api.cuEventCreate = event; + d.api.cuEventRecord = record; + d.api.cuEventQuery = query; + d.api.cuEventDestroy_v2 = destroy; + d.api.cuLaunchKernel = launch; + d.api.cuMemcpyDtoH_v2 = down; + d.api.cuCtxGetDevice = device; + d.api.cuDeviceGetAttribute = attribute; + return d; + } +}; + +const tiny: Geometry = .{ .hidden = 256, .layers = 2, .heads = 4, .kv_heads = 2, .head_dim = 64, .vocab = 512, .eps = 1e-6, .context = 4096, .acc = 512, .wide = 1024, .mid = 512, .theta = 1e4 }; + +fn lane() cuda.qmmf.Lane { + var l: cuda.qmmf.Lane = undefined; + l.clusters = false; // partBytes reads the tile plan alone + return l; +} + +test "a failed scratch upload wait proves completion by the context before init frees its buffers" { + var d = Fake.driver(); + const s: cuda.Stream = .{ .d = &d, .handle = Fake.queue }; + const l = lane(); + Fake.fail_wait = true; + try t.expectError(error.NotReady, mdl.Model.init(t.allocator, &d, s, tiny, 4096, &l)); + try t.expectEqual(@as(usize, 0), Fake.freed_pending); + try t.expectEqual(@as(usize, 0), Fake.live); + try t.expect(Fake.proofs >= 1); + Fake.fail_wait = false; + try s.synchronize(); + cuda.copy_source.dropSpares(&d, cuda.copy_source.current()); +} + +test "a failed attach wait proves completion before its arena and tables go, and leaves the model attachable" { + var d = Fake.driver(); + const s: cuda.Stream = .{ .d = &d, .handle = Fake.queue }; + const l = lane(); + var m = try mdl.Model.init(t.allocator, &d, s, tiny, 4096, &l); + const held = Fake.live; + Fake.fail_wait = true; + try t.expectError(error.NotReady, m.attach(s, 1, 2)); + try t.expectEqual(@as(usize, 0), Fake.freed_pending); + try t.expectEqual(held, Fake.live); + try t.expect(m.pages == 0 and m.slots == 0 and m.ids.len == 0 and m.arena == 0); + Fake.fail_wait = false; + try m.attach(s, 1, 2); + try t.expect(m.pages == 2 and m.slots == 2); + m.deinit(); + try t.expectEqual(@as(usize, 0), Fake.live); + cuda.copy_source.dropSpares(&d, cuda.copy_source.current()); +} + +fn buffer(d: *const cuda.Driver, offset: u64, len: usize) cuda.DeviceBuffer { + return .{ .d = d, .ptr = Fake.base + offset, .len = len, .borrowed = true }; +} + +fn function(d: *const cuda.Driver, id: usize) cuda.Function { + return .{ .d = d, .handle = @ptrFromInt(id) }; +} + +test "a sampled row over NaN or an infinity is refused before any draw reads it, as a greedy one is" { + var d = Fake.driver(); + const s: cuda.Stream = .{ .d = &d, .handle = Fake.queue }; + var k: kern.Kernels = undefined; + k.round = .{ .argmax_fn = function(&d, 1), .match_fn = function(&d, 9), .keep_fn = function(&d, 9), .gather_fn = function(&d, 9), .head_fn = function(&d, 9), .ids_fn = function(&d, 9) }; + k.draws = .{ .draw_fn = function(&d, 2), .draw_ids_fn = function(&d, 9), .logprob_rows_fn = function(&d, 9), .topk_slices_fn = function(&d, 9), .topk_merge_fn = function(&d, 9) }; + var greedy: [draws.ring + mdl.round_rows]bool = @splat(true); + var x: draws.Draws = .{ .d = &d, .k = &k, .s = s, .vocab = 512, .tok = buffer(&d, 0, 64), .picks = buffer(&d, 64, 128), .greedy = &greedy, .meta = buffer(&d, 512, 64) }; + const Pick = cuda.tree_round.Pick; + var picks = [2]Pick{ .{ .token = 3, .nonfinite = 0 }, .{ .token = 4, .nonfinite = 0 } }; + @memcpy(Fake.mem[64..][0..16], std.mem.sliceAsBytes(&picks)); // what the argmax writes for finite rows + @memcpy(Fake.mem[0..8], std.mem.sliceAsBytes(&[2]u32{ 7, 8 })); + try x.rows(true, 0x9000, 0xa000, 0, 2, 5, 0, 0); + try t.expectEqual(@as(usize, 1), Fake.argmax); // the sampled rows' non-finite flags + try t.expectEqual(@as(usize, 1), Fake.drawn); + var out: [2]u32 = undefined; + try x.tokens(0, &out); + try t.expectEqualSlices(u32, &.{ 7, 8 }, &out); // drawn, not picked + for ([_]u32{ 1, 2, 4 }) |mask| { // NaN, +inf, -inf (tree_round.cu's mask) + picks[1] = .{ .token = 0xFFFFFFFF, .nonfinite = mask }; + @memcpy(Fake.mem[64..][0..16], std.mem.sliceAsBytes(&picks)); + try t.expectError(error.NonFiniteLogits, x.rows(true, 0x9000, 0xa000, 0, 2, 5, 0, 0)); + try t.expectEqual(@as(usize, 1), Fake.drawn); // tf_draw never saw the row + } + try x.rows(false, 0, 0xa000, 0, 2, 5, 0, 0); + try t.expectError(error.NonFiniteLogits, x.tokens(0, &out)); + try t.expectEqual(@as(usize, 1), Fake.drawn); + picks[1] = .{ .token = 4, .nonfinite = 0 }; + @memcpy(Fake.mem[64..][0..16], std.mem.sliceAsBytes(&picks)); + try x.tokens(0, &out); + try t.expectEqualSlices(u32, &.{ 3, 4 }, &out); + cuda.copy_source.dropSpares(&d, cuda.copy_source.current()); +} diff --git a/zig/src/cuda_dense/draws.zig b/zig/src/cuda_dense/draws.zig new file mode 100644 index 00000000000..770d2cdd5be --- /dev/null +++ b/zig/src/cuda_dense/draws.zig @@ -0,0 +1,106 @@ +//! The dense engine's draws: a ring of drawn tokens behind handles, greedy picks or keyed samples of bf16 logits. +const std = @import("std"); +const cuda = @import("cuda"); +const kern = @import("kernels.zig"); +const mdl = @import("model.zig"); + +const Pick = cuda.tree_round.Pick; +const round_rows = mdl.round_rows; + +/// Drawn tokens a handle names (a ring; read before the round loop could wrap it). +pub const ring = 1024; + +pub const Draws = struct { + d: *const cuda.Driver, + k: *const kern.Kernels, + s: cuda.Stream, + vocab: usize, + tok: cuda.DeviceBuffer, // drawn tokens (u32): the handle ring, then a round's rows + picks: cuda.DeviceBuffer, // argmax picks (Pick), the same entries + greedy: []bool, // entry e of the ring (and the round's rows) is a pick, not a draw + meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each + next: u64 = 0, + + /// The device bytes `init` allocates: the token and pick entries and the metas. + pub fn openBytes() usize { + return (ring + round_rows) * (4 + @sizeOf(Pick)) + round_rows * 16; + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, k: *const kern.Kernels, s: cuda.Stream, vocab: usize) !Draws { + const greedy = try gpa.alloc(bool, ring + round_rows); + errdefer gpa.free(greedy); + @memset(greedy, true); + var tok = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * 4); + errdefer tok.free(); + var picks = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * @sizeOf(Pick)); + errdefer picks.free(); + const meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16); + return .{ .d = d, .k = k, .s = s, .vocab = vocab, .tok = tok, .picks = picks, .greedy = greedy, .meta = meta }; + } + + pub fn deinit(x: *Draws, gpa: std.mem.Allocator) void { + gpa.free(x.greedy); + x.tok.free(); + x.picks.free(); + x.meta.free(); + } + + /// The next handle and its ring entry. + pub fn handle(x: *Draws) struct { h: u64, at: usize } { + const h = x.next; + x.next += 1; + return .{ .h = h, .at = @intCast(h % ring) }; + } + + /// Rows [row0, row0 + rows) of bf16 `logits` into entries [at, at + rows), row r at `position0 + r`. + pub fn rows(x: *Draws, sampled: bool, rule: u64, logits: u64, row0: usize, n: usize, position0: u64, at: usize, meta_at: usize) !void { + try x.pick(sampled, logits, row0, n, at); + if (!sampled) return; + try x.finite(at, n); + try x.draw(rule, logits, row0, n, position0, at, meta_at); + } + + /// Every row's argmax pick into entries [at, at + n): a greedy row's token, and any row's NaN or infinity flag. + pub fn pick(x: *Draws, sampled: bool, logits: u64, row0: usize, n: usize, at: usize) !void { + const v = x.vocab; + @memset(x.greedy[at..][0..n], !sampled); + try x.k.round.argmax(x.s, logits + row0 * v * 2, x.picks.ptr + at * @sizeOf(Pick), n, v, v); + } + + /// Entries [at, at + n)'s picks read back: a row of NaN or infinite logits is refused before any draw reads it. + pub fn finite(x: *Draws, at: usize, n: usize) !void { + try x.s.synchronize(); + var picks: [round_rows]Pick = undefined; + try x.picks.download(at * @sizeOf(Pick), std.mem.sliceAsBytes(picks[0..n])); + for (picks[0..n]) |p| if (p.nonfinite != 0) return error.NonFiniteLogits; + } + + /// Keyed samples of rows whose picks `finite` passed (tf_draw faults on a row with no finite logit). + pub fn draw(x: *Draws, rule: u64, logits: u64, row0: usize, n: usize, position0: u64, at: usize, meta_at: usize) !void { + const v = x.vocab; + const meta = [4]i32{ @intCast(position0 - 1), 0, 0, 0 }; + const mp = x.meta.ptr + meta_at * 16; + try mdl.put(x.d, x.s, mp, std.mem.asBytes(&meta)); + try x.k.draws.draw(x.s, logits + row0 * v * 2, v, rule, mp, 0, x.tok.ptr + at * 4, n, null, null); + } + + /// Entries [at, at + out.len) as tokens after the stream's work; non-finite logits are an error, drawn or picked. + pub fn tokens(x: *Draws, at: usize, out: []u32) !void { + try x.s.synchronize(); + var picks: [round_rows]Pick = undefined; + try x.tok.download(at * 4, std.mem.sliceAsBytes(out)); + try x.picks.download(at * @sizeOf(Pick), std.mem.sliceAsBytes(picks[0..out.len])); + for (out, picks[0..out.len], x.greedy[at..][0..out.len]) |*t, p, g| { + if (p.nonfinite != 0) return error.NonFiniteLogits; + if (g) t.* = p.token; + } + } + + /// The token handle `h` names, once drawn; a handle the ring has passed is refused. + pub fn read(x: *Draws, h: u64) !u32 { + if (h >= x.next or x.next - h > ring) return error.NoSuchToken; + var out: [1]u32 = undefined; + try x.tokens(@intCast(h % ring), &out); + return out[0]; + } +}; diff --git a/zig/src/cuda_dense/kernels.zig b/zig/src/cuda_dense/kernels.zig new file mode 100644 index 00000000000..32365e5cc62 --- /dev/null +++ b/zig/src/cuda_dense/kernels.zig @@ -0,0 +1,145 @@ +//! The dense engine's shared launches: the lane matmul, norms, embedding rows, bias, draws, paged attention, SwiGLU. +const std = @import("std"); +const cuda = @import("cuda"); +const recipe = @import("recipe.zig"); +const pages = @import("pages.zig"); + +const Args = cuda.Args; +const Function = cuda.Function; +const Stream = cuda.Stream; +const Geometry = recipe.Geometry; + +/// Streams one shared-window attention launch takes (paged_attn.cuh's MAXSEG). +pub const max_segs = 8; + +/// paged_attn.cuh's AttnSegs: each stream's caches, position, page table and first row. +pub const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, pt: [max_segs]u64, ps: i64, row0: [max_segs]c_int, n: c_int }; + +const images = .{ "sample", "tree_round", "prefill_attention", "nemotron_attention", "lfm2" }; + +fn int(x: usize) c_int { + return @intCast(x); +} + +pub const Kernels = struct { + mods: [images.len]cuda.Module, + g: Geometry, + shape: recipe.Shape, + lane: cuda.qmmf.Lane, + norms: cuda.norms.Norms, + embed: cuda.embed_rows.EmbedRows, + vit: cuda.vit.Kernels, + draws: cuda.sample.Draws, + round: cuda.tree_round.Kernels, + pattn: Function, + kv_write: Function, + chunk: Function, + merge: Function, + silu_mul: Function, + + /// The shared modules, and the attention launches `g`'s shape names; a geometry the kernels lack refuses by name. + pub fn load(ctx: *const cuda.Context, g: Geometry) !Kernels { + const d = ctx.d; + const shape = try recipe.attentionShape(g); + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + k.g = g; + k.shape = shape; + var loaded: usize = 0; + errdefer for (k.mods[0..loaded]) |*m| m.unload(); + inline for (images, 0..) |name, i| { + k.mods[i] = try cuda.Module.load(d, @field(cuda.kernels, name)); + loaded += 1; + } + const m = &k.mods; + k.draws = try cuda.sample.Draws.resolve(m[0]); + k.round = try cuda.tree_round.Kernels.resolve(m[1]); + k.pattn = try m[2].function(shape.pattn); + try k.pattn.allowDynamicShared(shape.shared); + k.kv_write = try m[3].function("tf_nemo_kv_write_segs"); + k.chunk = try m[3].function(shape.chunk); + k.merge = try m[3].function(shape.merge); + k.silu_mul = try m[4].function("tf_lfm2_silu_mul"); + k.lane = try cuda.qmmf.Lane.load(d, (try ctx.capability()) / 10); + errdefer k.lane.unload(); + k.norms = try cuda.norms.Norms.load(d); + errdefer k.norms.unload(); + k.embed = try cuda.embed_rows.EmbedRows.load(d); + errdefer k.embed.unload(); + k.vit = try cuda.vit.Kernels.load(d); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.vit.unload(); + k.embed.unload(); + k.norms.unload(); + k.lane.unload(); + for (&k.mods) |*m| m.unload(); + } + + fn go(f: Function, s: Stream, grid: [3]usize, block: usize, smem: u32, a: *Args) !void { + const g: cuda.Dim3 = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }; + try cuda.launch.launch(f, .{ .grid = g, .block = .{ .x = @intCast(block) }, .shared = smem }, s, a); + } + + /// out [rows, f] = bf16(bf16(silu(gate)) * up) from rows gate | up [rows, 2 f] (lfm2.cu). + pub fn siluMul(k: *const Kernels, s: Stream, gu: u64, out: u64, rows: usize, f: usize) !void { + var a: Args = .{}; + a.add(gu); + a.add(out); + a.add(int(f)); + try go(k.silu_mul, s, .{ rows, (f + 255) / 256, 1 }, 256, 0, &a); + } + + /// Rows k, v (kv [rows, 2 kv width]) into each stream's pages at its position on. + pub fn kvWrite(k: *const Kernels, s: Stream, sg: AttnSegs, kv: u64, rows: usize) !void { + const kvw = k.g.kvWidth(); + var a: Args = .{}; + a.add(sg); + a.add(kv); + a.add(int(2 * kvw)); + a.add(@as(c_int, 0)); + a.add(int(kvw)); + try go(k.kv_write, s, .{ rows, 1, 1 }, 256, 0, &a); + } + + fn scale(k: *const Kernels) f32 { + return 1.0 / @sqrt(@as(f32, @floatFromInt(k.g.head_dim))); + } + + /// Decode rows: q [rows, heads, dim] over each stream's pages in 512-key chunks, merged in order into out and xs. + pub fn attend(k: *const Kernels, s: Stream, sg: AttnSegs, q: u64, po: u64, pm: u64, pl: u64, out: u64, xs: u64, rows: usize, nch: usize, z: usize) !void { + const g = k.g; + var a: Args = .{}; + a.add(sg); + a.add(q); + for ([_]u64{ po, pm, pl }) |v| a.add(v); + for ([_]usize{ g.qWidth(), g.heads, g.kv_heads, g.group(), pages.page_rows, nch }) |v| a.add(int(v)); + a.add(k.scale()); + try go(k.chunk, s, .{ rows, g.kv_heads, z }, 256, 0, &a); + var b: Args = .{}; + b.add(sg); + for ([_]u64{ po, pm, pl, out, xs }) |v| b.add(v); + for ([_]usize{ g.heads, g.kv_heads, g.group(), pages.page_rows, nch }) |v| b.add(int(v)); + try go(k.merge, s, .{ rows, g.kv_heads, 1 }, 256, 0, &b); + } + + /// A prompt chunk's rows at positions p0.. against its stream's pages (written through them): out [rows, H, dim]. + pub fn prompt(k: *const Kernels, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, rows: usize, pt: u64, ps: usize) !void { + const g = k.g; + var a: Args = .{}; + for ([_]u64{ q, kc, vc, out }) |v| a.add(v); + for ([_]usize{ p0, rows, g.heads, g.kv_heads, g.group() }) |v| a.add(int(v)); + a.add(k.scale()); + a.add(pt); + a.add(@as(i64, @intCast(ps))); + try go(k.pattn, s, .{ k.shape.blocks(rows), k.shape.grid_y, 1 }, 256, k.shape.shared, &a); + } +}; + +test "AttnSegs is paged_attn.cuh's struct, padded to 8 bytes: 304" { + try std.testing.expectEqual(@as(usize, 304), @sizeOf(AttnSegs)); + try std.testing.expectEqual(@as(usize, 256), @offsetOf(AttnSegs, "ps")); + try std.testing.expectEqual(@as(usize, 296), @offsetOf(AttnSegs, "n")); +} diff --git a/zig/src/cuda_dense/model.zig b/zig/src/cuda_dense/model.zig new file mode 100644 index 00000000000..06a293223cf --- /dev/null +++ b/zig/src/cuda_dense/model.zig @@ -0,0 +1,251 @@ +//! The dense engine's device state: the forward's scratch, the page arena of every layer's KV, each slot's table. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const recipe = @import("recipe.zig"); +const pages = @import("pages.zig"); +const rope = @import("rope.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const Geometry = recipe.Geometry; + +/// Rows a prompt chunk takes. +pub const prompt_rows = 1024; +/// Rows a round of windows takes over every stream. +pub const round_rows = 128; +/// Key chunks one decode attention launch walks in parallel at most. +const chunk_blocks = 64; + +/// The forward's buffers, `R` rows of each (a prompt chunk's); a recipe reads and writes them through its Pass. +pub const Scratch = struct { + ids: u64, // int32 [R]: tokens + pos: u64, // int32 [R]: each row's position + meta: u64, // int32 [round_rows]: each chain's first position (AttnSegs' META) + h: u64, // bf16 [R, hidden]: the residual + x: u64, // bf16 [R, hidden]: a norm's output + r: u64, // bf16 [R, hidden]: a block's output + acc: u64, // fp32 [R, acc width]: a projection's sums before its bias + wide: u64, // bf16 [R, wide]: gate|up + y: u64, // bf16 [R, max(q width, mid)]: attention out, the MLP's activation + q: u64, // bf16 [R, q width]: rotated queries + kv: u64, // bf16 [R, 2 kv width]: keys then values a row + xs: u64, // fp32 [round_rows, q width / 64]: decode attention's 64-value sums + po: u64, // decode attention partials: o, m, l + pm: u64, + pl: u64, + logits: u64, // bf16 [round_rows, vocab] + part: u64, // the lane matmul's slice partials +}; + +/// bytes to dst on `s`: a private pinned copy kept until the stream completes, so the source may change at once. +pub fn put(d: *const cuda.Driver, s: cuda.Stream, dst: u64, bytes: []const u8) !void { + try cuda.copy_source.enqueue(d, dst, bytes, s); +} + +pub const Model = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + g: Geometry, + context: usize, // the served window: a request's prompt and reply fit in it + slots: usize = 0, // streams with a table row: none until attach + pages: usize = 0, // pages in the arena + arena: u64 = 0, // [page][layer][K, V][512][kv width] bf16 + table: u64 = 0, // int32 [slots][perSlot()]: slot s's page ids in position order + ids: []u32 = &.{}, // the tables on the host + caps: []usize = &.{}, // positions slot s's pages cover + inv: u64 = 0, // fp32 [head_dim / 2]: RoPE's inverse frequencies (a geometry with theta) + acc_width: usize, + part_bytes: usize, + s: Scratch = undefined, + buffers: std.ArrayList(DeviceBuffer) = .empty, + planned: ?usize = null, // a plan: the bytes counted, nothing allocated + + pub fn perSlot(m: *const Model) usize { + return pages.perSlot(m.context); + } + + /// Positions a slot's table reaches: the window and its verify rows, in whole pages. + pub fn capacity(m: *const Model) usize { + return m.perSlot() * pages.page_rows; + } + + pub fn pageElems(m: *const Model) usize { + return pages.pageElems(m.g); + } + + pub fn pageBytes(m: *const Model) u64 { + return pages.pageBytes(m.g); + } + + pub fn layout(m: *const Model) kv.Layout { + return pages.layout(m.g, m.context); + } + + pub fn streamBytes(m: *const Model) u64 { + return pages.streamBytes(m.g, m.context); + } + + /// Decode attention's key chunks a row walks at most, and the blocks that walk them. + pub fn chunks(m: *const Model) usize { + return m.perSlot(); + } + + pub fn chunkBlocks(m: *const Model) usize { + return @min(m.chunks(), chunk_blocks); + } + + fn alloc(m: *Model, bytes: usize) !u64 { + if (m.planned) |*p| { + p.* += (@max(bytes, 256) + 255) / 256 * 256; + return 0; + } + const b = try DeviceBuffer.alloc(m.d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try m.buffers.append(m.gpa, b); + return b.ptr; + } + + /// The f32 sums width every biased projection fits, and the lane's largest slice partials over the shapes. + fn sizes(g: Geometry, lane: *const cuda.qmmf.Lane) struct { acc: usize, part: usize } { + const shapes = [_][2]usize{ .{ g.acc, g.hidden }, .{ g.hidden, g.qWidth() }, .{ g.wide, g.hidden }, .{ g.hidden, g.mid }, .{ g.vocab, g.hidden } }; + var part: usize = 0; + for (shapes) |nk| for (1..prompt_rows + 1) |rows| { + part = @max(part, lane.partBytes(rows, cuda.modelopt.bf16Weight(0, nk[0], nk[1]))); + }; + return .{ .acc = @max(g.acc, g.wide, g.hidden), .part = part }; + } + + fn blank(gpa: std.mem.Allocator, d: *const cuda.Driver, g: Geometry, context: usize, lane: *const cuda.qmmf.Lane) Model { + const z = sizes(g, lane); + return .{ .gpa = gpa, .d = d, .g = g, .context = context, .acc_width = z.acc, .part_bytes = z.part }; + } + + /// The scratch and the RoPE table; the arena and tables come with attach, once the host has planned the pages. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, g: Geometry, context: usize, lane: *const cuda.qmmf.Lane) !Model { + var m = blank(gpa, d, g, context, lane); + errdefer m.deinit(); + errdefer cuda.completion.drain(d, s, "the dense engine's scratch uploads did not drain before their buffers were freed"); + try m.build(s); + return m; + } + + /// The device bytes `init` allocates, counted by the same path without allocating. + pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, g: Geometry, context: usize, lane: *const cuda.qmmf.Lane) !usize { + var m = blank(gpa, d, g, context, lane); + m.planned = 0; + defer m.deinit(); + try m.build(s); + return m.planned.?; + } + + fn build(m: *Model, s: cuda.Stream) !void { + const g = m.g; + const r = prompt_rows; + const parts = round_rows * m.chunks() * g.heads; + m.s = .{ + .ids = try m.alloc(r * 4), + .pos = try m.alloc(r * 4), + .meta = try m.alloc(round_rows * 4), + .h = try m.alloc(r * g.hidden * 2), + .x = try m.alloc(r * g.hidden * 2), + .r = try m.alloc(r * g.hidden * 2), + .acc = try m.alloc(r * m.acc_width * 4), + .wide = try m.alloc(r * g.wide * 2), + .y = try m.alloc(r * @max(g.qWidth(), g.mid) * 2), + .q = try m.alloc(r * g.qWidth() * 2), + .kv = try m.alloc(r * 2 * g.kvWidth() * 2), + .xs = try m.alloc(round_rows * (g.qWidth() / 64) * 4), + .po = try m.alloc(parts * g.head_dim * 4), + .pm = try m.alloc(parts * 4), + .pl = try m.alloc(parts * 4), + .logits = try m.alloc(round_rows * g.vocab * 2), + .part = try m.alloc(m.part_bytes), + }; + const theta = g.theta orelse return; + m.inv = try m.alloc(rope.count(g.head_dim) * 4); + if (m.planned != null) return; + const inv = try m.gpa.alloc(f32, rope.count(g.head_dim)); + defer m.gpa.free(inv); + if (g.inv_freq) |table| { + if (table.len != inv.len) return error.BadRopeTable; + @memcpy(inv, table); + } else rope.inverse(theta, g.head_dim, inv); + try m.upload(s, m.inv, std.mem.sliceAsBytes(inv)); + try s.synchronize(); + } + + fn upload(m: *const Model, s: cuda.Stream, dst: u64, bytes: []const u8) !void { + try put(m.d, s, dst, bytes); + } + + pub fn deinit(m: *Model) void { + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + m.gpa.free(m.ids); + m.gpa.free(m.caps); + } + + /// The arena of `pool` pages (zeroed) and the tables of `extra + 1` slots, no page mapped. + pub fn attach(m: *Model, s: cuda.Stream, extra: usize, pool: usize) !void { + if (m.pages != 0 or pool == 0 or extra >= round_rows or m.planned != null) return error.BadAttach; + const mark = m.buffers.items.len; // the errdefer reads it at failure, when the list has grown + errdefer m.undo(s, mark); + const bytes = pool * m.pageBytes(); + m.arena = try m.alloc(bytes); + try m.d.check(m.d.api.cuMemsetD8Async(m.arena, 0, bytes, s.handle), "cuMemsetD8Async"); + const slots = extra + 1; + m.table = try m.alloc(slots * m.perSlot() * 4); + m.ids = try m.gpa.alloc(u32, slots * m.perSlot()); + @memset(m.ids, 0); + m.caps = try m.gpa.alloc(usize, slots); + @memset(m.caps, 0); + try m.upload(s, m.table, std.mem.sliceAsBytes(m.ids)); + try s.synchronize(); + m.slots = slots; + m.pages = pool; + } + + /// A failed attach: its fill and table upload proved complete, then its arena, tables and host rows gone. + fn undo(m: *Model, s: cuda.Stream, mark: usize) void { + cuda.completion.drain(m.d, s, "the dense engine's page arena fill did not drain before it was freed"); + while (m.buffers.items.len > mark) { + var b = m.buffers.pop().?; + b.free(); + } + m.gpa.free(m.ids); + m.gpa.free(m.caps); + m.ids = &.{}; + m.caps = &.{}; + m.arena = 0; + m.table = 0; + } + + /// Slot `slot`'s pages from table entry `from` on are `ids` (arena page numbers in position order). + pub fn mapPages(m: *Model, s: cuda.Stream, slot: usize, ids: []const u32, from: usize) !void { + const per = m.perSlot(); + if (slot >= m.slots or from + ids.len > per) return error.BadChain; + for (ids) |id| if (id >= m.pages) return error.BadChain; + const at = slot * per + from; + @memcpy(m.ids[at..][0..ids.len], ids); + m.caps[slot] = (from + ids.len) * pages.page_rows; + try m.upload(s, m.table + at * 4, std.mem.sliceAsBytes(m.ids[at..][0..ids.len])); + } + + /// Slot `slot` holds no pages: its table entries stay, no forward may reach them. + pub fn unmapPages(m: *Model, slot: usize) void { + if (slot < m.slots) m.caps[slot] = 0; + } + + /// Layer `layer`'s key (or value) cache as the kernels take it: its part of page 0. + pub fn cache(m: *const Model, layer: usize, value: bool) u64 { + return m.arena + ((2 * layer + @intFromBool(value)) * pages.page_rows * m.g.kvWidth()) * 2; + } + + /// Slot s's page table row. + pub fn tableOf(m: *const Model, slot: usize) u64 { + return m.table + slot * m.perSlot() * 4; + } +}; diff --git a/zig/src/cuda_dense/native.zig b/zig/src/cuda_dense/native.zig new file mode 100644 index 00000000000..818ec0f2bf9 --- /dev/null +++ b/zig/src/cuda_dense/native.zig @@ -0,0 +1,152 @@ +//! Native(R): a dense recipe as the native CUDA host's family, its decls generated once for every recipe. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const kv = @import("core").kv_pool; +const recipe = @import("recipe.zig"); +const kern = @import("kernels.zig"); +const mdl = @import("model.zig"); +const pages = @import("pages.zig"); +const backend = @import("backend.zig"); + +/// The refusals the engine words itself, for every recipe; a recipe's own come from its `explain`. +pub const engine_errors = [_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.DenseTextOnly, error.StructuresUnsupported, error.NonFiniteLogits, error.CapturedKernels, error.ParallelPastRound, error.NoStreamFits, error.UnsupportedDenseHeadDim, error.UnsupportedDenseGroup }; + +pub fn Native(comptime R: type) type { + comptime recipe.check(R); + return struct { + const Dense = backend.Dense(R); + + pub const model_type: []const u8 = R.model_type; + pub const formats: []const []const u8 = R.formats; + pub const default_context: i64 = R.default_context; + pub const max_segments: u32 = 1; + pub const prompt_rows: u32 = backend.prompt_rows; + /// A request's verify reserve: its pages hold that many rows past its prompt and reply. + pub const window_rows: u32 = backend.max_window; + /// `explain` words this family's open refusals too. + pub const explains_open = true; + /// No drafter: --drafter is refused before the load. + pub const takes_drafter = false; + pub const drafter_refusal = std.fmt.comptimePrint("--drafter: {s} on CUDA loads no drafter; serve it without --drafter", .{R.name}); + /// No prompt states: the host plans no prompt cache for this family and refuses --prompt-cache-gib. + pub const keeps_no_prompts = true; + /// Whether replies open with a reasoning block (the recipe's word, else the host's default). + pub const reasoning: bool = if (@hasDecl(R, "reasoning")) R.reasoning else true; + + /// The server's reading of the room left after the weights; null if unreadable. + pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 }; + /// No vision, precision, kv8, prompt_int8, loops or rank field: the host refuses those flags by name. + pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null }; + + pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32 = 1, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + target: *Dense, + /// Every stream's keys and values are pages of the core pool, the first stream's too. + layout: ?kv.Layout, + engine_paged: bool = true, + }; + + const Owned = struct { gpa: std.mem.Allocator, s: cuda.Stream, w: R.Weights, k: kern.Kernels, own: R.Kernels, m: mdl.Model, b: Dense }; + + /// The checkpoint, the shared and the recipe's kernels, the scratch; streams' pages come with attach. + 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.CapturedKernels; + if (o.lanes_fixed and o.lanes > backend.round_rows) return error.ParallelPastRound; + const g = try R.geometry(gpa, io, dir); + _ = try recipe.attentionShape(g); // refused by name before the weights load + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.s = try cuda.Stream.init(ctx.d, false); + errdefer own.s.deinit(); + own.w = try R.load(gpa, io, ctx.d, own.s, dir); + errdefer own.w.deinit(); + own.k = try kern.Kernels.load(ctx, g); + errdefer own.k.deinit(); + own.own = try R.Kernels.load(ctx); + errdefer own.own.deinit(); + const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); + const bytes = try mdl.Model.openBytes(gpa, ctx.d, own.s, g, o.context, &own.k.lane) + Dense.openBytes(g.vocab); + if (bytes > room) return error.NoStreamFits; + own.m = try mdl.Model.init(gpa, ctx.d, own.s, g, o.context, &own.k.lane); + errdefer own.m.deinit(); + own.b = try Dense.init(gpa, ctx.d, &own.m, &own.k, &own.own, &own.w, own.s); + return .{ + .backend = own.b.backend(), + .facts = own.b.facts(), + .rows = backend.max_window, + .stream_bytes = own.m.streamBytes(), + .ctx = own, + .deinit = release, + .target = &own.b, + .layout = own.m.layout(), + }; + } + + /// The checkpoint's weight bytes, counted from its tensors before the load. + pub fn textWeightBytes(a: std.mem.Allocator, io: std.Io, dir: []const u8) !?u64 { + return try R.weightBytes(a, io, dir); + } + + fn owned(p: *anyopaque) *Owned { + return @ptrCast(@alignCast(p)); + } + + /// The pool: `streams` more slots beside the first and an arena of `n` pages. + pub fn attach(p: *anyopaque, streams: u32, n: usize) !void { + try owned(p).b.attach(streams, n); + } + + /// Whether a request's pages could be promised now (the lane host's admission, on its engine thread). + pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { + return owned(p).b.admits(prompt, max_tokens); + } + + /// The pool the lane host's ledger admits this family's promises by. + pub fn ledgerPool(p: *anyopaque) ?*kv.Pool { + const b = &owned(p).b; + return if (b.pool) |*pool| pool else null; + } + + /// A request or load this engine refuses, in words; then the recipe's own; null: none. + pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return words(err) orelse if (@hasDecl(R, "explain")) R.explain(err) else null; + } + + fn words(err: anyerror) ?[]const u8 { + const name = R.name; + return switch (err) { + error.PastContext => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.NoFreeSlot => "every stream's caches are in use: retry once another request ends", + error.NoDraftHead => name ++ " on CUDA has no draft head: it drafts from the prompt's own text", + error.DenseTextOnly => name ++ " on CUDA reads text prompts only so far: this request has images", + error.StructuresUnsupported => "structured output is not built for " ++ name ++ " on CUDA yet", + error.NonFiniteLogits => "the model's logits went non-finite at this row: the request is stopped rather than answered from them", + error.CapturedKernels => name ++ " runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS", + error.ParallelPastRound => std.fmt.comptimePrint("--parallel: {s} on CUDA holds at most {d} streams (a round's rows): serve --parallel {d} or fewer, or auto", .{ name, backend.round_rows, backend.round_rows }), + error.NoStreamFits => "the CUDA memory budget fits no " ++ name ++ " stream after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + error.UnsupportedDenseHeadDim => "this " ++ name ++ " checkpoint's attention head dim is not 64 or 128, the two the CUDA dense engine's attention kernels take", + error.UnsupportedDenseGroup => "this " ++ name ++ " checkpoint's query heads a KV head are not an even count of at most 16, the groups the CUDA dense engine's attention kernels take", + else => null, + }; + } + + fn release(p: *anyopaque) void { + const own = owned(p); + own.b.deinit(); + own.m.deinit(); + own.own.deinit(); + own.k.deinit(); + own.w.deinit(); + own.s.deinit(); + own.gpa.destroy(own); + } + }; +} diff --git a/zig/src/cuda_dense/native_test.zig b/zig/src/cuda_dense/native_test.zig new file mode 100644 index 00000000000..6788f5675a0 --- /dev/null +++ b/zig/src/cuda_dense/native_test.zig @@ -0,0 +1,105 @@ +//! Native(R)'s host contract with a tiny recipe: the decls it emits and omits, its refusals' words, a compiled path. +const std = @import("std"); +const cuda = @import("cuda"); +const native = @import("native.zig"); +const Pass = @import("pass.zig").Pass; +const Geometry = @import("recipe.zig").Geometry; +const Allocator = std.mem.Allocator; + +const tiny: Geometry = .{ .hidden = 256, .layers = 2, .heads = 4, .kv_heads = 2, .head_dim = 64, .vocab = 512, .eps = 1e-6, .context = 4096, .acc = 512, .wide = 1024, .mid = 512, .theta = 1e4 }; + +/// A recipe of the dots shape on tiny sizes: biased q|k|v sums, attention, o, SwiGLU, an untied head. +const Tiny = struct { + pub const name = "Tiny"; + pub const model_type = "tiny_dense"; + pub const formats: []const []const u8 = &.{"bf16"}; + pub const default_context: i64 = 4096; + pub const reasoning = false; + pub const post_embed: ?fn (*const Pass(Tiny)) anyerror!void = null; + pub const Weights = struct { + embed: u64 = 0, + norm: u64 = 0, + bias: u64 = 0, + proj: cuda.qmmf.Weight = cuda.modelopt.bf16Weight(0, 512, 256), + pub fn deinit(_: *Weights) void {} + }; + pub const Kernels = struct { + pub fn load(_: *const cuda.Context) !Kernels { + return .{}; + } + pub fn deinit(_: *Kernels) void {} + }; + pub fn geometry(_: Allocator, _: std.Io, _: []const u8) !Geometry { + return tiny; + } + pub fn load(_: Allocator, _: std.Io, _: *const cuda.Driver, _: cuda.Stream, _: []const u8) !Weights { + return .{}; + } + pub fn weightBytes(_: Allocator, _: std.Io, _: []const u8) !u64 { + return 1 << 20; + } + pub fn embed(p: *const Pass(Tiny)) !void { + try p.k.embed.bf16(p.s, p.m.s.ids, p.w.embed, p.m.s.h, p.rows, p.g.hidden); + } + pub fn layer(p: *const Pass(Tiny), i: usize) !void { + const x = p.m.s; + try p.norm(if (i == 0) null else x.r, p.w.norm); + try p.projectBias(p.w.proj, p.w.bias, x.x, x.wide); + try p.projectF32(p.w.proj, x.x, x.acc); + try p.attend(i, x.q, x.kv, x.y); + try p.project(p.w.proj, x.y, x.r); + try p.siluMul(x.wide, x.y, p.g.mid); + } + pub fn final(p: *const Pass(Tiny)) !void { + try p.norm(p.m.s.r, p.w.norm); + } + pub fn head(w: *const Weights) cuda.qmmf.Weight { + return w.proj; + } + pub fn explain(err: anyerror) ?[]const u8 { + return if (err == error.TinyConfig) "this Tiny checkpoint is not the one the recipe serves" else null; + } +}; + +const N = native.Native(Tiny); + +test "Native emits the host decls a dense family takes and no others" { + inline for (.{ "model_type", "formats", "default_context", "max_segments", "prompt_rows", "window_rows", "explains_open", "takes_drafter", "drafter_refusal", "Room", "Options", "Loaded", "open", "textWeightBytes", "attach", "fits", "ledgerPool", "explain", "keeps_no_prompts", "reasoning" }) |d| { + if (!@hasDecl(N, d)) return error.MissingDecl; + } + // each of these switches a host path on (prompt cache, KV sharing, learning, own planning, drafters, ranks) + inline for (.{ "snap_save", "snap_restore", "snap_drop", "snap_bytes", "snap_slots", "snap_grid", "snap_lookahead", "share_attach", "share_pool", "share_states", "kvPool", "slide", "learnModel", "linked", "contextCeiling", "pagedLayout", "rankRoom", "request_starts", "request_starts_switch", "openDrafter", "warm_quiet" }) |d| { + if (@hasDecl(N, d)) return error.UnexpectedDecl; + } + inline for (.{ "context", "drafts", "segments", "slide", "lanes", "lanes_fixed", "room" }) |f| try std.testing.expect(@hasField(N.Options, f)); + inline for (.{ "vision", "precision", "kv8", "prompt_int8", "loops", "rank", "drafter" }) |f| try std.testing.expect(!@hasField(N.Options, f)); + inline for (.{ "layout", "engine_paged", "free_streams", "target" }) |f| try std.testing.expect(@hasField(N.Loaded, f)); + try std.testing.expectEqualStrings("tiny_dense", N.model_type); + try std.testing.expect(N.keeps_no_prompts and !N.reasoning and !N.takes_drafter); + try std.testing.expectEqual(@as(u32, 16), N.window_rows); + try std.testing.expectEqual(@as(u32, 1024), N.prompt_rows); + try std.testing.expectEqual(@as(u32, 1), N.max_segments); + try std.testing.expectStringStartsWith(N.drafter_refusal, "--drafter: Tiny on CUDA"); +} + +test "every engine refusal is worded, the recipe's own after it, the rest left to the server" { + for (native.engine_errors) |e| try std.testing.expect(N.explain(null, e) != null); + try std.testing.expectStringStartsWith(N.explain(null, error.TinyConfig).?, "this Tiny checkpoint"); + try std.testing.expectStringStartsWith(N.explain(null, error.DenseTextOnly).?, "Tiny on CUDA reads text prompts only"); + try std.testing.expect(N.explain(null, error.OutOfMemory) == null); +} + +test "open refuses captured kernels and a fixed --parallel past a round before it reads the checkpoint" { + const ctx: *const cuda.Context = undefined; + const base: N.Options = .{ .context = 4096, .drafts = true }; + try std.testing.expectError(error.CapturedKernels, N.open(std.testing.allocator, undefined, ctx, "", "captured", base)); + var wide = base; + wide.lanes = 129; + wide.lanes_fixed = true; + try std.testing.expectError(error.ParallelPastRound, N.open(std.testing.allocator, undefined, ctx, "", null, wide)); +} + +test "a recipe's whole path compiles against the engine: open, the lane backend, the forward" { + const hooks = .{ &N.open, &N.attach, &N.fits, &N.ledgerPool, &N.textWeightBytes, &N.explain }; + inline for (hooks) |h| try std.testing.expect(@intFromPtr(h) != 0); +} diff --git a/zig/src/cuda_dense/pages.zig b/zig/src/cuda_dense/pages.zig new file mode 100644 index 00000000000..83698cae7ea --- /dev/null +++ b/zig/src/cuda_dense/pages.zig @@ -0,0 +1,66 @@ +//! KV pages of the dense engine: 512 rows of every layer's keys and values, the host pool's layout, a request's pages. +const std = @import("std"); +const kv = @import("core").kv_pool; +const Geometry = @import("recipe.zig").Geometry; + +/// Rows a page holds: one decode attention chunk (nemotron_attention.cu's PAGE). +pub const page_rows = 512; +/// The widest window a stream verifies: a request's pages hold that many rows past its prompt and reply. +pub const window = 16; + +/// Pages that hold positions [0, positions). +pub fn pagesOf(positions: usize) usize { + return (positions + page_rows - 1) / page_rows; +} + +/// Pages a request holds at most: its prompt, its reply and the verify window. +pub fn pagesFor(prompt: usize, max_new: usize) usize { + return pagesOf(prompt + max_new + window); +} + +/// Table entries a stream takes at a `context`-token window, its verify window included. +pub fn perSlot(context: usize) usize { + return pagesOf(context + window); +} + +/// bf16 elements a page holds: [layer][K, V][512][kv width]. +pub fn pageElems(g: Geometry) usize { + return g.layers * 2 * page_rows * g.kvWidth(); +} + +pub fn pageBytes(g: Geometry) u64 { + return pageElems(g) * 2; +} + +/// Bytes of one stream at a whole `context`-token window: its pages and its table row. +pub fn streamBytes(g: Geometry, context: usize) u64 { + return perSlot(context) * (pageBytes(g) + 4); +} + +/// What the host's pool plans: a token's keys and values over every layer, a stream's table row fixed. +pub fn layout(g: Geometry, context: usize) kv.Layout { + const row: u64 = 2 * g.kvWidth() * 2; + return kv.Layout.ofRows(.{ .kind = .full, .row = row, .layers = @intCast(g.layers) }, page_rows, perSlot(context) * 4, 0); +} + +const dots: Geometry = .{ .hidden = 1536, .layers = 28, .heads = 12, .kv_heads = 2, .head_dim = 128, .vocab = 151936, .eps = 1e-6, .context = 131072, .acc = 2048, .wide = 17920, .mid = 8960, .theta = 1e6 }; + +test "dots' page: 512 rows of 28 layers' keys and values at 2 x 128 dims, 14 MiB" { + try std.testing.expectEqual(@as(u64, 28 * 2 * 512 * 256 * 2), pageBytes(dots)); + try std.testing.expectEqual(@as(u64, 14680064), pageBytes(dots)); + const l = layout(dots, 32768); + try std.testing.expectEqual(pageBytes(dots), l.pageBytes()); + try std.testing.expectEqual(@as(u64, 65 * 4), l.fixed_bytes); + try std.testing.expectEqual(@as(usize, 65), l.pages(32768 + window)); + try std.testing.expectEqual(@as(u64, 65 * (14680064 + 4)), streamBytes(dots, 32768)); +} + +test "a request holds its prompt, its reply and the verify window in whole pages" { + try std.testing.expectEqual(@as(usize, 1), pagesFor(1, 1)); + try std.testing.expectEqual(@as(usize, 1), pagesFor(400, 96)); + try std.testing.expectEqual(@as(usize, 2), pagesFor(400, 97)); + try std.testing.expectEqual(@as(usize, 65), pagesFor(32768, 0)); + try std.testing.expectEqual(@as(usize, 65), perSlot(32768)); + try std.testing.expectEqual(@as(usize, 1), perSlot(496)); + try std.testing.expectEqual(@as(usize, 2), perSlot(497)); +} diff --git a/zig/src/cuda_dense/pass.zig b/zig/src/cuda_dense/pass.zig new file mode 100644 index 00000000000..63d9a277094 --- /dev/null +++ b/zig/src/cuda_dense/pass.zig @@ -0,0 +1,150 @@ +//! One dense forward: the engine's operators a recipe's layers call, and the pass from token ids to bf16 logits. +const std = @import("std"); +const cuda = @import("cuda"); +const recipe = @import("recipe.zig"); +const kern = @import("kernels.zig"); +const mdl = @import("model.zig"); + +const Model = mdl.Model; +const Kernels = kern.Kernels; +const Stream = cuda.Stream; +const Weight = cuda.qmmf.Weight; +const tail: cuda.norms.Tail = .{ .round = true }; // the norms' torch points: bf16(w * bf16(h * inv)) + +/// A stream's rows in one forward: its slot, first position, tokens. +pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32 }; + +/// What a forward runs on: the model's state, the shared and the recipe's kernels, the weights, the CUDA stream. +pub fn Engine(comptime R: type) type { + return struct { m: *const Model, k: *const Kernels, own: *const R.Kernels, w: *const R.Weights, s: Stream }; +} + +/// A forward's rows (a prompt chunk, or a round of windows) and the operators a recipe's layers call on them. +pub fn Pass(comptime R: type) type { + return struct { + const P = @This(); + m: *const Model, + k: *const Kernels, + own: *const R.Kernels, + w: *const R.Weights, + s: Stream, + g: recipe.Geometry, + chains: []const Chain, + rows: usize, + prompt: bool, + + fn part(p: *const P, w: Weight) !u64 { + if (p.k.lane.partBytes(p.rows, w) > p.m.part_bytes) return error.LanePartTooSmall; + return p.m.s.part; + } + + /// out [rows, n] bf16 = x [rows, k] times `w`, fp32 sums rounded once. + pub fn project(p: *const P, w: Weight, x: u64, out: u64) !void { + try p.k.lane.matmul(p.s, x, w.k, p.rows, w, out, try p.part(w)); + } + + /// out [rows, n] fp32 = x [rows, k] times `w`: the sums a recipe's own kernel rounds. + pub fn projectF32(p: *const P, w: Weight, x: u64, out: u64) !void { + try p.k.lane.matmulF32(p.s, x, w.k, p.rows, w, out, try p.part(w)); + } + + /// out [rows, n] = bf16(x times `w` + b): the fp32 sums and the bias rounded once (vit.cu's bias). + pub fn projectBias(p: *const P, w: Weight, b: u64, x: u64, out: u64) !void { + if (w.n > p.m.acc_width) return error.AccTooNarrow; + try p.projectF32(w, x, p.m.s.acc); + try p.k.vit.bias(p.s, p.m.s.acc, b, out, p.rows, w.n); + } + + /// h = bf16(h + r) (r null: as is), x = bf16(w * bf16(h * inv_rms)). + pub fn norm(p: *const P, r: ?u64, w: u64) !void { + const x = p.m.s; + try p.k.norms.addRmsTail(p.s, x.h, r, w, x.h, x.x, null, p.rows, p.g.hidden, p.g.eps, tail); + } + + /// out [rows, f] = bf16(bf16(silu(gate)) * up) from gate|up [rows, 2 f]. + pub fn siluMul(p: *const P, gu: u64, out: u64, f: usize) !void { + try p.k.siluMul(p.s, gu, out, p.rows, f); + } + + /// Layer `layer`'s keys and values (kv [rows, 2 kv width]) into the pages, then q's attention into out. + pub fn attend(p: *const P, layer: usize, q: u64, kv: u64, out: u64) !void { + if (layer >= p.g.layers) return error.BadLayer; + const m = p.m; + const x = m.s; + const kc = m.cache(layer, false); + const vc = m.cache(layer, true); + const qw = p.g.qWidth(); + var row0: usize = 0; + var at: usize = 0; + while (at < p.chains.len) : (at += kern.max_segs) { + const group = p.chains[at..@min(p.chains.len, at + kern.max_segs)]; + var sg = std.mem.zeroes(kern.AttnSegs); + var rows: usize = 0; + for (group, 0..) |ch, j| { + sg.kc[j] = kc; + sg.vc[j] = vc; + sg.meta[j] = x.meta + (at + j) * 4; + sg.pt[j] = m.tableOf(ch.slot); + sg.row0[j] = @intCast(rows); + rows += ch.tokens.len; + } + sg.ps = @intCast(m.pageElems()); + sg.n = @intCast(group.len); + try p.k.kvWrite(p.s, sg, kv + row0 * 2 * p.g.kvWidth() * 2, rows); + const qa = q + row0 * qw * 2; + const oa = out + row0 * qw * 2; + if (p.prompt) { + try p.k.prompt(p.s, qa, kc, vc, oa, group[0].p0, rows, m.tableOf(group[0].slot), m.pageElems()); + } else { + try p.k.attend(p.s, sg, qa, x.po, x.pm, x.pl, oa, x.xs + row0 * (qw / 64) * 4, rows, m.chunks(), m.chunkBlocks()); + } + row0 += rows; + } + } + }; +} + +/// Every chain's rows (a prompt chunk, or windows): the recipe's layers, then logits rows [head0, head0 + heads). +pub fn forward(comptime R: type, e: Engine(R), chains: []const Chain, prompt: bool, head0: usize, heads: usize) !void { + const m = e.m; + var rows: usize = 0; + for (chains) |ch| { + if (ch.tokens.len == 0 or ch.slot >= m.slots or ch.p0 + ch.tokens.len > m.caps[ch.slot]) return error.BadChain; + rows += ch.tokens.len; + } + if (rows > @as(usize, if (prompt) mdl.prompt_rows else mdl.round_rows) or (prompt and chains.len != 1)) return error.BadWindow; + if (chains.len > mdl.round_rows or heads > mdl.round_rows or head0 + heads > rows) return error.BadWindow; + try upload(m, e.s, chains, rows); + const p: Pass(R) = .{ .m = m, .k = e.k, .own = e.own, .w = e.w, .s = e.s, .g = m.g, .chains = chains, .rows = rows, .prompt = prompt }; + try R.embed(&p); + for (0..m.g.layers) |i| try R.layer(&p, i); + try R.final(&p); + if (heads == 0) return; + const head: Weight = R.head(e.w); + const x = m.s.x + head0 * m.g.hidden * 2; + if (e.k.lane.partBytes(heads, head) > m.part_bytes) return error.LanePartTooSmall; + try e.k.lane.matmul(e.s, x, m.g.hidden, heads, head, m.s.logits, m.s.part); +} + +fn put(m: *const Model, s: Stream, dst: u64, bytes: []const u8) !void { + try mdl.put(m.d, s, dst, bytes); +} + +/// Tokens and positions a row, each chain's first position, staged for the launches. +fn upload(m: *const Model, s: Stream, chains: []const Chain, rows: usize) !void { + const ids = try m.gpa.alloc(i32, 2 * rows + chains.len); + defer m.gpa.free(ids); + var at: usize = 0; + for (chains, 0..) |ch, j| { + for (ch.tokens, 0..) |t, i| { + ids[at + i] = @intCast(t); + ids[rows + at + i] = @intCast(ch.p0 + i); + } + ids[2 * rows + j] = @intCast(ch.p0); + at += ch.tokens.len; + } + const x = m.s; + try put(m, s, x.ids, std.mem.sliceAsBytes(ids[0..rows])); + try put(m, s, x.pos, std.mem.sliceAsBytes(ids[rows .. 2 * rows])); + try put(m, s, x.meta, std.mem.sliceAsBytes(ids[2 * rows ..])); +} diff --git a/zig/src/cuda_dense/put_test.zig b/zig/src/cuda_dense/put_test.zig new file mode 100644 index 00000000000..236f2303c1d --- /dev/null +++ b/zig/src/cuda_dense/put_test.zig @@ -0,0 +1,72 @@ +//! The engine's uploads own their source: a fake driver reads it at the stream's sync, after the caller changed it. +const std = @import("std"); +const cuda = @import("cuda"); +const put = @import("model.zig").put; +const t = std.testing; + +const Fake = struct { + var pinned: ?[]align(16) u8 = null; + var pending: ?[]const u8 = null; + var result: [16]u8 = undefined; + fn host(out: *?*anyopaque, len: usize, _: c_uint) callconv(.c) cuda.abi.Result { + pinned = t.allocator.alignedAlloc(u8, .@"16", len) catch return cuda.abi.error_not_ready; + out.* = pinned.?.ptr; + return 0; + } + fn free(_: ?*anyopaque) callconv(.c) cuda.abi.Result { + std.debug.assert(pending == null); + t.allocator.free(pinned.?); + pinned = null; + return 0; + } + fn copy(_: u64, ptr: ?*const anyopaque, len: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + pending = @as([*]const u8, @ptrCast(ptr.?))[0..len]; + return 0; + } + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + if (pending) |bytes| @memcpy(result[0..bytes.len], bytes); + pending = null; + return 0; + } + fn event(out: *cuda.abi.Event, _: c_uint) callconv(.c) cuda.abi.Result { + out.* = @ptrFromInt(1); + return 0; + } + fn record(_: cuda.abi.Event, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return 0; + } + fn query(_: cuda.abi.Event) callconv(.c) cuda.abi.Result { + return if (pending != null) cuda.abi.error_not_ready else 0; + } + fn destroy(_: cuda.abi.Event) callconv(.c) cuda.abi.Result { + return 0; + } +}; + +test "an upload's source may change as soon as put returns: the device gets the bytes it was given" { + var d: cuda.Driver = undefined; + d.api.cuMemHostAlloc = Fake.host; + d.api.cuMemFreeHost = Fake.free; + d.api.cuMemcpyHtoDAsync_v2 = Fake.copy; + d.api.cuStreamSynchronize = Fake.sync; + d.api.cuEventCreate = Fake.event; + d.api.cuEventRecord = Fake.record; + d.api.cuEventQuery = Fake.query; + d.api.cuEventDestroy_v2 = Fake.destroy; + const s: cuda.Stream = .{ .d = &d, .handle = @ptrFromInt(0x1800) }; + Fake.pending = null; + Fake.result = @splat(0); + const host_before = cuda.memory.usage(false).host; + defer { + s.synchronize() catch unreachable; + cuda.copy_source.dropSpares(&d, cuda.copy_source.current()); + } + var meta = [4]i32{ 41, 0, 0, 0 }; + try put(&d, s, 0x1000, std.mem.asBytes(&meta)); + meta[0] = -7; // a draw's next meta, written into the same stack slot before the stream ran + try s.synchronize(); + cuda.copy_source.dropSpares(&d, cuda.copy_source.current()); + try t.expectEqual(@as(i32, 41), std.mem.bytesToValue(i32, Fake.result[0..4])); + try t.expect(Fake.pinned == null); + try t.expectEqual(host_before, cuda.memory.usage(false).host); +} diff --git a/zig/src/cuda_dense/recipe.zig b/zig/src/cuda_dense/recipe.zig new file mode 100644 index 00000000000..b0f35964a52 --- /dev/null +++ b/zig/src/cuda_dense/recipe.zig @@ -0,0 +1,136 @@ +//! A dense decoder's recipe contract: its geometry, the decls the engine reads, and the attention launch shape. +const std = @import("std"); + +/// A dense GQA decoder's sizes, read from its checkpoint's config before the weights load. +pub const Geometry = struct { + hidden: usize, + layers: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, // 64 or 128 + vocab: usize, + eps: f32, + context: usize, // the checkpoint's own window + acc: usize, // q|k|v width: a fused projection's f32 sums + wide: usize, // gate|up width + mid: usize, // the MLP's inner width + theta: ?f64 = null, // 1-D RoPE base (rope.zig's inverse frequencies on the device); null: no rotary positions + inv_freq: ?[]const f32 = null, // the reference's own inverse frequencies, when its pow is not correctly rounded + + pub fn qWidth(g: Geometry) usize { + return g.heads * g.head_dim; + } + + pub fn kvWidth(g: Geometry) usize { + return g.kv_heads * g.head_dim; + } + + pub fn group(g: Geometry) usize { + return g.heads / g.kv_heads; + } +}; + +/// Query heads a KV head the decode kernels take at most (paged_attn.cuh's GMAX). +pub const max_group = 16; + +/// The attention launches of one geometry: prompt pattn and its block, decode chunk and merge. +pub const Shape = struct { + pattn: [:0]const u8, + chunk: [:0]const u8, + merge: [:0]const u8, + hpc: usize, // query heads a pattn block takes from one KV head + rows: usize, // prompt rows a pattn block takes + grid_y: usize, // pattn blocks over the heads + shared: u32, // pattn's dynamic shared bytes + + pub fn blocks(s: Shape, rows: usize) usize { + return (rows + s.rows - 1) / s.rows; + } +}; + +fn pattnSymbol(comptime d: usize, comptime hpc: usize) [:0]const u8 { + return std.fmt.comptimePrint("_ZN20tf_prefill_attention12pattn_kernelILi{d}ELi8ELi{d}ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix", .{ d, hpc }); +} + +/// The pattn instances prefill_attention.cu builds at eight warps and eight slots: head dim, query heads a block. +const Instance = struct { d: usize, hpc: usize, symbol: [:0]const u8 }; +const instances = [_]Instance{ + .{ .d = 128, .hpc = 8, .symbol = pattnSymbol(128, 8) }, + .{ .d = 128, .hpc = 2, .symbol = pattnSymbol(128, 2) }, + .{ .d = 64, .hpc = 4, .symbol = pattnSymbol(64, 4) }, + .{ .d = 64, .hpc = 2, .symbol = pattnSymbol(64, 2) }, +}; + +/// pattn's shared bytes: eight staging slots of 32 keys, a key row in 16-byte chunks (prefill_attention_rows.cuh). +fn pattnShared(d: usize) u32 { + const chunks = (2 * ((d + 15) / 16) + 7) / 8 * 8; + return @intCast(8 * 32 * chunks * 16); +} + +/// The launches for `g`: the largest built pattn instance whose heads a block divide the group; refusals by name. +pub fn attentionShape(g: Geometry) error{ UnsupportedDenseHeadDim, UnsupportedDenseGroup }!Shape { + if (g.head_dim != 64 and g.head_dim != 128) return error.UnsupportedDenseHeadDim; + if (g.kv_heads == 0 or g.heads % g.kv_heads != 0 or g.group() > max_group) return error.UnsupportedDenseGroup; + const group_n = g.group(); + for (instances) |in| if (in.d == g.head_dim and group_n % in.hpc == 0) return .{ + .pattn = in.symbol, + .chunk = if (g.head_dim == 128) "tf_nemo_attn_chunk_segs" else "tf_paged_attn_chunk_segs_64", + .merge = if (g.head_dim == 128) "tf_nemo_attn_merge_segs" else "tf_paged_attn_merge_segs_64", + .hpc = in.hpc, + .rows = 16 * (8 / in.hpc), + .grid_y = g.kv_heads * (group_n / in.hpc), + .shared = pattnShared(g.head_dim), + }; + return error.UnsupportedDenseGroup; +} + +fn need(comptime R: type, comptime name: []const u8) void { + if (!@hasDecl(R, name)) @compileError("a cuda_dense recipe declares " ++ name); +} + +fn absent(comptime R: type, comptime name: []const u8) void { + if (@hasDecl(R, name) and @field(R, name) != null) @compileError(name ++ " is not built in the dense engine yet: set it null"); +} + +/// The decls a recipe gives the engine, at compile time; the capabilities not built yet must stay null. +pub fn check(comptime R: type) void { + const required = .{ "name", "model_type", "formats", "default_context", "Weights", "Kernels", "geometry", "load", "weightBytes", "embed", "layer", "final", "head" }; + inline for (required) |name| need(R, name); + inline for (.{ "post_embed", "Tower", "scores", "propose" }) |name| absent(R, name); +} + +fn shapeOf(heads: usize, kv_heads: usize, head_dim: usize) !Shape { + return attentionShape(.{ .hidden = heads * head_dim, .layers = 1, .heads = heads, .kv_heads = kv_heads, .head_dim = head_dim, .vocab = 1, .eps = 1e-6, .context = 1, .acc = 1, .wide = 1, .mid = 1 }); +} + +test "attention shapes: dots 12/2 x 128, 16/8 x 128, 16/8 x 64 and 32/8 x 64" { + const dots = try shapeOf(12, 2, 128); + try std.testing.expectEqual(@as(usize, 2), dots.hpc); + try std.testing.expectEqual(@as(usize, 64), dots.rows); + try std.testing.expectEqual(@as(usize, 6), dots.grid_y); + try std.testing.expectEqual(@as(u32, 65536), dots.shared); + try std.testing.expectEqualStrings("_ZN20tf_prefill_attention12pattn_kernelILi128ELi8ELi2ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix", dots.pattn); + try std.testing.expectEqualStrings("tf_nemo_attn_chunk_segs", dots.chunk); + try std.testing.expectEqual(@as(usize, 2), dots.blocks(65)); + const glm = try shapeOf(16, 8, 128); + try std.testing.expect(glm.hpc == 2 and glm.rows == 64 and glm.grid_y == 8 and glm.shared == 65536); + const parse = try shapeOf(16, 8, 64); + try std.testing.expect(parse.hpc == 2 and parse.rows == 64 and parse.grid_y == 8 and parse.shared == 32768); + try std.testing.expectEqualStrings("_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi2ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix", parse.pattn); + try std.testing.expectEqualStrings("tf_paged_attn_merge_segs_64", parse.merge); + const lfm = try shapeOf(32, 8, 64); + try std.testing.expect(lfm.hpc == 4 and lfm.rows == 32 and lfm.grid_y == 8 and lfm.shared == 32768); + try std.testing.expectEqualStrings("_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix", lfm.pattn); +} + +test "attention shapes refuse groups past 16, odd groups, uneven heads and head dims other than 64 and 128" { + try std.testing.expectError(error.UnsupportedDenseGroup, shapeOf(34, 2, 128)); + try std.testing.expectError(error.UnsupportedDenseGroup, shapeOf(12, 4, 128)); + try std.testing.expectError(error.UnsupportedDenseGroup, shapeOf(12, 5, 64)); + try std.testing.expectError(error.UnsupportedDenseGroup, shapeOf(8, 8, 64)); + try std.testing.expectError(error.UnsupportedDenseHeadDim, shapeOf(16, 8, 80)); + try std.testing.expectError(error.UnsupportedDenseHeadDim, shapeOf(16, 8, 256)); + try std.testing.expectEqual(@as(usize, 2), (try shapeOf(32, 8, 128)).hpc); // group 4 at head 128: no 4-head block + try std.testing.expectEqual(@as(usize, 8), (try shapeOf(32, 2, 128)).hpc); + try std.testing.expectEqual(@as(usize, 8), (try shapeOf(64, 8, 128)).hpc); +} diff --git a/zig/src/cuda_dense/root.zig b/zig/src/cuda_dense/root.zig new file mode 100644 index 00000000000..c20f45e0e5e --- /dev/null +++ b/zig/src/cuda_dense/root.zig @@ -0,0 +1,30 @@ +//! The shared CUDA dense-decoder engine: a recipe's layers over paged GQA attention, lanes, draws and the host decls. +pub const recipe = @import("recipe.zig"); +pub const pages = @import("pages.zig"); +pub const rope = @import("rope.zig"); +pub const kernels = @import("kernels.zig"); +pub const model = @import("model.zig"); +pub const pass = @import("pass.zig"); +pub const draws = @import("draws.zig"); +pub const backend = @import("backend.zig"); +pub const native = @import("native.zig"); + +pub const Native = native.Native; +pub const Pass = pass.Pass; +pub const Geometry = recipe.Geometry; +pub const check = recipe.check; + +test { + _ = recipe; + _ = pages; + _ = rope; + _ = kernels; + _ = model; + _ = pass; + _ = draws; + _ = backend; + _ = native; + _ = @import("native_test.zig"); + _ = @import("put_test.zig"); + _ = @import("device_fail_test.zig"); +} diff --git a/zig/src/cuda_dense/rope.zig b/zig/src/cuda_dense/rope.zig new file mode 100644 index 00000000000..e49633f0ef4 --- /dev/null +++ b/zig/src/cuda_dense/rope.zig @@ -0,0 +1,31 @@ +//! 1-D rotary inverse frequencies as HF's default rope computes them in float32; angles and cos/sin are device work. +const std = @import("std"); + +/// Frequencies a head of `dim` dims takes: one a rotated pair. +pub fn count(dim: usize) usize { + return dim / 2; +} + +/// out[i] = 1 / theta^(2i / dim): the exponent an f32 quotient, the power of the f32 operands rounded once to f32. +pub fn inverse(theta: f64, dim: usize, out: []f32) void { + std.debug.assert(out.len == count(dim)); + const base: f32 = @floatCast(theta); + for (out, 0..) |*v, i| { + const e: f32 = @as(f32, @floatFromInt(2 * i)) / @as(f32, @floatFromInt(dim)); + const pw: f32 = @floatCast(std.math.pow(f64, base, e)); + v.* = @as(f32, 1.0) / pw; + } +} + +test "dots' inverse frequencies (theta 1e6, head dim 128) in float32" { + var got: [64]f32 = undefined; + inverse(1e6, 128, &got); + for (got, 0..) |v, i| { + const e = @as(f32, @floatFromInt(i * 2)) / 128.0; + const want = @as(f32, 1.0) / @as(f32, @floatCast(std.math.pow(f64, @as(f32, 1e6), e))); + try std.testing.expectEqual(@as(u32, @bitCast(want)), @as(u32, @bitCast(v))); + } + try std.testing.expectEqual(@as(f32, 1), got[0]); + try std.testing.expectEqual(@as(f32, 1e-3), got[32]); // theta^(1/2): 1000 exactly, its f32 reciprocal + try std.testing.expect(got[63] < got[62] and got[63] > 1e-6); +} From 3c6b9aa11c34d5ac4d18881534ceea3b9cf36c93 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 9 Nov 2026 13:58:59 +0000 Subject: [PATCH 02/12] feat(dots): dots.mocr's text decoder as a CUDA dense recipe dots.mocr's Qwen2 decoder runs on the dense engine: q|k|v and gate|up fused into packed BF16 lanes, the q|k|v biases and rotary in its own Torch-order kernel with Transformers' own float32 frequency table, attention through the engine, the untied head, and no reasoning block. The loader drains its streams before freeing on a failure. It is not registered yet; its host tests check the geometry, text bytes, frequencies, declarations and refusals. --- zig/build/cuda_small_hosts.zig | 3 +- zig/build/dots_mocr_cuda.zig | 12 + zig/src/families/dots_mocr/cuda.zig | 60 +++++ zig/src/families/dots_mocr/cuda_kernels.zig | 31 +++ zig/src/families/dots_mocr/cuda_recipe.zig | 97 ++++++++ zig/src/families/dots_mocr/cuda_weights.zig | 231 ++++++++++++++++++ zig/src/families/dots_mocr/fixtures/README.md | 6 + .../dots_mocr/fixtures/rope_inv_freq.txt | 1 + 8 files changed, 440 insertions(+), 1 deletion(-) create mode 100644 zig/build/dots_mocr_cuda.zig create mode 100644 zig/src/families/dots_mocr/cuda.zig create mode 100644 zig/src/families/dots_mocr/cuda_kernels.zig create mode 100644 zig/src/families/dots_mocr/cuda_recipe.zig create mode 100644 zig/src/families/dots_mocr/cuda_weights.zig create mode 100644 zig/src/families/dots_mocr/fixtures/rope_inv_freq.txt diff --git a/zig/build/cuda_small_hosts.zig b/zig/build/cuda_small_hosts.zig index 0e6b160af17..c73fd287ceb 100644 --- a/zig/build/cuda_small_hosts.zig +++ b/zig/build/cuda_small_hosts.zig @@ -1,4 +1,4 @@ -//! Host test steps of the small CUDA families (Spark 2.5 and LFM2 MoE) and the dense engine. +//! Host test steps of the small CUDA families (Spark 2.5, LFM2 MoE, dots.mocr) and the dense engine. const std = @import("std"); pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, all: *std.Build.Step) void { @@ -6,4 +6,5 @@ pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Mod @import("lfm2_moe_cuda.zig").host(b, @import("lfm2_moe_cuda.zig").module(b, host, .debug, cuda, core, lanes), all); const dense = @import("cuda_dense.zig").module(b, host, .debug, cuda, core, lanes); @import("cuda_dense.zig").host(b, dense, all); + @import("dots_mocr_cuda.zig").host(b, @import("dots_mocr_cuda.zig").module(b, host, .debug, cuda, core, lanes, dense), all); } diff --git a/zig/build/dots_mocr_cuda.zig b/zig/build/dots_mocr_cuda.zig new file mode 100644 index 00000000000..b11667c5d44 --- /dev/null +++ b/zig/build/dots_mocr_cuda.zig @@ -0,0 +1,12 @@ +//! dots.mocr's CUDA text family over the runtime, the core, the lane core and the dense engine; its host tests. +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, dense: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/families/dots_mocr/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes }, .{ .name = "cuda_dense", .module = dense } } }); +} + +pub fn host(b: *std.Build, m: *std.Build.Module, all: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = m })); + b.step("test-dots-mocr-cuda-host", "dots.mocr on CUDA: geometry, text bytes, Native decls and refusals without a GPU").dependOn(&run.step); + all.dependOn(&run.step); +} diff --git a/zig/src/families/dots_mocr/cuda.zig b/zig/src/families/dots_mocr/cuda.zig new file mode 100644 index 00000000000..0311c5827dc --- /dev/null +++ b/zig/src/families/dots_mocr/cuda.zig @@ -0,0 +1,60 @@ +//! dots.mocr's text decoder on CUDA: the dots recipe behind the shared dense engine (native/cuda.zig's family decls). +const std = @import("std"); +const dense = @import("cuda_dense"); +const cfg = @import("config.zig"); +const fixtures = @import("config_test.zig"); + +pub const recipe = @import("cuda_recipe.zig"); +pub const weights = @import("cuda_weights.zig"); +pub const kernels = @import("cuda_kernels.zig"); +pub const native = dense.Native(recipe); + +test "dots geometry from the released config: 12 query heads over 2 KV heads of 128, two query heads a pattn block" { + const c = try cfg.Config.parse(std.testing.allocator, fixtures.sample, fixtures.generation); + const g = recipe.geometryOf(c); + var plain = g; + plain.inv_freq = null; + try std.testing.expectEqual(dense.Geometry{ .hidden = 1536, .layers = 28, .heads = 12, .kv_heads = 2, .head_dim = 128, .vocab = 151936, .eps = 1e-6, .context = 131072, .acc = 2048, .wide = 17920, .mid = 8960, .theta = 1e6 }, plain); + try std.testing.expectEqual(@as(usize, 64), g.inv_freq.?.len); // Transformers' own table (cuda_recipe.zig) + const shape = try dense.recipe.attentionShape(g); + try std.testing.expect(shape.hpc == 2 and shape.grid_y == 6 and shape.rows == 64); +} + +test "dots text bytes: embedding, untied head, final norm and 28 layers, the tower left out" { + const c = try cfg.Config.parse(std.testing.allocator, fixtures.sample, fixtures.generation); + const h: u64 = 1536; + const kv: u64 = 256; + const m: u64 = 8960; + const layer = 2 * h + (h + 2 * kv) * h + (h + 2 * kv) + h * h + 3 * m * h; + const want = (2 * 151936 * h + h + 28 * layer) * 2; + try std.testing.expectEqual(@as(u64, 3554176000), want); + try std.testing.expectEqual(want, try weights.textBytes(std.testing.allocator, c)); +} + +test "dots on CUDA: text only, no reasoning block, its config and tensor refusals worded" { + try std.testing.expect(!native.reasoning and native.keeps_no_prompts); + const hooks = .{ &native.open, &native.attach, &native.fits, &native.ledgerPool, &native.textWeightBytes }; // compiles the load and the layer + inline for (hooks) |h| try std.testing.expect(@intFromPtr(h) != 0); + try std.testing.expectEqualStrings("dots_ocr", native.model_type); + try std.testing.expectStringStartsWith(native.explain(null, error.DenseTextOnly).?, "dots.mocr on CUDA reads text prompts only"); + for ([_]anyerror{ error.UnsupportedDotsConfig, error.BadDotsConfig, error.BadEos, error.DotsShape, error.DotsMissingTensor, error.DotsTensorCount, error.DotsIndexMismatch }) |e| + try std.testing.expect(native.explain(null, e) != null); +} + +test "dots on CUDA reads its geometry from the checkpoint's files and refuses another model's config by name" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(io, .{ .sub_path = "generation_config.json", .data = fixtures.generation }); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = fixtures.sample }); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(dir); + try std.testing.expectEqual(@as(usize, 2048), (try recipe.geometry(a, io, dir)).acc); + try std.testing.expectError(error.FileNotFound, recipe.weightBytes(a, io, dir)); // no shards: nothing to count + const other = try std.mem.replaceOwned(u8, a, fixtures.sample, "\"dots_ocr\"", "\"qwen3_5\""); + defer a.free(other); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = other }); + try std.testing.expectError(error.UnsupportedDotsConfig, recipe.geometry(a, io, dir)); + try std.testing.expect(native.explain(null, error.UnsupportedDotsConfig) != null); +} diff --git a/zig/src/families/dots_mocr/cuda_kernels.zig b/zig/src/families/dots_mocr/cuda_kernels.zig new file mode 100644 index 00000000000..39a101a2e0c --- /dev/null +++ b/zig/src/families/dots_mocr/cuda_kernels.zig @@ -0,0 +1,31 @@ +//! dots.mocr's own CUDA launch beside the dense engine's: the q|k|v bias and Torch-order RoPE (dots_mocr.cu). +const cuda = @import("cuda"); + +/// Threads a (row, head) block takes: dims j and j + 64 of a 128-dim head. +const threads = 64; + +pub const Kernels = struct { + mod: cuda.Module, + rope: cuda.Function, + + pub fn load(ctx: *const cuda.Context) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(ctx.d, cuda.kernels.dots_mocr); + errdefer mod.unload(); + return .{ .mod = mod, .rope = try mod.function("tf_dots_qkv_rope") }; + } + + pub fn deinit(k: *Kernels) void { + k.mod.unload(); + } + + /// q [rows, h, 128] and kv [rows, k hk | v hk] from fp32 acc [rows, (h + 2 hk) 128] plus bias, rotated at pos. + pub fn qkvRope(k: *const Kernels, s: cuda.Stream, acc: u64, bias: u64, inv: u64, pos: u64, q: u64, kv: u64, rows: usize, h: usize, hk: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + for ([_]u64{ acc, bias, inv, pos, q, kv }) |v| a.add(v); + a.add(@as(c_int, @intCast(h))); + a.add(@as(c_int, @intCast(hk))); + try cuda.launch.launch(k.rope, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(h + 2 * hk), .z = 1 }, .block = .{ .x = threads } }, s, &a); + } +}; diff --git a/zig/src/families/dots_mocr/cuda_recipe.zig b/zig/src/families/dots_mocr/cuda_recipe.zig new file mode 100644 index 00000000000..3052c4d69a1 --- /dev/null +++ b/zig/src/families/dots_mocr/cuda_recipe.zig @@ -0,0 +1,97 @@ +//! dots.mocr's Qwen2 text decoder as a cuda_dense recipe: HIP's rounding points, its tower refused until it is built. +const std = @import("std"); +const cuda = @import("cuda"); +const dense = @import("cuda_dense"); +const cfg = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); + +const Recipe = @This(); +const P = *const dense.Pass(Recipe); + +pub const name = "dots.mocr"; +pub const model_type = "dots_ocr"; +pub const formats: []const []const u8 = &.{"bf16"}; +pub const default_context: i64 = 32768; +/// Replies carry no reasoning block: the text is content. +pub const reasoning = false; + +pub const Weights = wts.Weights; +pub const Kernels = kern.Kernels; + +/// The engine's geometry of an admitted dots.mocr config. +pub fn geometryOf(c: cfg.Config) dense.Geometry { + const h: usize = c.hidden; + return .{ .hidden = h, .layers = c.layers, .heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.headDim(), .vocab = c.vocab, .eps = c.eps, .context = c.context, .acc = h + 2 * @as(usize, c.kvWidth()), .wide = 2 * @as(usize, c.intermediate), .mid = c.intermediate, .theta = c.theta, .inv_freq = if (c.theta == 1e6 and c.headDim() == 128) &hf_inv_freq else null }; +} + +/// config.json and generation_config.json admitted as dots.mocr's, before any tensor is read. +pub fn geometry(a: std.mem.Allocator, io: std.Io, dir: []const u8) !dense.Geometry { + return geometryOf(try cfg.Config.read(a, io, dir)); +} + +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, dir: []const u8) !Weights { + return wts.load(gpa, io, d, s, dir); +} + +/// The text decoder's bytes on the device, once every shard's header is checked. +pub fn weightBytes(a: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { + const c = try cfg.Config.read(a, io, dir); + _ = try @import("weights.zig").inspect(a, io, dir, c); + return wts.textBytes(a, c); +} + +pub fn embed(p: P) !void { + try p.k.embed.bf16(p.s, p.m.s.ids, p.w.embed, p.m.s.h, p.rows, p.g.hidden); +} + +/// One Qwen2 layer: RMS, q|k|v sums, bias and RoPE in Torch's order, attention, o, RMS, SwiGLU. +pub fn layer(p: P, i: usize) !void { + const x = p.m.s; + const l = p.w.layers[i]; + try p.norm(if (i == 0) null else x.r, l.ln1); + try p.projectF32(l.qkv, x.x, x.acc); + try p.own.qkvRope(p.s, x.acc, l.bias, p.m.inv, x.pos, x.q, x.kv, p.rows, p.g.heads, p.g.kv_heads); + try p.attend(i, x.q, x.kv, x.y); + try p.project(l.o, x.y, x.r); + try p.norm(x.r, l.ln2); + try p.project(l.gate_up, x.x, x.wide); + try p.siluMul(x.wide, x.y, p.g.mid); + try p.project(l.down, x.y, x.r); +} + +pub fn final(p: P) !void { + try p.norm(p.m.s.r, p.w.norm); +} + +pub fn head(w: *const Weights) cuda.qmmf.Weight { + return w.head; +} + +pub fn explain(err: anyerror) ?[]const u8 { + return switch (err) { + error.UnsupportedDotsConfig, error.BadDotsConfig, error.BadEos => "this checkpoint's config is not dots.mocr's (DotsOCRForCausalLM, Qwen2 1.5B text decoder in BF16), the one this engine serves", + error.DotsDType, error.DotsShape, error.DotsExtent, error.DotsMissingTensor, error.DotsTensorCount, error.DotsIndexMismatch => "this dots.mocr checkpoint's tensors are not the released BF16 ones: download it again", + else => null, + }; +} + +/// Transformers' own float32 inv_freq for theta 1e6 at head dim 128 (fixtures/README.md): Torch's CPU pow, bit for bit. +const hf_inv_freq: [64]f32 = blk: { + @setEvalBranchQuota(20_000); + const text = @embedFile("fixtures/rope_inv_freq.txt"); + var out: [64]f32 = undefined; + var it = std.mem.tokenizeAny(u8, text, ", \n"); + for (&out) |*v| v.* = @bitCast(std.fmt.parseInt(u32, it.next().?, 0) catch unreachable); + if (it.next() != null) @compileError("fixtures/rope_inv_freq.txt holds more than 64 values"); + break :blk out; +}; + +test "dots' rotary frequencies are Transformers' own: the correctly rounded pow differs at entry 37 only" { + var mine: [64]f32 = undefined; + @import("cuda_dense").rope.inverse(1e6, 128, &mine); + for (hf_inv_freq, mine, 0..) |hf, own, i| { + const one_ulp = @as(u32, @bitCast(hf)) -% @as(u32, @bitCast(own)) <= 1 or @as(u32, @bitCast(own)) -% @as(u32, @bitCast(hf)) <= 1; + try std.testing.expect(if (i == 37) hf != own and one_ulp else hf == own); + } +} diff --git a/zig/src/families/dots_mocr/cuda_weights.zig b/zig/src/families/dots_mocr/cuda_weights.zig new file mode 100644 index 00000000000..edda9806954 --- /dev/null +++ b/zig/src/families/dots_mocr/cuda_weights.zig @@ -0,0 +1,231 @@ +//! dots.mocr's text decoder on CUDA: BF16 tensors on the device, q|k|v and gate|up fused into packed lane weights. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const cfg = @import("config.zig"); +const inventory = @import("weights.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const Weight = cuda.qmmf.Weight; + +/// The text decoder's names: the tower's (vision_tower.) stay on the host until the tower is built on CUDA. +const text = [_][]const u8{ "model.", "lm_head." }; + +pub const Layer = struct { + ln1: u64, // bf16 [hidden]: input_layernorm + ln2: u64, // bf16 [hidden]: post_attention_layernorm + qkv: Weight, // q|k|v [hidden + 2 kv width, hidden] + bias: u64, // bf16 [hidden + 2 kv width]: q|k|v biases + o: Weight, + gate_up: Weight, // gate|up [2 intermediate, hidden] + down: Weight, +}; + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: cfg.Config, + buffers: std.ArrayList(DeviceBuffer) = .empty, + layers: []Layer = &.{}, + embed: u64 = 0, // bf16 [vocab, hidden] + norm: u64 = 0, // bf16 [hidden]: model.norm + head: Weight = undefined, // the untied lm_head + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.layers); + } + + fn alloc(w: *Weights, d: *const cuda.Driver, bytes: usize) !DeviceBuffer { + try w.buffers.ensureUnusedCapacity(w.gpa, 1); + const b = try DeviceBuffer.alloc(d, bytes); + w.buffers.appendAssumeCapacity(b); + return b; + } +}; + +/// The text tensors' bytes as stored, which their device copies take too (every lane weight's rows a multiple of 128). +pub fn textBytes(gpa: std.mem.Allocator, c: cfg.Config) !u64 { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + var total: u64 = 0; + for (try inventory.inventory(arena.allocator(), c)) |s| { + for (text) |p| if (std.mem.startsWith(u8, s.name, p)) { + total += s.bytes(); + }; + } + return total; +} + +/// Reads the checkpoint's text tensors into the device on `s`, fusing and packing the projections. +const Loader = struct { + w: *Weights, + ck: *core.Checkpoint, + d: *const cuda.Driver, + s: cuda.Stream, + pack: *const cuda.modelopt.Packer, + stage: DeviceBuffer, + + /// A bf16 tensor as stored. + fn raw(l: *Loader, name: []const u8, shape: []const usize) !u64 { + return l.joined(&.{name}, &.{shape[0]}, if (shape.len == 2) shape[1] else 1); + } + + /// bf16 tensors [rows[i], k] one after another in a buffer of their own. + fn joined(l: *Loader, names: []const []const u8, rows: []const usize, k: usize) !u64 { + var n: usize = 0; + for (rows) |r| n += r; + const b = try l.w.alloc(l.d, n * k * 2); + var at: usize = 0; + for (names, rows) |name, r| { + const t = try l.ck.expect(name, .bf16, if (k == 1) &.{r} else &.{ r, k }); + try b.upload(at, t.bytes); + at += t.bytes.len; + } + return b.ptr; + } + + /// bf16 projections [rows[i], k] stacked into one [sum rows, k] lane weight, packed on the device. + fn lane(l: *Loader, names: []const []const u8, rows: []const usize, k: usize) !Weight { + var n: usize = 0; + var at: usize = 0; + for (names, rows) |name, r| { + const t = try l.ck.expect(name, .bf16, &.{ r, k }); + try l.stage.upload(at, t.bytes); + at += t.bytes.len; + n += r; + } + const b = try l.w.alloc(l.d, cuda.modelopt.padded(n) * k * 2); + try l.pack.bf16Lane(l.s, l.stage.ptr, b.ptr, n, k); + try l.s.synchronize(); // the next projection's upload reuses the stage + return cuda.modelopt.bf16Weight(b.ptr, n, k); + } + + /// A failed load's work proved complete (the packing stream, then the uploads' null stream) before buffers go. + fn settle(l: *const Loader) void { + cuda.completion.drain(l.d, l.s, "a dots.mocr weight load did not drain before its buffers were freed"); + cuda.completion.drain(l.d, .{ .d = l.d, .handle = null }, "dots.mocr weight uploads did not drain before their buffers were freed"); + } +}; + +/// The text decoder of the dots.mocr checkpoint at `dir`, every header checked first; a tensor left over refuses. +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, dir: []const u8) !Weights { + const c = try cfg.Config.read(gpa, io, dir); + _ = try inventory.inspect(gpa, io, dir, c); + var ck = try core.Checkpoint.openModelPrefixes(gpa, io, dir, &text); + defer ck.close(); + var pack = try cuda.modelopt.Packer.load(d); + defer pack.unload(); + var w: Weights = .{ .gpa = gpa, .config = c }; + errdefer w.deinit(); + w.layers = try gpa.alloc(Layer, c.layers); + const h: usize = c.hidden; + const kvw: usize = c.kvWidth(); + const m: usize = c.intermediate; + var stage = try DeviceBuffer.alloc(d, @as(usize, c.vocab) * h * 2); + defer stage.free(); + var l: Loader = .{ .w = &w, .ck = &ck, .d = d, .s = s, .pack = &pack, .stage = stage }; + errdefer l.settle(); // runs first on failure: the stage, the weights and the packer go after it + w.embed = try l.raw("model.embed_tokens.weight", &.{ c.vocab, h }); + w.norm = try l.raw("model.norm.weight", &.{h}); + w.head = try l.lane(&.{"lm_head.weight"}, &.{c.vocab}, h); + var names: [7][96]u8 = undefined; + for (w.layers, 0..) |*y, i| { + var n: [7][]const u8 = undefined; + for (&n, &names, [_][]const u8{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj" }) |*out, *buf, part| + out.* = try std.fmt.bufPrint(buf, "model.layers.{d}.{s}.weight", .{ i, part }); + var bias: [3][96]u8 = undefined; + var b: [3][]const u8 = undefined; + for (&b, &bias, n[0..3]) |*out, *buf, weight| out.* = try std.fmt.bufPrint(buf, "{s}.bias", .{weight[0 .. weight.len - ".weight".len]}); + var norms: [2][96]u8 = undefined; + y.* = .{ + .ln1 = try l.raw(try std.fmt.bufPrint(&norms[0], "model.layers.{d}.input_layernorm.weight", .{i}), &.{h}), + .ln2 = try l.raw(try std.fmt.bufPrint(&norms[1], "model.layers.{d}.post_attention_layernorm.weight", .{i}), &.{h}), + .qkv = try l.lane(n[0..3], &.{ h, kvw, kvw }, h), + .bias = try l.joined(&b, &.{ h, kvw, kvw }, 1), + .o = try l.lane(n[3..4], &.{h}, h), + .gate_up = try l.lane(n[4..6], &.{ m, m }, h), + .down = try l.lane(n[6..7], &.{h}, m), + }; + } + if (ck.unused() != 0) return error.DotsTensorCount; + try s.synchronize(); + return w; +} + +/// A device whose packing stream's wait fails: it counts buffers freed while a launch may still write them. +const Fake = struct { + const queue: cuda.abi.Stream = @ptrFromInt(0x1800); + var next: u64 = 0x10000; + var live: usize = 0; + var pending = false; + var freed_pending: usize = 0; + var proofs: usize = 0; + fn alloc(out: *u64, len: usize) callconv(.c) cuda.abi.Result { + out.* = next; + next += len + 256; + live += 1; + return 0; + } + fn free(_: u64) callconv(.c) cuda.abi.Result { + freed_pending += @intFromBool(pending); + live -= 1; + return 0; + } + fn copy(_: u64, _: ?*const anyopaque, _: usize) callconv(.c) cuda.abi.Result { + return 0; + } + fn launch(_: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + pending = true; + return 0; + } + fn wait(s: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return if (s == queue) cuda.abi.error_not_ready else 0; + } + fn context() callconv(.c) cuda.abi.Result { + proofs += 1; + pending = false; + return 0; + } + fn driver() cuda.Driver { + var d: cuda.Driver = .{ .lib = undefined, .api = undefined }; + d.api.cuMemAlloc_v2 = alloc; + d.api.cuMemFree_v2 = free; + d.api.cuMemcpyHtoD_v2 = copy; + d.api.cuLaunchKernel = launch; + d.api.cuStreamSynchronize = wait; + d.api.cuCtxSynchronize = context; + return d; + } +}; + +test "a failed pack wait proves completion by the context before the stage or any weight is freed" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = "{\"x.weight\":{\"dtype\":\"BF16\",\"shape\":[2,64],\"data_offsets\":[0,256]}}"; + var file: [8 + header.len + 256]u8 = @splat(0); + std.mem.writeInt(u64, file[0..8], header.len, .little); + @memcpy(file[8..][0..header.len], header); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &file }); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(dir); + var ck = try core.Checkpoint.openModelPrefixes(a, io, dir, &.{}); + defer ck.close(); + var d = Fake.driver(); + var pack: cuda.modelopt.Packer = undefined; + pack.bf16_lane = .{ .d = &d, .handle = @ptrFromInt(1) }; + var w: Weights = .{ .gpa = a, .config = undefined }; + var stage = try DeviceBuffer.alloc(&d, 256); + const l: Loader = .{ .w = &w, .ck = &ck, .d = &d, .s = .{ .d = &d, .handle = Fake.queue }, .pack = &pack, .stage = stage }; + var run = l; + try std.testing.expectError(error.NotReady, run.lane(&.{"x.weight"}, &.{2}, 64)); + try std.testing.expect(Fake.pending and Fake.live == 2); + l.settle(); + stage.free(); + w.deinit(); + try std.testing.expectEqual(@as(usize, 0), Fake.freed_pending); + try std.testing.expectEqual(@as(usize, 0), Fake.live); + try std.testing.expectEqual(@as(usize, 1), Fake.proofs); +} diff --git a/zig/src/families/dots_mocr/fixtures/README.md b/zig/src/families/dots_mocr/fixtures/README.md index f75f365ad29..75fbffe9e66 100644 --- a/zig/src/families/dots_mocr/fixtures/README.md +++ b/zig/src/families/dots_mocr/fixtures/README.md @@ -50,3 +50,9 @@ FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE \ OUTPUT_DIR --pci PCI_ID --max-new-tokens 64 python tools/zig/dots_mocr_reference.py compare OUTPUT_DIR NATIVE_PREFIX ``` + +`rope_inv_freq.txt` holds the 64 float32 bit patterns of the checkpoint's +rotary `inv_freq` as Transformers 4.57.6 computes them on the CPU (theta 1e6, +head dim 128), dumped by the oracle's `text` mode (`rope-inv_freq.npy`). The +CUDA recipe uploads them as they are, because Torch's CPU `powf` is not +correctly rounded: entry 37 is one ulp from the correctly rounded value. diff --git a/zig/src/families/dots_mocr/fixtures/rope_inv_freq.txt b/zig/src/families/dots_mocr/fixtures/rope_inv_freq.txt new file mode 100644 index 00000000000..8d31a9bc1db --- /dev/null +++ b/zig/src/families/dots_mocr/fixtures/rope_inv_freq.txt @@ -0,0 +1 @@ +0x3f800000,0x3f4e4bad,0x3f263de0,0x3f05f6ef,0x3ed7e89b,0x3eadfcff,0x3e8c3504,0x3e61f835,0x3e361887,0x3e12bd91,0x3dec7fd6,0x3dbe94c6,0x3d99940d,0x3d778513,0x3d47763f,0x3d20bc1d,0x3d0186e3,0x3cd0c1a8,0x3ca8398b,0x3c879008,0x3c5a7bf2,0x3c301052,0x3c0de12d,0x3be4aa46,0x3bb8449c,0x3b947dae,0x3b6f520e,0x3b40dac5,0x3b1b690d,0x3afa78f0,0x3ac9d75c,0x3aa2a6f7,0x3a83126f,0x3a533f28,0x3a2a3b44,0x3a092e02,0x39dd1725,0x39b229fb,0x398f9272,0x39676492,0x393a7753,0x39164324,0x38f22ce2,0x38c327b5,0x389d43a4,0x387d75d4,0x384c3fbe,0x382497ab,0x3804a2b2,0x37d5c441,0x37ac431d,0x378ad0ed,0x375fba4f,0x37344a0e,0x371148e3,0x36ea2732,0x36bcb0c1,0x36980e02,0x36751070,0x36457bab,0x361f23e4,0x36003dec,0x35ceaf79,0x35a68e4c From 7c74a58d60ce01060f2b1f4cf259a08aefc74bc9 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Sun, 8 Nov 2026 22:53:28 +0000 Subject: [PATCH 03/12] feat(native): serve dots.mocr's text decoder on CUDA The native CUDA server now lists dots_ocr with the publisher's BF16 weights and opens it through the dense engine, built once for the server's engines. Its capability entry is checked, and the host's reasoning test now expects only dots.mocr to declare replies without a reasoning block. Images stay refused until its tower lands. --- zig/build/cuda_engines.zig | 1 + zig/src/native/cuda.zig | 4 ++-- zig/src/native/cuda_families.zig | 10 +++++++++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/zig/build/cuda_engines.zig b/zig/build/cuda_engines.zig index b04a7aef562..80f037e3d22 100644 --- a/zig/build/cuda_engines.zig +++ b/zig/build/cuda_engines.zig @@ -89,6 +89,7 @@ pub fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu .{ .name = "spark25", .module = @import("spark25.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "lfm2_moe_cuda", .module = @import("lfm2_moe_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "dory", .module = @import("dory.zig").module(b, target, optimize, cuda, core, lanes) }, + .{ .name = "dots_mocr_cuda", .module = @import("dots_mocr_cuda.zig").module(b, target, optimize, cuda, core, lanes, @import("cuda_dense.zig").module(b, target, optimize, cuda, core, lanes)) }, }, }); mod.addImport("qwen4", @import("qwen4.zig").module(b, target, optimize, core)); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 10d955392e3..57c34a6bac1 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -490,12 +490,12 @@ fn reasoning(comptime F: type) bool { return if (@hasDecl(F, "reasoning")) F.reasoning else true; } -test "a CUDA family declares replies without reasoning with `pub const reasoning = false`; every served one keeps it" { +test "a CUDA family declares replies without reasoning with `pub const reasoning = false`; only dots.mocr does" { try std.testing.expect(reasoning(struct {})); try std.testing.expect(!reasoning(struct { pub const reasoning = false; })); - inline for (registry) |F| try std.testing.expect(reasoning(F)); + inline for (registry) |F| try std.testing.expectEqual(!std.mem.eql(u8, F.model_type, "dots_ocr"), reasoning(F)); } test "the Info openWith serves carries each family's reasoning declaration over the host's" { diff --git a/zig/src/native/cuda_families.zig b/zig/src/native/cuda_families.zig index e4bc3c2a493..de591b1b621 100644 --- a/zig/src/native/cuda_families.zig +++ b/zig/src/native/cuda_families.zig @@ -14,7 +14,7 @@ const dory = @import("dory"); const qwen3 = @import("qwen3_cuda"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -pub const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native, @import("lfm2_moe_cuda").native }; +pub const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native, @import("dots_mocr_cuda").native, @import("lfm2_moe_cuda").native }; pub const families: []const api.Family = blk: { var out: [registry.len + 5]api.Family = undefined; for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats, .decision_formats = if (@hasDecl(F, "decision_formats")) F.decision_formats else &.{} }; @@ -37,6 +37,14 @@ test "every registered family and d1's decision engine are listed for capabiliti try std.testing.expectEqualStrings("qwen3", families[registry.len + 4].model_type); } +test "dots.mocr serves its publisher's BF16 weights on CUDA through the dense engine" { + for (families) |f| if (std.mem.eql(u8, f.model_type, "dots_ocr")) { + try std.testing.expectEqual(@as(usize, 1), f.formats.len); + return std.testing.expectEqualStrings("bf16", f.formats[0]); + }; + return error.NoDotsFamily; +} + test "GLM capabilities: CUDA offers the EXL3 format its loader accepts" { const config_json = "{\"quantization_config\":{\"quant_method\":\"exl3\",\"bits\":4,\"codebook\":\"mcg\"}}"; var why: @import("core").modelopt.Why = .{}; From 059a489eb8d7b67e43cc8d654b2927c13c9123bd Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Sun, 8 Nov 2026 23:06:09 +0000 Subject: [PATCH 04/12] feat(tools): dots oracle text references and scoring on any GPU The dots.mocr oracle only ran the model card's image example on one HIP GPU. It now runs text-only prompts on a HIP or CUDA GPU at the pinned Transformers, with the rotary values the native kernel's checks read, and a score mode that teacher-forces native ids through HF and classifies each as exact, one BF16 step or a failure, exiting nonzero on any failure or on inputs that do not line up. Without flash-attn, a refusing stand-in lets the tower's import load for text. --- tools/zig/dots_mocr_reference.py | 213 +++++++++++++++++++++++++------ 1 file changed, 177 insertions(+), 36 deletions(-) diff --git a/tools/zig/dots_mocr_reference.py b/tools/zig/dots_mocr_reference.py index 6af268c9734..4c0038cfdb9 100644 --- a/tools/zig/dots_mocr_reference.py +++ b/tools/zig/dots_mocr_reference.py @@ -14,6 +14,7 @@ REVISION = "e539fbb52280393adc081b289ec597430a0f9031" ORACLE_VERSIONS = {"transformers": "4.57.6"} +ROPE_POSITIONS = [0, 1, 4096, 131071] # the rotary values the native kernel's contract compares, short and long PROCESSOR_VERSIONS = {"transformers": "4.51.0", "pillow": "11.2.1", "numpy": "2.2.5"} # The revision's non-weight sources: code, configuration, tokenizer, template, processor and the model card. SOURCES = { @@ -67,8 +68,11 @@ def example_prompt(checkpoint): def selected_pci(torch): - if torch.version.hip is None or not torch.cuda.is_available() or torch.cuda.device_count() != 1: - raise RuntimeError("the oracle requires exactly one visible HIP GPU") + if not torch.cuda.is_available() or torch.cuda.device_count() != 1: + raise RuntimeError("the oracle requires exactly one visible GPU") + if torch.version.hip is None: + p = torch.cuda.get_device_properties(0) + return f"{p.pci_domain_id:04x}:{p.pci_bus_id:02x}:{p.pci_device_id:02x}.0" lib = ctypes.CDLL("libamdhip64.so") function = lib.hipDeviceGetPCIBusId function.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int] @@ -79,6 +83,40 @@ def selected_pci(torch): return text.value.decode().lower() +class Scores: + """Retain the scores used for greedy selection without a growing device tensor history (a LogitsProcessor).""" + + def __init__(self, path, vocab): + self.file = path.open("w+b") + self.vocab, self.count = vocab, 0 + self.header() + + def header(self): + text = repr({"descr": " Date: Mon, 9 Nov 2026 13:58:05 +0000 Subject: [PATCH 05/12] feat(contract): records for the dense engine's other two launches The dense engine also launches the BF16 embedding rows and the Torch-order SwiGLU, and neither had a launch record. Their register and shared-memory usage, read from nvcc 13.0 images for sm_80, sm_86, sm_89, sm_120 and sm_121, now sits with dots' rotary record; the contract's image table gains the lfm2 image the SwiGLU lives in. --- zig/src/core/contract/dots_mocr_cuda_records.zig | 4 +++- zig/tests/cuda/contract_limits.zig | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/zig/src/core/contract/dots_mocr_cuda_records.zig b/zig/src/core/contract/dots_mocr_cuda_records.zig index 16a194dc45b..24500deb427 100644 --- a/zig/src/core/contract/dots_mocr_cuda_records.zig +++ b/zig/src/core/contract/dots_mocr_cuda_records.zig @@ -1,10 +1,12 @@ -//! dots.mocr's CUDA glue: tf_dots_qkv_rope's usage from nvcc 13.0 images for every SM. +//! dots.mocr's CUDA glue and the dense engine's other launches: usage from nvcc 13.0 images for every SM. const std = @import("std"); const registry = @import("registry.zig"); const R = registry.CudaRecord; pub const cuda = [_]R{ .{ .launch = .{ .name = "dots_mocr tf_dots_qkv_rope", .image = "dots_mocr", .symbol = "tf_dots_qkv_rope", .threads = 64 }, .usage = .{ .{ .regs = 28, .static = 0, .bound = 64 }, .{ .regs = 24, .static = 0, .bound = 64 }, .{ .regs = 24, .static = 0, .bound = 64 }, .{ .regs = 25, .static = 0, .bound = 64 }, .{ .regs = 25, .static = 0, .bound = 64 } } }, + .{ .launch = .{ .name = "embed_rows tf_embed_rows_bf16", .image = "embed_rows", .symbol = "tf_embed_rows_bf16", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "lfm2 tf_lfm2_silu_mul", .image = "lfm2", .symbol = "tf_lfm2_silu_mul", .threads = 256 }, .usage = .{ .{ .regs = 19, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 20, .static = 0, .bound = 256 }, .{ .regs = 20, .static = 0, .bound = 256 } } }, }; test "dots.mocr records: every launch fits every SM by its measured record" { diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index fa601d6bbd2..02f9ef56dc2 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -30,7 +30,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple", "gemma4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "lfm2_moe", "dots_mocr" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "lfm2_moe", "lfm2", "dots_mocr" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; return error.UnknownContractImage; } From 66baeb5e3493dd570fa357bd51dd634230f1b7d4 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 9 Nov 2026 13:58:42 +0000 Subject: [PATCH 06/12] docs(dots): a CUDA page for dots.mocr's text decoder dots.mocr's text decoder now serves on CUDA through the shared dense engine, so it gets a recipe page of its own beside the HIP one: how to build and serve it, which flags it refuses and why, where each part lives, and how the oracle's text and score modes check its replies. The HIP page, the recipe index and the README row point to it. --- README.md | 2 +- docs/recipes/README.md | 1 + docs/recipes/dots-mocr-cuda.md | 37 ++++++++++++++++++++++++++++++++++ docs/recipes/dots-mocr-hip.md | 3 ++- 4 files changed, 41 insertions(+), 2 deletions(-) create mode 100644 docs/recipes/dots-mocr-cuda.md diff --git a/README.md b/README.md index 4469c700e83..33840f0b8dc 100644 --- a/README.md +++ b/README.md @@ -177,7 +177,7 @@ Conversions are explicitly marked: their upstream weights need the linked conver | Model / format | Download | Machine and memory scope | Native recipe / limits | | --- | --- | --- | --- | -| dots.mocr, BF16 OCR | Use the local `dots-studio/dots.mocr` checkpoint. | Native HIP; AMD R9700 (`gfx1201`) checked, one GPU. | [Native OCR recipe](docs/recipes/dots-mocr-hip.md); full model-card demo matches HF tokens. No video, MTP head, CUDA, or Metal qualification. | +| dots.mocr, BF16 OCR | Use the local `dots-studio/dots.mocr` checkpoint. | Native HIP; AMD R9700 (`gfx1201`) checked, one GPU. | [Native OCR recipe](docs/recipes/dots-mocr-hip.md); full model-card demo matches HF tokens. Native CUDA serves text prompts only ([CUDA recipe](docs/recipes/dots-mocr-cuda.md)). No video, MTP head, or Metal qualification. | | Qwen3.5-2B, MLX 4-bit | `tensorfold pull mlx-community/Qwen3.5-2B-MLX-4bit@93760be4f1f69842a46bc13dbdc0f19e291392a3` | One Mac (M5 Max checked), or one GB10; no RAM minimum published. | [2B](docs/recipes/qwen3.5-2b-native.md); [CUDA release checks](CHANGELOG.md#110). Text only. | | Qwen3.8-27B, MLX 4-bit | `tensorfold pull TensorFold/Qwen3.8-27B-MLX-4bit`
`tensorfold pull z-lab/Qwen3.8-27B-DFlash2` | One Mac (M5 Max/M3 Ultra), or one GB10; a 32 GB Mac needs `iogpu.wired_limit_mb` raised (the startup refusal names the value). | [27B](docs/recipes/qwen3.8-27b.md). Add `--drafter z-lab/Qwen3.8-27B-DFlash2`; without it decoding is plain. | | Nemotron 3.5 Lightning 30B-A3B, MLX 4-bit | `tensorfold pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit` | One Mac (M1-M5), or one GB10; also checked on RTX 3090 / RTX PRO 6000. | [Lightning](docs/recipes/nemotron-3.5.md). Includes MTP. | diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 6abb258dffa..a84de785b19 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -8,6 +8,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | Family | Recipe | | --- | --- | | dots.mocr | [Native HIP OCR, BF16, one GPU](dots-mocr-hip.md) | +| dots.mocr | [Native CUDA text decoder, BF16, one GPU](dots-mocr-cuda.md) | | Kimi K3 / K3.1 | [Native Metal bring-up, synthetic only, resident budget refusal](kimi-k3-metal.md) | | Nemotron 3.5 Lightning | [MLX](nemotron-3.5.md), [Intel Arc, experimental](xpu-arc.md) | | Nemotron 3.5 Parse | [Native HIP OCR, validation and qualification limits](nemotron-parse.md) | diff --git a/docs/recipes/dots-mocr-cuda.md b/docs/recipes/dots-mocr-cuda.md new file mode 100644 index 00000000000..3a218382fec --- /dev/null +++ b/docs/recipes/dots-mocr-cuda.md @@ -0,0 +1,37 @@ +# dots.mocr on native CUDA + +The native CUDA engine serves dots.mocr's Qwen2 decoder for text prompts +through the shared dense engine (`zig/src/cuda_dense/`), from the BF16 +`dots-studio/dots.mocr` checkpoint at revision +`e539fbb52280393adc081b289ec597430a0f9031`. It is qualified on a GB10 +(sm_121) only. Images are not served on CUDA yet: `--vision` is refused +by name. The HIP engine, with the image tower, is in +[dots.mocr on native HIP](dots-mocr-hip.md). + +```sh +zig build native -Dnvcc=/usr/local/cuda/bin/nvcc +zig-out/native/bin/tensorfold-native serve MODEL_DIRECTORY \ + --context 8192 --parallel 8 --name dots --host 127.0.0.1 --port 8090 +``` + +Replies carry no reasoning block: all text is answer content. The engine +keeps no prompt states, so `--prompt-cache-gib` is refused; so are +`--drafter`, `--slide`, `--kv-cache fp8` and `--prompt-precision int8`. +The engine's own copy drafts and `--no-drafts` give the same tokens. + +| Responsibility | Files | +|---|---| +| Lanes, pages, attention, draws, the host's declarations | `zig/src/cuda_dense/` | +| Geometry, weights, one layer's arithmetic | `zig/src/families/dots_mocr/cuda_recipe.zig`, `cuda_weights.zig` | +| q, k and v biases with the rotary in Torch's rounding | `zig/kernels/cuda/dots_rope.cuh` (`tf_dots_qkv_rope`) | + +The text reference runs the pinned oracle's `text` mode on any HIP or CUDA +GPU, and `score` classifies every native id by HF's teacher-forced logits +(exact, one BF16 step, or a failure, which exits nonzero): + +```sh +python3 tools/zig/dots_mocr_reference.py text MODEL_DIRECTORY REFERENCE \ + --prompt-file PROMPT.txt --pci PCI_ADDRESS +python3 tools/zig/dots_mocr_reference.py score MODEL_DIRECTORY REFERENCE \ + NATIVE_IDS.json DIFFERENCES.json +``` diff --git a/docs/recipes/dots-mocr-hip.md b/docs/recipes/dots-mocr-hip.md index 221ff743ea7..68d5701c9e9 100644 --- a/docs/recipes/dots-mocr-hip.md +++ b/docs/recipes/dots-mocr-hip.md @@ -3,7 +3,8 @@ This implementation serves the BF16 `dots-studio/dots.mocr` checkpoint at revision `e539fbb52280393adc081b289ec597430a0f9031`. It runs the image tower and Qwen2 decoder in native Zig/HIP. Python is -used only for the validation reference. CUDA and Metal are not qualified. +used only for the validation reference. CUDA serves the text decoder only +([dots.mocr on native CUDA](dots-mocr-cuda.md)); Metal is not qualified. Build for the local GPU architecture. The checked host uses `gfx1201`: From f0ffe921c078b4178f9d14054becb82d35110a11 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 15:28:45 +0000 Subject: [PATCH 07/12] style(dots): move a trailing comment above its line in the CUDA test The declaration test's hook list carried its comment at the end of the line, past the column limit; the comment now sits on the line above it, and nothing else changes. --- zig/src/families/dots_mocr/cuda.zig | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/zig/src/families/dots_mocr/cuda.zig b/zig/src/families/dots_mocr/cuda.zig index 0311c5827dc..837f383434b 100644 --- a/zig/src/families/dots_mocr/cuda.zig +++ b/zig/src/families/dots_mocr/cuda.zig @@ -33,7 +33,8 @@ test "dots text bytes: embedding, untied head, final norm and 28 layers, the tow test "dots on CUDA: text only, no reasoning block, its config and tensor refusals worded" { try std.testing.expect(!native.reasoning and native.keeps_no_prompts); - const hooks = .{ &native.open, &native.attach, &native.fits, &native.ledgerPool, &native.textWeightBytes }; // compiles the load and the layer + // compiles the load and the layer + const hooks = .{ &native.open, &native.attach, &native.fits, &native.ledgerPool, &native.textWeightBytes }; inline for (hooks) |h| try std.testing.expect(@intFromPtr(h) != 0); try std.testing.expectEqualStrings("dots_ocr", native.model_type); try std.testing.expectStringStartsWith(native.explain(null, error.DenseTextOnly).?, "dots.mocr on CUDA reads text prompts only"); From 89f98feb6ea61bc4b06c4474afd516877e652c9b Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 02:12:28 +0000 Subject: [PATCH 08/12] changelog: #2673 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d3f34c0cc3..9449975c4e2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. + - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. - GLM-5.3-Flash on CUDA: with `--vision` and no video library, rank 0 says at startup how much each stream sets aside for `--vision-video-tokens` rows that no request can fill (128 MiB a stream at the default 16,384), and that `--vision-video-tokens 64` on every rank cuts it to 512 KiB (#2813). Thanks to @Bizuayeu. From eed039531e717bfca85c52a35a2c605e3975923f Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 11:48:28 +0000 Subject: [PATCH 09/12] changelog: #2673's line from lines/2673.txt in place of its own --- CHANGELOG.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9449975c4e2..2c5452db10a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,8 +10,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - GLM-5.3-Flash on CUDA: with `--vision` and no video library, rank 0 says at startup how much each stream sets aside for `--vision-video-tokens` rows that no request can fill (128 MiB a stream at the default 16,384), and that `--vision-video-tokens 64` on every rank cuts it to 512 KiB (#2813). Thanks to @Bizuayeu. - - - CUDA: `Qwen/Qwen3-Embedding-0.6B`, `-4B` and `-8B` load on the native engine from the published BF16 weights and serve `/v1/embeddings` (last-token pooling, unit-normed, `dimensions` crops). On one GB10 the worst cosine distance to float32 is 1.55e-4, 1.34e-4 and 1.08e-4 (transformers' own BF16: 1.88e-4, 1.79e-4 and 1.77e-4), a request's vectors equal each input sent alone, and a reranker folder or a flag the family does not serve is refused by name (#2285). Thanks to @Arminova. - Qwen3.8-27B on Metal: the startup line says how many streams draft ("for both streams", "for all 4 streams", "for 8 of 16 streams"), the smaller of the stream count and `TF_Q27_DRAFT_SLOTS`, instead of "for one of them"; log text only (#2806). From df5f50850676be0f1e07f296700ed66317a60643 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 02:33:33 +0000 Subject: [PATCH 10/12] cuda exl3 grouped: launch and size storage at the caller's tile and pair count (#2823) Grouped.run and bytes take the real pair count and the caller's tile (16 or 64; any other tile is refused with TileNotBuilt). maxm is the tile, grid.z is mats * sk * (tile / 16), and x0, x1 and z are spaced by pairs instead of items * 64. Every caller passes its own tile: decode and verify 16, wide prompts 64, MiniMax-M3's prefill 64. Qwen sizes its grouped storage once at max_rows * slots. No kernel source changes. --- zig/src/cuda/exl3_grouped.zig | 126 +++++++++++++++--- zig/src/cuda/experts.zig | 16 +-- zig/src/cuda/experts_host_test.zig | 18 +-- zig/src/families/glm/cuda_experts.zig | 4 +- zig/src/families/kolibri1/forward.zig | 8 +- zig/src/families/lfm2_moe/cuda_forward.zig | 8 +- zig/src/families/minimax_m3/cuda_forward.zig | 11 +- zig/src/families/nemotron/cuda_forward.zig | 33 ++--- zig/src/families/nemotron/cuda_latent.zig | 8 +- zig/src/families/nemotron/cuda_train_fwd.zig | 4 +- .../families/qwen3_5/cuda_exl3_weights.zig | 13 +- zig/src/families/qwen3_5/cuda_routed.zig | 8 +- zig/tests/cuda/exl3_grouped_oai_tests.zig | 2 +- zig/tests/cuda/exl3_grouped_tile_tests.zig | 100 ++++++++++++++ zig/tests/cuda/experts3_tests.zig | 14 +- zig/tests/cuda/fp8_experts_tests.zig | 4 +- zig/tests/cuda/ggml_tests.zig | 12 +- zig/tests/cuda/main.zig | 1 + zig/tests/cuda/modelopt_tests.zig | 8 +- zig/tests/cuda/moe_live_tests.zig | 4 +- zig/tests/cuda/moe_tests.zig | 4 +- zig/tests/cuda/nvfp4_experts_tests.zig | 4 +- zig/tests/cuda/quant_ref_experts.zig | 16 +-- zig/tests/exl3_compile.zig | 6 +- 24 files changed, 314 insertions(+), 118 deletions(-) create mode 100644 zig/tests/cuda/exl3_grouped_tile_tests.zig diff --git a/zig/src/cuda/exl3_grouped.zig b/zig/src/cuda/exl3_grouped.zig index d1abcceda32..9ced70012c7 100644 --- a/zig/src/cuda/exl3_grouped.zig +++ b/zig/src/cuda/exl3_grouped.zig @@ -48,16 +48,17 @@ pub const Grouped = struct { g.storage = null; } - pub fn bytes(k: usize, n: usize, items: usize, sk: usize, mats: usize) !usize { - const p = try std.math.mul(usize, items, 64); - if (p > std.math.maxInt(i32) or k > std.math.maxInt(i32) or n > std.math.maxInt(i32)) + pub fn bytes(k: usize, n: usize, pairs: usize, tile: usize, sk: usize, mats: usize) !usize { + _ = try tileBlocks(tile); + if (pairs > std.math.maxInt(i32) or k > std.math.maxInt(i32) or n > std.math.maxInt(i32)) return error.Exl3Shape; const x = try std.math.mul(usize, k, 4); - const z = try std.math.mul(usize, try std.math.mul(usize, n, sk * mats), 4); - return std.math.mul(usize, p, try std.math.add(usize, x, z)); + const z = try std.math.mul(usize, try std.math.mul(usize, n, try std.math.mul(usize, sk, mats)), 4); + return std.math.mul(usize, pairs, try std.math.add(usize, x, z)); } - pub fn run(g: *Grouped, s: Stream, epi: experts.Epi, in: experts.Rows, l: experts.Layer, p: grouped.Plan, out: u64, items: usize, skip: i32, _: bool) !void { + pub fn run(g: *Grouped, s: Stream, epi: experts.Epi, in: experts.Rows, l: experts.Layer, p: grouped.Plan, out: u64, items: usize, pairs: usize, tile: usize, skip: i32) !void { + const mt = try tileBlocks(tile); if (epi == .up and g.up_contract != .bf16) return error.UnsupportedEpilogue; const w = l.exl3 orelse return error.BadExperts; if (w.dims != l.dims or w.width != l.width or w.experts != l.experts or l.relu2 or @@ -66,17 +67,17 @@ pub const Grouped = struct { const up = epi == .up; const k = if (up) w.dims else w.width; const n = if (up) w.width else w.dims; - if (k == 0 or n == 0 or k % 128 != 0 or n % 128 != 0 or in.stride < k) return error.Exl3Shape; + if (k == 0 or n == 0 or k % 128 != 0 or n % 128 != 0 or in.stride < k or items > std.math.maxInt(i32)) + return error.Exl3Shape; const sk = try ex.plan(k, n, up); const mats: usize = if (up) 2 else 1; - const size = try bytes(k, n, items, sk, mats); - if (items == 0) return; + const size = try bytes(k, n, pairs, tile, sk, mats); + if (items == 0 or pairs == 0) return; if (g.storage == null or g.storage.?.len < size) { const replacement = try memory.DeviceBuffer.alloc(g.d, size); if (g.storage) |*old| old.free(); g.storage = replacement; } - const pairs = items * 64; const x0 = g.storage.?.ptr; const x1 = x0 + pairs * k * 2; const z = x1 + pairs * k * 2; @@ -90,16 +91,25 @@ pub const Grouped = struct { try launch.launch(g.kernels.plan_rotate, .{ .grid = .{ .x = @intCast(items), .y = @intCast(k / 128), .z = @intCast(mats) }, .block = .{ .x = 32 } }, s, &rot); var args: launch.Args = .{}; for ([_]u64{ x0, x1, if (up) w.gate else w.down, if (up) w.up else w.down, if (up) w.gate_k2 else w.down_k2, if (up) w.up_k2 else w.down_k2, p.items, p.counts, p.members, z }) |v| args.add(v); - for ([_]usize{ k, n, pairs, sk, 64, in.slots, items, mats, if (up) w.gu_stride else 0, 1 }) |v| + for ([_]usize{ k, n, pairs, sk, tile, in.slots, items, mats, if (up) w.gu_stride else 0, 1 }) |v| args.add(@as(i32, @intCast(v))); args.add(skip); - try launch.launch(g.kernels.grouped[@backingInt(w.codebook)], .{ .grid = .{ .x = @intCast(items), .y = @intCast(n / 128), .z = @intCast(mats * sk * 4) }, .block = .{ .x = 128 } }, s, &args); + try launch.launch(g.kernels.grouped[@backingInt(w.codebook)], .{ .grid = .{ .x = @intCast(items), .y = @intCast(n / 128), .z = @intCast(mats * sk * mt) }, .block = .{ .x = 128 } }, s, &args); var finish: launch.Args = .{}; finish.add(a); try launch.launch(if (l.after) g.kernels.plan_finish_after else g.kernels.plan_finish, .{ .grid = .{ .x = @intCast(items), .y = @intCast(n / 128) }, .block = .{ .x = 32 } }, s, &finish); } }; +/// Member-tile blocks a grouped launch spans: 16-member decode items or 64-member prompt items. +fn tileBlocks(tile: usize) !usize { + return switch (tile) { + 16 => 1, + 64 => 4, + else => error.TileNotBuilt, + }; +} + /// The activations plan_finish runs (0 bf16, 1 fp32, 2 SwiGLU-OAI); SwiGLU-OAI needs its clamp and alpha. fn checkAct(w: ex.Layer, after: bool) !void { if (w.act_mode < 0 or w.act_mode > 2) return error.BadExperts; @@ -135,16 +145,96 @@ test "Grouped takes the clamp-after order at fp32 activations only: not at bf16, try std.testing.expectError(error.OptionNotBuilt, checkAct(w, true)); } -test "Grouped scratch includes every pair and rejects overflowing geometry" { - try std.testing.expectEqual(@as(usize, 64 * (128 * 4 + 256 * 4 * 2 * 4)), try Grouped.bytes(128, 256, 1, 4, 2)); - try std.testing.expectError(error.Exl3Shape, Grouped.bytes(std.math.maxInt(usize), 256, 1, 4, 2)); - try std.testing.expectError(error.Exl3Shape, Grouped.bytes(128, 256, std.math.maxInt(i32), 4, 2)); +test "Grouped scratch scales linearly with pairs, refuses unsupported tiles, and rejects overflowing geometry" { + const one = try Grouped.bytes(2048, 512, 1, 16, 4, 2); + try std.testing.expectEqual(@as(usize, 2048 * 4 + 512 * 4 * 2 * 4), one); + try std.testing.expectEqual(144 * one, try Grouped.bytes(2048, 512, 144, 16, 4, 2)); + try std.testing.expectEqual(144 * one, try Grouped.bytes(2048, 512, 144, 64, 4, 2)); + for ([_]usize{ 0, 8, 15, 17, 32, 48, 128 }) |bad| { + try std.testing.expectError(error.TileNotBuilt, Grouped.bytes(2048, 512, 144, bad, 4, 2)); + } + try std.testing.expectError(error.Exl3Shape, Grouped.bytes(std.math.maxInt(usize), 256, 1, 16, 4, 2)); + try std.testing.expectError(error.Exl3Shape, Grouped.bytes(128, 256, @as(usize, std.math.maxInt(i32)) + 1, 16, 4, 2)); } test "Grouped's bf16 up refuses a call that asks for the combined fp32 activation pipeline" { var g: Grouped = .{ .d = undefined, .kernels = undefined, .up_contract = .combined_fp32 }; const layer: experts.Layer = .{ .format = .exl3, .up = 0, .down = 0, .dims = 128, .width = 128, .experts = 1 }; - try std.testing.expectError(error.UnsupportedEpilogue, g.run(undefined, .up, undefined, layer, undefined, 0, 0, -1, false)); + for ([_]usize{ 0, 8, 32, 128 }) |bad| { + try std.testing.expectError(error.TileNotBuilt, g.run(undefined, .up, undefined, layer, undefined, 0, 1, 1, bad, -1)); + } + try std.testing.expectError(error.UnsupportedEpilogue, g.run(undefined, .up, undefined, layer, undefined, 0, 0, 0, 16, -1)); g.up_contract = .bf16; - try std.testing.expectError(error.BadExperts, g.run(undefined, .up, undefined, layer, undefined, 0, 0, -1, false)); + try std.testing.expectError(error.BadExperts, g.run(undefined, .up, undefined, layer, undefined, 0, 0, 0, 16, -1)); +} + +test "Grouped.run launches the plan tile's grid and compact pair stride, dropping no member at tile 64" { + const abi = @import("abi.zig"); + const Trace = struct { + var n: usize = 0; + var grids: [3][3]u32 = undefined; + var rot: Stage = undefined; + var fin: Stage = undefined; + var ptrs: [10]u64 = undefined; + var ints: [11]i32 = undefined; + fn launchFn(_: abi.Function, gx: c_uint, gy: c_uint, gz: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) abi.Result { + grids[n] = .{ gx, gy, gz }; + if (n == 0) rot = @as(*const Stage, @ptrCast(@alignCast(args.?[0].?))).*; + if (n == 1) { + for (&ptrs, 0..) |*p, i| p.* = @as(*const u64, @ptrCast(@alignCast(args.?[i].?))).*; + for (&ints, 0..) |*v, i| v.* = @as(*const i32, @ptrCast(@alignCast(args.?[10 + i].?))).*; + } + if (n == 2) fin = @as(*const Stage, @ptrCast(@alignCast(args.?[0].?))).*; + n += 1; + return 0; + } + }; + var d: Driver = undefined; + d.api.cuLaunchKernel = Trace.launchFn; + var ks: ex.Experts = undefined; + ks.plan_rotate = .{ .d = &d, .handle = @ptrFromInt(1) }; + ks.plan_finish = .{ .d = &d, .handle = @ptrFromInt(2) }; + ks.plan_finish_after = .{ .d = &d, .handle = @ptrFromInt(3) }; + for (&ks.grouped, 0..) |*f, i| f.* = .{ .d = &d, .handle = @ptrFromInt(10 + i) }; + const base: u64 = 0x4000_0000; + var g: Grouped = .{ .d = &d, .kernels = &ks, .storage = .{ .d = &d, .ptr = base, .len = 1 << 28 } }; + const w: ex.Layer = .{ .gate = 0x10, .up = 0x20, .down = 0x30, .gate_k2 = 0x40, .up_k2 = 0x50, .down_k2 = 0x60, .suh_g = 1, .suh_u = 2, .suh_d = 3, .svh_g = 4, .svh_u = 5, .svh_d = 6, .dims = 2048, .width = 512, .experts = 257, .codebook = .mul1 }; + const l: experts.Layer = .{ .format = .exl3, .exl3 = w, .up = 0, .down = 0, .dims = 2048, .width = 512, .experts = 257 }; + const pl: grouped.Plan = .{ .members = 0x100, .items = 0x200, .counts = 0x300, .rank = 0, .hist = 0 }; + const s: Stream = .{ .d = &d, .handle = null }; + const grp: experts.Grouped = .{ .exl3 = &g }; + for ([_]struct { tile: usize, pairs: usize, items: usize, epi: experts.Epi, prompt: bool }{ + .{ .tile = 16, .pairs = 9, .items = 9, .epi = .up, .prompt = false }, + .{ .tile = 16, .pairs = 144, .items = 153, .epi = .down_f32, .prompt = false }, + .{ .tile = 64, .pairs = 512, .items = 136, .epi = .up, .prompt = false }, + .{ .tile = 64, .pairs = 576, .items = 266, .epi = .down_bf16, .prompt = true }, + }) |c| { + Trace.n = 0; + const up = c.epi == .up; + const k: usize = if (up) 2048 else 512; + const n: usize = if (up) 512 else 2048; + const sk: usize = try ex.plan(k, n, up); + const mats: usize = if (up) 2 else 1; + const in: experts.Rows = .{ .x = 0x9000, .stride = k, .slots = if (up) 9 else 0 }; + if (c.prompt) try grp.prompt(s, c.epi, in, l, pl, 0xa000, c.items, c.pairs, c.tile, -1) else try grp.decode(s, c.epi, in, l, pl, 0xa000, c.items, c.pairs, c.tile, -1); + try std.testing.expectEqual(@as(usize, 3), Trace.n); + try std.testing.expectEqual(@as(u32, @intCast(mats * sk * (c.tile / 16))), Trace.grids[1][2]); + try std.testing.expectEqual(@as(i32, @intCast(c.pairs)), Trace.rot.p); + try std.testing.expectEqual(@as(i32, @intCast(c.pairs)), Trace.fin.p); + try std.testing.expectEqual(@as(i32, @intCast(c.pairs)), Trace.ints[2]); + try std.testing.expectEqual(@as(i32, @intCast(c.tile)), Trace.ints[4]); + try std.testing.expectEqual(base + c.pairs * k * 2, Trace.rot.x1); + try std.testing.expectEqual(base + c.pairs * k * 4, Trace.rot.z); + try std.testing.expectEqual(Trace.rot.x1, Trace.ptrs[1]); + try std.testing.expectEqual(Trace.rot.z, Trace.ptrs[9]); + if (c.tile == 64 and up) { + var seen: [64]usize = @splat(0); + const mt: u32 = @intCast((@as(usize, @intCast(Trace.ints[4])) + 15) / 16); + for (0..Trace.grids[1][2]) |zz| { + const mtile = zz % mt; + for (0..16) |tx| seen[mtile * 16 + tx] += 1; + } + for (seen) |hits| try std.testing.expectEqual(mats * sk, hits); + } + } } diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 3c5828bea36..3d70527c1d8 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -53,11 +53,11 @@ pub const Grouped = struct { nvfp4_csf0: ?*const @import("experts32.zig").Nvfp4b32Experts = null, iq2xxs: ?*const Iq2xxsExperts = null, - /// Decode rows: `items` bounds the plan's items (grouped.maxItems at tile 16); `skip` an expert left out, or -1. - pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { + /// Decode rows: `items` bounds the plan, `pairs` its picks; `tile` 16 or 64 (EXL3 only); `skip` an expert or -1. + pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, pairs: usize, tile: usize, skip: i32) !void { if (l.after and l.format != .exl3) return error.OptionNotBuilt; // the clamp-after order is built for EXL3 only switch (l.format) { - .exl3 => return (g.exl3 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip, false), + .exl3 => return (g.exl3 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, pairs, tile, skip), .affine4 => { const a = g.affine4 orelse return error.FormatNotLoaded; if (epi == .down_bf16) return error.EpilogueNotBuilt; @@ -80,11 +80,11 @@ pub const Grouped = struct { } } - /// Prompt rows: `items` bounds the plan's items at the format's prompt tile (prompt_tile). - pub fn prompt(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { + /// Prompt rows: `items` bounds the plan, `pairs` its picks; `tile` 16 or 64 (EXL3 only); `skip` an expert or -1. + pub fn prompt(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, pairs: usize, tile: usize, skip: i32) !void { if (l.after and l.format != .exl3) return error.OptionNotBuilt; // the clamp-after order is built for EXL3 only switch (l.format) { - .exl3 => return (g.exl3 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip, true), + .exl3 => return (g.exl3 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, pairs, tile, skip), .affine4 => { const a = g.affine4 orelse return error.FormatNotLoaded; if (epi == .down_f32) return error.EpilogueNotBuilt; @@ -556,8 +556,8 @@ test "nvfp4 rows of either kind are refused when the format is not loaded" { const l: Layer = .{ .format = .nvfp4, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .relu2 = true }; const s: Stream = undefined; // never reached: the refusal comes first const p: Plan = .{ .members = 0, .items = 0, .counts = 0, .rank = 0, .hist = 0 }; - try std.testing.expectError(error.FormatNotLoaded, g.prompt(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, -1)); - try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, l, p, 0, 1, -1)); + try std.testing.expectError(error.FormatNotLoaded, g.prompt(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, 1, 64, -1)); + try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, l, p, 0, 1, 1, 16, -1)); try std.testing.expectEqual(@as(usize, 2 * 2 * 3 * 144 * 4), Nvfp4Experts.bytes(2, 64, 96, 1)); try std.testing.expectEqual(@as(u32, 2 * 17408), Nvfp4Experts.promptSmem(1, 8)); // 64 rows x 64 inputs, 8 x 2 blocks try std.testing.expectEqual(@as(u32, 2 * 17408), Nvfp4Experts.promptSmem(2, 4)); diff --git a/zig/src/cuda/experts_host_test.zig b/zig/src/cuda/experts_host_test.zig index 1b645b64f84..635eea49d2f 100644 --- a/zig/src/cuda/experts_host_test.zig +++ b/zig/src/cuda/experts_host_test.zig @@ -12,15 +12,15 @@ test "a format that is not loaded, or an option it was not built with, is refuse const l: Layer = .{ .format = .fp8g, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2 }; const s: Stream = undefined; // never reached: the refusal comes first const p: Plan = .{ .members = 0, .items = 0, .counts = 0, .rank = 0, .hist = 0 }; - try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, -1)); + try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, 1, 16, -1)); const a4: experts.Affine4Experts = undefined; const h: Grouped = .{ .affine4 = &a4 }; var la = l; la.format = .affine4; - try std.testing.expectError(error.EpilogueNotBuilt, h.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1)); - try std.testing.expectError(error.EpilogueNotBuilt, h.prompt(s, .down_f32, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1)); - try std.testing.expectError(error.OptionNotBuilt, h.decode(s, .up, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 3)); - try std.testing.expectError(error.OptionNotBuilt, h.prompt(s, .up, .{ .x = 0, .stride = 64 }, .{ .format = .affine4, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .after = true }, p, 0, 1, -1)); + try std.testing.expectError(error.EpilogueNotBuilt, h.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 1, 16, -1)); + try std.testing.expectError(error.EpilogueNotBuilt, h.prompt(s, .down_f32, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 1, 64, -1)); + try std.testing.expectError(error.OptionNotBuilt, h.decode(s, .up, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 1, 16, 3)); + try std.testing.expectError(error.OptionNotBuilt, h.prompt(s, .up, .{ .x = 0, .stride = 64 }, .{ .format = .affine4, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .after = true }, p, 0, 1, 1, 64, -1)); } test "the clamp-after order: exl3 takes it on both paths, every other format refuses it on both" { @@ -30,11 +30,11 @@ test "the clamp-after order: exl3 takes it on both paths, every other format ref const in: experts.Rows = .{ .x = 0, .stride = 64 }; var l: Layer = .{ .format = .exl3, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .limit = 10, .after = true }; // exl3 accepts the order, then finds its kernels not loaded - try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, in, l, p, 0, 1, -1)); - try std.testing.expectError(error.FormatNotLoaded, g.prompt(s, .up, in, l, p, 0, 1, -1)); + try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, in, l, p, 0, 1, 1, 16, -1)); + try std.testing.expectError(error.FormatNotLoaded, g.prompt(s, .up, in, l, p, 0, 1, 1, 64, -1)); for (std.enums.values(experts.Format)) |f| if (f != .exl3) { l.format = f; // nvfp4 among them; not loaded either, so the refusal comes before the format's - try std.testing.expectError(error.OptionNotBuilt, g.decode(s, .up, in, l, p, 0, 1, -1)); - try std.testing.expectError(error.OptionNotBuilt, g.prompt(s, .up, in, l, p, 0, 1, -1)); + try std.testing.expectError(error.OptionNotBuilt, g.decode(s, .up, in, l, p, 0, 1, 1, 16, -1)); + try std.testing.expectError(error.OptionNotBuilt, g.prompt(s, .up, in, l, p, 0, 1, 1, 64, -1)); }; } diff --git a/zig/src/families/glm/cuda_experts.zig b/zig/src/families/glm/cuda_experts.zig index 8eca3815e3b..48f4c2b5c71 100644 --- a/zig/src/families/glm/cuda_experts.zig +++ b/zig/src/families/glm/cuda_experts.zig @@ -60,13 +60,13 @@ pub fn run(ks: *const Kernels, stream: cuda.Stream, which: Which, l: cuda.expert }; const in: cuda.experts.Rows = .{ .x = x, .stride = x_stride, .slots = plan.slots }; const pairs = rows * plan.slots; - if (plan.prompt) return g.prompt(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, staged_tile), skip); + if (plan.prompt) return g.prompt(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, staged_tile), pairs, staged_tile, skip); if (plan.fact) { // gate-up reads two matrices a K tile and gains from the factored tiles; down is short and keeps the waves if (epi == .up and !l.relu2) return cuda.experts.Nvfp4Experts.runFact(&ks.nvfp4, stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, staged_tile), skip); return cuda.experts_waves.run(&ks.nvfp4, stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, tile), skip); } - return g.decode(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, tile), skip); + return g.decode(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, tile), pairs, tile, skip); } test "block 32's kernels are wanted for a csf0_v1 module alone, not for block 16's recipes or a single-format NVFP4" { diff --git a/zig/src/families/kolibri1/forward.zig b/zig/src/families/kolibri1/forward.zig index dcfa7d3abb8..28b504135cc 100644 --- a/zig/src/families/kolibri1/forward.zig +++ b/zig/src/families/kolibri1/forward.zig @@ -464,11 +464,11 @@ pub const Model = struct { const up_in: experts.Rows = .{ .x = m.s.x, .stride = c.hidden, .slots = slots_e }; const down_in: experts.Rows = .{ .x = m.s.act, .stride = c.moe_width }; if (prompt) { - try gx.prompt(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, -1); - try gx.prompt(s, .down_bf16, down_in, L.experts, m.s.plan, m.s.y, items, -1); + try gx.prompt(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, pairs, tile, -1); + try gx.prompt(s, .down_bf16, down_in, L.experts, m.s.plan, m.s.y, items, pairs, tile, -1); } else { - try gx.decode(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, -1); - try gx.decode(s, .down_f32, down_in, L.experts, m.s.plan, m.s.y, items, -1); + try gx.decode(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, pairs, tile, -1); + try gx.decode(s, .down_f32, down_in, L.experts, m.s.plan, m.s.y, items, pairs, tile, -1); } try o.combine(m.s.y, prompt, m.s.wts, m.s.mo, n, slots_e, c.hidden); const after = if (i + 1 < w.layers.len) w.layers[i + 1].input_norm else w.norm; diff --git a/zig/src/families/lfm2_moe/cuda_forward.zig b/zig/src/families/lfm2_moe/cuda_forward.zig index 6c2b7ef3577..b0ac16a93af 100644 --- a/zig/src/families/lfm2_moe/cuda_forward.zig +++ b/zig/src/families/lfm2_moe/cuda_forward.zig @@ -133,11 +133,11 @@ const Pass = struct { const rows_in: cuda.experts.Rows = .{ .x = x.x, .stride = c.hidden, .slots = c.topk }; const act_in: cuda.experts.Rows = .{ .x = x.eact, .stride = c.moe_intermediate }; if (p.prompt) { - try grouped.prompt(p.s, .up, rows_in, e.experts, plan, x.eact, items, -1); - try grouped.prompt(p.s, .down_bf16, act_in, e.experts, plan, x.ey, items, -1); + try grouped.prompt(p.s, .up, rows_in, e.experts, plan, x.eact, items, pairs, tile, -1); + try grouped.prompt(p.s, .down_bf16, act_in, e.experts, plan, x.ey, items, pairs, tile, -1); } else { - try grouped.decode(p.s, .up, rows_in, e.experts, plan, x.eact, items, -1); - try grouped.decode(p.s, .down_f32, act_in, e.experts, plan, x.ey, items, -1); + try grouped.decode(p.s, .up, rows_in, e.experts, plan, x.eact, items, pairs, tile, -1); + try grouped.decode(p.s, .down_f32, act_in, e.experts, plan, x.ey, items, pairs, tile, -1); } try k.moe.combine(p.s, x.ey, !p.prompt, x.wts, x.r, p.rows, c.topk, c.hidden, 0); } diff --git a/zig/src/families/minimax_m3/cuda_forward.zig b/zig/src/families/minimax_m3/cuda_forward.zig index 1f1298cd30c..42fbbb2a084 100644 --- a/zig/src/families/minimax_m3/cuda_forward.zig +++ b/zig/src/families/minimax_m3/cuda_forward.zig @@ -61,10 +61,11 @@ pub const Modules = struct { /// The grouped experts' scratch for a window of `rows` (allocated before serving, never mid-forward). pub fn reserve(m: *const Modules, c: *const cfg.Config, share: anytype, rows: usize, prefill: bool) !void { - const items = cuda.grouped.maxItems(rows * c.topk, c.experts, if (prefill) 64 else 16); + const pairs = rows * c.topk; + const tile: usize = if (prefill) 64 else 16; const ni = share.moe.len(); - const up = try cuda.exl3_grouped.Grouped.bytes(c.hidden, ni, items, try cuda.exl3_experts.plan(c.hidden, ni, true), 2); - const down = try cuda.exl3_grouped.Grouped.bytes(ni, c.hidden, items, try cuda.exl3_experts.plan(ni, c.hidden, false), 1); + const up = try cuda.exl3_grouped.Grouped.bytes(c.hidden, ni, pairs, tile, try cuda.exl3_experts.plan(c.hidden, ni, true), 2); + const down = try cuda.exl3_grouped.Grouped.bytes(ni, c.hidden, pairs, tile, try cuda.exl3_experts.plan(ni, c.hidden, false), 1); const need = @max(up, down); const g = &m.x3.grouped; if (g.storage != null and g.storage.?.len >= need) return; @@ -243,8 +244,8 @@ pub const Forward = struct { const in_up: cuda.experts.Rows = .{ .x = b.x, .stride = h, .slots = topk }; const in_down: cuda.experts.Rows = .{ .x = b.eact, .stride = m.routed.width }; // one path for a prompt chunk and a window: every pick's down rows in fp32, so a row's bits never depend on it - try g.decode(f.stream, .up, in_up, m.routed, b.plan(), b.eact, b.max_items, -1); - try g.decode(f.stream, .down_f32, in_down, m.routed, b.plan(), b.ey, b.max_items, -1); + try g.decode(f.stream, .up, in_up, m.routed, b.plan(), b.eact, b.max_items, pairs, tile, -1); + try g.decode(f.stream, .down_f32, in_down, m.routed, b.plan(), b.ey, b.max_items, pairs, tile, -1); try f.mlp(m.shared, f.w.share.shared.len(), R, b.sy, .f32); try kern.moePartial(&f.m.k, f.stream, b.ey, b.wts, b.sy, f.slot(R, 2), R, topk, h); } diff --git a/zig/src/families/nemotron/cuda_forward.zig b/zig/src/families/nemotron/cuda_forward.zig index 02b01c2c5b5..0bf7a6cc629 100644 --- a/zig/src/families/nemotron/cuda_forward.zig +++ b/zig/src/families/nemotron/cuda_forward.zig @@ -274,13 +274,13 @@ pub const Forward = struct { const g = o.grouped(); const in = f.rowsIn(ex); if (prompt) { - try g.prompt(o.s, .up, in[0], ex, b.plan, b.act, items, -1); - try g.prompt(o.s, .down_bf16, in[1], ex, b.plan, b.ymoe, items, -1); + try g.prompt(o.s, .up, in[0], ex, b.plan, b.act, items, pairs, tile, -1); + try g.prompt(o.s, .down_bf16, in[1], ex, b.plan, b.ymoe, items, pairs, tile, -1); if (f.dump) |d| for ([_][]const u8{ "pick", "act", "ymoe" }, [_]u64{ b.pick, b.act, b.ymoe }, [_]usize{ 4, ex.width * 2, ex.dims * 2 }) |l, p, w| try d.input(o, l, p, pairs * w); } else { - try g.decode(o.s, .up, in[0], ex, b.plan, b.act, items, -1); + try g.decode(o.s, .up, in[0], ex, b.plan, b.act, items, pairs, tile, -1); try f.mark(.experts_up); - try g.decode(o.s, .down_f32, in[1], ex, b.plan, b.ymoe, items, -1); + try g.decode(o.s, .down_f32, in[1], ex, b.plan, b.ymoe, items, pairs, tile, -1); try f.mark(.experts_down); } try f.sharedRest(o, m, rows, !prompt); @@ -305,12 +305,12 @@ pub const Forward = struct { const up_in: cuda.experts.Rows = .{ .x = b.lat, .stride = c.latent, .slots = K }; const down_in: cuda.experts.Rows = .{ .x = b.act, .stride = ex.width }; if (prompt) { - try g.prompt(o.s, .up, up_in, ex, b.plan, b.act, items, -1); - try g.prompt(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, -1); + try g.prompt(o.s, .up, up_in, ex, b.plan, b.act, items, pairs, tile, -1); + try g.prompt(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, pairs, tile, -1); } else { - try g.decode(o.s, .up, up_in, ex, b.plan, b.act, items, -1); + try g.decode(o.s, .up, up_in, ex, b.plan, b.act, items, pairs, tile, -1); try f.mark(.experts_up); - try g.decode(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, -1); + try g.decode(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, pairs, tile, -1); try f.mark(.experts_down); } try o.latentSum(b.ymoe, b.wts, b.comb, rows, K, c.latent); @@ -322,11 +322,11 @@ pub const Forward = struct { const sh_up: cuda.experts.Rows = .{ .x = b.y, .stride = c.hidden, .slots = 1 }; const sh_down: cuda.experts.Rows = .{ .x = b.act, .stride = sh.width }; if (prompt) { - try g.prompt(o.s, .up, sh_up, sh, b.plan, b.act, items, -1); - try g.prompt(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, -1); + try g.prompt(o.s, .up, sh_up, sh, b.plan, b.act, items, rows, tile, -1); + try g.prompt(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, rows, tile, -1); } else { - try g.decode(o.s, .up, sh_up, sh, b.plan, b.act, items, -1); - try g.decode(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, -1); + try g.decode(o.s, .up, sh_up, sh, b.plan, b.act, items, rows, tile, -1); + try g.decode(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, rows, tile, -1); } try o.addBf16(b.delta, b.delta, b.ymoe, true, rows * c.hidden); try f.mark(.shared_wait); @@ -373,8 +373,9 @@ pub const Forward = struct { try sd.fork.record(o.s); try sd.s.wait(sd.fork); const sp = b.sharedPlan(rows); - try g.decode(sd.s, .up, in[0], ex, sp, b.act, 2, -1); - try g.decode(sd.s, .down_f32, in[1], ex, sp, b.ymoe, 2, -1); + const pairs = rows * c.slots(); + try g.decode(sd.s, .up, in[0], ex, sp, b.act, 2, pairs, 16, -1); + try g.decode(sd.s, .down_f32, in[1], ex, sp, b.ymoe, 2, pairs, 16, -1); try f.sharedRest(so, m, rows, true); try sd.join.record(sd.s); try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.slots(), c.routed_scaling, c.norm_topk); @@ -382,9 +383,9 @@ pub const Forward = struct { try o.planRouted(b.pick, rows, c.slots(), c.top_k, c.experts, 16, b.plan); try f.mark(.plan); const items = kern.maxItems(rows * c.top_k, c.experts, 16); - try g.decode(o.s, .up, in[0], ex, b.plan, b.act, items, -1); + try g.decode(o.s, .up, in[0], ex, b.plan, b.act, items, pairs, 16, -1); try f.mark(.experts_up); - try g.decode(o.s, .down_f32, in[1], ex, b.plan, b.ymoe, items, -1); + try g.decode(o.s, .down_f32, in[1], ex, b.plan, b.ymoe, items, pairs, 16, -1); try f.mark(.experts_down); try o.s.wait(sd.join); try f.mark(.shared_wait); diff --git a/zig/src/families/nemotron/cuda_latent.zig b/zig/src/families/nemotron/cuda_latent.zig index ffdb4ded492..a119d6cbf5b 100644 --- a/zig/src/families/nemotron/cuda_latent.zig +++ b/zig/src/families/nemotron/cuda_latent.zig @@ -339,8 +339,8 @@ fn routed(o: kern.Ops, t: Tables, s: Scratch, picks: u64, lat: u64, rows: usize, const g = o.grouped(); try o.plan(picks, pairs, t.routed.experts, 16, s.plan); const items = kern.maxItems(pairs, t.routed.experts, 16); - try g.decode(o.s, .up, .{ .x = lat, .stride = c.latent, .slots = c.top_k * own }, t.routed, s.plan, s.act, items, -1); - try g.decode(o.s, .down_f32, .{ .x = s.act, .stride = t.routed.width }, t.routed, s.plan, out, items, -1); + try g.decode(o.s, .up, .{ .x = lat, .stride = c.latent, .slots = c.top_k * own }, t.routed, s.plan, s.act, items, pairs, 16, -1); + try g.decode(o.s, .down_f32, .{ .x = s.act, .stride = t.routed.width }, t.routed, s.plan, out, items, pairs, 16, -1); } /// The shared expert's (row, share) pairs into `out` (partials' second region), on `s`'s fixed plan where it has one. @@ -350,8 +350,8 @@ fn sharedRows(o: kern.Ops, t: Tables, s: Scratch, hidden: u64, rows: usize, own: const sp = s.shared_plan orelse s.plan; if (s.shared_plan == null) try o.plan(s.sids, rows * own, own, 16, sp); const sh_items = kern.maxItems(rows * own, own, 16); - try g.decode(o.s, .up, .{ .x = hidden, .stride = c.hidden, .slots = own }, t.shared, sp, s.sact, sh_items, -1); - try g.decode(o.s, .down_f32, .{ .x = s.sact, .stride = t.shared.width }, t.shared, sp, out, sh_items, -1); + try g.decode(o.s, .up, .{ .x = hidden, .stride = c.hidden, .slots = own }, t.shared, sp, s.sact, sh_items, rows * own, 16, -1); + try g.decode(o.s, .down_f32, .{ .x = s.sact, .stride = t.shared.width }, t.shared, sp, out, sh_items, rows * own, 16, -1); } /// The MoE block's delta into b.delta: the other ranks' shares asked for, ours run, all summed in share order. diff --git a/zig/src/families/nemotron/cuda_train_fwd.zig b/zig/src/families/nemotron/cuda_train_fwd.zig index cb4ad780e00..7c0c4d30449 100644 --- a/zig/src/families/nemotron/cuda_train_fwd.zig +++ b/zig/src/families/nemotron/cuda_train_fwd.zig @@ -92,8 +92,8 @@ pub fn experts(f: *const fwd.Forward, m: weights.MoE, rows: usize, scale: Scale, const items = kern.maxItems(pairs, ex.experts, tile); const g = o.grouped(); const in = f.rowsIn(ex); - try g.prompt(o.s, .up, in[0], ex, b.plan, b.act, items, -1); - try g.prompt(o.s, .down_bf16, in[1], ex, b.plan, b.ymoe, items, -1); + try g.prompt(o.s, .up, in[0], ex, b.plan, b.act, items, pairs, tile, -1); + try g.prompt(o.s, .down_bf16, in[1], ex, b.plan, b.ymoe, items, pairs, tile, -1); for (m.rest, 0..) |r, h| if (r != 0) try o.restRows(b.act, ex.width, r, b.ymoe, ex.dims, false, rows, c.slots(), c.top_k + h, ex.dims, ex.width); const t: train_ops.Ops = .{ .k = &o.k.train, .s = o.s }; if (m.pocket) |p| try t.pocketRows(b.y, ex.dims, p.up, p.down, p.n, b.ymoe, false, rows, c.slots(), c.top_k, scale.rows, scale.from[layer]); diff --git a/zig/src/families/qwen3_5/cuda_exl3_weights.zig b/zig/src/families/qwen3_5/cuda_exl3_weights.zig index bf1f1c6b1d5..8e001a8ddde 100644 --- a/zig/src/families/qwen3_5/cuda_exl3_weights.zig +++ b/zig/src/families/qwen3_5/cuda_exl3_weights.zig @@ -33,14 +33,17 @@ pub const Kit = struct { k.grouped = .{ .d = d, .kernels = &k.experts }; if (!g.moe()) return k; // a dense checkpoint's MLP is three Dense projections // the grouped experts' scratch grows on first use: allocate it before any graph capture, for the widest call - const pairs = engine.max_rows * g.slots(); - const items = cuda.grouped.maxItems(pairs, g.experts + 1, 16); - const up = try cuda.exl3_grouped.Grouped.bytes(g.hidden, g.intermediate, items, try ex.plan(g.hidden, g.intermediate, true), 2); - const down = try cuda.exl3_grouped.Grouped.bytes(g.intermediate, g.hidden, items, try ex.plan(g.intermediate, g.hidden, false), 1); - k.grouped.storage = try cuda.DeviceBuffer.alloc(d, @max(up, down)); + k.grouped.storage = try cuda.DeviceBuffer.alloc(d, try storageBytes(g)); return k; } + pub fn storageBytes(g: geometry.Geometry) !usize { + const pairs = engine.max_rows * g.slots(); + const up = try cuda.exl3_grouped.Grouped.bytes(g.hidden, g.intermediate, pairs, 64, try ex.plan(g.hidden, g.intermediate, true), 2); + const down = try cuda.exl3_grouped.Grouped.bytes(g.intermediate, g.hidden, pairs, 64, try ex.plan(g.intermediate, g.hidden, false), 1); + return @max(up, down); + } + pub fn deinit(k: *Kit, gpa: std.mem.Allocator) void { k.grouped.deinit(); k.b16.unload(); diff --git a/zig/src/families/qwen3_5/cuda_routed.zig b/zig/src/families/qwen3_5/cuda_routed.zig index 79c31dcaa16..636d1bfc673 100644 --- a/zig/src/families/qwen3_5/cuda_routed.zig +++ b/zig/src/families/qwen3_5/cuda_routed.zig @@ -57,11 +57,11 @@ pub fn rows(e: *E.Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u try launch(k, e.stream, .up, rows_in, layer, plan, s.act, items, -1); try launch(k, e.stream, .down_f32, act_in, layer, plan, s.y, items, -1); } else if (wide) { - try grouped.prompt(e.stream, .up, rows_in, layer, plan, s.act, items, -1); - try grouped.prompt(e.stream, .down_bf16, act_in, layer, plan, s.y, items, -1); + try grouped.prompt(e.stream, .up, rows_in, layer, plan, s.act, items, pairs, tile, -1); + try grouped.prompt(e.stream, .down_bf16, act_in, layer, plan, s.y, items, pairs, tile, -1); } else { - try grouped.decode(e.stream, .up, rows_in, layer, plan, s.act, items, -1); - try grouped.decode(e.stream, .down_f32, act_in, layer, plan, s.y, items, -1); + try grouped.decode(e.stream, .up, rows_in, layer, plan, s.act, items, pairs, tile, -1); + try grouped.decode(e.stream, .down_f32, act_in, layer, plan, s.y, items, pairs, tile, -1); } try e.moe.combine(e.stream, s.y, !wide or nvfp4, s.wts, out, total, slots, g.hidden, params); } diff --git a/zig/tests/cuda/exl3_grouped_oai_tests.zig b/zig/tests/cuda/exl3_grouped_oai_tests.zig index f08573d4885..063e27c416d 100644 --- a/zig/tests/cuda/exl3_grouped_oai_tests.zig +++ b/zig/tests/cuda/exl3_grouped_oai_tests.zig @@ -76,7 +76,7 @@ fn up(r: *Rig, k: *Kit, l: cuda.experts.Layer, x: u64, pick: u64, m: usize, tile }; try k.router.route(r.s, pick, pairs, E, tile, plan); const out = try r.zeros(pairs * I * 2); - try k.g.run(r.s, .up, .{ .x = x, .stride = D, .slots = SLOTS }, l, plan, out, items, -1, false); + try k.g.run(r.s, .up, .{ .x = x, .stride = D, .slots = SLOTS }, l, plan, out, items, pairs, tile, -1); return r.back(out, pairs * I * 2); } diff --git a/zig/tests/cuda/exl3_grouped_tile_tests.zig b/zig/tests/cuda/exl3_grouped_tile_tests.zig new file mode 100644 index 00000000000..39bb79cc44a --- /dev/null +++ b/zig/tests/cuda/exl3_grouped_tile_tests.zig @@ -0,0 +1,100 @@ +//! Grouped EXL3: maxm 16 and the old maxm 64 are bit-equal; maxm 16 on a tile-64 plan drops members. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const quant_ref_tests = @import("quant_ref_tests.zig"); +const f64t = @import("exl3_experts_f64_tests.zig"); + +const Rig = quant_ref_tests.RigT; +const E = f64t.E; +const D = f64t.D; +const I = f64t.I; +const SLOTS = f64t.SLOTS; + +const Plan = struct { p: cuda.grouped.Plan, items: usize, pairs: usize }; +const Out = struct { dev: u64, bytes: []u8 }; + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(61); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try cuda.grouped.Router.resolve(mod); + var exp = try cuda.exl3_experts.Experts.load(gpu.d); + defer exp.unload(); + const b = try f64t.buildLayer(&r, &.{ 3, 4, 5, 6, 3, 4, 5, 6, 3 }); + const l: cuda.experts.Layer = .{ .format = .exl3, .exl3 = b.layer, .up = 0, .down = 0, .width = I, .dims = D, .experts = E }; + var g: cuda.exl3_grouped.Grouped = .{ .d = gpu.d, .kernels = &exp }; + defer g.deinit(); + const grp: cuda.experts.Grouped = .{ .exl3 = &g }; + + for ([_]usize{ 1, 4, 16 }) |rows| { + const pairs = rows * SLOTS; + const p16 = try plan(&r, router, b.dpick, pairs, 16); + const p64 = try plan(&r, router, b.dpick, pairs, 64); + const up16 = try upBytes(&r, grp, l, b.dx, p16, 16); + const upOld = try upBytes(&r, grp, l, b.dx, p16, 64); + try check.sameBytes("up: maxm 16 vs the old maxm 64 on a tile-16 plan", upOld.bytes, up16.bytes); + const upT64 = try upBytes(&r, grp, l, b.dx, p64, 64); + try check.sameBytes("up: a tile-64 plan through decode()", upT64.bytes, up16.bytes); + const down16 = try downBytes(&r, grp, l, up16.dev, p16, 16); + const downOld = try downBytes(&r, grp, l, up16.dev, p16, 64); + try check.sameBytes("down: maxm 16 vs the old maxm 64 on a tile-16 plan", downOld, down16); + const downT64 = try downBytes(&r, grp, l, upT64.dev, p64, 64); + try check.sameBytes("down: a tile-64 plan through decode()", downT64, down16); + check.pass("exl3-grouped-tile {d} rows: maxm 16 and the old 64 bit-equal, tile-64 plan bit-equal", .{rows}); + } + + // 16 rows' 128 pairs all pick expert 0: two 64-member items, so the maxm-16 mutant drops members 16..63. + const skew_host = try r.a.alloc(i32, 16 * SLOTS); + @memset(skew_host, 0); + const skew = try r.dev(std.mem.sliceAsBytes(skew_host)); + const mp = try plan(&r, router, skew, 16 * SLOTS, 64); + const poison_x = try r.bfs(16 * SLOTS * D, 0.5); + const poison_x_dev = try r.dev(std.mem.sliceAsBytes(poison_x)); + var gm: cuda.exl3_grouped.Grouped = .{ .d = gpu.d, .kernels = &exp }; + defer gm.deinit(); + const gmgrp: cuda.experts.Grouped = .{ .exl3 = &gm }; + const poisoned = try upBytes(&r, gmgrp, l, poison_x_dev, mp, 64); + const mutant = try upBytes(&r, gmgrp, l, b.dx, mp, 16); + var gc: cuda.exl3_grouped.Grouped = .{ .d = gpu.d, .kernels = &exp }; + defer gc.deinit(); + const real = try upBytes(&r, .{ .exl3 = &gc }, l, b.dx, mp, 64); + const row: usize = 16 * I * 2; + try check.sameBytes("planted mutant keeps the poisoned members 16..31", mutant.bytes[row..][0..row], poisoned.bytes[row..][0..row]); + try check.expect(!std.mem.eql(u8, mutant.bytes, real.bytes), "planted maxm 16 on a tile-64 plan dropped members and must differ from maxm 64", .{}); + check.pass("exl3-grouped-tile planted mutant (maxm 16 on a tile-64 plan) drops members 16..63", .{}); +} + +fn plan(r: *Rig, router: cuda.grouped.Router, pick: u64, pairs: usize, tile: usize) !Plan { + const items = cuda.grouped.maxItems(pairs, E, tile); + const p: cuda.grouped.Plan = .{ + .members = try r.zeros(pairs * 4 + 16), + .items = try r.zeros(items * 3 * 4 + 16), + .counts = try r.zeros(16), + .rank = try r.zeros(pairs * 4 + 16), + .hist = try r.zeros((E + 1) * 4 * ((pairs + 1023) / 1024 + 1)), + }; + try router.route(r.s, pick, pairs, E, tile, p); + return .{ .p = p, .items = items, .pairs = pairs }; +} + +fn upBytes(r: *Rig, grp: cuda.experts.Grouped, l: cuda.experts.Layer, x: u64, p: Plan, tile: usize) !Out { + const out = try r.zeros(p.pairs * I * 2); + try grp.decode(r.s, .up, .{ .x = x, .stride = D, .slots = SLOTS }, l, p.p, out, p.items, p.pairs, tile, -1); + return .{ .dev = out, .bytes = try r.back(out, p.pairs * I * 2) }; +} + +fn downBytes(r: *Rig, grp: cuda.experts.Grouped, l: cuda.experts.Layer, act: u64, p: Plan, tile: usize) ![]u8 { + const out = try r.zeros(p.pairs * D * 4); + try grp.decode(r.s, .down_f32, .{ .x = act, .stride = I }, l, p.p, out, p.items, p.pairs, tile, -1); + return r.back(out, p.pairs * D * 4); +} diff --git a/zig/tests/cuda/experts3_tests.zig b/zig/tests/cuda/experts3_tests.zig index 3c61bdeef20..7e3497976a4 100644 --- a/zig/tests/cuda/experts3_tests.zig +++ b/zig/tests/cuda/experts3_tests.zig @@ -195,11 +195,11 @@ fn sameBits(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random, sh: Shape) !voi const rows_in: cuda.experts.Rows = .{ .x = dx.ptr, .stride = sh.dims, .slots = sh.slots }; const act_in: cuda.experts.Rows = .{ .x = act.ptr, .stride = sh.ni }; if (prompt) { - try g.prompt(s, .up, rows_in, l, pb.plan(), act.ptr, pb.items_n, -1); - try g.prompt(s, .down_bf16, act_in, l, pb.plan(), y.ptr, pb.items_n, -1); + try g.prompt(s, .up, rows_in, l, pb.plan(), act.ptr, pb.items_n, pairs, 64, -1); + try g.prompt(s, .down_bf16, act_in, l, pb.plan(), y.ptr, pb.items_n, pairs, 64, -1); } else { - try g.decode(s, .up, rows_in, l, pb.plan(), act.ptr, pb.items_n, -1); - try g.decode(s, .down_f32, act_in, l, pb.plan(), y.ptr, pb.items_n, -1); + try g.decode(s, .up, rows_in, l, pb.plan(), act.ptr, pb.items_n, pairs, 16, -1); + try g.decode(s, .down_f32, act_in, l, pb.plan(), y.ptr, pb.items_n, pairs, 16, -1); } try s.synchronize(); outs[li] = .{ try check.download(gpu, act), try check.download(gpu, y) }; @@ -256,7 +256,7 @@ fn projections(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random) !void { var y = try cuda.DeviceBuffer.alloc(gpu.d, pairs * n * 4); defer y.free(); const in: cuda.experts.Rows = .{ .x = dx.ptr, .stride = kk }; - if (prompt) try g.prompt(s, .down_bf16, in, l, pb.plan(), y.ptr, pb.items_n, -1) else try g.decode(s, .down_f32, in, l, pb.plan(), y.ptr, pb.items_n, -1); + if (prompt) try g.prompt(s, .down_bf16, in, l, pb.plan(), y.ptr, pb.items_n, pairs, 64, -1) else try g.decode(s, .down_f32, in, l, pb.plan(), y.ptr, pb.items_n, pairs, 16, -1); try s.synchronize(); outs[li] = try check.download(gpu, y); for (0..pairs * n) |i| { @@ -334,8 +334,8 @@ fn timeDecode(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random, reps: usize) for (0..reps) |i| { const p = plans[i % plans.len]; if (form < 2) { - try g.decode(s, .up, rows_in, l3, p.plan(), act.ptr, p.items_n, -1); - } else try g.decode(s, .down_f32, act_in, l3, p.plan(), y.ptr, p.items_n, -1); + try g.decode(s, .up, rows_in, l3, p.plan(), act.ptr, p.items_n, rows * sh.slots, 16, -1); + } else try g.decode(s, .down_f32, act_in, l3, p.plan(), y.ptr, p.items_n, rows * sh.slots, 16, -1); } var graph = try cuda.graph.endCapture(s); defer graph.deinit(); diff --git a/zig/tests/cuda/fp8_experts_tests.zig b/zig/tests/cuda/fp8_experts_tests.zig index aa74443dee0..3e50068162d 100644 --- a/zig/tests/cuda/fp8_experts_tests.zig +++ b/zig/tests/cuda/fp8_experts_tests.zig @@ -140,7 +140,7 @@ fn oneCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, s: cud const bound = grouped.maxItems(pairs, l.experts, t); try act.fill8(0xff, s.handle); // a call that writes nothing fails rather than pass on old bytes const up_in: ex_.Rows = .{ .x = dx.ptr, .stride = l.dims, .slots = slots }; - if (prompt) try g.prompt(s, .up, up_in, l, plan, act.ptr, bound, -1) else try g.decode(s, .up, up_in, l, plan, act.ptr, bound, -1); + if (prompt) try g.prompt(s, .up, up_in, l, plan, act.ptr, bound, pairs, t, -1) else try g.decode(s, .up, up_in, l, plan, act.ptr, bound, pairs, t, -1); try s.synchronize(); const got_act = try check.download(gpu, act); defer a.free(got_act); @@ -153,7 +153,7 @@ fn oneCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, s: cud defer y.free(); try y.fill8(0xff, s.handle); const down_in: ex_.Rows = .{ .x = act.ptr, .stride = l.width }; - if (prompt) try g.prompt(s, .down_bf16, down_in, l, plan, y.ptr, bound, -1) else try g.decode(s, .down_f32, down_in, l, plan, y.ptr, bound, -1); + if (prompt) try g.prompt(s, .down_bf16, down_in, l, plan, y.ptr, bound, pairs, t, -1) else try g.decode(s, .down_f32, down_in, l, plan, y.ptr, bound, pairs, t, -1); try s.synchronize(); const got_y = try check.download(gpu, y); defer a.free(got_y); diff --git a/zig/tests/cuda/ggml_tests.zig b/zig/tests/cuda/ggml_tests.zig index 849349b7074..a5cd2b28d42 100644 --- a/zig/tests/cuda/ggml_tests.zig +++ b/zig/tests/cuda/ggml_tests.zig @@ -150,7 +150,7 @@ fn expertCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, k: // up: relu^2 through experts.Grouped, then fp32 sums by the same kernels with no activation try y.fill8(0xff, s.handle); // a call that writes nothing fails rather than pass on old bytes const call = if (prompt) &ex_.Grouped.prompt else &ex_.Grouped.decode; - try call(g, s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, y.ptr, bound, -1); + try call(g, s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, y.ptr, bound, pairs, tile, -1); try s.synchronize(); try compare(gpu, try std.fmt.bufPrint(&name, "iq2 up relu^2 {s} {d} rows x {d}", .{ form, rows, slots }), y, ref_up, .relu2, tolerance(.relu2)); try y.fill8(0xff, s.handle); @@ -160,7 +160,7 @@ fn expertCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, k: // down from pair rows, fp32 and bf16 inline for (.{ .{ ex_.Epi.down_f32, ggml.Out.f32 }, .{ ex_.Epi.down_bf16, ggml.Out.bf16 } }) |c| { try y.fill8(0xff, s.handle); - try call(g, s, c[0], .{ .x = dh.ptr, .stride = l.width }, l, plan, y.ptr, bound, -1); + try call(g, s, c[0], .{ .x = dh.ptr, .stride = l.width }, l, plan, y.ptr, bound, pairs, tile, -1); try s.synchronize(); try compare(gpu, try std.fmt.bufPrint(&name, "iq2 down {t} {s} {d} rows x {d}", .{ c[1], form, rows, slots }), y, ref_down, c[1], tolerance(c[1])); } @@ -256,8 +256,8 @@ pub fn bench(gpu: Gpu) !void { for (0..2) |warm| { try t0.record(s); for (0..reps) |_| for (ps) |p| { - if (pi != 1) try g.decode(s, .up, .{ .x = x.ptr, .stride = dims, .slots = 9 }, l, p, h.ptr, bound, -1); - if (pi != 0) try g.decode(s, .down_f32, .{ .x = h.ptr, .stride = width }, l, p, y.ptr, bound, -1); + if (pi != 1) try g.decode(s, .up, .{ .x = x.ptr, .stride = dims, .slots = 9 }, l, p, h.ptr, bound, 9, 16, -1); + if (pi != 0) try g.decode(s, .down_f32, .{ .x = h.ptr, .stride = width }, l, p, y.ptr, bound, 9, 16, -1); }; try t1.record(s); try t1.synchronize(); @@ -299,9 +299,9 @@ pub fn bench(gpu: Gpu) !void { defer t2.deinit(); for (0..2) |warm| { try t0.record(s); - try g.prompt(s, .up, .{ .x = xp.ptr, .stride = dims, .slots = 9 }, l, pp, hp.ptr, pb, -1); + try g.prompt(s, .up, .{ .x = xp.ptr, .stride = dims, .slots = 9 }, l, pp, hp.ptr, pb, pairs, 64, -1); try t1.record(s); - try g.prompt(s, .down_bf16, .{ .x = hp.ptr, .stride = width }, l, pp, yp.ptr, pb, -1); + try g.prompt(s, .down_bf16, .{ .x = hp.ptr, .stride = width }, l, pp, yp.ptr, pb, pairs, 64, -1); try t2.record(s); try t2.synchronize(); if (warm == 0) continue; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 941b0c5a80b..0f4d6c266d6 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -457,6 +457,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else "none"); + if (std.mem.eql(u8, cmd, "exl3-grouped-tile")) return @import("exl3_grouped_tile_tests.zig").run(gpu); if (std.mem.eql(u8, cmd, "exl3-tiles")) return exl3_candidate_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-shapes")) return exl3_plan_bench.profileShapes(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-midm")) return exl3_plan_bench.midm(gpu); diff --git a/zig/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig index 42ec3e8748d..b848dcbd255 100644 --- a/zig/tests/cuda/modelopt_tests.zig +++ b/zig/tests/cuda/modelopt_tests.zig @@ -354,11 +354,11 @@ fn experts(gpu: Gpu, rng: std.Random, p: *const modelopt.Packer, s: cuda.Stream) const up_in: cuda.experts.Rows = .{ .x = dx.ptr, .stride = dims, .slots = slots }; const down_in: cuda.experts.Rows = .{ .x = act.ptr, .stride = width }; if (arm == 0) { - try g.decode(s, .up, up_in, l, plan, act.ptr, items, -1); - try g.decode(s, .down_f32, down_in, l, plan, y.ptr, items, -1); + try g.decode(s, .up, up_in, l, plan, act.ptr, items, np, tile, -1); + try g.decode(s, .down_f32, down_in, l, plan, y.ptr, items, np, tile, -1); } else { - try g.prompt(s, .up, up_in, l, plan, act.ptr, items, -1); - try g.prompt(s, .down_bf16, down_in, l, plan, y.ptr, items, -1); + try g.prompt(s, .up, up_in, l, plan, act.ptr, items, np, tile, -1); + try g.prompt(s, .down_bf16, down_in, l, plan, y.ptr, items, np, tile, -1); } try s.synchronize(); outs[arm] = .{ try check.download(gpu, act), try check.download(gpu, y) }; diff --git a/zig/tests/cuda/moe_live_tests.zig b/zig/tests/cuda/moe_live_tests.zig index ab15cd4c9a5..29d0e4847d8 100644 --- a/zig/tests/cuda/moe_live_tests.zig +++ b/zig/tests/cuda/moe_live_tests.zig @@ -75,8 +75,8 @@ const Chain = struct { try c.f.route(c.s, x, dims, rows, c.w, c.logits, dims, count, params); try c.f.pick(c.s, c.logits, c.picks, c.wts, rows, routed, top_k, params); try c.router.routeRows(c.s, c.picks, pairs, count, 16, c.plan, if (params == 0) 0 else params + 8, slots); - try c.grouped.decode(c.s, .up, .{ .x = x, .stride = dims, .slots = slots }, c.layer, c.plan, c.act, items, -1); - try c.grouped.decode(c.s, .down_f32, .{ .x = c.act, .stride = ni }, c.layer, c.plan, c.y, items, -1); + try c.grouped.decode(c.s, .up, .{ .x = x, .stride = dims, .slots = slots }, c.layer, c.plan, c.act, items, pairs, 16, -1); + try c.grouped.decode(c.s, .down_f32, .{ .x = c.act, .stride = ni }, c.layer, c.plan, c.y, items, pairs, 16, -1); try c.f.combine(c.s, c.y, true, c.wts, out, rows, slots, dims, params); } diff --git a/zig/tests/cuda/moe_tests.zig b/zig/tests/cuda/moe_tests.zig index b3301a13b4f..809cc2f97a5 100644 --- a/zig/tests/cuda/moe_tests.zig +++ b/zig/tests/cuda/moe_tests.zig @@ -339,8 +339,8 @@ fn pairExperts(gpu: Gpu, s: cuda.Stream, r: std.Random) !void { 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 grouped.decode(s, .up, .{ .x = dx.ptr, .stride = dims, .slots = slots }, layer, plan, dact.ptr, items_n, pairs, 16, -1); + try grouped.decode(s, .down_f32, .{ .x = dact.ptr, .stride = ni }, layer, plan, dy.ptr, items_n, pairs, 16, -1); try s.synchronize(); const got = try check.download(gpu, dy); defer a.free(got); diff --git a/zig/tests/cuda/nvfp4_experts_tests.zig b/zig/tests/cuda/nvfp4_experts_tests.zig index 20310178366..10f2e0a2108 100644 --- a/zig/tests/cuda/nvfp4_experts_tests.zig +++ b/zig/tests/cuda/nvfp4_experts_tests.zig @@ -138,7 +138,7 @@ fn oneCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, s: cud var act = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.width * 2); defer act.free(); try act.fill8(0xff, s.handle); // a call that writes nothing fails rather than pass on old bytes - try g.decode(s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, act.ptr, bound, -1); + try g.decode(s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, act.ptr, bound, pairs, 16, -1); try s.synchronize(); const got_act = try check.download(gpu, act); defer a.free(got_act); @@ -150,7 +150,7 @@ fn oneCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, s: cud var y = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.dims * width); defer y.free(); try y.fill8(0xff, s.handle); - try g.decode(s, if (wide) .down_bf16 else .down_f32, .{ .x = act.ptr, .stride = l.width }, l, plan, y.ptr, bound, -1); + try g.decode(s, if (wide) .down_bf16 else .down_f32, .{ .x = act.ptr, .stride = l.width }, l, plan, y.ptr, bound, pairs, 16, -1); try s.synchronize(); const got_y = try check.download(gpu, y); defer a.free(got_y); diff --git a/zig/tests/cuda/quant_ref_experts.zig b/zig/tests/cuda/quant_ref_experts.zig index f841aaf3bbf..98dae8897cb 100644 --- a/zig/tests/cuda/quant_ref_experts.zig +++ b/zig/tests/cuda/quant_ref_experts.zig @@ -67,11 +67,11 @@ fn batch(r: *Rig, comptime label: []const u8, g: experts.Grouped, router: groupe const up_in: experts.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(x[0 .. rows * D])), .stride = D, .slots = slots }; const down_in: experts.Rows = .{ .x = act, .stride = NI }; if (prompt) { - try g.prompt(r.s, .up, up_in, l, pl.p, act, pl.items, -1); - try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y, pl.items, -1); + try g.prompt(r.s, .up, up_in, l, pl.p, act, pl.items, pairs, tile, -1); + try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y, pl.items, pairs, tile, -1); } else { - try g.decode(r.s, .up, up_in, l, pl.p, act, pl.items, -1); - try g.decode(r.s, .down_f32, down_in, l, pl.p, y, pl.items, -1); + try g.decode(r.s, .up, up_in, l, pl.p, act, pl.items, pairs, tile, -1); + try g.decode(r.s, .down_f32, down_in, l, pl.p, y, pl.items, pairs, tile, -1); } // gate-up: each pair's row of x through its expert, SwiGLU with the kernels' roundings const want = try r.a.alloc(f64, pairs * NI); @@ -140,15 +140,15 @@ fn skipped( const up_in: experts.Rows = .{ .x = xd, .stride = D, .slots = slots }; const a0 = try r.zeros(pairs * NI * 2); const a1 = try r.dev(fill[0 .. pairs * NI * 2]); - try g.prompt(r.s, .up, up_in, l, pl.p, a0, pl.items, -1); - try g.prompt(r.s, .up, up_in, l, pl.p, a1, pl.items, skip); + try g.prompt(r.s, .up, up_in, l, pl.p, a0, pl.items, pairs, experts.promptTile(l.format), -1); + try g.prompt(r.s, .up, up_in, l, pl.p, a1, pl.items, pairs, experts.promptTile(l.format), skip); const up_plain = try r.back(a0, pairs * NI * 2); try leftAlone(label ++ " prompt gate-up", pk, skip, up_plain, try r.back(a1, pairs * NI * 2), NI * 2, fill); const down_in: experts.Rows = .{ .x = a0, .stride = NI }; const y0 = try r.zeros(pairs * D * 2); const y1 = try r.dev(fill[0 .. pairs * D * 2]); - try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y0, pl.items, -1); - try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y1, pl.items, skip); + try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y0, pl.items, pairs, experts.promptTile(l.format), -1); + try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y1, pl.items, pairs, experts.promptTile(l.format), skip); const down_plain = try r.back(y0, pairs * D * 2); try leftAlone(label ++ " prompt down", pk, skip, down_plain, try r.back(y1, pairs * D * 2), D * 2, fill); } diff --git a/zig/tests/exl3_compile.zig b/zig/tests/exl3_compile.zig index a3abe2e3790..e5e8468f983 100644 --- a/zig/tests/exl3_compile.zig +++ b/zig/tests/exl3_compile.zig @@ -21,9 +21,9 @@ export fn exl3Compile(d: *const cuda.Driver, dense: *cuda.exl3.Dense, group: *cu linear.decodeF32(s.*, weight, .{ .x = 0 }, 0, 1) catch {}; linear.promptF32(s.*, weight, .{ .x = 0 }, 0, 1) catch {}; const grouped: cuda.experts.Grouped = .{ .exl3 = group }; - grouped.decode(s.*, .up, .{ .x = 0, .stride = l.dims }, l.*, p.*, 0, 1, -1) catch {}; - grouped.prompt(s.*, .down_f32, .{ .x = 0, .stride = l.width }, l.*, p.*, 0, 1, -1) catch {}; - grouped.prompt(s.*, .down_bf16, .{ .x = 0, .stride = l.width }, l.*, p.*, 0, 1, -1) catch {}; + grouped.decode(s.*, .up, .{ .x = 0, .stride = l.dims }, l.*, p.*, 0, 1, 1, 16, -1) catch {}; + grouped.prompt(s.*, .down_f32, .{ .x = 0, .stride = l.width }, l.*, p.*, 0, 1, 1, 64, -1) catch {}; + grouped.prompt(s.*, .down_bf16, .{ .x = 0, .stride = l.width }, l.*, p.*, 0, 1, 1, 64, -1) catch {}; } export fn candidateCompile(d: *const cuda.Driver, m: *const cuda.Module, s: *const cuda.Stream) void { From fe3ac1e748cc7920202f085536156af5e149b7d0 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 02:14:23 +0000 Subject: [PATCH 11/12] changelog: #2843 under Unreleased --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c5452db10a..89c6addd9b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. +- CUDA EXL3 grouped experts launch and size their storage at the caller's tile and pair count, so decode no longer launches three quarters of its grid empty, and Qwen3.6-35B-A3B EXL3 sizes its grouped storage for 36,864 pair rows instead of 163,904 (about 3 GB less); outputs are bit-equal (#2843, for #2823). - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. From 6158d0cc217d05b41afdfac2da3863ce5e61f67b Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 11:48:30 +0000 Subject: [PATCH 12/12] changelog: #2843's line from lines/2843.txt in place of its own --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 89c6addd9b9..ae965c88f54 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,9 +4,10 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. - CUDA EXL3 grouped experts launch and size their storage at the caller's tile and pair count, so decode no longer launches three quarters of its grid empty, and Qwen3.6-35B-A3B EXL3 sizes its grouped storage for 36,864 pair rows instead of 163,904 (about 3 GB less); outputs are bit-equal (#2843, for #2823). +- CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. + - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. - GLM-5.3-Flash on CUDA: with `--vision` and no video library, rank 0 says at startup how much each stream sets aside for `--vision-video-tokens` rows that no request can fill (128 MiB a stream at the default 16,384), and that `--vision-video-tokens 64` on every rank cuts it to 512 KiB (#2813). Thanks to @Bizuayeu.