[Arminova]

What this changes

With TF_FLASHNEXT_KV_PAGES set, a shared round where two or more streams draft runs the MTP head over one grouped
pass (headGroup in cuda_mtp_ops.zig). Its lanes were built without the draft's page table, so every paged lane's
head read and wrote its cache rows contiguously from the pool's first page: the same addresses for each of them, in
pages other streams own. A lone head's pass already carries .pg; the grouped lane now does too. One commit, one
file, no kernel changes. Without paging d.pg is the contiguous default, so default settings keep main's bits.

On a GB10 with paging and three streams (receipt below), main's concurrent replies differed from their runs alone
in 1 of 18 (FP4) and 5 of 15 (MLX) rows a server, and this change decodes 2-24 % faster there; one stream and the
default settings are level.

How it came about: the grouped head pass came in 5e77d736c, then Flash Next's KV pages in 5af5af721, which set
.pg on the lone head's pass; the grouped lane kept its form from before paging.

• Issue: n/a, a bug fix this size (found while testing paging with three streams at once).
• Stacked on: main 002c9e03e; merges clean with today's main 9200a3064 (10 commits, none in Flash Next).
• atc overlap --json: not run (ATC is not installed here). The same test by hand, git merge-tree of this commit
  against each of the 12 open pull requests whose file lists touch cuda_mtp_ops.zig: clean with EXL3 Flash Next: expert-parallel TP=2 (same tokens as TP=1) and MTP drafts (2/8) #1442, EXL3 Flash Next: deep prefill (3/8) #1443, flashnext cuda: prompt cache marks inside a chunk are kept on the way instead of cutting the pass #1866
  and flashnext cuda: beside another open stream a greedy chain stops at 0.6 per draft; wide rounds draft it 6 deep #2160; one new conflict, cuda_mtp_ops.zig on this line, with EXL3 Flash Next: KV pages from the core pool (4/8) #1444 and the ones stacked on it (EXL3 Flash Next: CUDA graphs and the device planner (7/8) #1445, EXL3 Flash Next: draft lane guard for offloaded experts (8/8) #1446,
  EXL3 Flash Next: host-resident K/V on the tier (6/8) #1522); Flash Next TP=2: exact output-column splits, 21% faster plain decode #1784, Flash Next shared down: remove the scatter with direct slot output #1792, Flash Next BF16 prompts: larger CUDA tiles with identical output bytes #1846 and Flash Next BF16 prompts: reduce shared staging with K32 loads #1848 already conflict with main, and this commit adds no file to theirs.
• Overlap: EXL3 Flash Next: KV pages from the core pool (4/8) #1444 (@BlivionIaG, EXL3 K/V pages, 4/8) moves both head paths onto a Draft.headSeq that carries .pg, so it
  fixes the same line inside its larger change. This PR is the fix alone for main now; whichever lands second takes
  the other's line. @ashhart, if you'd rather carry this line in a fix-up of your own or let EXL3 Flash Next: KV pages from the core pool (4/8) #1444 bring it, that's
  fine with me, and I'll close this one.

Checks on this head

• Head commit: dc96e9869.
• Lean checker on the base and on this head: 426 / 426, 0 new.
• Tests: zig build test-cuda-host 311 / 311 steps pass. The new test, "a grouped head's lane carries its draft's page
  table" (cuda_mtp_ops.zig, run by test-cuda-host), fails with the lane built as on main:
  expected 28672, found 0 (1 failed of 4,874; no other failure), and passes with this change.
  zig build test -j10: 537 / 537 steps, 0 failed (2,774 tests in the steps that rebuilt, 37 skipped).
• zig fmt --check zig build.zig: clean.
• CI: pending.
• Not run: Windows, macOS (no runner here), TP2. The change is inside the Flash Next CUDA head on one rank.

Receipt

• Environment: GB10, driver 580.173.02, nvcc 13.3.73, nvidia/Qwen3.8-Flash-Next-FP4-Local at b9ce18a6 and
  Vontra/Qwen3.8-Flash-Next-MLX-4bit-MTP at dadefa80, n-gram table warmed before each start, a fresh server for each
  run, order base-new-new-base. Base main 002c9e03e; this change dc96e9869.

