[Arminova]

What this changes

Seven changes to prompt passes on one GPU. Each keeps the bits of main:

• BF16 projections' prompt passes of 1,024 rows or more run on prompt16's 128 x 128 tile with the lane's sums (TF_FLASHNEXT_PROMPT_ROWS sets the row count; TF_FLASHNEXT_PROMPT_BF16=0: the lane). Below 1,024 rows the lane is faster on a GB10 (cold single passes, three rounds: 512 rows +3.6 % on the tile, 768 level, 1,024 -0.8 %, 1,536 -2.0 %).
• A layer's hyper-connection writeback and its read-out's normalize run in one launch (TF_FLASHNEXT_HC_FUSE=0: two).
• A native projection asked for K slices returns its reduced slice and one zero slice, which keeps the old sums' bits (a -0 becomes +0 as before); the other slices are never zeroed.
• tf_route_logits_wide walks the expert blocks first, so a row tile's input rows are read once (its contract test launches the new grid).
• fn_qsa_select keeps 8 keys in flight a pass and places its picks by warp runs: the same ids, ties included.
• The fused prompt DeltaNet kernel runs its batch's gated norm on the chain warps.
• The prompt scratch sizes head_local, native_part, hc_part and part by what reads them; the prompt cache gets the rest.

The default chunk stays 1,536. Five new kernels (fn_hc_writeback_norm0 / _moe / _slices, tf_prompt_bf16 / _f32): the writeback-and-normalize ones registered as bit-equal replacements of the two launches, the BF16 tile under the projection contract's FP64 reference. fn_qsa_select and fn_gdn_fused take more registers; their nvcc 13.3 records are updated.

#1763 is in main. main 088cfbd35 is merged into this branch (a merge commit on top, nothing rewritten), so the diff is this PR's own change, commit 69c4cc49e. 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), on main af9d7cde7. Base: flashnext cuda: pending DeltaNet rows, reused lane buffers and the shared expert on a side stream #1763's commit (890c4e6b1); this change 69c4cc49e. Every server takes one warm-up prompt before the timed ones.

• Exactness: every reply's token_sha256 equals the base's at 1, 4 and 16 streams, greedy and seeded sampled, code and chat, and on every cold prompt below (with the prompt cache on and off, and with TF_FLASHNEXT_CHUNK=8192 TF_FLASHNEXT_DRAFT_HEAD=affine4). bench_concurrent --alone --serial at 1 and 4 streams (3 reps) and at 16 streams (--context 32768 --parallel 16): every reply equals its solo and serial run. Drafted replies equal "draft": false on four prompts (300 to 15K tokens). A 32K prompt repeated (resumed from 32,767 cached tokens) and sent concurrently equals the first.

