From e4062e6ac51984d475b276912b2a01b63a9e6dd2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:29:32 +0200 Subject: [PATCH] Flash Next PLE gather: parked worker pool instead of a thread spawn per call Ngram.rows spawned min(128, items/4) threads per call and joined them: about 150 us for one decode token, 1 ms for four and 10-11 ms for a prefill chunk. The table now keeps a pool for its life; workers are spawned on first need and park on a futex. A call's parts (the same [lo, hi) ranges as before) are claimed from one counter by the caller and the workers, and each claim that leaves parts unclaimed wakes one more worker, so quick parts stay mostly on the caller and slow (faulting) parts fan out. The first failed part's error is returned, in part order, as before. rows() on the FP4-Local table, p50: 16 items 164 -> 20 us, 64 items 0.99 ms -> 31 us, 4096 items 10.8 -> 0.96 ms, 8192 items 11.0 -> 1.66 ms; output bytes identical. Served (fresh server per arm, A B A B): c=4 decode -1.0 ms/step (61.5 -> 60.5), c=1 and 8k TTFT inside noise; greedy, seeded and resumed-cache replies identical. --- zig/src/families/flashnext/cuda_ple.zig | 44 +++--- zig/src/families/flashnext/cuda_ple_pool.zig | 139 +++++++++++++++++++ 2 files changed, 161 insertions(+), 22 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_ple_pool.zig 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)); +}