[cshintov]

What this changes

A row gives probability 0 to a key it must not see, and 0 times an infinite or NaN value in P V gave NaN in that row. Now such a value adds nothing to the rows that must not see it. Rows that do see it keep the IEEE result, the same class float64 gives, as the decode path already did. Inputs with finite values keep their bytes. This covers the three kernels that mask keys row by row: attention_prompt_sg (prompt rows before M5), attention_prompt (prompt rows on M5) and attention_mpp's q27_attn_prefix (decode rows on M5).

• Issue: Refs Qwen3.6-35B-A3B on Metal: routed experts in the 27B's engine (slice 1 of #931) #1009 (the finding from 28 Oct, and the fixes requested on 4 Nov).
• Stacked on: main (002c9e03e). Eight commits: the contract registration with the finding marked known, the fix, the registration moved into its own suite so metal_main.zig stays at 600 lines, the changelog, then four changes from review: the masked read moved off the steps every row sees whole, gate 1's decode round made a chain on an M5, gate 1's prompt chunk made to hold a whole step past its first row's keys, and a per-launch pick so a finite prompt launch runs main's kernel.
• atc overlap --json: not run. ATC isn't installed here, so whoever opens the PR needs to run it.

Design

A finite prompt must run the kernel main runs. Two earlier heads did not:

• The first version read every value fragment of every 32-key step through a mask that is 0 when nothing is non-finite. Measured on an M5 Ultra: -0.8% at 2,048 tokens, -11.0% at 65,536.
• 36bea7342 ran main's step on the steps every row sees, but in a kernel that also held the masked step, the vote and the add-back. Measured served on the 27B on an M5 Ultra: -1.2% at 2K to -4.3% at 32K. On M4, a timing-only build with the masked code removed cut the same kind of cost from +2.2% to +0.6% at 8,192 keys, so the code's presence costs time even when it never runs. Why it grows with prompt length on the M5 is inferred, not measured: attention is a larger share of a longer prompt, so a fixed slowdown of its kernel shows more.

Now each prompt launch picks one of two kernels:

• The pick. q27_attn_prompt_scan reads the launch's own new value rows, keys start + 1 to start + rows - 1, all 256 dims of every KV head, one simdgroup a key. Those are the only keys any row of the launch can be masked from: row r sees start + r + 1 keys, so every key below start + 1 is seen by every row. Any non-finite bf16 sets one flag word. q27_attn_prompt_pick (one thread) reads and clears the flag and writes two indirect grids, the prompt kernel's grid into one and zeros into the other. Both kernels are then dispatched indirectly; the empty one does nothing. Cost: two small dispatches and two barriers per launch, about 0.03 ms on M4.
• A finite launch runs q27_attn_prompt_8 (M5) or q27_attn_prompt_sg_8 (before M5), the MASKED = false instance of the kernel: main's loop, no vote, no add-back. Its IR is identical to main's, both the kernel and all four step instances, once value names, callee names and the unused mask parameter are normalised.
• A launch with a non-finite value runs q27_attn_prompt_masked_8 or q27_attn_prompt_sg_masked_8, which is the previous head's kernel, described below.

Inside the masked kernel:

• Each simdgroup votes once per chunk on whether any value at a key some live row must not see is non-finite. "Live" means not padding; seen is the fewest keys any live row sees. The scan covers only keys from seen to the end of the chunk, so it is constant per chunk, one 16-byte load per lane per two keys.
• A step whose 32 keys all lie below seen (every live row sees every key) runs the step with MASKED false at compile time: plain value fragments, no mask argument, no add-back. Only the rest of the chunk, from the step that holds seen on, runs the second copy: values read under the vote's mask, then the add-back.
• The masked read and add-back are unchanged. A non-finite value is replaced by a zero with the same sign (exponent tested with integer ops on the bf16 bits, finite values untouched). After each step the probabilities pass through the simdgroup's own xs row, and each lane adds value times probability back for every non-finite value whose key its row sees, using the same probability the matrix step used (fp32 before M5, half on M5).
  • Rows that don't see the key get only the zero, which is what a finite value of that sign would contribute.
  • Rows that see it get infinity or NaN, as in float64 and the decode path.