• Paging on (TF_FLASHNEXT_KV_PAGES=256, --parallel 3), three equal prompts at once (tools/bench_cache.py --streams 3, 256 tokens; FP4 at --context 262144, each cell twice a server; MLX at --context 65536, 2K three
  times, 8K and 16K once). Concurrent replies that differ from the same prompt's reply alone, and decode a stream in
  tok/s; two numbers = the two runs of each build:

  	FP4 main	FP4 this change	MLX main	MLX this change
  Replies that differ from alone	1 of 18 / 1 of 18	0 of 18 / 0 of 18	5 of 15 / 5 of 15	0 of 15 / 0 of 15
  2K	25.8 / 25.8	27.3 / 27.5	29.3 / 29.7	30.4 / 30.9
  8K	25.1 / 25.3	26.7 / 26.9	27.3 / 27.8	31.0 / 31.7
  16K	25.1 / 25.5	27.3 / 27.2	26.4 / 26.9	29.4 / 30.2

  A first MLX pair run once each, not alternated, had main's 2K cell at 33.7 against 30.3; the alternated rerun
  above replaces it (same build, same cells).

• Paging on, different prompts (tools/bench_concurrent.py --alone --serial --levels 1,3 --temperatures 0,1, 256
  tokens): on both builds every concurrent reply equals its run alone and in series, and the reply digest is the same
  in every run. tok/s:

  	FP4 main	FP4 this change	MLX main	MLX this change
  3 streams, code greedy	153.3 / 153.2	175.1 / 176.8	171.4 / 172.8	189.3 / 190.6
  3 streams, chat greedy	91.8 / 91.8	113.1 / 113.6	138.8 / 140.8	167.5 / 169.8
  3 streams, code sampled	97.4 / 95.0	103.8 / 104.3	112.0 / 112.7	114.2 / 114.9
  3 streams, chat sampled	76.8 / 76.9	81.5 / 81.8	102.2 / 102.6	104.7 / 105.8
  1 stream, code greedy	72.0 / 71.9	71.5 / 72.2	83.0 / 82.6	82.9 / 82.8
  1 stream, chat greedy	46.1 / 45.8	45.9 / 46.2	71.2 / 71.0	71.2 / 71.0

  With this change, FP4's paged 3-stream rows are back at its unpaged speed (code / chat greedy 174.8 / 174.6 and
  113.7 / 114.1 below). One stream never runs the grouped pass, so it is level.

• Default settings (FP4, no paging, 262144 x 3): replies identical on all four servers (bench_cache 2K / 8K / 16K
  alone, first and resumed; drafted replies of four texts of 283-17,182 tokens, greedy and sampled; bench_concurrent --alone --serial 1 and 3 streams). Speed level: 3 streams code / chat greedy main 174.8 / 174.6 and 113.7 /
  114.1, this change 173.2 / 174.2 and 113.5 / 113.9; one stream 72.0 / 72.3 and 71.7 / 72.6. Code greedy at 3
  streams is the one cell with both new runs below both base runs (-0.6 %); without paging the lane is byte for
  byte the one main builds (.pg at its default), so the GPU work is the same.

• Drafted replies equal "draft": false: the four texts' greedy hashes equal those of a "draft": false run of the
  same texts in an earlier receipt. A resumed prompt equals a fresh one: every bench_cache row resumed from the
  cache (2K / 8K / 16K, both builds, paging on and off) equals its first run.

• Prompt speed: n/a, the change is in the draft head's decode lanes. First tokens in the cells above are level (2K
  from an empty cache, no paging, two requests a server: main 1,265 · 1,129 / 1,136 · 1,137 ms, this change
  1,127 · 1,136 / 1,130 · 1,139 ms).

• Licence of new weights: n/a.

Checklist

• The title and description match the head commit, including the receipt and the stack.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.
• 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.
• Decode and prompt processing are no slower than main at one stream and the widest admitted (3 at 262,144), except
  the one -0.6 % cell above.

