Skip to content

perf: decode only the span chains' coefficients in the Apple7/8 register decode tile - #20

Merged
paperniuk merged 1 commit into
paperniuk:apple7-m1-1.3.2from
ghisguth:apple7-m1-mlx-span-coefs
Oct 21, 2026
Merged

paperniuk merged 1 commit into
paperniuk:apple7-m1-1.3.2from
ghisguth:apple7-m1-mlx-span-coefs

Conversation

@ghisguth

@ghisguth ghisguth commented Oct 21, 2026 •

Copy link
Copy Markdown

This one is for MLX 4-bit, the default model. In the Apple7/8 F32 register decode tile, MLX affine 4-bit in groups of
64 (and 128) runs span chains: one chain per 64-input span, closed with one coefficient. The tile still decoded a
coefficient for every 32-input group, so half of them were decoded and never read, and each had its meta load and
threadgroup store. Removing them makes MLX 4-bit decode 2.8% faster at one lane on my M2 Max. As before, I used
Claude Code to author it and tested it by hand.

What changed

Only the Apple7/8 F32 instances of the register decode tile (linear_gguf_sgmatrix.metal) change. Apple9 and host
code are untouched.

  • For formats with span chains (SpanChain: MLX affine in groups of 64 or 128), the F32 instances decode only every
    other coefficient of a unit, the ones the span chains read. The coefficient values and the chains are unchanged, so
    the outputs are bit for bit the same.
  • No test changes. gguf-projection already runs every format through the F32 instances: natively on Apple7/8, and on
    Apple9+ through PR perf: zero points as mins and nibble operands in the Apple7/8 register decode tile #15's pass.

Results (M2 Max, 38-core GPU)

Kernel time, gguf-projection-benchmark, at 1 / 2 / 3 / 4 lanes, ABBA order; Δ: − is faster.

Format, shape Before (ms) After (ms) Δ (%)
af4g64 17408×5120 gate/up 0.415 / 0.732 / 1.011 / 1.383 0.404 / 0.725 / 1.004 / 1.361 −2.7 / −1.0 / −0.6 / −1.6
af4g64 5120×17408 down +res 0.218 / 0.382 / 0.531 / 0.718 0.211 / 0.376 / 0.529 / 0.711 −3.1 / −1.5 / −0.4 / −1.0
af4g64 5120×6144 out +res 0.088 / 0.160 / 0.220 / 0.297 0.086 / 0.157 / 0.219 / 0.293 −2.3 / −1.5 / −0.5 / −1.3
af4g128 17408×5120 gate/up 0.405 / 0.726 / 1.012 / 1.366 0.401 / 0.717 / 1.001 / 1.369 −1.0 / −1.3 / −1.1 / +0.3
q4k (unchanged code) 0.435 / 0.772 / 1.089 / 1.405 0.434 / 0.771 / 1.091 / 1.403 −0.1 / −0.2 / +0.2 / −0.1

MLX 4-bit on the M2 Max, before and after

make test-performance-real, mlx-community/Qwen3.8-27B-4bit, ANE share 0.35. Times in seconds; decode in tok/s.

14K cold 10K hit + 4K 2K Decode 1 Decode 2 Decode 4
Before 60.64 18.91 8.07 86.9 103.0 112.0
After 60.64 18.91 8.07 89.3 104.2 113.1
Δ (%) 0.0 0.0 0.0 +2.8 +1.2 +1.0

Prefill doesn't change: the change is decode-only, and prefill runs the staged MMA tile. Three MLX runs of this branch
agree within 0.3%.

Testing

  • gguf-projection passes on the M2 Max (all 43 formats, F32 instances natively).
  • make check passes on the M2 Max.
  • make test-performance-real with MLX 4-bit, before and after.

…ter decode tile

The F32 span chains (MLX affine in groups of 64 or 128) read one coefficient
per 64-input span, but the tile decoded one per 32-input group, two per span,
each with its meta load and threadgroup store. The F32 instances of those
formats now decode the coefficients their chains read alone, half the
unit's. The Apple9 instances are unchanged.

M2 Max, ms at one to four lanes:
  af4g64 17408x5120 gate/up  0.416 0.732 1.012 1.383 -> 0.404 0.724 1.004 1.360
  af4g64 5120x17408 down     0.218 0.382 0.531 0.718 -> 0.212 0.377 0.528 0.710
  af4g64 5120x6144 out       0.088 0.159 0.220 0.297 -> 0.086 0.158 0.219 0.294
  af4g128 17408x5120 gate/up 0.405 0.726 1.010 1.364 -> 0.401 0.718 1.000 1.367
  q4k (unchanged)            within 0.2%
@paperniuk
paperniuk merged commit ae00021 into paperniuk:apple7-m1-1.3.2 Oct 21, 2026
8 of 10 checks passed
airiclenz pushed a commit to airiclenz/splash that referenced this pull request Oct 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants