Pull requests / #2655
#2655 Metal: a tile for gr_down on windows of more than one token (STRATA_METAL_GR_TILE)
open · @alvarorsouza-arch · 0 commentaires · Sur GitHub
BenchmarksSetup & installAMD / HIPNVIDIA / CUDAModels & quantsWindows
Description
## 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.
Sur le site
Liens install, modèles, releases.