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

[BobClawblaw]

What this changes

The shared sessions sample. The batching host stops reporting greedy_only. A stream with a temperature draws each token on the host from the head's bf16 logits with the lane core's keyed sampler, the same draw the Qwen3.5 Metal backend makes, so temperature, top_k, top_p, min_p and seed work.

• First token: drawn from the prompt pass's last row, right after the last chunk.
• Each round: each sampled stream's window rows are drawn from t.logits at that stream's positions, overwriting the GPU's argmax picks, so a draft is kept only when it equals the drawn token. Greedy streams keep the GPU picks untouched, so greedy and sampled streams share rounds.
• families/flashnext/sample.zig: the same file as in Flash Next on Metal: sampled decoding (temperature, top_k, top_p, min_p, seed) on the host ashhart/TensorFold#719. It does a one-pass top-k over bf16 keys, then lanes.sampling.choose, and has a host test against qwen3_5's draw. Whichever lands first, the other merges it unchanged.

One commit, 1874ffe0, on 77ebc3fc; it merges cleanly with the load-limit PR in either order.

Second commit, a wide top-k. top_k 0 (or above 256) went through the reference draw, which sorts every candidate in f64: 13.7 to 16.8 ms a row at this vocabulary, 23 to 30 tok/s served. Now a counting sort of the bf16 keys gives choose's order. Top-p and min-p cut it as choose does, with libm's exp and the same pairwise sum. The keyed Gumbel draw runs in that order and passes over candidates whose noise is provably a nat short of the best score: one hash, no logs. It stops once none can reach the best. That's 0.33 to 0.65 ms a row and 91 to 119 tok/s served. Five sampled replies (top_k 0 with top-p 1 and 0.95, min-p, top_k 500, the default) equal their hashes from before the change. The host test now also compares top_k 0, 300 and the whole row, top-p 1 and min-p, on narrow and wide logits.

Receipt

Mac Studio M5 Ultra (256 GB), macOS 27.0.1, TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP at 45032a19, --parallel 4 --context 16384 --prompt-cache-gib 0 --temperature 0 --no-thinking:

• Greedy: token_sha on four prompts equals the release build's.
• Seeded, at 0.7: the same seed gives the same reply, another seed differs, and drafted equals "draft": false (124 and 129 against 90 tok/s).
• Mixed streams: four sampled streams sent together each equal their solo replies. A round with two greedy and two sampled streams equals their solo replies.
• Speed: a 300-token reply on the shared path ran greedy 123.1 tok/s, sampled at 0.7 122.6 tok/s.
• Tests: zig build test-flashnext-host 53/53; zig fmt --check clean.
• One expected failure, in my own check: 20 seeds at temperature 1 asked for a random number from 1 to 10 all answer 7. The model puts nearly all its weight there.

The same change, on a branch that also carries vision and video, has served a 262k window with four streams and a 24 GiB prompt cache since this morning. At eight streams (131k window) on that build, eight sampled streams sent together each equal their solo replies, twice. Four greedy and four sampled mixed did too, and drafted equals undrafted.

Not run: M1 to M4 (including your M3 Ultra), CUDA (unchanged), and speed-up mode (the host keeps it greedy only).

[cwschroeder]

The new wide-top-k optimization at b7372ac847224ca90180c310da9d479d2757bc0a has a reproducible mismatch against reference(). This finding applies to the added optimization, not the original 1874ffe0 commit.

lanes.sampling.uniform can round to exactly 1.0 in f64. That makes the reference Gumbel score positive infinity, so the early d > 38 break can discard the winning candidate. I reproduced this in Zig 0.17.0 after integrating the current PR head onto upstream 0efbd447 (local integration cf4cf125): the existing 58 host tests pass; adding this case gives 58 passed, 1 failed, expected 4, found 0.

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

Please preserve the current keyed sampler's output when handling this endpoint; changing uniform globally would change seeded replies. Falling back to the reference for the affected case, or proving the pruning safe for the actual floating-point endpoints, would keep the contract. This sampler is shared with ashhart#719, so the fix should stay identical there. The integration check was CPU-only; no new M3 model qualification is claimed.

[BobClawblaw]

@cwschroeder thanks, that's a real one, and a nice catch. (2^53 - 1) · 2^-53 + 2^-54 sits exactly halfway between 1 - 2^-53 and 1.0, rounds to even, and lands on 1.0. The noise is then +inf, and my "noise is at most 37.43" stopping rule skipped right past the winner.

Fixed in d519c085 on this branch, with uniform untouched, so seeded replies are unchanged:

• Past a 38-nat shortfall the draw no longer stops outright. It still asks each later candidate whether its uniform is exactly 1.0 (one hash, no logs). It stops only once the best score is itself +inf, because a later +inf can only tie, and the reference keeps the first.
• The pruning bound now reads 1 - u from the f64 uniform the draw actually uses, instead of rebuilding it from the bits. That subtraction is exact (Sterbenz) wherever the bound can apply, and a u of 1 gives 0, so it is never passed over. The largest finite noise is -log(-log(1 - 2^-52)) = 36.04, since 1 - 2^-53 is never produced, so 38 still bounds it.

Your test is in sample.zig verbatim, plus the same candidate inside the pruning bound (leader at 20.0) and behind a different leader. On the previous commit it fails as you saw (expected 4, found 0); now test-flashnext-host passes 54/54. Speed is unchanged: 0.33 to 0.63 ms a row at the 248k vocabulary, against 13.4 to 16.4 ms for the reference. Five served top_k 0 / top_k 500 / default replies keep their token shas.

The same commit is on ashhart#719 as 6809ccb8, so sample.zig stays byte-identical between the two PRs.