[alvarorsouza-arch]

What this changes

gr_down, one of the three dispatches of the hyper-connection read, is built with a fixed tile: 8 simdgroups per threadgroup, 8 of them splitting one row (gr_down_q8_8_8, metal/src/engine/window.cpp:638 and the same line in gr_read_m). gr.metal instantiates six tiles (GRD(8,8) (8,4) (4,4) (4,2) (16,16) (16,8)) and the engine only ever asks for one of them.

On an Apple M3 Max that tile is the slowest of the six as soon as a window carries more than one token, and it gets worse fast. metal/tests/gr_parity.cpp --bench with GR_Q8=1 (the form this pack uses) and GR_DOWN=S_K, median of three runs, per call:

Tile	1 token	2 tokens	3 tokens
8_8 (the engine's)	14.7 us	27.5 us	52.3 us
8_4	12.9 us	20.5 us	37.0 us
4_4	12.4 us	19.5 us	34.0 us
4_2	13.1 us	15.6 us	20.7 us

8_8 costs 1.88x for the second token and 3.56x for the third; 4_2 costs 1.19x and 1.58x for the same work, reading the same bytes. The two variants that split the row more ways are also the two that hold up.

The extra token is not weight traffic. One gr_down call reads 324 rows (hc_lr + hc) of hc * n_embd = 10,240 Q8_0 elements, 3.52 MB, and the tile change does not touch that. Its arithmetic is 3.32M MAC, which at the 3,729.6 Gops/s fp32 FMA rate metal/tests/alu_bench.cpp measures on this GPU is 0.89 us. So the marginal token, measured as the difference against one token in the table above, costs 12.8 us in 8_8 (7% of that rate) and 2.5 us in 4_2 (36%). What scales with the token count is the re-read of xn (10,240 floats per row, so about 13 MB per call, from cache) and the per-thread accumulation: grd_acc decodes each weight unit once and then loops the tokens inside it (metal/kernels/gr.metal:62-71), which is the right shape, but with KS = 8 each thread holds eight live accumulators over a shorter unit run and the eight partials meet in threadgroup memory behind a barrier (metal/kernels/gr.metal:82,94-102). Halving KS to 2 and putting two rows in each threadgroup is where the second token stops costing 1.88x.

The same scaling is what the engine pays: a --spec 2 window issues 96 of these calls, so the tile is worth 4.9% of the window's GPU time (below).

The patch reads the tile from STRATA_METAL_GR_TILE=S_K, keeps 8_8 when the variable is unset, and fixes the launch with it: the grid is ceil((hc_lr + hc) / RS) with RS = SG / KS and the threadgroup is 32 * SG, the same math gr_parity uses (metal/tests/gr_parity.cpp:106). Unset, the engine builds the same pipeline with the same grid and threadgroup as before.

What it leaves alone

gr_norm and gr_up are untouched. The CUDA, HIP and SYCL backends are untouched: metal/ is its own CMake project and does not enter those builds. No kernel source changes, so the six tiles already in gr.metal are the only ones reachable. The variable accepts only those six and says so on stderr if given another.

With the variable unset the engine asks for the same pipeline with the same grid and threadgroup as before, and the rebuilt binary printed the same 120 output ids as a build of main at 296b59d on the same prompt.

What it buys, measured

In the engine, STRATA_METAL_KPROF=2 on the --spec 2 window (48 layers, 96 gr_down dispatches per window, one boot each):

Tile	gr_down classes	whole window, GPU
8_8	2.656 ms	27.144 ms
4_2	1.522 ms	25.817 ms

That is 1.13 ms, or 4.9% of the window's GPU time. What the user sees, on the same machine, one boot per row, the fixed 449 token harness prompt with 256 out, three runs each, median first:

Drafter	tile 8_8	tile 4_2	tile 4_4
--spec 2 plus suffix (what setup writes)	57.3 (56.1 / 57.3 / 57.5)	59.1 (57.9 / 59.1 / 59.2)	59.1 (57.8 / 59.1 / 59.3)
drafter off	48.7 (47.8 / 48.7 / 48.7)		49.7 (49.2 / 49.8 / 49.7)
--spec 3 plus suffix	49.9 (49.1 / 49.9 / 49.9)	53.0 (52.4 / 53.1 / 53.0)

So +3.1% at the defaults, +6.2% at --spec 3 (where the tile table's 52.3 us against 20.7 us per call is the whole story), and +2.1% with no drafter using 4_4, which is the tile the bench prefers at one token. Draft acceptance did not move (113 of 146 offered with 8_8, 113 of 146 with 4_2). A second boot of the unchanged default in a later session gave 57.1 (56.1 / 57.2 / 57.1), so the default's own boot-to-boot spread is 0.2 tok/s against a 1.8 tok/s difference here. A separate 200 token greedy one-shot, same session, gave 62.88 tok/s for 8_8 and 64.63 for 4_2.

How parity looks

GR_Q8=1 GR_DOWN=4_2 build-metal/gr_parity --ntok N against the double precision reference, for N = 1, 2, 3, 8, both with and without the folded write: PASS, with the same relative errors as the current tile (lo 1.2e-7 against 1.4e-7, inject 1.1e-7 against 1.2e-7, prepared planes bitwise equal to xprep). The tiles are not bitwise identical to each other, since the row is summed in a different order; a greedy one-shot of 200 tokens on this machine came out with the same 200 ids for 8_8, 4_4 and 4_2, which is agreement, not proof.

Related

#2070 found the same shape of problem on the CUDA side and fixed it for sm_75: the hyper-connection read as about 19% of an RTX 2080 Ti's decode GPU time, 96 reads per window, with a dedicated gr_up kernel for 3-plus tokens and the staged down kernel taking a token group, opt-in behind STRATA_GR_UP75. My numbers say the Metal engine has the same token scaling, and there the fix so far is only the tile choice, which is why this PR keeps it opt-in. #1775 (SYCL, three kernels instead of four for the same read) is the other thread in this area.

What I did not test

Only the M3 Max, one pack (IQ3_S with the Q8_0 hyper-connection copies at load), one model. The M1 and M2 families and the M5 (which gets the M5 cooperative layout) are unmeasured here, which is why the tile is opt-in rather than a new default. The BF16 and F32-inject forms of the read (gr_down_8_8, gr_down_q8f_*) were benched only in their BF16 form in gr_parity, not in the engine. gr_up scales the same way with the token count (13.0 us at one token, 24.4 us at two, 35.3 us at three, and it is the larger half of the read at 2.35 ms per --spec 2 window) and has no tile knob in gr.metal; that is where I would go next, and it is what #2070 built for Turing.

[alvarorsouza-arch]

One more measurement in this file, since the same scaling is in gr_up and it is the larger half of the read. I tried the cheap version of the fix there too: stage the window's lo rows in threadgroup memory once per threadgroup instead of re-reading them from the device for each of the hc * n_embd rows. New kernel gr_up_sh_f in gr.metal, picked by STRATA_METAL_GR_UP=sh, dots and epilogue copied line for line.

gr_parity --bench with GR_Q8=1, median of three, per call:

Tokens per window	gr_up_q8	gr_up_q8_sh
1	13.1 us	13.6 us
2	24.5 us	20.6 us
3	35.3 us	28.7 us
4	41.7 us	34.4 us

In the engine the class goes from 2.344 ms to 2.025 ms per --spec 2 window (STRATA_METAL_KPROF=2), which is 1.2% of the window, and gr_parity reports the same parity figures for both (mixed 7.12e-8 at two tokens in each, prepared planes bitwise equal to xprep in each).

End to end I cannot see it: on the fixed 449 token harness prompt, 256 out, three runs each, one boot each, --spec 2 gives 57.3 / 57.6 / 57.5 tok/s for default, sh, default again, and --spec 3 gives 49.7 for both. The boot-to-boot spread of the unchanged default is the same size as the effect. So 1.13 ms of kernel time (this PR) shows up as +3.1% decode and 0.32 ms (this one) does not show up at all, which is a useful note about what is worth changing here: the gr_down tile clears the noise, gr_up as I left it does not.

I am keeping the gr_up_sh branch local rather than sending a second PR on these numbers. If someone can measure where the window dominates, an M5 with the cooperative layout or a batch of concurrent requests, the same change is worth another look there.

[Niko1221]

Thanks for the pull request.

We're resetting the tracker for 0.2.0-beta. There are more than 1,300 open issues and pull requests, too many for us to review properly, so we're closing everything that isn't being worked on right now. Closed here does not mean rejected.

To reopen it takes about two minutes:

1. Rebase on main if it no longer applies, and re-run your checks on 0.2.0-beta.
2. Start the title with its area: [Engine], [GUI], [TUI], [Benchmark], [Docs] or [Other].
3. Fill in the pull request template: the problem and result, and the validation (what you ran, on which revision, what you saw).
4. Comment /reopen. A check runs and either reopens the pull request or lists what is still missing.

The rules are in CONTRIBUTING.md.