diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index 9aea70bba3..2a938bf22d 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const core = @import("core"); const embed = @import("cuda_embed.zig"); const hc = @import("cuda_hc.zig"); +const ple_pool = @import("cuda_ple_pool.zig"); const golden: u64 = 0x9E3779B97F4A7C15; const mix1: u64 = 0xBF58476D1CE4E5B9; @@ -343,6 +344,7 @@ pub const Ngram = struct { b: [max_shards][]const u8 = undefined, format: enum { mlx, nvfp4 } = .mlx, tables: [max_shards]core.nvfp4_rows.Table = undefined, + pool: ?*ple_pool.Pool = null, // the gather's parked workers (null: every row on the caller's thread) pub fn init(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !Ngram { const cfg = try readSpec(gpa, io, model_dir); @@ -410,10 +412,12 @@ pub const Ngram = struct { g.b[shard] = bias.bytes; } } + g.pool = try ple_pool.Pool.create(gpa, io); return g; } pub fn deinit(g: *Ngram) void { + if (g.pool) |p| p.destroy(); for (g.files.items) |*o| o.close(g.gpa); g.files.deinit(g.gpa); } @@ -436,11 +440,19 @@ pub const Ngram = struct { } } - fn worker(g: *const Ngram, ids: []const i64, out: []u16, lo: usize, hi: usize, failed: *?anyerror) void { - g.fill(ids, out, lo, hi) catch |err| { - failed.* = err; - }; - } + /// One call's rows split into `parts` equal ranges, part i filling [i * per, (i + 1) * per). + const Split = struct { + g: *const Ngram, + ids: []const i64, + out: []u16, + per: usize, + + fn part(ctx: *const anyopaque, i: usize) anyerror!void { + const s: *const Split = @ptrCast(@alignCast(ctx)); + const lo = @min(s.ids.len, i * s.per); + try s.g.fill(s.ids, s.out, lo, @min(s.ids.len, lo + s.per)); + } + }; /// `embedding` for each token after `history` (moved past the chunk) into out and xs, read by up to 128 threads. pub fn rows(g: *const Ngram, history: []i64, tokens: []const u32, out: []u16, xs: []f32) !void { @@ -457,25 +469,12 @@ pub const Ngram = struct { for (0..history.len - 1) |i| hist[i] = hist[i + 1]; hist[history.len - 1] = token; } - const max_threads = 128; - const n_threads = @min(max_threads, items / 4); - if (n_threads <= 1) { + const n_threads = @min(ple_pool.max_parts, items / 4); + if (n_threads <= 1 or g.pool == null) { try g.fill(ids, out, 0, items); } else { - var threads: [max_threads]std.Thread = undefined; - var failed: [max_threads]?anyerror = @splat(null); - var spawned: usize = 0; - defer for (threads[0..spawned]) |th| th.join(); - const per = (items + n_threads - 1) / n_threads; - for (0..n_threads) |i| { - const lo = @min(items, i * per); - const hi = @min(items, lo + per); - threads[i] = try std.Thread.spawn(.{}, worker, .{ g, ids, out, lo, hi, &failed[i] }); - spawned += 1; - } - for (threads[0..spawned]) |th| th.join(); - spawned = 0; - for (failed[0..n_threads]) |f| if (f) |err| return err; + const split: Split = .{ .g = g, .ids = ids, .out = out, .per = (items + n_threads - 1) / n_threads }; + try g.pool.?.run(n_threads, .{ .ctx = &split, .run = Split.part }); } for (0..tokens.len) |t| { const row = out[t * w ..][0..w]; @@ -517,4 +516,5 @@ test "a fresh window hashes token 7 into sixteen rows" { test { _ = @import("cuda_ple_nvfp4_test.zig"); + _ = @import("cuda_ple_pool.zig"); } diff --git a/zig/src/families/flashnext/cuda_ple_pool.zig b/zig/src/families/flashnext/cuda_ple_pool.zig new file mode 100644 index 0000000000..a9e6ea95c2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple_pool.zig @@ -0,0 +1,139 @@ +//! The n-gram gather's parked host workers; parts claimed from one counter, each claim waking one more worker. +const std = @import("std"); + +/// The most parts a call may have (the caller beside up to 127 workers). +pub const max_parts = 128; + +/// One part of a call: `run(ctx, part)`. +pub const Task = struct { + ctx: *const anyopaque, + run: *const fn (ctx: *const anyopaque, part: usize) anyerror!void, +}; + +pub const Pool = struct { + gpa: std.mem.Allocator, + io: std.Io, + mutex: std.Io.Mutex = .init, // one call at a time + threads: [max_parts - 1]std.Thread = undefined, + go: [max_parts - 1]std.atomic.Value(u32) = undefined, // bumped to wake a worker: to claim parts, or to stop + spawned: usize = 0, + stopping: std.atomic.Value(bool) = .init(false), + // the current call: generation << 32 | parts << 16 | parts claimed (a claim is a compare-exchange on all three) + state: std.atomic.Value(u64) = .init(0), + done: std.atomic.Value(u32) = .init(0), // parts finished + waking: std.atomic.Value(u32) = .init(0), // workers woken + task: Task = undefined, + failed: [max_parts]?anyerror = undefined, + + pub fn create(gpa: std.mem.Allocator, io: std.Io) !*Pool { + const p = try gpa.create(Pool); + p.* = .{ .gpa = gpa, .io = io }; + return p; + } + + /// Wakes every worker to exit, joins them and frees the pool. + pub fn destroy(p: *Pool) void { + p.stopping.store(true, .release); + for (p.go[0..p.spawned]) |*g| { + _ = g.fetchAdd(1, .release); + p.io.futexWake(u32, &g.raw, 1); + } + for (p.threads[0..p.spawned]) |th| th.join(); + p.gpa.destroy(p); + } + + /// Runs parts 0..parts-1 of `task` and returns when all are done: the first failed part's error, in part order. + pub fn run(p: *Pool, parts: usize, task: Task) !void { + if (parts == 0) return; + if (parts > max_parts) return error.TooManyParts; + p.mutex.lockUncancelable(p.io); + defer p.mutex.unlock(p.io); + while (p.spawned < parts - 1) { + p.go[p.spawned] = .init(0); + p.threads[p.spawned] = try std.Thread.spawn(.{}, work, .{ p, p.spawned }); + p.spawned += 1; + } + p.task = task; + @memset(p.failed[0..parts], null); + p.done.store(0, .monotonic); + p.waking.store(0, .monotonic); + const gen = (p.state.load(.monotonic) >> 32) +% 1; + p.state.store(gen << 32 | @as(u64, parts) << 16, .release); + p.drain(); + while (true) { + const done = p.done.load(.acquire); + if (done == parts) break; + p.io.futexWaitUncancelable(u32, &p.done.raw, done); + } + for (p.failed[0..parts]) |f| if (f) |err| return err; + } + + /// Claims and runs parts of the current call until none is left unclaimed. + fn drain(p: *Pool) void { + var s = p.state.load(.acquire); + while (true) { + const parts: u32 = @intCast(s >> 16 & 0xffff); + const part: u32 = @intCast(s & 0xffff); + if (part >= parts) return; + if (p.state.cmpxchgWeak(s, s + 1, .acq_rel, .acquire)) |now| { + s = now; + continue; + } + if (part + 1 < parts) p.wakeOne(parts - 1); + p.failed[part] = if (p.task.run(p.task.ctx, part)) null else |err| err; + if (p.done.fetchAdd(1, .acq_rel) + 1 == parts) p.io.futexWake(u32, &p.done.raw, 1); + s = p.state.load(.acquire); + } + } + + /// Wakes the next parked worker: a call of quick parts stays mostly on the caller, one of slow parts fans out. + fn wakeOne(p: *Pool, workers: u32) void { + const w = p.waking.fetchAdd(1, .monotonic); + if (w >= workers) return; + _ = p.go[w].fetchAdd(1, .release); + p.io.futexWake(u32, &p.go[w].raw, 1); + } + + fn work(p: *Pool, w: usize) void { + var seen: u32 = 0; + while (true) { + const now = p.go[w].load(.acquire); + if (now == seen) { + p.io.futexWaitUncancelable(u32, &p.go[w].raw, seen); + continue; + } + seen = now; + if (p.stopping.load(.acquire)) return; + p.drain(); + } + } +}; + +test "parts run once each across calls of changing size, and the first failed part's error wins" { + const Job = struct { + hits: [max_parts]std.atomic.Value(u32) = @splat(.init(0)), + fail_from: usize = max_parts, + + fn run(ctx: *const anyopaque, part: usize) anyerror!void { + const j: *@This() = @ptrCast(@alignCast(@constCast(ctx))); + _ = j.hits[part].fetchAdd(1, .monotonic); + if (part >= j.fail_from) return if (part == j.fail_from) error.First else error.Later; + } + }; + const p = try Pool.create(std.testing.allocator, std.testing.io); + defer p.destroy(); + var job: Job = .{}; + const task: Task = .{ .ctx = &job, .run = Job.run }; + const sizes = [_]usize{ 4, 16, 1, 128, 16, 4 }; + for (sizes) |n| try p.run(n, task); + try std.testing.expectEqual(@as(usize, 127), p.spawned); + for (job.hits, 0..) |h, part| { + var calls: u32 = 0; + for (sizes) |n| calls += @intFromBool(part < n); + try std.testing.expectEqual(calls, h.load(.monotonic)); + } + job.fail_from = 5; + try std.testing.expectError(error.First, p.run(64, task)); + try p.run(5, task); + try std.testing.expectError(error.TooManyParts, p.run(max_parts + 1, task)); +}