• q27_attn_prefix (decode on M5) is chosen by the host. A round where every stream has one row runs q27_attn_prefix_N, main's kernel by IR for all 16 tile counts: that row sees every key its tile reads, so nothing in its range is masked. A round with a stream of rows runs q27_attn_prefix_masked_N, the kernel described next. A test-qwen27 test walks the kernel's ranges over one-row plans and finds no masked key, and finds one in a two-row chain.
• The masked q27_attn_prefix reads values straight from device memory through the tensor op, so there is no point at which to mask them. In a pass that voted yes, each step that holds a key some row must not see sums its P V in plain code and skips each row's unseen keys. Steps that every row fully sees still use the tensor op, as they did on main. Its fast steps differ from main's only by one uniform compare and branch per 64-key step (!masked || kt + TK <= seen). Its steps aren't split into two copies within the kernel, because they are one long body with the tensor ops in it.
• Which decode rows reach that kernel on an M5: a chain of two or more rows (fused in the plan). Its prefix kernel runs to the last row's key, so row r sees start + r + 1 keys and the chunk holds keys up to start + rows: every row but the last has keys it must not see. Any other round (a tree, or one row) sends its own keys to the tail kernel, which reads only a row's ancestors, and stops the prefix at the tile before them. A plan of 5 rows at start 2 gives, as a tree, prefix_chunks=0, tail_chunks=1, and as a chain prefix_chunks=1, tail_chunks=0, limit 7 with rows seeing 3, 4, 5, 6 and 7 keys.

Rejected alternatives:

• Selecting p == 0 ? 0 : p * v for every product. It turns a -0 product into +0, so finite outputs would change bytes. It also costs a select on every product of every step.
• Flushing non-finite values to 0 with no add-back. It hides an infinity from the rows that do see it. Gate 1 fails ("finite where float64 gives +Inf").
• A vote at every step instead of once per chunk. Measured on an M4 Pro, keeping a non-finite flag live through the step's loads and voting on it made the prompt kernel 4% to 6% slower. Voting once before the steps costs nothing measurable.
• The mask as a runtime argument in every step (the first version). The cost above. It scales with the key count because every step pays it, and the vote is not the cause: the vote reads only keys from seen on.
• Three copies of the step: unmasked, masked whole, masked partial. On M4's prompt kernel it ran 6.1% slower than main at 8,192 keys and 2.5% slower at 32,768, against 2.2% and 0.5% for the two-copy form. Each copy is a long fully unrolled body.
• Two copies of the whole step loop, picked by the vote (four step copies). On M4 with finite inputs, 2.1% slower than main at 8,192 keys and 2.4% at 32,768 (the chosen form: 2.2% and 0.5%), and 5% slower than the chosen form when the planted value makes the vote yes. An earlier measurement here gave "about 9 times slower"; I could not reproduce that on these kernels and no longer claim it. What did reproduce is that more step copies cost more.
• The masked copy in the same kernel as the finite path, behind the vote (36bea7342). The M5 cost above: the finite path's code was main's, the kernel's was not.
• A runtime branch per launch or per tile inside one kernel. Same problem: the masked code is still in the kernel. On M4 the four-copy form above is that design.
• The flag fused into the kernel that appends the new values. It would skip the scan's read of the new rows, but that kernel is shared with decode, and the scan costs about 0.03 ms a launch on M4.
• Host-side choice. The host can't read the values without waiting for the GPU between the append and the attention kernel.
• Staging values through threadgroup memory in q27_attn_prefix. It needs 16 KiB per simdgroup, more than the 32 KiB budget allows from 2 tiles (4 simdgroups) up.

Checks on this head

