Repository navigation
perf: Apple7/8 MoE gate and up in one decode pass - #45
Conversation
Apple7/8's MMA expert tile ran a decode step's gate and up projections as two passes, the up pass reading the gate's bf16 results back. One pass now runs both (moe_expert_gguf_mma_m8_gate_up[_mixed], the staged tile's bindings and grid): a gate and an up simdgroup on the same 32 columns, each projection on the tile its own pass takes, the gate parked in its stage. Where both take the column tile (Q2_0 tiles of one or two live rows), one lane runs a column's gate and up chains interleaved (column_dot_pair), so their FMAs overlap and each activation load serves both. Each chain keeps its order, so the intermediate is bitwise the two passes'; gguf-moe checks every MMA format against them. Prefill keeps two passes. M2 Max 38 cores, against 1.5.0-m1 (A B B A): gguf-moe-benchmark, ms a layer, two passes -> one, 1-row verify / 8-32 rows: Q2_0 (Flash-Next shape) -24.0% / -1.5..-15.8%; Qwen3.6-35B-A3B shape with Q4_K/Q5_K/Q6_K/IQ4_XS/IQ4_NL/Q8_0/MLX 4-bit gate/up -27..-34% / -0.7..-11.2%; Flint (two shared slices) Q4_K and MLX 4-bit -30% / +0.2..-20% (+0.2% inside its 0.3% spread). decode-profile, Qwen3.8-Flash-Next 2K, int8 KV, eight verify rows a lane, fused ms B1/B2/B3/B4: Q2_0 60.05/90.71/116.72/144.03 -> 57.38/86.88/113.38/140.55 (-4.5/-4.2/-2.9/-2.4%) Coder 70.87/103.11/134.06/161.91 -> 68.72/101.38/133.50/160.82 (-3.0/-1.7/-0.4/-0.7%) IQ3_XXS 69.90/104.38/133.28/162.24 -> 68.50/102.91/132.18/161.30 (-2.0/-1.4/-0.8/-0.6%) Q2_0 32K 61.57/94.49/121.76/149.73 -> 59.14/90.57/117.86/146.64 (-3.9/-4.1/-3.2/-2.1%) Server, Flash-Next Q2_0, 5 greedy 400-token requests: 18.5 -> 18.9 cycles/s. Agent chat replay (78 turns, sent twice, 32-128K context), cycles/s: Q2_0 +3.4..+4.0%, Coder +0.4..+2.5%; prefill unchanged.
|
Ran it on my M1 Max (the #29/#38/#40 machine). Same recipe: the ISTA-DASLab The end-to-end medians of 3 interleaved runs each sit inside this machine's
−7.2% on the experts, ≈ −2.8% of the whole B1 cycle — right in your 2.4–4.5% |
|
@abssoul thank you for running it on your M1 Max! Good to see the fused kernels win there too, even if the cycle is inside the noise. ±8% run to run is a lot. On my M2 Max I get ~1% with the display off, If you have time, #46 is another small one for the same model. The cycle effect is ~1%, so on your machine it would only show in the profile too: |
|
Thanks, merged. On my M1 Max together with #46, Flash-Next Q2_0 through the server, four runs each in ABBA order: 1 lane 73.9 to 75.8 tok/s, 2 lanes 93.1 to 94.4, prefill unchanged. |
@paperniuk this one is for MoE decode on M1/M2. I profiled Flash-Next Q2_0 decode on my M2 Max and the routed experts
are about a third of a one-lane cycle, ~40% at eight verify rows. On Apple7/8 the MMA expert tile runs gate and up
as two passes, the up pass reading the gate back. This runs both in one pass, like the staged tile does on Apple9+.
Outputs are bitwise the same. Flash-Next Q2_0 decode cycle is 2.4–4.5% faster, the Coder and IQ3_XXS
checkpoints a bit less, and the expert kernel is faster for every GGUF format I tried, MLX 4-bit too. As before, I
used Claude Code to author it and tested it by hand.
What changed
runtime/ops/MoE.cpp:oneGateUpPassis also true for the Apple7/8 MMA tile in decode. Prefill keeps two passes(I didn't measure it). The rest is the existing one-pass path: bindings, 32-column grid,
_mixedtables.kernels/shared/moe_gguf_mma.metal:moe_expert_gguf_mma_m8_gate_up[_mixed][_s2]. A gate simdgroup and an upsimdgroup on the same 32 columns, each projection on the tile its own pass takes (the column tile for a Q2_0 tile
of one or two rows, the MMA tile otherwise). The gate parks bf16(gate) in its own stage, one barrier, then the up
writes bf16(up) * silu(gate) like the up pass. No new threadgroup memory (6 KB, 7 KB for
_mixed).(
column_dot_pairingguf_column_tile.h), so their FMAs overlap and the activation loads serve both. Each chainkeeps
column_dot's order. The even threadgroup's two simdgroups take the odd one's columns, so the pairs run on asmany simdgroups as before.
gguf_moe_test.mm: expects one pass on the MMA tile in decode. The existing check then splits the one-pass dispatchback into the
_a/_gpasses and compares the output bit for bit on every run, every format the MMA tile takes.What I tried first and dropped: the staged one-pass kernel on Apple7/8 is 8–14% slower in the cycle (its multi-row
tiles go to the MPP matmul). Gate then up serially in each simdgroup made the Coder up to 1.4% slower.
Results (M2 Max, 38-core GPU)
Expert kernel alone,
gguf-moe-benchmark(one MoE layer, random weights in each model's shape and format, no modelloaded), ms per layer, Δ two passes → one pass (%), A B B A, 1 = a verify of one row,
r / p = 8–32 rows with random / pooled routes:
The only plus (Flint MLX 4-bit, 16 pooled) is inside that cell's run-to-run noise (0.3%).
decode-profile, Flash-Next, 2048-token prompt, int8 KV, fused GPU ms per cycle at 1 / 2 / 3 / 4 lanes, A B B A,eight verify rows per lane (Q2_0 four runs per side, the others two):
With the verify budget's own rows (noisier, the rows differ between runs):
The IQ3_XXS one-lane +1.4% is not noise in these runs, both runs per side agree within 1%. With the budget the two
builds verify different rows. A second A B B A gave +1.0% again. With a debug build that logs the rows, the same
kernel at the same rows is 2.5% faster at one lane, and at a fixed eight rows (table above) it is 2.0% faster. So I think
it is the rows the budget picks, but I can't show the rows of those exact runs.
Server, Flash-Next Q2_0 with my usual flags, median decode tok/s of 5 greedy 400-token requests per run, two runs
each: 55.83 / 55.50 before and 57.64 / 56.66 after (+2.6%). Per request in cycles per second, which
doesn't depend on how many draft tokens get accepted, 18.5 → 18.9 (+2.4%), faster on every request.
Long agent chat replay in one of my tests (78 turns up to 125K context, each sent twice), server log, base → this
PR. Cycles per second is the decode speed without the effect of how many draft tokens get accepted:
Prefill is the same (paired requests +0.2% and −0.1%). IQ3_XXS didn't fit 156K of context at my memory setting, so it
has no replay numbers.
Testing
gguf-moeandgguf-moe-e512pass on the M2 Max with and withoutMTL_SHADER_VALIDATION. I also checked thatthey catch a broken mixed tile (gate and up on different tiles, which a shared expert with Q2_0 gate and another up
format hits): a build that zeroes those tiles fails both tests, in both directions.
linear-planpasses with and without validation.make checkpasses on the M2 Max up to the app tests, all Metal tests under validation included (the MoE onesdon't time out anymore).
swift testfor the app hung downloading Sparkle here, so I stopped it. It doesn't touchthis code.