Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits

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:

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");
 }

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)

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));
+}

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,
     };

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");