From 1874ffe0b29f849909fd143e7a38983ca2b2d014 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:17:05 -0500 Subject: [PATCH 1/3] Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits The first token from the prompt pass's last row, each verify row from the window's logits, with the lane core's keyed sampler (top-k in one pass over bf16 keys, equal to qwen3_5's draw; families/flashnext/sample.zig, the same file as #719). Sampled replies are the same drafted and undrafted, alone and beside other streams; greedy replies are unchanged. --- docs/recipes/qwen3.8-flash-next.md | 4 +- zig/flashnext_host.zig | 1 + zig/src/families/flashnext/backend.zig | 21 ++++- zig/src/families/flashnext/sample.zig | 107 ++++++++++++++++++++++++ zig/src/native/flashnext_batch_host.zig | 3 +- zig/src/tensorfold.zig | 1 + 6 files changed, 133 insertions(+), 4 deletions(-) create mode 100644 zig/src/families/flashnext/sample.zig 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..3fbe8e1981 --- /dev/null +++ b/zig/src/families/flashnext/sample.zig @@ -0,0 +1,107 @@ +//! 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"); + +/// k at most this takes the one-pass path; a wider (or unlimited) top-k takes the reference's. +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 reference(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)); +} + +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) * 3.0); + 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{ 1, 2, 20, 64, 256 }) |k| for ([_]f64{ 0.6, 1.0, 1.5 }) |t| { + const p: lanes.Sampling = .{ .seed = trial * 31 + k, .temperature = t, .top_k = k, .top_p = 0.95 }; + const pos: u64 = trial * 7 + k; + try std.testing.expectEqual(try reference(gpa, words, p, pos), try draw(gpa, words, p, pos)); + }; + } +} 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"); From b7372ac847224ca90180c310da9d479d2757bc0a Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:02:28 -0500 Subject: [PATCH 2/3] Flash Next's host sampler: an unlimited or wide top-k by a counting sort of bf16 keys, not a full sort, the same draws top_k 0 (or above 256) went through the reference: every candidate in f64, sorted, then choose. Now a counting sort of the 65,536 bf16 keys gives choose's order (ids ascending within a key), top-p and min-p cut it as choose does with libm's exp and the same pairwise sum, and the keyed Gumbel draw runs in that order, passing over candidates whose noise is provably a nat short of the best score (one hash, no logs) and stopping once none can reach it. 248k vocabulary: 13.7-16.8 ms a row to 0.33-0.65 ms; the host test now also compares top_k 0, 300 and the whole row, top-p 1, min-p and wide logits. --- zig/src/families/flashnext/sample.zig | 102 ++++++++++++++++++++++++-- 1 file changed, 95 insertions(+), 7 deletions(-) diff --git a/zig/src/families/flashnext/sample.zig b/zig/src/families/flashnext/sample.zig index 3fbe8e1981..7763842f90 100644 --- a/zig/src/families/flashnext/sample.zig +++ b/zig/src/families/flashnext/sample.zig @@ -4,7 +4,10 @@ const std = @import("std"); const lanes = @import("lanes"); -/// k at most this takes the one-pass path; a wider (or unlimited) top-k takes the reference's. +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. @@ -62,7 +65,7 @@ pub fn row(base: [*]const u8, vocab: usize, at: usize) []const u16 { /// 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 reference(gpa, logits, p, position); + 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; @@ -84,6 +87,89 @@ pub fn draw(gpa: std.mem.Allocator, logits: []const u16, p: lanes.Sampling, posi 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) break; // a draw's noise is at most -log(-log(1 - 2^-54)) = 37.43, and x only falls from here + 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, read from the uniform's integer bits against exp of a whole nat at or below d - 1. +fn beyondReach(seed: u64, position: u64, id: u64, d: f64) bool { + const m = lanes.sampling.keyBits(seed, position, id) >> 11; + const one_minus_u = @as(f64, @floatFromInt((@as(u64, 1) << 53) - m)) * 0x1p-53 - 0x1p-54; + 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); @@ -92,16 +178,18 @@ test "the one-pass draw equals the reference on random rows with ties, signed ze defer gpa.free(words); for (0..40) |trial| { for (words) |*w| { - const f: f32 = @floatCast(r.floatNorm(f64) * 3.0); + 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{ 1, 2, 20, 64, 256 }) |k| for ([_]f64{ 0.6, 1.0, 1.5 }) |t| { - const p: lanes.Sampling = .{ .seed = trial * 31 + k, .temperature = t, .top_k = k, .top_p = 0.95 }; - const pos: u64 = trial * 7 + k; - try std.testing.expectEqual(try reference(gpa, words, p, pos), try draw(gpa, words, p, pos)); + 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)); + } }; } } From d519c08573e1e089d9655460bb99bf7fb3b4c74c Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:35:59 -0500 Subject: [PATCH 3/3] Flash Next's wide top-k draw keeps a uniform that rounds to 1: its +inf noise wins as in the reference lanes.sampling.uniform's largest value, (2^53 - 1) 2^-53 + 2^-54, rounds to exactly 1.0 in f64, so the reference's Gumbel score is +inf there; the early stop assumed noise of at most 37.43 and could pass that candidate over (a seed that hits it: expected 4, found 0). Past a 38-nat shortfall the draw now still asks each candidate whether its uniform is 1 (one hash, no logs) and stops only once the best score is +inf; the pruning bound reads 1 - u from the f64 uniform itself (exact), so a u of 1 is never passed over. The reviewer's case is a host test, with the same candidate inside the bound and behind another leader. Speed unchanged (0.33-0.63 ms a row at 248k). --- zig/src/families/flashnext/sample.zig | 29 ++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/zig/src/families/flashnext/sample.zig b/zig/src/families/flashnext/sample.zig index 7763842f90..c388bb4fc2 100644 --- a/zig/src/families/flashnext/sample.zig +++ b/zig/src/families/flashnext/sample.zig @@ -142,8 +142,10 @@ fn wide(gpa: std.mem.Allocator, logits: []const u16, p: lanes.Sampling, position const x = value(logits[id]) / t; if (j > 0) { const d = best_score - x; - if (d > 38) break; // a draw's noise is at most -log(-log(1 - 2^-54)) = 37.43, and x only falls from here - if (d > 1 and beyondReach(p.seed, position, id, d)) continue; + 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) { @@ -162,10 +164,10 @@ const below_e = blk: { }; /// 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, read from the uniform's integer bits against exp of a whole nat at or below 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 m = lanes.sampling.keyBits(seed, position, id) >> 11; - const one_minus_u = @as(f64, @floatFromInt((@as(u64, 1) << 53) - m)) * 0x1p-53 - 0x1p-54; + 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); } @@ -193,3 +195,20 @@ test "the one-pass draw equals the reference on random rows with ties, signed ze }; } } + +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)); +}