• Prompt speed, base → this change, tok/s, cold prompts (bench_cache --require-cold, 0 cached tokens; exact token arrays: main's docs/**/*.md tokenized by the server), mean of two alternated rounds:

  --prompt-cache-gib 0	2,048	8,192	32,768	65,536
  default settings otherwise	1,592 → 1,692 (+6.3 %)	1,695 → 1,811 (+6.8 %)	1,681 → 1,791 (+6.5 %)	1,660 → 1,761 (+6.1 %)

  prompt cache at its default size	512	2,048	16,384
  default settings	1,270 → 1,281 (+0.8 %)	1,586 → 1,642 (+3.5 %)	1,681 → 1,769 (+5.2 %)

  	8K	32K	128K
  TF_FLASHNEXT_CHUNK=8192 TF_FLASHNEXT_DRAFT_HEAD=affine4 (base once, this change twice)	1,910 → 2,044 (+7.0 %)	1,873 → 2,015 (+7.6 %)	1,811 → 1,934 (+6.8 %)

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

  	greedy	seeded sampled
  bench_openai code / chat (256 tokens, 3 reps)	67.1 → 67.2 (0 %) / 42.2 → 42.2 (0 %)
  bench_concurrent code x1 / x4	66.9 → 66.9 (0 %) / 120.7 → 120.0 (-0.5 %)	65.0 → 64.9 (0 %) / 89.3 → 89.3 (0 %)
  bench_concurrent chat x1 / x4	42.1 → 42.2 (0 %) / 78.7 → 78.7 (0 %)	40.1 → 40.0 (0 %) / 66.2 → 65.6 (-1 %)
  16 streams (--context 32768 --parallel 16) code / chat	284.4 → 279.9 (-1.6 %) / 247.6 → 246.2 (-0.5 %)	200.7 → 200.1 (0 %) / 167.4 → 167.1 (0 %)

  The 16-stream greedy code cell moves 272-286 tok/s between runs of one build (this change alone: 272.2, 278.2, 281.6 and 285.1 in four runs); the other cells stay within 1 %.

• Memory: with 8,192-row chunks the prompt scratch is 3.2 GiB smaller, so the prompt cache gets 933 → 4,221 MiB; at default settings the difference stays within the run-to-run spread of the host's free memory.

• Tests on the head: zig build test 456 / 456 steps and test-cuda-host 276 / 276 steps, 0 failures. tools/zig/lean_check.py: 424 problems, the same as main af9d7cde7 (0 new).

• Kernel contract (nvcc 13.3.73, GB10), on this change: every gate line of this change's kernels passes.

  SUMMARY 535 kernels: 535 pass gates 1-4, 0 fail, 6 known findings; 344 controls: 344 red as planted, 0 missed; 229 limit problems

  kernel-contract-limits on fatbins for sm_80 / 86 / 89 / 120 / 121: none of this change's records stale, no launch over a limit; the 274 stale records and the 229 limit problems above are main's own 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 kernels do 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.

[ashhart]

Thanks @Arminova. Our GB10 checks are in, base, new, new, base on each:

• FP4 (Qwen3.8-Flash-Next-FP4-Local): GREEN. Cold prompts are faster at every size and start:

  prompt	main	flashnext cuda: prompt passes with the same bits and less work, and a prompt scratch sized by use #1764	change
  2,048	1,645.3	1,702.8	+3.5%
  2,048, repeated	1,667.8	1,740.7	+4.4%
  8,192	1,744.6	1,825.3	+4.6%
  8,192, repeated	1,740.7	1,827.1	+5.0%
  16,384	1,748.7	1,839.0	+5.2%

  Both new runs are above both base runs in every cell, and no cell is slower.

• MLX 4-bit + MTP: GREEN. The prompt cells are flat, decode after a repeated 64K prompt is 43.7 → 45.4 (+3.7%), and nothing is slower.

This lands now. Your branch carries #1763's pre-squash commit under its merge, so we landed your own change, 69c4cc49e, cherry-picked onto current main with you as its author. It applied cleanly.

On the overlap with @laudney's #1846 and #1847 (the same 128 x 128 BF16 prompt tile and the same writeback-and-normalize fuse): we're landing this one first. It's measured on both checkpoints, it runs on any SM behind its switches, and it covers six parts theirs doesn't. We'll ask them to rebase what's left onto it.

[Arminova]

Thanks, Ash, and for checking it on both checkpoints. One heads-up for the cherry-pick: with #1739 in main, 69c4cc49e takes cuda_batch.zig to 604 lines, so the lean check will flag it. One commit moves the two fused read-out branches into cuda_batch_hc.zig (attnReadout, mlpReadout: the same calls in the same order) and keeps it at 592: Arminova/TensorFold@5c5228cea on our fork's flashnext/prefill-lean branch. It applies on top of your cherry-pick as is (lean 0 new, compiles), and on our GB10 every reply equals main's at 1, 4 and 16 streams.

[ashhart]

Thank you @Arminova. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

[ashhart]

Landed as 9693706. Thanks @Arminova!