Commit Graph

1 Commits

Author SHA1 Message Date
John
6fcf677cef Fuse the K^3 taps into one gather + matmul (2.1x on M3 Ultra)
Fleet benchmarking exposed the problem: throughput was ANTI-correlated with GPU
core count. The 80-core M3 Ultra came in slowest at 128^3/128ch (81.6ms) behind a
38-core M2 Max (58.3ms), M1 Ultra (66.7ms) and even a 32-core M1 Max (69.2ms).
That ordering only makes sense if the op is bound by dispatch latency rather than
compute - the per-offset loop issued 2*K^3 = 54 tiny GPU ops per layer, none big
enough to occupy the machine, and the Ultra's fused-die design penalises exactly
that.

Concatenating the K^3 neighbour taps along the channel axis collapses it to a
single [N, K^3*Cin] x [K^3*Cin, Cout] matmul. Chunked over rows so peak memory
stays ~256MB (the unchunked buffer is ~2.9GB at 128^3/128ch - fine on a Studio,
not fine on an 8GB mini).

m3ultra 128^3/128ch: 81.6ms -> 39.3ms (2.08x), 2.57 -> 5.34 Mvox/s
m3ultra  64^3/128ch: 21.9ms ->  6.4ms (3.4x)

7/7 tests still pass against the naive reference.
2026-08-02 10:07:53 +10:00