[jrb-ai]

What this changes

A shared round's Flash Next streams with a temperature above zero now draft in the batched MTP head pass on CUDA instead of one at a time: each step, after the shared best-column download, a sampled lane's score row is downloaded and its draft drawn on the host under its own keyed rule at the drafted position, the draw its lone pick makes, so drafts, probabilities, the confidence stop and replies are unchanged. Greedy lanes keep their best column. TF_FLASHNEXT_SAMPLED_HEADS=0 restores the previous behaviour; TF_FLASHNEXT_BATCH_HEADS=0 still drafts every stream alone. The parity checker gains sampled and mixed shapes with per-lane seeds and positions. Fixes #1608.

Receipt

Full receipt in docs/receipts/flashnext-cuda-sampled-heads/ (README, the four bench_concurrent reports, the parity output, commit.txt).

• Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision: TensorFold 1.1.8 (03d08f1e) plus this change (measured there; the branch is rebased onto main at c9c3aa03, which does not touch these files, with the source byte-identical), built from source with Zig 0.17.0 and nvcc 13.3 in nvcr.io/nvidia/pytorch:26.07-py3 (PyTorch unused by the native server); one DGX Spark (GB10, sm_121, 128 GB), kernel driver 580.178.04, CUDA 13.3 forward-compatibility driver 610.43.02; nvidia/Qwen3.8-Flash-Next-FP4-Local at b9ce18a6887e18c089fdea8509a33e63429a3300; --context 131072 --parallel 5 --kv-cache bf16 --temperature 1.0 --top-p 0.95 --top-k 20 --thinking.
• Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false:

check	result
bench_concurrent.py --alone (every concurrent reply against its lone run), baseline and this PR, two runs each, 64 cells	0 unequal, 0 failed
bench_concurrent.py --serial (every lone run against "draft": false), baseline and this PR, two runs each, 64 cells	0 unequal, 0 failed
tf-cuda-test flashnext-mtp-parity: 6 greedy shapes (confidence 0, 2) and 12 sampled and mixed shapes (confidence 0, 0.5)	18 pass: logits byte-equal, proposals, probabilities and the stop equal; planted control caught

• Decode speed before and after, from tools/bench_openai.py: measured with bench_concurrent.py on the same prompts and settings, one binary. Baseline is TF_FLASHNEXT_SAMPLED_HEADS=0, the behaviour main has today; this PR is the change (on by default). Two runs each, alternated baseline, this PR, baseline, this PR; each figure is the mean of the two runs (per-run values are in the receipt README). Aggregate tok/s:

output	temperature	streams	baseline tok/s	this PR tok/s	change
code	1	1	67.1	67.4	+0%
code	1	2	77.7	82.5	+6%
code	1	4	99.9	118.8	+19%
code	1	5	107.3	132.8	+24%
code	0	1	67.2	67.5	+0%
code	0	2	123.9	124.5	+0%
code	0	4	211.1	211.4	+0%
code	0	5	242.4	242.5	+0%
chat	1	1	42.1	42.5	+1%
chat	1	2	61.1	65.3	+7%
chat	1	4	82.2	94.6	+15%
chat	1	5	87.0	103.3	+19%
chat	0	1	42.6	42.7	+0%
chat	0	2	80.8	81.2	+0%
chat	0	4	139.4	141.4	+1%
chat	0	5	162.7	164.4	+1%

Time to first token, the median request's ttft_s in the same runs, mean of the two runs:

output	temperature	streams	baseline TTFT	this PR TTFT
code	1	1	96 ms	96 ms
code	1	2	142 ms	144 ms
code	1	4	239 ms	237 ms
code	1	5	283 ms	286 ms
code	0	1	96 ms	96 ms
code	0	2	143 ms	144 ms
code	0	4	238 ms	240 ms
code	0	5	286 ms	287 ms
chat	1	1	83 ms	84 ms
chat	1	2	122 ms	122 ms
chat	1	4	210 ms	213 ms
chat	1	5	256 ms	252 ms
chat	0	1	86 ms	86 ms
chat	0	2	128 ms	128 ms
chat	0	4	212 ms	214 ms
chat	0	5	263 ms	254 ms

• Prompt speed before and after: tools/bench_cache.py --tokens 64 --repeats 2 --streams 2 --require-cold on the 12,001-token fixture prompt (zig/tests/cuda/flashnext/lr7/long12k.ids), a fresh server for each side, both qualified cold (prefill-baseline.json, prefill-this-pr.json in the receipt):

request (12,001-token prompt, 64 tokens)	baseline	this PR
cold first request, ttft_s (mean of the two cold requests)	7.22 s (7.39, 7.05)	7.22 s (7.29, 7.14)
prompt tokens over that time	1,662 tok/s	1,663 tok/s
repeated request on the kept prompt, ttft_s	0.088 s	0.092 s
two concurrent requests on the kept prompt, ttft_s	0.084 s, 0.258 s	0.091 s, 0.238 s
decode, first request	39.8 tok/s	39.5 tok/s
repeated and concurrent replies equal to the first	4 of 4	4 of 4

• Tests run, with pass, fail and skip counts, and not run, and why:

step	result
zig build test-cuda-host (the new drawPick test included)	191 pass, 0 fail, 2 skip
tf-cuda-test flashnext-mtp-parity on the GB10	18 pass, 0 fail
zig fmt --check on the changed files	clean
tools/zig/lean_check.py	425 on the branch, the same 425 as main at c9c3aa03 (428 on both at 1.1.8, when measured); none added
zig build compile-flashnext-mtp-parity-host	builds
Metal, HIP, XPU	not run: the change touches CUDA Flash Next files only
kernel-contract	not run: no kernel added or changed
TP=2	not run: one Spark

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
• No precision traded for speed. If the bits change, the description says which sums change and why.
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. (No kernel changed.)
• 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. (CUDA GB10 only.)
• New tests fail before the change, pass after it, and skip cleanly without their dependency. (The sampled and mixed parity shapes cannot run against 1.1.8's draftMany, which has no sampling input; the planted control fails as before.)
• 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.

Prior art: Metal's shared rounds draw a sampled stream's drafts under its keyed rule since 1.1.7 (#1340, #1343, @gilby); draftMany for greedy streams on CUDA is #1155 (@jschmied); MiaAI-Lab's Flash Next recipe for one Spark carries the same idea for CUDA on v1.1.4 as its patch 0007 (TF_FLASHNEXT_SAMPLED_HEADS=1). This is an independent implementation against 1.1.8.

Overlap (atc check 1608: free, claimed; git merge-tree against every open PR that touches these files): #1446, the last part of the EXL3 Flash Next stack, inserts its deferred-draft check right after the line in batched that this changes, and adds EXL3 early returns beside pick and draft; the hunks are adjacent, not the same logic, and both apply (keep this skip line, then the EXL3 check). I will rebase onto whichever lands first. CHANGELOG.md conflicts with every open PR, as usual.

A pre-existing observation, untouched here: headGroup (the batched head's pass) does not pass Draft.pg where forwardHead does, so batched heads and paged KV (TF_FLASHNEXT_KV_PAGES, off by default) have not met.

Authored by JRB AI in a private capacity. This contribution does not represent the views or positions of any organisation.