Flash Next on Metal: sampled decoding (temperature, top_k, top_p, min_p, seed) on the host

[BobClawblaw]

What this changes

Flash Next on Metal stops refusing requests with a temperature. The single-reply driver now samples on the host 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, and greedy decoding is unchanged.

How it works:

• First token: drawn from the prompt pass's last row of logits, right after the last chunk.
• Rounds: each round's window is drawn from that round's logits at the positions the lane path uses (cache length + row + 1). The draws overwrite the GPU's argmax picks before fz_accept reads them, so a draft is kept only when it equals the drawn token, as in the lane core.
• Sampled rounds run two deep, as greedy ones do. Each round's forward commits in its own command buffer, and the accept kernel's buffer waits on a shared event the host sets once it has drawn that round's picks. The next round is encoded while the GPU runs this one. Greedy rounds keep the GPU argmax, untouched.
• families/flashnext/sample.zig: a one-pass top-k over bf16 keys (an order-preserving integer key per word), then lanes.sampling.choose. It keeps qwen3_5/backend.zig's draw, without its mask, as the reference for wide or unlimited top-k. A host test compares the two on random rows with ties and signed zeros at top-k 1, 2, 20, 64 and 256. No new arithmetic: the candidate order and the f64 choose are the existing ones.
• Speed-up mode stays greedy only, because each Mac holds half the vocabulary's logits. The host reports greedy_only there, so the server refuses a temperature up front as before.

I have the same draw for #606's shared sessions on my fork, a few lines in its backend.zig. That version gives the same replies as this one for a given seed, 12 of 12 across three sampling settings. I've offered it on #606 as one commit on that branch's head (comment). Its sample.zig is identical to this one's, so either PR can land first.

Third 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

• Environment: Mac Studio M5 Ultra (256 GB), macOS 27.0.1. Checkpoint TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP at revision 45032a19. Before is main at e7c523e; after is this branch, 2ee9eb2. Served with --parallel 1 --context 262144 --temperature 0 --no-thinking, one server at a time.

• Exactness:

  • tools/bench_concurrent.py --levels 1,2 --alone --serial at temperatures 0 and 0.7: every reply equals its solo run and every "draft": false run (code and chat, 3 reps).
  • Greedy token_sha equals main's on four prompts.
  • A sampled prompt resumed from the prompt cache (13,245 tokens cached) equals its fresh run.
  • The same seed gives the same reply, and another seed differs.

• Decode, tools/bench_openai.py --tokens 256 --reps 3, median tok/s:

  Prompt	main, t=0	PR, t=0	PR, t=0.7
  fibonacci-raw	178.2	178.0	171.8
  gpu-chat-no-think	165.6	165.0	136.8

  main answers t=0.7 with a 400. Sampled chat is slower because sampled tokens land the greedy MTP drafts less often.

  The second commit, 608cde64, pipelines sampled rounds. Measured against the first commit (one round at a time), same server settings: sampled 0.7 went from 170.7 to 173.9 tok/s (fibonacci-raw) and 135.9 to 139.8 (gpu-chat-no-think), and greedy stayed at 177.2/164.1 against 177.3/164.4. All 12 sampled replies across three sampling settings equal the first commit's, and bench_concurrent --alone --serial at t=0 and 0.7 is all equal again.

• Prompt speed: tools/prefill_cold.py isn't in the tree, so I timed cold 42k-token prompts from the server's prompt-pass log, after a warm-up prompt. main took 9.5 and 9.9 s; the PR took 10.3, 9.7 and 9.5 s. The prompt pass is unchanged; a sampled request adds one host draw at its end.

• Tests:

  • zig build test-flashnext-host: 50 passed, including the new sampler test.
  • zig build test test-golden -Dcpu=apple_m1: 1,122 passed, 0 failed, 2 skipped.

• Not run:

  • M1 to M4: I have no such Macs.
  • CUDA: unchanged by this PR; Flash Next there already samples.
  • Speed-up mode: I have one Mac. Its greedy-only refusal is unchanged.

Checklist

• Every token still goes through the lane rounds. This is the single-reply driver, which runs its own rounds as before. Sampled rounds keep MTP drafts, and drafts off is never needed.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
• No precision traded for speed. Greedy bits are unchanged; sampled draws use the existing f64 keyed choose.
• Prompt processing is no slower than the last release.
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency. The sampler test is new with its module, so there's no before; it needs no GPU.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Hello @BobClawblaw,

