You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Flash Next CUDA: gather sampled top-k candidates on device, retain keyed host draws #1624
Follow-up to #1608 / #1609, stacked on flashnext-cuda-sampled-heads while open.
The #1609 receipt reports five-stream sampled code at 132.8 tok/s versus 242.5 greedy, and sampled chat at 103.3 versus 164.4 greedy. It identifies the remaining host work: downloading and scanning 79,591-column draft rows and 248,320-column verify rows.
The paths are zig/src/families/flashnext/cuda_mtp.zig (drawRow, drawPick) and zig/src/families/flashnext/cuda_mtp_ops.zig (Ops.pick, Ops.draftMany, Ops.draws). download synchronizes the stream and makes a blocking device-to-host copy. Batched draft steps still do this separately for each sampled lane; verify downloads the window's full score rows before host scans.
Proposed change: a fresh CUDA implementation selects exact top-k BF16 candidates for 0 < top_k <= 64, ordered by descending value and ascending original column. Batch eligible draft lanes in one launch/download per step and verify rows in one launch/download per window. Map draft columns through h.ids, then keep lanes.sampling.choose and its 53-bit keyed host draw unchanged. Keep the existing log-sum-exp and probability arithmetic. Greedy rows keep fn_row_best. A NaN anywhere in a row, masks, top_k=0 or top_k>64 keep the whole-row path. TF_FLASHNEXT_SAMPLE_CAND=0 restores the baseline; default on.
Gate: standalone CUDA candidate tests against an independent brute-force host sort across widths including 79,591 and 248,320, k=1..64, ties, signed zeros, infinities and NaNs; sampled/mixed MTP parity with the switch on/off for tokens, probability bits, logits and stopping; kernel-contract registration, controls and backend SUMMARY lines; CUDA host tests, formatting and zero new lean-check findings. Benchmark one binary in A/B/A/B order with #1609's concurrent/alone/serial protocol, plus fresh-server cold 12,001-id prompt receipts. Report exactness failures and unmeasured cells, throughput and TTFT, and unavailable backends explicitly.
Prior art: MiaAI-Lab/Qwen3.8-Flash-Next-Single-DGX-Spark-TensorFold patch 0008 and patch 0011 (MIT); Metal sampled selection/shared rounds #1338 and #1343 by @gilby; Nemotron's CUDA device draws as the existing launch/buffering pattern. This implementation will be written fresh. Nemotron's device random draw is not reused because its precision differs from the host rule.
Overlap review will cover the adjacent EXL3 stack #1441-#1446 and other touching PRs before publication.
Authored by JRB AI in a private capacity. This contribution does not represent the views or positions of any organisation.
Activity
jrb-ai commentedon Oct 29, 2026
/atc claim
jrb-ai commentedon Oct 29, 2026
/atc presence c33da245d7b971bccd4a0027b36a0bd9eff7e521dadd2e8ce7d3c4fac95fcb97