• Head commit: d89ce2b9c, on main cde4fe7d9. The newest commit picks q27_attn_prefix per round; lean, test-qwen27 and the contract were rerun on this head.
• IR, xcrun metal -O3 -S -emit-llvm: q27_attn_prefix_1 to _16 are main's kernels, and q27_attn_prefix_masked_1 to _16 are the previous head's (54ff80b87), all 32 identical.
• Lean checker, python3 tools/zig/lean_check.py: 425 problems on the base (cde4fe7d9), 425 on this head, 0 new (tools/zig/lean_new.py). metal_main.zig is at 600 lines.
• Tests run, with the command and pass, fail and skip counts:
  • zig fmt --check build.zig zig: clean.
  • zig build test-qwen27: 46/46 steps, 152/152 tests passed. a stream of one row has no key in the prefix kernel's range it must not see walks the prefix kernel's per-tile ranges over one-row plans (1 to 3 streams, 14 starts across the tile and 8,192 boundaries, shared prefix and fused chains on and off) and finds no masked key; it finds one in a two-row chain. The other new one is prompt pick: one non-finite value at every key, head and kind; masked iff start < key < start + rows. It runs Attention.encodePick back to back on one flag: for 1, 2, 3, 8, 33 and 40 rows at starts 0 and 5, one value (+Inf, -Inf, NaN or a negative NaN) at every key from start - 1 to start + rows on each of 4 KV heads, each followed by a finite launch. It checks the grid picked, that the other grid is zero, and that the flag is clear after each launch.
  • zig build test-kernel-contract: 22/22 steps, 173/173 tests passed.
  • tf-kernel-contract --only qwen27: SUMMARY 30 kernels: 30 pass gates 1-4, 0 fail, 0 known findings; 3 controls: 3 red as planted, 0 missed; 0 limit problems, exit 0.
  • On the first commit (main's kernels): 22 pass gates 1-4, 0 fail, 3 known findings. The three are gate 1 on prompt rows at 24 over 4, 32 over 2 and 16 over 4, each "nan where float64 gives pos_inf".
• CI: not run yet.
• Mutations of the pick (each run alone, then reverted). Gate 1 alone misses the first two: it plants values at both start + 1 and the last key, so a scan short at either end still finds the other. The new test catches all five:
  • Scan from start + 2: "rows 2 start 0 key 1 head 0: masked false".
  • Scan to start + rows - 2: the same line.
  • Scan from start: "rows 1 start 0 key 0 head 0: masked true".
  • The pick never clears the flag: the next finite launch picks the masked grid.
  • The scan reads KV head 0 only: "rows 2 start 0 key 1 head 1: masked false".
• Mutations of the masked kernel's split, on the prompt rows' registration (measured on 36bea7342's kernel, which is this head's masked kernel):
  • Add-back off in the masked copy: gate 1 fails, "47 of 49152 outputs differ in class; first at 6147: pos_inf where float64 gives nan". On the earlier chunk it failed with "16 of 49152 ... finite where float64 gives pos_inf".
  • The split taken at the chunk's end instead of at seen, so steps that rows must not see read unmasked: gate 1 fails, "6 of 49152 outputs differ in class; first at 3: nan where float64 gives pos_inf". Before gate 1's prompt chunk held a whole step past its first row's keys, this mutant passed. That is why the chunk is now k = 57.
  • The planted control now turns the scan off, so every launch runs the finite kernel as main does: 46 of 49,152 outputs differ, "nan where float64 gives pos_inf".
• Not run, and why:
  • Anything on an M5. The attention_prompt kernels were built here and compared as IR, but none of their GPU paths ran: those kernels assume the M5's tensor-op element layout.
  • The M5 cases are yours to run: the one-stream decode cells on this head, and tf-qwen27-attention-precision --prompt --plant on its tensor-unit paths. On 54ff80b87 the M5 contract passed with all four controls red, and the prompt cost was gone.

Receipt

• Environment: TensorFold d89ce2b9c, Zig 0.17.0, Xcode Metal toolchain, an M4 Pro (no tensor units). No checkpoint: these are synthetic kernel checks.

• Paths and platforms the change reaches, which I ran and which I could not:

  • Ran: prompt rows before M5 (attention_prompt_sg) through the contract and the precision tool.
  • Could not run: prompt rows on M5 (attention_prompt) and decode rows on M5 (q27_attn_prefix).
  • Not touched: decode rows before M5. They take the tail kernel, whose rows see every key it reads, and they already passed gate 1.

• Exactness: tf-qwen27-attention-precision output digests are byte-equal across main, the earlier heads of this fix and this head:

  Case	Main, earlier heads and this head
  --prompt (24 over 4), prompt kernel / decode path, queries width 1	ddeecd15f434b130 / d126ea575c5b560a
  --prompt (24 over 4), queries width 4	bd7e7beac1aaa2dc / acceedb02042af55
  --prompt --heads 32 --kv-heads 2	48b07607b1f53adb df5b68aac294eb05 c3355635363e142d dfaac0400478a740
  --prompt --heads 16 --kv-heads 4	8dcf3c049caaf93a e46988c141212968 02b4e9c2c174e005 139544178cb7cc85
  decode (24 over 4)	d56632e65c16e02f 60c1f0342f93123a 0591e9cba21ae225 e38e42c9532d7e28

  With --plant (-Inf in a value of each case's last key), main's prompt kernel writes 3,300 non-finite values (width 1) into rows that never see that key; this head writes 0 at all three head counts, and gives the same bytes as the earlier heads (2c34835aec7d2f56 fe3104dea4d43b66 2522dafd65de43ec 700d248ac565b389 at 24 over 4). Those rows hash the same as the unplanted rows, the same rows the decode path gives with and without the plant.

• Decode speed before and after: the decode path before M5 is unchanged (byte-equal digests above). Decode on M5 is not measured.

• Prompt speed before and after: not measured end to end (no checkpoint, and no M5). The attention call's GPU time (scan and pick, prompt kernel, merge) of one 128-row chunk at the end of N cached keys, 24 over 4, on an M4 Pro through its pre-M5 prompt kernel. Median of 5 interleaved runs, each the median of 5 timed repetitions after a dropped warm-up. In ms; the key count is set by a flag I added to the tool locally and have not included:

  Cached keys	main	36bea7342	this head, finite	this head vs main	this head, planted
  8,192	5.98	6.13	6.01	+0.5%	6.12
  16,384	12.80	12.88	12.88	+0.7%	not run
  32,768	26.76	26.97	26.87	+0.4%	26.90
  65,536	53.38	53.89	53.62	+0.4%	not run

  The 36bea7342 column is from the earlier run, on main c2aa0fed3. With finite values this head runs main's kernel, and what is left is the scan and pick (about 0.03 ms a launch). At 16,384 and 65,536 keys the runs of main and this head overlap. Only a launch with a non-finite value it must mask pays the masked kernel: +1.5% over a finite launch at 8,192 keys and +0.5% at 32,768 (median of 3 runs). The M5 speed is yours to measure.

• Licence of any new weights or drafter: n/a.

Kernel contract, only when a GPU kernel is new or changed

• Kernels registered with the contract, and the kernel each one replaces: metal qwen27 attention {27B, GQA16, 9B} {decode rows, prompt rows}, which run Attention.encode's whole calls against float64. They are new registrations of existing kernels, at one Mac's head counts: 24 over 4, 32 over 2 and 16 over 4. The rank's 12 over 2 merges as a pair, so it is not one call.
• New kernels: q27_attn_prompt_scan and q27_attn_prompt_pick, covered by the test-qwen27 test above and run by the prompt rows' registrations. q27_attn_prompt_masked_8 and q27_attn_prompt_sg_masked_8 are the previous head's prompt kernels under new names.
• A planted control for each new kernel, and each one fails as planted: control: qwen27 attention 27B prompt rows with non-finite values unmasked turns the scan off, so every launch runs the finite kernel, which sends values to P V as main does. It is red at gate 1, as planted. On M5 the same control covers attention_prompt, plus a decode-rows control that disables q27_attn_prefix's vote. Before M5, decode rows have no masked path, so that control isn't registered there.
  • The decode control did not bite on an M5 Ultra on the earlier head. Cause: gate 1 runs one decode round, the first served shape with 5 to 8 rows, which was a tree (start 2, even, 3 rows or more). A tree's keys go to the tail kernel, so the control's edit of the prefix kernel was never on the path. On an M5 gate 1's decode round is now a chain, and the registration refuses to run the control on a round the plan does not fuse (PlantedRoundMissesPrefix). Before M5 the round is still a tree: that kernel is the tail, and the sibling that must not see another's key is tested there.
  • The prompt control now differs in 46 of 49,152 outputs (6 before), because gate 1's prompt chunk now holds a whole step past its first row's keys and the last row's value is planted too.
  • Mutations on this head: main's three kernels under these tests fail gate 1 at all three head counts. The three mutations of the split are listed under Checks.
• Each backend's SUMMARY line and exit status: Metal (M4 Pro), 30 kernels: 30 pass gates 1-4, 0 fail, 0 known findings; 3 controls: 3 red as planted, 0 missed; 0 limit problems, exit 0.
• Chips I could not run: M5 (tensor units), M1, M2.

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. Finite inputs are byte-equal, so these are unchanged.
• No precision traded for speed. Finite inputs keep their bits.
• Decode and prompt processing are no slower than main at one stream and the widest admitted. Measured only on M4's prompt kernel: within +0.4% to +0.7% of main from 8,192 to 65,536 keys, the scan and pick. A finite launch on the M5 runs main's kernel by IR, but its speed there is not measured.

[ashhart]

Thanks @cshintov, this is careful work: one vote per chunk keeps main's instructions whenever no masked value is non-finite, the rejected alternatives are measured, and the M4 Pro digests are byte-equal with the 3,300 stray non-finite values gone. Two things before it lands. On our side, the M5 runs you couldn't do: the contract on the tensor-unit paths with the decode control, the precision tool with --plant, and the 27B served against main for tokens and speed, on an M5 Ultra and an M3 Ultra. On yours, a CHANGELOG line: replies change only where a NaN or infinity already sits in a value a row can't see, and finite inputs keep their bytes. We'll post the M5 results here.

[ashhart]

Thanks @cshintov. Here's the first of the runs we owed you, on an M5 Ultra: the 27B served at your head against its base, base and new alternated, four runs each.

Tokens: identical in all 20 rows.

Prompt speed is lower, and it grows with prompt length:

Prompt	Base, tok/s	This PR, tok/s	Change
2,048	1,884	1,869	-0.8%
8,192	1,856	1,814	-2.3%
16,384	1,790	1,725	-3.6%
32,768	1,670	1,568	-6.1%
65,536	1,466	1,304	-11.0%

The repetitions agree to about a tenth of a percent, so this is a steady cost and not noise. Finite outputs don't change, so this is a speed question only, but we can't take a default-path slowdown of this size for inputs that already hold a NaN or infinity.

Since it grows with length, the vote looks like the likely cost on these tensor-unit paths: each simdgroup checks every key a live row can't see, on a path your M4 Pro doesn't run. That's our reading, not a measurement. Could you look for a form that doesn't cost finite prompts, such as a cheaper check, or one that runs only when a non-finite value could already be present? We can run any candidate on the M5 Ultra for you.

The contract with the decode control and the precision tool's --plant on the tensor-unit paths are still to come.

[ashhart]

@cshintov, the contract on the M5 Ultra, on your head (tf-kernel-contract --only qwen27, exit 1):

SUMMARY 30 kernels: 30 pass gates 1-4, 0 fail, 0 known findings; 4 controls: 3 red as planted, 1 missed; 0 limit problems

All 30 kernels pass. The one miss is the decode control:

PASS nonfinite control: qwen27 attention 27B decode rows with non-finite values unmasked: 3358 non-finite and 27362 finite outputs match the float64 classes
CONTROL MISSED control: qwen27 attention 27B decode rows with non-finite values unmasked: nonfinite did not go red

The prompt control does go red here (6 of 49,152 outputs differ in class; first at 3: nan where float64 gives pos_inf). So on tensor-unit decode, the control doesn't show what it plants: either this path doesn't need the mask, or the control doesn't reach it on an M5. Please make the decode control bite there, or tell us the M5's decode isn't affected and leave that path alone. If it isn't affected, that may also be a place to save the prompt cost we posted above.

[cshintov]

Two changes on the head (36bea7342), both for your M5 run.

Cost: it was not the vote, which reads only keys from seen to the end of the chunk. It was the mask argument in every 32-key step of the prefix. Steps whose 32 keys every live row sees now run main's step (compile-time unmasked). The M5 step's IR is identical to main's; I could not see M5 machine code. On the M4 prompt kernel the new head is within +0.5% to +2.2% of main from 8K to 64K keys, with no growth with length. M4 cannot show your effect, so please rerun the M5 Ultra A/B on this head at 2,048 to 65,536 tokens.

Decode control: gate 1's decode round was a tree (5 rows at start 2). A tree's keys go to the tail kernel, so the prefix-kernel edit never ran. Only a chain of 2 or more rows reaches masked keys in q27_attn_prefix. Gate 1's decode round is now a chain on M5, and the control refuses to run on a round that is not one. The q27_attn_prefix change stays, since that path is reachable. If the control is still green on your M5, 0 times an infinity does not poison that tensor op, and I will drop the change.

I also moved gate 1's prompt chunk so it holds a whole step past the first row's keys: a wrongly placed split passed gate 1 before. The CHANGELOG line is in too.

[ashhart]

Thanks @cshintov, that's a sharp diagnosis: the mask on every fully seen step, and a tree that never reached the prefix kernel. We'll rerun the M5 Ultra A/B on 36bea7342 at 2,048 to 65,536 tokens, the contract with the chain decode control, and the precision tool with --plant, and post all three here.

[ashhart]

@cshintov, our M5 Ultra rerun on 36bea7342. The contract is right now, and the prompt cost is real.

Contract (--only qwen27): 30 kernels pass gates 1-4, 0 fail, and all 4 controls go red as planted. The decode control you fixed is caught.

Served, on the 27B, alternating builds, tokens identical in all 20 rows:

Prompt	Base	This PR	Change
2K	1,903 tok/s	1,880	-1.2%
8K	1,856	1,821	-1.9%
16K	1,790	1,742	-2.7%
32K	1,670	1,598	-4.3%

These don't overlap: every repetition sits apart, so it isn't noise, and the cost grows with prompt length. Decode is level, and four streams are level too (182.1 -> 183.0).

The fix is a correctness one, so we want it, but a slower default prompt on every M5 is a lot to pay for values that are rarely non-finite. Is there a way to keep the masked path's cost off rows that are all finite, for example a cheap check per launch or per block? Could you look at where the time goes? I haven't run --plant yet, and I'll post that separately.

[cshintov]

The time went to the masked code being in the kernel at all: on 36bea7342 the finite steps were main's by IR, but the kernel also held the masked step, the vote and the add-back. On M4, taking that code out cut the same kind of cost from +2.2% to +0.6% at 8,192 keys.

54ff80b87 does it per launch. A scan reads the launch's new value rows (keys start + 1 to start + rows - 1, the only ones a row can be masked from), and a one-thread pick writes two indirect grids. A finite launch runs q27_attn_prompt_8, which is main's kernel by IR. Only a launch with a non-finite value it must mask runs q27_attn_prompt_masked_8, the previous kernel.

On M4, finite launches are +0.4% to +0.7% over main from 8K to 64K keys (the scan and pick), and planted ones +1.5% at 8K. The prompt control now turns the scan off. Gate 1 plants at both ends of the window, so it can't see a scan that is short at one end. A new test-qwen27 test covers that by planting at every key around the window on each KV head.

Could you rerun the M5 A/B on 54ff80b87?

[ashhart]

Thanks @cshintov. Finding that the masked code's presence in the kernel was the cost, and picking the kernel per launch so a finite prompt runs main's own, is the right fix, and the window test that plants at every key around start + 1 .. start + rows - 1 makes the scan's range easy to trust. The open question is the M5: your timings are M4's, where finite launches are within +0.4% to +0.7%. So the next step is ours: an M5 Ultra A/B on 54ff80b87, with prompts at 2K to 64K against main and decode beside them, the contract, and the precision tool with --plant, which we said we'd run last time and didn't. Then an M3 Ultra run. We'll post both here.

[ashhart]

The M5 Ultra run on 54ff80b87, main against this head, 27B, alternating builds. Tokens are identical in every row (0 of 20 differ), and the prompt cost is gone.

Cell	Change
Prompts at 32,768 tokens	+0.1%, +0.1%, -0.1%
4 streams, sweep	-0.2% to +0.2%
1 stream, sweep	-0.7% to -0.1%
1 stream decode after a long prompt	-1.2%, -1.0%, -0.1%
1 stream with the DFlash2 drafter	-0.1% to +0.9%

The --only qwen27 contract passes: 30 kernels through gates 1-4, 4 of 4 controls red. So picking the kernel per launch worked, @cshintov. That's a clear fix for the prompt side.

One thing to look at: one-stream decode leans slightly negative, while four streams don't. Your note said q27_attn_prefix keeps the masked code in the kernel with a compare per 64-key step, and it reaches decode chains on M5. Could it take the same treatment, or could you show that isn't the cause? It's small, but we hold "never slower". The M3 Ultra run and the precision tool with --plant are still ours to do.

[cshintov]

It can take the same treatment, and here it needs no scan. A stream of one row sees every key its tile reads (a decode row's keys run to its own position, and every query row of the tile is that one row), so no value in its range is masked. The host now picks q27_attn_prefix_N when every stream has one row, and that is main's kernel by IR for all 16 tile counts. Only a stream of rows (a drafted chain) runs q27_attn_prefix_masked_N, which is the previous kernel by IR. A new test-qwen27 test walks the prefix kernel's ranges over one-row plans (1 to 3 streams, starts across the tile and 8,192 boundaries, shared prefix and fused chains on and off) and finds no masked key, and it does find one in a two-row chain.

So one-stream decode without the drafter now runs exactly main's kernels. Four streams of one row each do too.

The M5 decode control still goes through the masked kernel, since gate 1's decode round is a chain. Could you rerun the one-stream decode cells on d89ce2b9c?

[ashhart]

Thanks @cshintov. Picking q27_attn_prefix_N whenever every stream has one row, with the range-walk test showing no masked key across one-row plans and one in a two-row chain, is a clean argument that one-stream decode runs main's kernel. And the prompt kernels are unchanged since 54ff80b87, so the prompt numbers we posted still stand. What's left is the one-stream decode cell on an M5 Ultra on d89ce2b9c (256 tokens after 2K and 16K prompts, against main), plus the same with the DFlash2 drafter, since drafted chains still take the masked kernel. That run is the next step.

[ashhart]

Our M5 Ultra runs on this head (d89ce2b9c), against its merge base with main, Qwen3.8-27B:

• Plain: no token differences in 20 prompt rows. Prefill at 32K is level (+0.1% to +0.2%). Decode after a prompt is -1.1% to +0.8% across 2K to 64K prompts, with no consistent direction. One-stream decode is 48.6 to 48.7 tok/s for greedy chat and 51.3 to 51.4 for sampled chat, and level for code.
• With the DFlash2 drafter: one-stream decode is 161.4 to 160.6 tok/s for greedy chat, 128.3 to 128.2 for sampled chat, 230.2 to 230.1 for greedy code and 197.1 to 197.0 for sampled code.

Both are level. The M3 Ultra run, which is the chip this change is for, is running now.

[ashhart]

Our M3 Ultra run, the chip this is for, on this head (d89ce2b9c) against main at cde4fe7d9, Qwen3.8-27B MLX 4-bit, one stream, greedy, 256 reply tokens, a fresh server for each arm. Every reply's tokens are equal to the first plain main run at that length, warm-ups and repeats included, plain and drafted.

Mode	Prompt	Prefill tok/s, main → this	Decode tok/s, main → this
plain	2,048	369.2 → 369.5 (+0.1%)	33.99 → 33.92 (-0.2%)
plain	16,384	343.3 → 343.4 (+0.0%)	29.19 → 29.22 (+0.1%)
DFlash2	2,048	364.9 → 365.4 (+0.1%)	78.29 → 78.13 (-0.2%)
DFlash2	16,384	342.3 → 342.4 (+0.0%)	51.52 → 51.53 (+0.0%)

Six measured samples a side per cell, in alternating order. Level, with equal tokens. With the M5 Ultra runs above, nothing else is owed from our side, so this goes in the landing queue. Thanks @cshintov.

[ashhart]

Thank you @cshintov. 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.