diff --git a/docs/recipes/qwen3.8-flash-next.md b/docs/recipes/qwen3.8-flash-next.md index 613ae04d5e..9dc40db99c 100644 --- a/docs/recipes/qwen3.8-flash-next.md +++ b/docs/recipes/qwen3.8-flash-next.md @@ -77,7 +77,9 @@ New requests join between rounds, and cancelling a request releases its slot for Additional slots allocate caches for the selected context. Startup reserves their metadata and the prompt-cache budget, then checks the physical footprint before accepting requests. A fixed slot count that does not fit is refused; lower the count, context or prompt-cache budget. `--prompt-cache-over-cap` overrides only the prompt-cache -budget. Greedy decoding remains required, so clients must send temperature 0. +budget. Shared sessions sample: a request with a temperature draws each token on the host from the head's logits +with the server's keyed sampler (top-k, top-p, min-p, a `seed` or one derived from the prompt), so a sampled reply is +the same with drafting on and off and alone or beside other streams. `--parallel 1` still decodes greedily only. The checkpoint gate compares shared replies with fresh legacy replies, with drafting off and on, then checks prompt admission, cancellation, slot reuse and a restored prefix while another request remains active: diff --git a/zig/flashnext_host.zig b/zig/flashnext_host.zig index 2091e630e6..328f8b04f1 100644 --- a/zig/flashnext_host.zig +++ b/zig/flashnext_host.zig @@ -2,5 +2,6 @@ test { _ = @import("src/families/flashnext/host_test.zig"); _ = @import("src/families/flashnext/pack.zig"); + _ = @import("src/families/flashnext/sample.zig"); _ = @import("src/families/flashnext/batch_meta_test.zig"); } diff --git a/zig/src/families/flashnext/backend.zig b/zig/src/families/flashnext/backend.zig index e4c02f2591..e5b2c74b58 100644 --- a/zig/src/families/flashnext/backend.zig +++ b/zig/src/families/flashnext/backend.zig @@ -12,6 +12,19 @@ const Stream = lanes.Stream; const Buf = fz.Buf; const mtl = @import("metal"); const segments = @import("../../core/segments.zig"); +const sample = @import("sample.zig"); + +/// A stream's sampling when it samples (a temperature above 0); null: the GPU's argmax stands. +fn sampling(stream: *const Stream) ?lanes.Sampling { + const p = stream.sampling orelse return null; + return if (p.temperature > 0) p else null; +} + +/// Row `row` of a [rows, VOCAB] bf16 logits buffer, drawn with `p` keyed at `position`. +fn drawRow(gpa: std.mem.Allocator, logits: Buf, row: usize, p: lanes.Sampling, position: u64) !u32 { + const bytes = logits.b.contents()[logits.off + row * fz.VOCAB * 2 ..][0 .. fz.VOCAB * 2]; + return sample.draw(gpa, @as([*]const u16, @ptrCast(@alignCast(bytes.ptr)))[0..fz.VOCAB], p, position); +} pub const Backend = struct { gpa: std.mem.Allocator, @@ -120,7 +133,6 @@ pub const Backend = struct { const pool = mtl.objc.Pool.push(); defer pool.pop(); const b = self(ptr); - if (stream.sampling) |p| if (p.temperature > 0) return error.GreedyOnly; const prompt = stream.prompt(); if (prompt.len == 0) return error.EmptyPrompt; try b.settle(); @@ -162,6 +174,10 @@ pub const Backend = struct { end = at + call.rows; const ps = [2]*fz.Prompt{ b.e.pr, b.e.pr2 }; s.first = try fz.Prompt.chunkN(ps[0..call.parts], &s.m, b.gpa, prompt[at..end]); + // the pass's last row's logits are in the model's head buffer until the next forward: a sampled first token now + if (end == prompt.len) if (sampling(stream)) |p| { + s.first = try drawRow(b.gpa, s.m.t.logits, 0, p, prompt.len); + }; for (ps[0..call.parts], 0..) |p, k| { const start = at + segments.start(call.rows, call.parts, k); const rows = segments.rows(call.rows, call.parts, k); @@ -243,6 +259,9 @@ pub const Backend = struct { try b.primary.finish(cb); Session.copy(b.primary.last, b.hidden, plan.total * fz.WIDE * 2); for (out, 0..) |o, i| @memcpy(o.sampled, t.picks.b.slice(u32, plan.total)[plan.starts[i]..][0..o.sampled.len]); + for (windows, out, 0..) |w, o, i| if (sampling(w.stream)) |p| { + for (o.sampled, w.positions, 0..) |*tok, pos, j| tok.* = try drawRow(b.gpa, t.logits, plan.starts[i] + j, p, pos); + }; b.window_count += 1; if (windows.len > 1) b.shared_count += 1; if (std.c.getenv("FZ_BATCH_LOG") != null) diff --git a/zig/src/families/flashnext/sample.zig b/zig/src/families/flashnext/sample.zig new file mode 100644 index 0000000000..c388bb4fc2 --- /dev/null +++ b/zig/src/families/flashnext/sample.zig @@ -0,0 +1,214 @@ +//! The keyed host sampler over a bf16 logits row, in one pass: the top-k candidates by an order-preserving integer key +//! of each bf16 word, then lanes.sampling.choose. The candidates, their order and the draw equal `reference`, which is +//! qwen3_5/backend.zig's draw without its mask (two passes in f64), kept here for wide top-k and the tests. +const std = @import("std"); +const lanes = @import("lanes"); + +extern "c" fn exp(x: f64) f64; // libm's, as lanes.sampling's +extern "c" fn log(x: f64) f64; + +/// k at most this takes the one-pass path; a wider (or unlimited) top-k takes `wide`. +pub const MAX_K = 256; + +/// A bf16 word's order as an unsigned key: larger floats give larger keys, -0 and +0 the same. +fn key(word: u16) u16 { + const w = if (word == 0x8000) 0 else word; + return if (w & 0x8000 != 0) ~w else w | 0x8000; +} + +fn value(word: u16) f64 { + return @as(f32, @bitCast(@as(u32, word) << 16)); +} + +/// qwen3_5/backend.zig's draw without a mask: the unchanged top-k candidates for the keyed sampler. +pub fn reference(gpa: std.mem.Allocator, logits: []const u16, p: lanes.Sampling, position: u64) !u32 { + var best: ?usize = null; + for (logits, 0..) |word, id| { + const x = value(word); + if (!std.math.isFinite(x)) return error.NonfiniteQwenLogits; + if (best == null or x > value(logits[best.?])) best = id; + } + if (p.temperature <= 0) return @intCast(best.?); + const allowed = logits.len; + const k = if (p.top_k == 0) allowed else @min(allowed, p.top_k); + const values = try gpa.alloc(f64, k); + defer gpa.free(values); + const ids = try gpa.alloc(u64, k); + defer gpa.free(ids); + var filled: usize = 0; + for (logits, 0..) |word, id| { + const x = value(word); + if (k == allowed) { + values[filled] = x; + ids[filled] = id; + filled += 1; + continue; + } + if (filled == k and x <= values[k - 1]) continue; + var at = @min(filled, k - 1); + while (at > 0 and x > values[at - 1]) : (at -= 1) { + values[at] = values[at - 1]; + ids[at] = ids[at - 1]; + } + values[at] = x; + ids[at] = id; + filled = @min(filled + 1, k); + } + return @intCast(try lanes.sampling.choose(gpa, values, ids, position, p)); +} + +/// Row `row` of a [rows, vocab] bf16 logits buffer whose bytes start at `base`. +pub fn row(base: [*]const u8, vocab: usize, at: usize) []const u16 { + return @as([*]const u16, @ptrCast(@alignCast(base + at * vocab * 2)))[0..vocab]; +} + +/// The draw for `logits` with `p` keyed at `position` (p.temperature above 0). +pub fn draw(gpa: std.mem.Allocator, logits: []const u16, p: lanes.Sampling, position: u64) !u32 { + const k: usize = p.top_k; + if (k == 0 or k > MAX_K or k >= logits.len) return wide(gpa, logits, p, position); + var keys: [MAX_K]u16 = undefined; + var ids: [MAX_K]u64 = undefined; + var filled: usize = 0; + for (logits, 0..) |word, id| { + if (word & 0x7f80 == 0x7f80) return error.NonfiniteQwenLogits; // exponent all ones: inf or nan + const x = key(word); + if (filled == k and x <= keys[k - 1]) continue; + var at = @min(filled, k - 1); + while (at > 0 and x > keys[at - 1]) : (at -= 1) { + keys[at] = keys[at - 1]; + ids[at] = ids[at - 1]; + } + keys[at] = x; + ids[at] = id; + filled = @min(filled + 1, k); + } + var values: [MAX_K]f64 = undefined; + for (ids[0..filled], values[0..filled]) |id, *v| v.* = value(logits[id]); + return @intCast(try lanes.sampling.choose(gpa, values[0..filled], ids[0..filled], position, p)); +} + +/// A wide or unlimited top-k as `reference` draws it: the candidates in choose's order by a counting sort of their +/// keys (ids ascending within one), top-p and min-p cut as choose cuts them, then the keyed Gumbel draw in that order, +/// passing over candidates whose noise cannot lift them past the best score so far (most, without their two logs). +fn wide(gpa: std.mem.Allocator, logits: []const u16, p: lanes.Sampling, position: u64) !u32 { + const starts = try gpa.alloc(u32, 1 << 16); + defer gpa.free(starts); + @memset(starts, 0); + for (logits) |word| { + if (word & 0x7f80 == 0x7f80) return error.NonfiniteQwenLogits; // exponent all ones: inf or nan + starts[0xffff - key(word)] += 1; // larger keys first + } + var at: u32 = 0; + for (starts) |*c| { + const n = c.*; + c.* = at; + at += n; + } + const order = try gpa.alloc(u32, logits.len); + defer gpa.free(order); + for (logits, 0..) |word, id| { + const b = 0xffff - @as(usize, key(word)); + order[starts[b]] = @intCast(id); + starts[b] += 1; + } + const t = @max(p.temperature, 1e-6); + const k = @max(1, @min(if (p.top_k != 0) @as(usize, p.top_k) else logits.len, logits.len)); + const x0 = value(logits[order[0]]) / t; + var keep = k; + if (p.top_p > 0.0 and p.top_p < 1.0) { + const probs = try gpa.alloc(f64, k); + defer gpa.free(probs); + for (probs, order[0..k]) |*q, id| q.* = exp(value(logits[id]) / t - x0); + const total = lanes.sampling.pairwiseSum(probs); + var cum: f64 = 0.0; + var below: usize = 0; + for (probs) |q| { // cum only grows: past top_p no later one counts + cum += q / total; + if (cum < p.top_p) below += 1 else break; + } + keep = @min(keep, below + 1); + } + if (p.min_p > 0.0) { + const floor = x0 + p.minLog(); + var n: usize = 0; + for (order[0..keep]) |id| { + if (value(logits[id]) / t >= floor) n += 1; + } + keep = n; + } + var best: usize = 0; + var best_score = -std.math.inf(f64); + for (order[0..keep], 0..) |id, j| { + const x = value(logits[id]) / t; + if (j > 0) { + const d = best_score - x; + if (d > 38) { // finite noise is at most -log(-log(1 - 2^-52)) = 36.04; a uniform that rounds to 1 gives +inf + if (best_score == std.math.inf(f64)) break; // the first +inf stands: a later one ties, never passes + if (lanes.sampling.uniform(p.seed, position, id) != 1.0) continue; + } else if (d > 1 and beyondReach(p.seed, position, id, d)) continue; + } + const score = x - log(-log(lanes.sampling.uniform(p.seed, position, id))); + if (j == 0 or score > best_score) { + best = j; + best_score = score; + } + } + return order[best]; +} + +/// exp(-j) for the whole nats below a candidate's shortfall, for beyondReach. +const below_e = blk: { + var e: [38]f64 = undefined; + for (&e, 0..) |*v, j| v.* = @exp(-@as(f64, @floatFromInt(j))); + break :blk e; +}; + +/// Whether the noise for `id` stays a nat short of `d`: -log(-log u) <= -log(1 - u), so 1 - u >= exp(-(d - 1)) +/// bounds it by d - 1, against exp of a whole nat at or below d - 1. 1 - u is exact for the f64 u the draw uses +/// (Sterbenz past 0.5; below it the bound cannot hold), and a u rounded to 1 gives 0: never passed over. +fn beyondReach(seed: u64, position: u64, id: u64, d: f64) bool { + const one_minus_u = 1.0 - lanes.sampling.uniform(seed, position, id); + const j: usize = @intFromFloat(@floor(d) - 1); + return one_minus_u >= below_e[j] * (1 + 1e-9); +} + +test "the one-pass draw equals the reference on random rows with ties, signed zeros and every top-k" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(7); + const r = rng.random(); + const words = try gpa.alloc(u16, 4096); + defer gpa.free(words); + for (0..40) |trial| { + for (words) |*w| { + const f: f32 = @floatCast(r.floatNorm(f64) * @as(f64, if (trial % 2 == 0) 3.0 else 20.0)); // narrow, then wide + w.* = @intCast(@as(u32, @bitCast(f)) >> 16); + } + for (0..64) |j| words[r.uintLessThan(usize, words.len)] = words[j]; // ties + words[r.uintLessThan(usize, words.len)] = 0x8000; + words[r.uintLessThan(usize, words.len)] = 0; + for ([_]u32{ 0, 1, 2, 20, 64, 256, 300, 4096 }) |k| for ([_]f64{ 0.6, 1.0, 1.5, 4.0 }) |t| { + for ([_][2]f64{ .{ 0.95, 0 }, .{ 1.0, 0 }, .{ 0.5, 0.05 }, .{ 1.0, 0.2 } }) |cut| { + const p: lanes.Sampling = .{ .seed = trial * 31 + k, .temperature = t, .top_k = k, .top_p = cut[0], .min_p = cut[1] }; + const pos: u64 = trial * 7 + k; + try std.testing.expectEqual(try reference(gpa, words, p, pos), try draw(gpa, words, p, pos)); + } + }; + } +} + +test "wide draw preserves reference when uniform rounds to one" { + const gpa = std.testing.allocator; + var words: [300]u16 = @splat(0); + words[0] = 0x4280; // 64.0; other logits are zero + const p: lanes.Sampling = .{ .seed = 2870505676250484375, .temperature = 1, .top_k = 0, .top_p = 1, .min_p = 0 }; + try std.testing.expectEqual(@as(f64, 1), lanes.sampling.uniform(p.seed, 0, 4)); + const expected = try reference(gpa, &words, p, 0); + try std.testing.expectEqual(@as(u32, 4), expected); + try std.testing.expectEqual(expected, try draw(gpa, &words, p, 0)); + // the same candidate within the pruning bound (shortfall under 38), and behind a different leader + words[0] = 0x41a0; // 20.0 + try std.testing.expectEqual(try reference(gpa, &words, p, 0), try draw(gpa, &words, p, 0)); + words[0] = 0; + words[2] = 0x4280; + try std.testing.expectEqual(try reference(gpa, &words, p, 0), try draw(gpa, &words, p, 0)); +} diff --git a/zig/src/native/flashnext_batch_host.zig b/zig/src/native/flashnext_batch_host.zig index 4956950c60..a54afd2030 100644 --- a/zig/src/native/flashnext_batch_host.zig +++ b/zig/src/native/flashnext_batch_host.zig @@ -54,7 +54,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, dump: ?[]const errdefer h.core.deinit(); const plan = try fit.budget(cache_gib, over, 0, a, why, ""); eng.snap_pool.max = plan.budget_bytes; - h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "flashnext-zig", .lanes = slots, .greedy_only = true, .context_window = @intCast(window), .prefill_step = @intCast(eng.pr.step), .prompt_cache_plan = plan }); + h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "flashnext-zig", .lanes = slots, .context_window = @intCast(window), .prefill_step = @intCast(eng.pr.step), .prompt_cache_plan = plan }); h.cache = null; if (plan.budget_bytes > 0) { h.cache = api.prompt_cache.Store.init(gpa, .{ .ptr = &h.back, .vtable = &.{ .bytes = Backend.snapBytes, .save = Backend.snapSave, .restore = Backend.snapRestore, .drop = Backend.snapDrop, .charged = Backend.snapCharged, .spare = Backend.snapSpare, .trim = Backend.snapTrim, .reuses = Backend.snapReuses } }, .{ .lookahead = 1, .planned = true }, plan.budget_bytes); @@ -72,7 +72,6 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, dump: ?[]const fn words(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { - error.GreedyOnly => "the native Flash Next engine decodes greedily only: send temperature 0", error.ContextFull => "the prompt and max_tokens do not fit this server's --context", else => null, }; diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index bdb31c39d7..1804603c1d 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -15,6 +15,7 @@ pub const segments = @import("core/segments.zig"); pub const row_projection = @import("core/row_projection.zig"); pub const nemotron = @import("families/nemotron/nemotron.zig"); pub const flashnext_replay = @import("families/flashnext/replay.zig"); +pub const flashnext_sample = @import("families/flashnext/sample.zig"); pub const flashnext_engine = @import("families/flashnext/engine.zig"); pub const flashnext_snapshot = @import("families/flashnext/snapshot.zig"); pub const flashnext_session = @import("families/flashnext/session.zig");