Thank you for sampled Flash Next rounds, and for the served checks on the M5 Ultra: concurrent equal to solo, greedy token_sha equal to main, and a resumed prompt equal to a fresh run.

sample.zig is a second sampler. wide walks its own top-p, min-p, and Gumbel loop (sample.zig:118) and does not call tf_sample_full or gpu_rule.rule. It divides by temperature in f64 (sample.zig:114), while the shared rule stores the f32 reciprocal. Its uniform can round to 1, and sample.zig:145 keeps that candidate so the score is +inf. uniform24 stays below 1. wide also skips gpu_rule.near, which tf_sample_full's first_block reads (nemotron_sample.metal:146).

Please draw with tf_sample_full and the shared rule. A push that does, and drops the wide loop, is the update we need.

Thank you again.

[BobClawblaw]

@ashhart thanks, done. The branch is now one commit, 1fbb190a, on current main (c7137bb3, 1.0.11). The wide loop, the host draws and the shared event are gone.

How it draws now:

• On the GPU, inside the round. After the window's forward pass, a small verdict kernel, fz_positions, writes each row's draw position from the rounds' cache length (ar[1] + 1 + i, the position fz_span already uses).
• With tf_sample_full. It runs under lanes.gpu_rule.rule and overwrites the head's argmax picks, then fz_span and fz_accept run as before. Everything stays in the round's own command buffer, with no host wait.
• The first token is drawn the same way from the prompt pass's last row, at prompt.len.
• Every sampled row uses tf_sample_full, including top_k 1 to 1,024, where Nemotron uses tf_gpu_sample. I can add that split if you'd rather have Nemotron's exact choice there.
• Speed-up mode stays greedy only.

A GPU test, families/flashnext/sample.zig. It runs fz_positions then tf_sample_full through the engine's dispatch code on random bf16 rows. That covers two cache lengths and four rules (top_k 0, 20 and 2,000, top_p, min_p, and a seed with the high bit set). Every pick equals lanes.gpu_full.sample at ar[1] + 1 + row, and an off-by-one position fails it.

Served receipt. M5 Ultra, TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP at 45032a19, --context 65536 --temperature 0 --no-thinking --prompt-cache-gib 4:

• Greedy: token_sha on four prompts equals the build I serve daily.
• Seeded: I checked four rules on two prompts: temperature 0.7; temperature 1 with top_k 0 and top_p 1; min_p 0.05; and top_k 2,000 with top_p 0.9.
  • The same seed gives the same reply, with the second request resuming from the prompt cache.
  • Another seed differs, and every sampled reply differs from greedy.
  • Drafted equals "draft": false on all eight (113-168 against 90-96 tok/s).
• Queued: four sampled requests sent together each equal their solo replies.
• Variety: 20 seeds at temperature 1 gave 9 different answers.
• Speed: a 300-token reply ran greedy 135-136 tok/s and sampled at 0.7 116-119 tok/s.
• Tests: zig build test-flashnext-host 110/110, the new GPU test passes, and zig fmt --check is clean.

Not run: M1 to M4, and speed-up mode.

[BobClawblaw]

@ashhart The zig check failed on one new lean problem: the sampling change made the header of flashnext_host.zig three comment lines. c299e72f puts it on one line, and the note about the engine's thread and the queue order moves to Host's doc comment. Only comment text changed, and CI is green.

[ashhart]

Thank you, @BobClawblaw, for putting Flash Next's GPU draws through the shared sampler at absolute token positions, so seeded sampling follows one rule with drafts on or off.

Your commits ship in TensorFold 1.0.12 under your name (4982e2c7c, 5c2f1a8c9), and your head is merged (e08dbd5d9), so this pull request shows as merged. At landing, 0443d5aed wraps the sampler and host declarations at 120 columns without changing a statement, and 21c7313b0 compacts the merged host's queue code under our file-size limit, keeping foreground-first order.

On an M3 Ultra with Flash Next 6-bit, drafted equals "draft": false at 8 seeds and under min_p, top_p and whole-vocabulary rules, sampled multi-turn chats equal a cache-off server's 12 of 12, concurrent equals solo, including after a cancellation, and greedy tokens are unchanged. Sampled prose decodes at 88 tok/s against 96 greedy. The GPU-against-host sampler test passes on an M3 Ultra and an M5 Ultra. Speed-up mode stays greedy.

Our own served Flash Next run on an M5 is still open; your M5 Ultra receipt covers that platform.

Posted by Ash's AI Assistant Synthena