Repository navigation
perf: gate and up in one decode pass for GGUF formats on Apple7/8 - #27
Conversation
On the Apple7/8 register tile a gate/up pair of one GGUF format now runs as
one dispatch instead of a gate pass and an up pass: grid {N/32, S}, each
simdgroup's fragment nf=0 is 8 gate columns and nf=1 the same 8 up columns,
so SwiGLU is applied in registers. With S splits the last of the S arrivals
on a 32-column tile sums both tensors' partials and applies the gate. The
outputs are the two passes' bit for bit (every register format at 768x2048,
1-4 lanes, S 1-8; Q4_K and MLX 4-bit at 17408x5120).
MLX affine formats (af4g64) keep two passes: with one pass their 17408x5120
gate/up was -7.6 / -5.3 / +2.3 / +4.4% at 1 / 2 / 3 / 4 lanes. The choice is
by format, not by lane count, so a model runs the same plan at every lane
count. In UD files only layers whose gate and up share a format take one
pass (36 of 64 in UD-Q4_K_M). gguf-projection-benchmark gets register_gu
rows for the one-pass plan.
M2 Max, 17408x5120 gate/up at 1 / 2 / 3 / 4 lanes (S4, ABBA): q4k -6.4 /
-3.3 / -3.2 / -3.1%, q5k -7.1 / -5.4 / -5.8 / -2.2%, q6k -6.4 / -0.5 / -3.3 /
-3.6%, q3k -6.8 / -4.0 / -4.0 / -3.8%, q40 -6.1 / -3.6 / -4.0 / -3.2%, q80
-3.6 / -3.3 / -4.7 / -4.0%, iq4xs -7.5 / -5.9 / -1.4 / -2.5%; af4g64 and the
draft's 6144x2048 af4g64 unchanged (within 0.3%).
decode-profile, 27B Q4_0, 2048-token prompt, mean of two ABBA runs, with the
GDN verify and norm loads changes under it: the cycle 82.19 -> 80.53 ms
(-2.0%), 145.41 -> 143.50 (-1.3%), 199.63 -> 196.59 (-1.5%), 263.67 ->
257.44 (-2.4%; one slow base run, -1.2% against the other). MLX 4-bit, which
keeps two passes, is flat within its runs' noise.
Check bit-exactness of one-pass gate/up against two passes across all 42 register formats at 768x2048 (1-4 lanes, S 1-8) and Q4_K and MLX 4-bit at 17408x5120 (1 and 4 lanes). Add linear-plan verification for Apple7/8 oneGateUpPass policy (GGUF formats one-pass, MLX af4g64 two-pass, mixed formats two-pass) and scratch sizing bounds.
ghisguth
changed the title
|
@paperniuk I finished going through the ideas I had for decode on M1/M2. #25, #26 and #27 are the ones that paid off. Together with the merged #18–#21 this is what changed since 1.3.2-m1.2 on my M2 Max (without #24). Decode t/s at 1 / 2 / 4 lanes,
Prefill is the same in all three. These are from runs on three different days, so take each number as ±0.5%. I am re-running all of it stacked in one session now (A, B, C, then A again) and will post the numbers here when done. I also tried a few more things that didn't pay off on M2 Max: wider column tile for the register tile, two groups per barrier in the prefill tile, splitting the ANE prefill tail past 32 rows, clamped loads in prefill GDN and norm. Also register-tile draft attention: the kernel is 36–67% faster, but it is only ~1% of the cycle and nothing shows end to end, so I didn't open a PR. If you want it anyway let me know. |
|
@paperniuk Ok, the stacked re-run is done. All builds in one session on my M2 Max, display off, the release tag measured at the start and at the end. The two tag runs agree within 0.5%, so the day didn't drift much. Decode t/s at 1 / 2 / 4 lanes,
Prefill is the same in all of them (60.6 / 18.9 / 8.1 s on MLX). It is close to the numbers from three different days I posted above, MLX at one lane came out a bit lower (+7.0 instead of +7.9).
I also replayed my agent chat on MLX 4-bit (157 requests, up to 125K context). Aggregated Metrics (B = apple7-m1-1.3.2 replay from Oct 21, C = the stacked branch). Cycles/s is decode t/s divided by tok/cycle from the server log, so it doesn't depend on how many draft tokens were accepted. The 32k row has only 3–5 requests.
All requests was successfull. At long context the gain is smaller since attention takes more of the cycle and these PRs don't touch it.
|
|
Thanks! Confirmed on M1 Max (32-core GPU).
Small on this file, as expected with only part of the layers in one pass, but consistent. Merged. |
…port) into the 1.4 port Conflicts: - paged_attention.metal, rope.metal: 1.4's own versions, which already hold #24 and the precise RoPE tables. - LinearGguf.cpp, linear_plan_test.mm: #27's one-pass check for the F32 register tile next to 1.4's 16-row tile check; #27's one-pass scratch with LinearPlan::table(). - gguf_projection_test.mm: #27's register one-pass test where the device is not Apple9, before 1.4's prefill. - gguf_projection_benchmark.mm: register_gu rows next to 1.4's staged16 rows; the staged rows take one pass only from a staged policy.
Base:
paperniuk/splashapple7-m1-1.3.2(branchapple7-m1-gate-up-pass-1.3.2: the change, then the test). Independent of the GDN verify and norm loads PRs (#25, #26). The same two commits also apply cleanly toapple7-m1-1.4.On the M1/M2 register tile each MLP layer runs gate and up as two decode passes, and the up pass applies SwiGLU in its epilogue. This change runs both as one dispatch for GGUF formats: in each 32-column threadgroup every simdgroup computes 8 gate columns and the same 8 up columns, so SwiGLU happens in registers. On my M2 Max the 27B's gate/up takes 2–8% less time on every GGUF format, and the Q4_0 decode cycle is 1–2% faster. Outputs are bitwise the same as the two passes. As before, I used Claude Code to author it and tested it by hand.
What changed
gguf_decode_sgf_<fmt>_l<lanes>_guinstances inlinear_gguf_sgmatrix.metal(the F32 table instances, so Apple7/8 only). Grid {N/32, S}. With one split a threadgroup writesup · silu(gate)directly. With S splits it writes both tensors' partials, and the last of the S arrivals on a 32-column tile sums them and applies the gate.Linear::ggufBaselinesetsoneGateUpPasson the Apple7/8 register plan when gate and up are single tensors of one format and that format is not MLX affine (apple7GateUpFormat). The scratch bound covers the one-pass plan: partials of both tensors and one counter per 32 columns.gguf-projection-benchmarkgetsregister_gurows. Its staged rows tookoneGateUpPassfrom the policy, which is now set on the register plan, so on Apple7/8 they would have run the staged one-pass kernel. Now they take it only from a staged policy, and they run what they ran before on every family.Results (M2 Max, 38-core GPU)
Measured on
apple7-m1-1.3.2with #25 and #26 under it, and onapple7-m1-1.4.Kernel time,
gguf-projection-benchmark, 17408×5120 gate/up, the policy's config (S4), ms at 1 / 2 / 3 / 4 lanes, ABBA. Δ: − is faster. In the same runs the two-pass rows of the new build match the base within 0.3%.decode-profile, 2048-token prompt, ms per cycle at 1 / 2 / 3 / 4 lanes, mean of two ABBA runs, onapple7-m1-1.3.2:The Q4_0 4-lane number is flattered by one slow base run (266.8 against 260.6 ms). Against the other one it is −1.2%. MLX doesn't run the new code, and its two base runs differ by up to 0.8%, so I read it as flat. On
apple7-m1-1.4: Q4_0 −1.9 / −1.7 / −1.2 / −1.0%, UD-Q4_K_M −1.0 / −0.6 / −1.1 / −0.0%, MLX 4-bit +0.3 / −0.2 / 0.0 / −0.1%.make test-performance-real(ANE share 0.35, MLX first and last), before and after back to back, decode in tok/s at 1 / 2 / 4 lanes, onapple7-m1-1.3.2. Δ is from the decode times of the two runs. Prefill doesn't change (GGUFs within 1.1%).The UD-Q5_K_M 4-lane number is one run, and I think it is noise, but I didn't rerun it.
Testing
gguf-projection: a new check runs the one pass against the two passes, bitwise, for all 42 register formats at 768×2048 (1–4 lanes, S 1–8) and for Q4_K and MLX 4-bit at 17408×5120 (1 and 4 lanes). Passes with and withoutMTL_SHADER_VALIDATION.linear-plan: Apple7/8 plans one pass for a Q4_K pair and two for a Q4_K/Q6_K pair and an af4g64 pair, and the scratch bound covers the one-pass plan. Passes with and without validation.make checkpasses on the M2 Max onapple7-m1-gate-up-pass-1.3.2.decode-profileandmake test-performance-realruns above.