Flash Next CUDA: a shared round's sampled streams draft one at a time, so sampled chat at 5 streams decodes at half the greedy rate on a GB10 #1608

What I see

On stock 1.1.8, a Flash Next stream with a temperature above zero is left out of the shared round's batched MTP head pass and drafts alone, one head forward per draft token. Draft acceptance is the same as greedy, but aggregate decode halves with concurrency. The 1.1.0 notes record this as a known limit ("Sampled streams, trees and media keep drafting one stream at a time", #1069, #1155). Agents and chat clients that send no temperature get the served default (the recipe serves --temperature 1.0), so this is the common case for served Flash Next on a GB10.

Where it happens:

• zig/src/families/flashnext/cuda_lanes_draft.zig:31: if (r.stream.sampling) |sm| if (sm.temperature > 0) continue; in batched.
• zig/src/families/flashnext/cuda_lanes.zig:510 and :513: the lone draft path those streams take.

Environment

• One DGX Spark (GB10, sm_121), kernel driver 580.178.04; CUDA 13.3 (forward-compat 610.43.02) inside nvcr.io/nvidia/pytorch:26.07-py3.
• TensorFold v1.1.8 (03d08f1e) built from source with Zig 0.17.0 (zig build -Doptimize=fast fatbins native), no patches, no TF_FLASHNEXT_* variables.
• nvidia/Qwen3.8-Flash-Next-FP4-Local at b9ce18a688, --context 147456 --parallel 5 --kv-cache bf16 --max-tokens 32768 --temperature 1.0 --top-p 0.95 --top-k 20 --thinking, no vision. Server: "model 73.29 GiB; 5 streams at once, 4.52 GiB each at a 147456-token window, of 27.4 GiB left after a 9.0 GiB reserve; own glue kernels".

Measurements

Streaming /v1/chat/completions, 256 completion tokens, chat_template_kwargs.enable_thinking set explicitly, prompt "Explain why the sky is blue in a few paragraphs." with thinking on. Aggregate is the sum of each request's tokens_per_second from the response's tensorfold stats, median of two repeats; tokens per round and accepted drafts are from the server's done chatcmpl lines. The only difference between the two rows is temperature.

For reference, a greedy code prompt with thinking off on the same server: 69.5, 135.9, 232.1, 272.4 tok/s at 1, 2, 4, 5 streams (5.2 tokens/round, 76–81% accepted). With TF_FLASHNEXT_CONFIDENCE=-0.4 (every chain drafted to full depth) the sampled row falls further, to 37.9 and 76.1 at 1 and 5 streams.

With the change below in place (PR to follow), tools/bench_concurrent.py --levels 1,2,4,5 --tokens 256 --reps 3 --temperatures 1.0,0 --alone --serial on the same Spark at 5 x 131,072, one binary, before/after alternated twice (TF_FLASHNEXT_SAMPLED_HEADS=0 against the default), aggregate tok/s:

Greedy cells and time to first token are unchanged within noise; every concurrent reply equalled its lone run and every lone run equalled "draft": false (0 unequal, 0 failed over 64 cells). The prose table above used my own prompts; these are the public ones.

Prior art

• Metal, 1.1.7: flashnext: shared rounds draft every stream in one packed MTP pass a step #1340 drafts every stream of a shared round in one packed MTP pass, and flashnext: a shared round draws a sampled stream's drafts under its keyed rule #1343 draws a sampled stream's drafts under its own keyed rule at the drafted position (@gilby). flashnext: a shared round draws a sampled stream's drafts under its keyed rule #1343 notes that CUDA's lone path already drafts this way.
• CUDA, 1.1.0: Flash Next CUDA: a shared round's greedy streams draft their MTP heads together #1155's draftMany absorbs every greedy stream's kept rows in one head pass and runs one pass a chain step for the streams still drafting (@jschmied).
• MiaAI-Lab's Flash Next recipe for one Spark carries patch 0007-sampled-streams-in-the-batched-head on v1.1.4, behind TF_FLASHNEXT_SAMPLED_HEADS=1: sampled streams join the batched pass and each draws its drafts on the host from its own row, the draws a lone head step makes.

The change (PR to follow)

A shared round's sampled streams (no trees, no media, as for greedy) join batched and draftMany. At each chain step a sampled lane's score row is downloaded after the shared best-column download and its draft drawn on the host under its own keyed rule at the drafted position, the same draw the lone pick makes (the tail of pick, now shared as drawPick), so drafts, rollbacks, the confidence stop and replies are unchanged; greedy lanes keep their best column. TF_FLASHNEXT_SAMPLED_HEADS=0 restores the previous behaviour and TF_FLASHNEXT_BATCH_HEADS=0 still drafts every stream alone. The parity checker gains sampled and mixed shapes with per-lane seeds and positions (18 shapes, all bit-identical, control caught); test-cuda-host, zig fmt and lean_check.py are clean against 1.1.8; the receipt has the four runs and the parity output.

The rest of the gap between sampled and greedy cells is the host-side sampling work still done a lane at a time: the row download and top-k scan for each draft (79,591 columns) and for each verify row (248,320 columns). Moving those candidates to the device and waiting once a window would be a second step; I would like your view on whether to draw with the shared keyed sampler GLM's target rows use, or to keep the host rule and only move the candidate selection.

I will run atc check and claim this only if the check says it is free.

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