[BobClawblaw]

What

Copy drafts for Flash Next on CUDA. When a stream's last 8 tokens occur earlier in its prompt or reply, the round drafts that text's continuation whole, up to the depth, instead of asking the MTP head. The verify step is unchanged, so every reply equals one-token decoding. Nemotron already does this (nemotron_h/cuda/decode.py: CopyIndex); this is the same index, in its own module for the family, wired into the four places Flash Next drafts: the serial MTP loop (decode.py: mtp_decode), the shared rounds (multi.py: _draft_all), the lone stream's graph slot (multi_solo.py: _solo_round) and the first drafts after a prompt (multi_fill.py: _joined). The kept rows still go into the MTP cache, so the head resumes cleanly on the next round that needs it.

Two ranks: both hold the same tokens (sampling is gathered), so both build the same index and propose the same chain, and the round plan digest (multi_tp.shape) already covers each stream's drafts. In the shared rounds the first chain level's picks still cover every absorbed stream, because the two-rank gather reads them in order; a copying stream's pick is dropped and it joins no later level.

On by default with MTP drafts; TENSORFOLD_COPY_DRAFTS=0 turns it off. Off without MTP drafts and for grammar-constrained replies. A 250k-token prompt indexes in 0.12 s at admission; a chain lookup is microseconds.

I read #286: the Python engine is frozen and new work goes to the Zig engine. Flash Next is not on CUDA there yet, so this is offered as the reference implementation and the measurements for when it is; rewrite it however fits. #444 describes a prompt-lookup drafter for this family on EXL3 with the same conclusion (copy-heavy replies gain, fresh text does not); this one reuses Nemotron's index and runs on two ranks.

Receipt

Environment. TensorFold cb2ebf0 (v0.6.6 = main) plus this branch, inside nvcr.io/nvidia/pytorch:26.07-py3 (PyTorch 2.13.0a0, CUDA 13.3 forward-compat on kernel driver 580.178.04). Two NVIDIA DGX Spark (GB10, 128 GB) at --tp 2 --parallel 16, --kv-dtype int8, --mtp-drafts 15 --mtp-confidence 0.70, context 262144, NCCL over one RoCE HCA. Checkpoint TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP at 2b170fa6309d5d1ee380b35636075fac7945f286. "Before" is the same image and session with TENSORFOLD_COPY_DRAFTS=0.

Exactness. tools/bench_concurrent.py ... --levels 1,4 --tokens 512 --alone --serial, before and after: every concurrent reply equal to its solo run (alone 1/1 and 4/4 in all cells) and every solo run equal to "draft": false (serial 1/1 in all cells). Five one-stream requests sent twice, the second with "draft": false, token_sha equal in all five, before and after (table below).

Decode speed, tools/bench_openai.py --reps 2 --temperatures 0:

prompt	before	after
fibonacci-raw	91.9 tok/s	92.9 tok/s
gpu-chat-no-think	96.5	97.3

tools/bench_concurrent.py, aggregate tok/s: code 1 user 108.4 -> 103.7, 4 users 323.5 -> 320.7; chat 1 user 89.1 -> 87.9, 4 users 277.6 -> 273.2 (within this pair's run-to-run noise; those prompts repeat nothing).

The cells the change is for, one stream, greedy, 900-token replies, tensorfold.decode_s and rounds from the reply:

cell	before	after	tokens per round	"draft": false
copy a 3 kB passage verbatim	127.8	337.1	6.3 -> 15.5	47.7
fix typos in the passage	129.5	321.7	6.0 -> 15.5	47.3
rename a variable across a 3 kB file	129.3	288.6	6.4 -> 13.9	47.7
fresh prose (500-word essay)	70.6	71.5	2.1	48.3
fresh code (module plus tests)	127.1	124.7	5.7	47.5

Prompt speed, tools/prefill_cold.py (2k to 64k cold prompts, 3 reps each), copy drafts off, on, and on again in the same session:

prompt tokens	before: TTFT, prompt tok/s	after	after again
2,048	0.85 s, 2,400	0.85 s, 2,398	0.86 s, 2,377
8,192	3.25 s, 2,517	3.27 s, 2,502	3.28 s, 2,499
16,383	6.51 s, 2,518	6.54 s, 2,505	6.54 s, 2,504
32,768	13.21 s, 2,480	13.29 s, 2,465	13.45 s, 2,436
65,536	27.42 s, 2,390	27.55 s, 2,379	27.66 s, 2,370

Within the pair's noise at every length (the index is built after the prompt's first token is sampled: 0.12 s for 250k tokens on the host, 1 ms at 2k).

Tests. tests/test_copy_drafts.py (6, new; fails before the change with ModuleNotFoundError on the new module) and the full upstream suite inside the image on the CPU: this branch 4096 passed, 127 failed, 11 errors, 175 skipped; pristine main in the same container 4090 passed, 127 failed, 11 errors, 175 skipped, and the failure and error sets are identical (MLX-only tests and Triton kernel tests that cannot run on the CPU, such as test_affine_rows.py and test_exact_sampling.py). The six extra passes are tests/test_copy_drafts.py. The first push had one extra failure, test_cuda_ignore_eos_engines.py::test_every_path_takes_stop_eos[True-False-flashnext]: that test builds the engine without __init__, so the new attribute was missing; the reads are now getattr with a False default, as the surrounding code does for optional attributes.

Not run. Metal: no Mac here, and the change is CUDA-only (families/qwen4_exp/cuda). One GPU (--tp 1) and the EXL3 and NVFP4 formats on one GPU: not served on this pair; the code paths are the same four draft sites. --parallel 1 at --tp 2 (the serial engine on two ranks) was served earlier the same night on the previous image for the serial-against-concurrent comparison, but the "draft": false check above ran on the concurrent profile only.

[BobClawblaw]

Receipt completed in the body: cold prefill before, after and after again (within noise at 2k to 64k), and the full suite on this branch against pristine main in the same container (identical failure and error sets; the six extra passes are the new tests). One fix since the first push: test_cuda_ignore_eos_engines.py builds the engine without __init__, so the new attribute reads are now getattr with a False default (6bf4caf).

[ashhart]

Thank you! Copy drafts for Flash Next on CUDA are merged into main with your commit as it is (6bf4caf). Copying a 3 kB passage going from 127.8 to 337.1 tok/s is a lovely result. Reusing Nemotron's CopyIndex in all four places Flash Next drafts, with the kept rows still feeding the MTP cache, is how we'd want it done.

Before merging we ran the Flash Next CUDA suite with real kernels on a Blackwell card, with and without your change. Everything that passes on main passes with it, and your six new tests pass too. After it, in ed78d6f, we folded the multi-line docstrings onto one line each, which is our house style. The code is unchanged.

The Python engine is frozen at 0.6.6, so this lives in main rather than a tagged release. In the Zig engine, copy drafts come from the shared core, so Flash Next on Zig CUDA (#476) gets them there.