[Arminova]

What this changes

Seven changes to drafted rounds on one GPU. Replies keep the bits of main; drafts change, accepted tokens do not. Each has its own switch:

• A drafted stream's first token is sent before its first drafts (TF_FLASHNEXT_FIRST_EARLY=0: after).
• The n-gram table is read ahead into the page cache, and loaded weight shards are dropped from it (TF_FLASHNEXT_TABLE_WARM=0, TF_FLASHNEXT_DROP_LOADED=0).
• The draft chain stops on the running product of its probabilities; the core gets the chain's chances.
• Greedy streams draft up to 8 deep; sampled streams keep 6.
• Sampled draws read a row's top_k candidates from the GPU instead of the whole row (TF_FLASHNEXT_GPU_TOPK=0: the whole row).
• A drafted verify window leaves its scores on the device (TF_FLASHNEXT_VERIFY_HOST_LOGITS=1: copied).
• Drafted windows, draft-head pieces and rollbacks run as CUDA graphs on one GPU (TF_FLASHNEXT_WINDOW_GRAPH=0: plain launches).

One new kernel, fn_row_topk (a row's K best in column order, ties to the lowest columns), registered with the kernel contract against a host oracle; its set joins the family sets. Off with tensor parallelism. New files cuda_mtp_draw, cuda_mtp_pieces, cuda_window_graph.

#1445 adds CUDA graphs on the EXL3 path; this PR's graphs are the FP4-Local path's drafted windows, in their own file. Glad to line the two up if you prefer one graph cache.

Depends on #1764: its first commit is #1763's, its second #1764's change rebased onto #1763's current head; this PR's own change is the third commit. Best reviewed after #1764. Part of #1762.

Receipt

• Environment: as flashnext cuda: pending DeltaNet rows, reused lane buffers and the shared expert on a side stream #1763 (GB10, driver 580.173.02, nvcc 13.3.73, nvidia/Qwen3.8-Flash-Next-FP4-Local at b9ce18a6, --context 262144 --parallel 3, bf16 KV, prompt cache at its default size, a fresh server for every run). Base: the second commit here (flashnext cuda: prompt passes with the same bits and less work, and a prompt scratch sized by use #1764 on flashnext cuda: pending DeltaNet rows, reused lane buffers and the shared expert on a side stream #1763's head). The runs used dcf28f53a; the head adds one line to draft_sources_cuda_test.zig (its Cuda sets the new depth_greedy) and fn_row_topk's launch record measured in the contract run below, neither part of the server build.

• Exactness: every reply's token_sha256 equals the base's at 1, 4 and 16 streams, greedy and seeded sampled, code and chat, at default settings, and at 1 and 4 streams with TF_FLASHNEXT_CHUNK=8192 TF_FLASHNEXT_DRAFT_HEAD=affine4. bench_concurrent --levels 1,4 --alone --serial: every reply equals its solo and serial run. Drafted replies equal "draft": false on four prompts (300 to 15K tokens). bench_cache on cold 512 / 2,048 / 16,384-token prompts gives the base's replies, and a 32K prompt repeated (resumed from 32,767 cached tokens) and sent concurrently equals the first.

• Decode speed at default settings, base → this change, tok/s, mean of two alternated rounds (base, this change, this change, base):

  	greedy	seeded sampled
  bench_openai code / chat (256 tokens, 3 reps)	67.5 → 70.5 (+4 %) / 42.5 → 45.6 (+7 %)
  bench_concurrent code x1 / x4	67.3 → 70.3 (+4 %) / 119.8 → 126.6 (+6 %)	65.2 → 68.8 (+5 %) / 90.5 → 98.0 (+8 %)
  bench_concurrent chat x1 / x4	42.5 → 45.5 (+7 %) / 79.2 → 82.1 (+4 %)	40.3 → 43.6 (+8 %) / 66.0 → 70.6 (+7 %)
  16 streams (--context 32768 --parallel 16) code / chat	272.9 → 330.8 (+21 %) / 245.5 → 308.4 (+26 %)	200.0 → 207.7 (+4 %) / 166.7 → 179.2 (+7 %)
  decode after a cold 512 / 2,048 / 16,384-token prompt	69.3 → 71.5 (+3 %) / 39.6 → 43.0 (+8 %) / 46.4 → 52.0 (+12 %)

  With TF_FLASHNEXT_CHUNK=8192 TF_FLASHNEXT_DRAFT_HEAD=affine4 (one round): bench_concurrent code x1 / x4 76.2 → 81.6 / 133.6 → 143.0, chat 47.1 → 49.7 / 85.9 → 90.4.

• First token, ms (lower is better), same rounds: a new ~90-token prompt, one stream, 233.8 → 229.4; the same prompt resumed 71.3 → 63.5; after cold 512 / 2,048 / 16,384-token prompts (bench_cache --require-cold) 420 → 379 / 1,260 → 1,165 / 9,379 → 9,028.

• Memory: unchanged. Both builds hold the same fixed use within the log's 0.1 GiB rounding; the prompt cache's size follows the host's free memory at start on both.

• Tests on the head: zig build test 449 / 449 steps and test-cuda-host 274 / 274 steps, 0 failures.

• Kernel contract (nvcc 13.3.73, GB10), run with upstream's fbbc20d0f applied (the brace in contract_controls.cu; main 20705d604 does not build the runner without it):

  SUMMARY 518 kernels: 518 pass gates 1-4, 0 fail, 6 known findings; 342 controls: 342 red as planted, 0 missed; 226 limit problems

  fn_row_topk's gate lines all pass. kernel-contract-limits: its record in this push is the run's RECORD line (sm_121: 31 registers, 9,228 static shared bytes); the run itself still had the estimate, which is the one limit problem of this PR's. The others are main's own records on this nvcc build.

• Not run: TP2 (one GB10; every change keeps today's path with tensor parallelism), Metal and HIP (CUDA-only change), CUDA chips other than GB10 (launch records only).

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.
• 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 new kernel does not exist before this change.)
• 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.