[gilby]

What this changes

After each prompt chunk, Flash Next's MTP head runs over the chunk's rows to fill its own caches (head.run from prefillStepFn). It ran as a decode window (o.window set, o.prefill off, moe.run(..., false)), so its projections and its MoE took the stream tiles for up to a whole chunk of rows (4,096 with #1651), reading each weight once per four rows. A kernel trace of a cold 7,869-token prompt shows about 0.54 s of affine_dot2_stream launches with 928- to 4,096-row grids, one set per chunk.

head.run now takes prompt: true from the prompt fill, which sets o.prefill and the MoE's prompt path and leaves o.window off; the decode windows pass false and keep their path. This is what the main layers already do with prompt rows.

The head's caches after a prompt are not the same bytes as before: the prompt tiles' sums differ from the stream tiles' in the last bits, and the head's cache hash after a cold 7,869-token prompt differs from main's. This does not change a served token, because the target model verifies every draft: replies are equal to main's, and only the drafted tokens (and so the acceptance) can move. Measurements below.

• Issue: n/a (a follow-up to hip: 4096-row prompt chunks by default on the RDNA 3.5 iGPUs #1651's 4,096-row chunks).
• Stacked on: main.
• atc overlap --json against main 000b534a2: incomplete (pullFilesLimit, timeLimit). The change touches CHANGELOG.md and engine.zig, head.zig and lanes.zig under zig/src/families/flashnext/hip/. hip flashnext: a shared round's streams draft in one head pass a step #2696 also changes head.run's signature; whichever lands second needs a small update. hip flashnext: size the scratch arena by the dry run's peak, with the MTP head and a verify round's logits #2670 is in main, and its dry run now passes prompt = true here, so it sizes the head's prompt path.

Checks on this head

• Head commit: 63914d6ae: the reviewed commit 7169ea5ac, a merge of main 000b534a2 (conflict in CHANGELOG.md only), head.run taking the prompt flag as an argument with hip flashnext: size the scratch arena by the dry run's peak, with the MTP head and a verify round's logits #2670's dry run passing true (6cd97ef66), and the changelog line (d885b0e71, 63914d6ae). The receipt ran on d885b0e71; the head differs from it only in the changelog figure.
• Lean checker, python3 tools/zig/lean_check.py: 345 problems on the base and 345 on this head, 0 new findings. python3 tools/pr_check.py --base origin/main --pr 1878: 0 errors, 0 warnings.
• Tests run: zig build test -Dhip=true --summary all: 889 of 891 steps succeeded, 1 test failed: dashboard_test "dashboard page polls the public stats snapshot without private details", which fails the same way on main 000b534a2. zig build hip-kernel-test on gfx1151: 20 of 20 steps, 159 of 165 tests passed, 6 skipped, the same as main.
• CI: not yet reported for this head; I will post each failure and whether this change caused it.
• Not run: Metal and CUDA (HIP host change only); tensor parallelism (one GPU); gfx1030, gfx1100 and gfx1201.

Receipt

• Environment: TensorFold main 000b534a2 against d885b0e71. Zig 0.17.0, zig build native-hip -Dhip=true -Dhip-arch=gfx1151 -Doptimize=ReleaseFast, system hipcc (HIP 7.1, clang 21), Linux 7.0. One Strix Halo iGPU (Radeon 8060S, gfx1151), 128 GB unified memory. Checkpoint unigilby/Qwen3.8-Flash-Next-oQ4e-HIP revision 5a0fb414, MTP 3 drafts, --context 131072 --parallel 4 --prompt-cache-gib 0. Arms in the order main, this, this, main, a fresh server with the page cache dropped for each; both arms of each side are shown.
• Paths and platforms: the head's prompt fill on HIP. I ran gfx1151 only.
• Cold first token, a fresh server for each size and arm, tools/bench_cache.py --require-cold, seconds: 2,048 tokens 4.19 / 4.07 on main vs 4.09 / 3.86 here (−3.7%); 7,377 tokens 12.28 / 12.28 vs 11.70 / 11.60 (−5.1%); 29,364 tokens 45.79 / 45.84 vs 42.98 / 43.31 (−5.8%); 58,728 tokens 92.48 / 92.90 vs 87.42 / 87.38 (−5.7%). Reply token_sha256 equal to main's at every size.
• Exactness of served tokens, greedy and sampled: tools/bench_concurrent.py --levels 1,2,4 --reps 3 --alone --serial --temperatures 1.0,0 --prompts tools/bench_prompts.json: every request's token_sha256 equals main's in each arm (42 of 42 warm, 8 of 8 on a cold four-stream start); concurrent equals alone (3/3, 6/6, 12/12) and alone equals "draft": false (4/4) in every cell of every arm, 0 unequal, 0 failed.
• Head caches and drafts are not bit-identical to before; replies are. The prompt tiles' sums differ from the stream tiles' in the last bits, so the head's caches after a prompt differ from main's and a draft can differ; the target verifies every draft, so served tokens do not. Acceptance (tokens a round, tools/bench_openai.py --tokens 256 --reps 3), code / chat: main 2.612 / 2.032 in both arms, here 2.612 / 2.065 in both arms. Earlier comments have the FP64 check of the head's two projections and the resumed multi-turn run.
• Decode speed, aggregate tok/s at 1 / 2 / 4 streams. Greedy: 51.5 / 67.1 / 66.5 and 51.6 / 67.5 / 67.3 (main) vs 51.5 / 67.5 / 67.0 and 51.5 / 67.5 / 67.1 (here). Sampled: 45.1 / 59.2 / 63.8 and 45.2 / 59.5 / 64.5 vs 45.2 / 59.5 / 64.3 and 45.2 / 60.1 / 64.6. bench_openai code 46.19 / 46.30 vs 46.14 / 46.17 tok/s (−0.2%, inside the spread between main's own arms), chat 40.03 / 40.12 vs 40.44 / 40.46.
• 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. (No new test: the token equality across the benches is the check.)
• 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. (Drafted and concurrent above, greedy and sampled; the resumed multi-turn run is in the comments, on the earlier head.)
• No precision traded for speed. The head's cache bits change (see above); served replies do not.
• Decode and prompt processing are no slower than main at one stream and four streams, within the noise above.

[ashhart]

Thanks @gilby. Routing the MTP head's chunk fill through the prompt path instead of a whole chunk of decode windows is a sensible fix for that 0.54 s, and the kernel trace makes the cost clear. What we need before it lands, since we have no AMD GPU to check it on:

1. Exactness of the head's caches: the MTP head's caches after a prompt must be the same bytes whether they were filled by the prompt tiles or the old decode windows. That holds if the prompt tiles give decode's bytes at every row count. Please show it directly: the head's K/V (or its cache hash) after a cold 7,869-token prompt, old path against new.
2. Served tokens: token_sha256 equal to main for drafted replies, plus the acceptance rate before and after, which should be unchanged if (1) holds.
3. Speed: cold first-token time at 8K and 32K, and decode at 1 stream, main against this head, alternating runs.

Then take it out of draft and we'll review it for landing.

[gilby]

Thanks. I ran the three checks on the rebased head, and the first one does not hold, so I'm leaving this in draft until you decide.

1. Head caches: they are not the same bytes. The hash of the head's whole cache buffer after a cold 7,869-token prompt differs between main and this head, so the prompt tiles do not give the decode windows' bytes for the head's projections and MoE. The new state is the same at prefill_step 1,024, 2,048 and 4,096. The head's draft logits differ from main's: at 8K, mean KL 7.8e-4, max 2.4e-2, top-1 differs on 1 of 49 draft calls; at 32K, mean KL 0.65, max 31.7, top-1 differs on 5 of 102. I have no float64 reference for the head, so I can't say which path is nearer the exact value.
2. Served tokens: token_sha256 equals main on every request at 2,048, 7,377, 29,364 and 63,728 prompt tokens and on both bench_openai prompts, and concurrent equals solo and serial equals "draft": false at 1 and 4 streams. Acceptance went from 2.753 to 2.783 tokens a round on the code prompt and stayed 2.37 on chat. The target's logits differ from main's only in verify rows that follow a different drafted token.
3. Speed, alternating main, this, this, main on fresh servers: cold first token 2.3 to 3.2% faster (7,377 tokens 18.46 -> 17.87 s, 29,364 tokens 82.50 -> 79.91 s, 63,728 tokens 189.35 -> 183.75 s), decode at one stream 55.96 against 56.15 tok/s on code, which is inside the spread between main's own two arms.

So the change buys about 3% on cold prompts, and it changes the head's cache bytes without changing a served token. If you want the head's cache to stay bit-equal to decode's, I'll drop this PR. If a draft-only difference is acceptable, I'll take it out of draft with the body updated to say so. Which do you prefer?

I also rebased both PRs onto current main: this one is now 7169ea5ac and #1877 is 4d7d418e8, the same patches. I ran the build, host suite and (for #1877) the contract on those heads; the timing arms ran on the pre-rebase heads, which differ only in the changelog.

[gilby]

Follow-up on the head's cache bytes, with the checks I could add. Head 7169ea5ac, same checkpoint and machine as the receipt in the body.

1. Where the bytes differ: I ran the head's forward three ways on each chunk of a cold 7,377-token prompt (local probe, not committed) and hashed every intermediate: the old path (the whole chunk on the stream tiles), the new path (prompt tiles), and decode's own shape (the same chunk in windows of four rows). The old path equals decode's bytes at every buffer (both projections, both hyper-connection inputs, the attention output and the MoE output), and the new path differs from decode from the first projection on. So your premise was right for the old path, and it does not hold for this change: the prompt tiles are built so a prompt's bytes do not depend on how it is cut, and they do not reproduce the stream tiles' bytes. Making them equal would need a stream-tile kernel that reads each weight once per many rows and sums in decode's order, which is a new kernel for about 0.5 s of an 18 s cold 8K prompt, so I did not write it.
2. Accuracy of the two paths: for the head's two 4-bit projections (fc_embedding and fc_hidden) on 16 rows of the first chunk, I computed a float64 product from the checkpoint's own dequantized weights and the identical bf16 inputs. Both paths are within one bf16 rounding of it: fc_embedding old 6 of 40,960 outputs differ from the correctly rounded float64 value and new 8; fc_hidden old 19 and new 13; the mean error is the same to four digits. Where the two paths differ (10 and 22 outputs), the old path is nearer in 6 and 8 cases and the new one in 4 and 14. This is two ops on 16 rows, not a whole-head reference, and I do not have one.
3. Acceptance and decode after long prompts, alternating main, this, this, main on fresh servers, --prompt-cache-gib 0, 256 tokens at temperature 0 (tokens a round, decode tok/s on the second request): 7,377-token prompt 2.462 on both, 45.25 and 45.29; 29,364 tokens 2.032 on both, 41.30 and 41.34; a 7,441-token prompt 3.938 on both, 65.26 and 65.27; a 2,048-token prompt 2.349 on main and 2.286 on this change (-2.7%), 44.15 and 43.61. token_sha256 is equal to main on every request. Earlier I only had bench_openai's 14-token prompts, where acceptance rose from 2.753 to 2.783; so across prompts the change moves acceptance both ways by at most 3%.
4. Resumed prompts: I could not check that a resumed prompt equals a fresh one on this build, because with --prompt-cache-gib 8 the cache kept 0 states (hits 0, misses on every request), so no request resumed. The head's state after a prompt is the same at prefill_step 1,024, 2,048 and 4,096.

So the change is about 3% faster on cold prompts, gives the same tokens and about the same acceptance, and does change the head's cache bytes from decode's. If bytes equal to decode's are the bar for this PR, I will close it; if the target's verification is enough, it is ready to review. Which do you prefer?

[gilby]

Follow-up on the resumed-prompt check, which I said I could not run: my earlier runs sent single-message prompts, and a single message has no chat history, so this family's prompt cache (planned chunk starts, a history mark only) never keeps a state. That was a flaw in my test, not in the build.

I reran it as a three-turn chat on a prompt of about 7,000 tokens (temperature 0, thinking off, 48 new tokens at most a turn), on main (this PR's parent) and on this head, each on a fresh server, with --prompt-cache-gib 8 and with --prompt-cache-gib 0.

• With the cache on, both builds kept a state each turn and resumed: turn 2 at 6,976 cached tokens and turn 3 at 7,040 (hits 0, 1, 2). First token was about 1 s on turns 2 and 3, against about 17 s fresh.
• The reply hash of every turn is identical in all four arms (main and this head, cache on and off), so a resumed prompt equals a fresh one on both builds, and this head's resumed replies equal main's.
• Drafts accepted are identical in all four arms (24/31, 10/18, 8/13). The replies are short, so this says little about acceptance differences.

This covers the head cache being saved and restored with the prompt state. It is one prompt, so it is a check, not a sweep. The question in my previous comment stands: if bytes equal to decode's are the bar for this PR, I will close it; if the target's verification is enough, it is ready to review.

[ashhart]

Thanks @gilby, and thanks for the careful follow-ups. You found exactly where the bytes part: the old path matched
decode's bytes and the prompt tiles don't. You checked both against FP64 on the head's two projections, and you ran
resumed prompts as a real multi-turn chat. That answers what we asked.

We'll take the draft-only difference. The target verifies every served token, replies equal main's at every length
you ran, and resumed equals fresh. The head's prompt rows now use the same prompt arithmetic as the target's own prompt
rows. Acceptance moves both ways by up to 3 %, which we read as noise around no change.

Before it comes out of draft:

1. Rerun on today's main. hip flashnext: sparse attention with a lane a key over several heads, the scalar kernel's bits #1877 landed after your runs, and with o.prefill set the head's sparse attention now
   takes its prompt kernel. Please merge or rebase onto main and post token_sha256 against main, plus the cold
   first-token times at 8K and 32K, alternating as before. The CHANGELOG line's 3 % should come from that run.
2. Description and CHANGELOG line. Please keep saying plainly that the head's caches and drafts aren't bit-identical
   to before and replies are, as the body already does, and end the line with "Thanks to @gilby." We can add the
   credit at landing if that's easier.

One heads-up: #2696 also changes head.run's signature, and #2670 calls head.run from its dry run
(engine.zig:306 there). Whichever of the three lands later will need a small update. With this PR in, #2670's dry
run should pass prompt = true so it sizes the head's prompt path.

[gilby]

Thanks @ashhart. Both points are done, on head 63914d6ae.

1. Rerun on today's main. The branch now carries a merge of main 000b534a2 (a plain merge commit, no force push; the conflict was CHANGELOG.md), so the head's sparse attention takes hip flashnext: sparse attention with a lane a key over several heads, the scalar kernel's bits #1877's prompt kernel when o.prefill is set. Main against this branch on one Strix Halo iGPU (gfx1151), arms main, this, this, main, a fresh server with the page cache dropped for each:

cold prompt	main, first token	this head	change	reply token_sha256
2,048 tokens	4.19 / 4.07 s	4.09 / 3.86 s	−3.7%	equal
7,377 tokens	12.28 / 12.28 s	11.70 / 11.60 s	−5.1%	equal
29,364 tokens	45.79 / 45.84 s	42.98 / 43.31 s	−5.8%	equal
58,728 tokens	92.48 / 92.90 s	87.42 / 87.38 s	−5.7%	equal

• Tokens against main, greedy and sampled (bench_concurrent --levels 1,2,4 --reps 3 --alone --serial --temperatures 1.0,0 --prompts tools/bench_prompts.json): every request's token_sha256 equals main's in each arm, 42 of 42 warm and 8 of 8 on a cold four-stream start. Concurrent equals alone (3/3, 6/6, 12/12) and alone equals "draft": false (4/4) in every cell of every arm, 0 unequal, 0 failed.
• Decode is level. Aggregate tok/s at 1 / 2 / 4 streams, greedy: 51.5 / 67.1 / 66.5 and 51.6 / 67.5 / 67.3 on main against 51.5 / 67.5 / 67.0 and 51.5 / 67.5 / 67.1 here; sampled: 45.1 / 59.2 / 63.8 and 45.2 / 59.5 / 64.5 against 45.2 / 59.5 / 64.3 and 45.2 / 60.1 / 64.6. bench_openai code 46.19 / 46.30 against 46.14 / 46.17, chat 40.03 / 40.12 against 40.44 / 40.46. Acceptance: code 2.612 tokens a round on both sides, chat 2.032 on main and 2.065 here.

2. Description and CHANGELOG line. The description still says plainly that the head's caches and drafts are not bit-identical to before and replies are. The line now reads "cold prompts finish about 5% sooner at 8K and 32K", from the table above, keeps that sentence, and ends "(hip flashnext: the MTP head absorbs a prompt chunk on the prompt tiles #1878). Thanks to @gilby."

On the heads-up: #2670 is in main, and its dry run now passes prompt = true through head.run's new argument (6cd97ef66), so it sizes the head's prompt path. #2696 still changes the same signature; whichever of the two lands second needs a one-line update, and I'll make it on mine.

Checks on 63914d6ae. python3 tools/pr_check.py --base origin/main --pr 1878: 0 errors, 0 warnings; lean 345 on the base and 345 here. zig build test -Dhip=true --summary all: 889 of 891 steps succeeded with one failing test, dashboard_test "dashboard page polls the public stats snapshot without private details", which fails the same way on main 000b534a2. zig build hip-kernel-test on the GPU: 159 of 165 passed and 6 skipped, the same as main. The tests and the receipt ran on d885b0e71; the head differs from it only in the changelog figure.

Environment: Zig 0.17.0, zig build native-hip -Dhip=true -Dhip-arch=gfx1151 -Doptimize=ReleaseFast, system hipcc (HIP 7.1, clang 21), Linux 7.0, unigilby/Qwen3.8-Flash-Next-oQ4e-HIP at revision 5a0fb414, --context 131072 --parallel 4 --prompt-cache-gib 0. I've marked the PR ready for review and will post the CI results for this head when the run finishes.

[ashhart]

Thanks @gilby, both are answered. On today's main with #1877's prompt kernel, cold first token is 3.7% to 5.8% shorter from 2K to 59K-token prompts, every reply's token_sha256 equals main's greedy and sampled (42/42 warm, 8/8 cold), concurrent equals alone, and decode is level. It merges cleanly with main, so it's queued to land with your changelog line.

• Synthena (Ash's Agent)

[gilby]

CI is green on 28ef18b24: 8 checks passed and the two nightly host jobs were skipped. This head is 63914d6ae, the head of the receipt above, plus a plain merge of main 6c5cef894 (no force push); main had moved and the branch conflicted in CHANGELOG.md, so CI could not start without it. The merge changes nothing else of this PR.

[ashhart]

Thank you @gilby. 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 @gilby! This is in TensorFold 1.2.8, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.8

• Synthena (Ash's Agent)