[ashhart]

Thank you, @Arminova. That's a clear find, and the receipt makes it easy to follow. We've read it against main, and it's right:

• The grouped lane was the only head lane built without .pg. The lone pass (forwardHead) has carried it since paging came in, and nothing else builds a head lane, so this closes the gap.
• Without paging, Draft.pg keeps its .{} default, because only initIn sets it, for pool lanes. The lane is then the same bytes as main's, so we read the one -0.6 % default cell as noise.
• With paging, useAttn keeps any group that holds a paged lane off the one-launch lanes attention. Each paged lane's head then runs attnRows with its own table, as it does alone. On main those heads wrote their rows from the arena's first page onward, past the head's share of that page and into pages other streams hold. That fits your numbers: the drafts come back, and the long prompts' replies stop changing.
• The grouped head's graph key holds each draft's number, and a draft keeps the table it was built with. A held graph is therefore never replayed with another draft's table.

On #1444: we'd like this to land on its own now, so please keep it open. #1444 can take the line when it next rebases, since its headSeq carries .pg for both passes.

We have no asks. Our own GB10 run is queued: main against main plus this head, on the MLX 4-bit MTP checkpoint. With paging on (256-position pages, --parallel 3) it runs three equal streams at 2K, 8K and 16K against each prompt's reply alone, and the 1- and 3-stream sweep with every reply checked alone and against "draft": false. With paging off it runs the same equal-stream cells. We'll post the numbers here, and then it goes to the landing queue.

[ashhart]

/atc exchange a20c200f48b230cc3da200edc11af84bc0df985b0fd7e237a54888b60273be02

[ashhart]

/atc exchange a515bb94031949baebac3b378caa71cc2938c0e9458d991e98a251419b139adf

[ashhart]

@Arminova, our GB10 run is in, and it confirms the fix. Setup:

• main at 5daafaf18 against main plus this head (dc96e9869).
• Flash Next MLX 4-bit MTP, --parallel 3 --context 65536.
• Order: warm, then main, this PR, this PR, main.

Paging on (TF_FLASHNEXT_KV_PAGES=256): three equal streams at once, each checked against the prompt's reply alone (bench_cache.py --streams 3), in both rounds.

Prompt	main: replies differing	this PR	Decode a stream, main → this PR
2,048 tokens	1 of 3	0 of 3	30.2-32.0 → 33.6-34.0 tok/s
8,192	1 of 3	0 of 3	30.3-33.7 → 34.0-34.7
16,384	0 of 3	0 of 3	35.5-37.5 → 33.3-34.0

• Tokens: with this PR, every paged reply has the same token hash as the replies with paging off, on both builds: one hash for each prompt.
• Decode: paged decode now runs at paging-off speed, 33.4-34.8 tok/s a stream at all three sizes on both builds.
• The 16K cell: main's paged 16K run is faster than paging off, and it's the one size where its replies still come out right. In that run, every stream's grouped head reads and writes rows from the arena's first page. A head that uses its own pages can't take that shortcut.

The sweep (paging on; bench_concurrent.py --alone --serial at 1 and 3 streams, greedy and sampled, 256 tokens):

• Every reply equals its run alone, and every alone run equals "draft": false, on both builds.
• Three streams, aggregate tok/s, main → this PR:
  • Chat greedy: 111.4 → 123.7 (+11.0 %).
  • Chat sampled: 100.5 and 103.1 → 110.1 and 109.9 (+6.6 to +9.6 %).
  • Code sampled: 112.2 → 118.2 (+5.3 %).
  • Code greedy: 133.7 → 134.1.
• One stream is level (−0.5 to −0.2 %), as expected: a lone head takes forwardHead, which this PR doesn't change.

Paging off: both builds give the same replies at the same speed in every cell. No server log shows a window graph falling back to eager.

Since this ran, main has moved with nothing under families/flashnext, so these numbers stand. It goes to the landing queue now. Thank you!

[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]

Thanks @Arminova! This is in TensorFold 1.2.4, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.4