Use the explicit gemm for wide convs on neural accelerators - #4214
Use the explicit gemm for wide convs on neural accelerators#4214erwinzhang7 wants to merge 1 commit into
Conversation
|
Following up with two real models since the synthetic Short version is this helps convolution-dominated models and is ResNet-50 forward, fp16, M5 Max. Every stage above layer1 crosses the 512 layer1 is the control: both builds run identical code for it, and it measures Whisper, where the shapes are real but the impact is not. Its encoder But the whole encoder is 20.4 ms per window, 41 ms across the two windows of a |
conv.cpppicks between the implicit gemm and unfolding into an explicit gemm. That choice wastuned when both paths ran on
BlockMMA.steel_matmulnow takes the NAX path on hardware withneural accelerators, so one side of the tradeoff got a lot faster and the boundary hasn't moved.
Measured on an M5 Max (
applegpu_g17s), before and after, over eight shape families thatcurrently reach the implicit gemm, so nothing winograd takes. n=8 per cell:
The 128 and 256 rows are a control: they are below the threshold, so both builds run the same
code for them and the true ratio is 1.00x. They measure 0.97x to 1.10x, which puts the noise
floor around 10% and leaves every real result well clear of it.
The axis is the output channel count, not the filter size or the spatial extent. Unfolding costs
an extra pass over the input, roughly
2*M*Kelements, while the gemm that follows does2*M*N*K, so the work per materialized element scales withN. HoldingC_outfixed and movingKfrom 1152 to 9216 leaves the ratio flat (2.62x, 2.47x, 2.58x, 2.58x); holdingKfixed andmoving
C_outwalks it across the whole range in the table above. The same behaviour shows upfor 3x3, 5x5 and 7x7 filters, stride 1 and 2, dilation 1 and 2, aligned and unaligned channels,
and
Mfrom 1024 to 65536, all within about 0.1x.Gating on
is_nax_available()because without the accelerators the explicit path isworse. On an M4 Pro (
applegpu_g16s) it loses at every size, reaching 1.01x only atC_out = 4096:Restricted to float16 and bfloat16 because for float32 the gain is precision, not speed. NAX runs
float32 at TF32 mantissa, and conv doesn't do that today. With
MLX_ENABLE_TF32=0so both pathsrun true float32, the explicit path loses everywhere: 0.52x, 0.67x, 0.78x, 0.88x for the same
four sizes.
Verified
M5 Max, macOS 26.6, and M4 Pro, both against a CPU reference: 144 cases and 36 cases over filter
1/3/5/7, stride 1/2, padding 0/1/2, dilation 1/2, aligned and unaligned channels, two batch
shapes.
float16 and bfloat16 sit at dtype rounding. float32 is two orders tighter, which is the check
that it stays off the TF32 path.
Separately, 504 cases per dtype comparing the two dispatch paths directly against each other:
worst relative error 7.79e-04 for float16, 6.06e-03 for bfloat16, no shape mismatches.
On an M4 Pro the branch is unreachable, so it runs exactly what it does today.
Notes
Winograd, depthwise, grouped and the small-channel paths are untouched; the branch sits after
those and only catches what would otherwise go to the implicit gemm.
The explicit path already bounds its own memory.
max_unfold_rowstiles the unfold againstmaxBufferLengthand reuses one buffer, so this doesn't introduce an unbounded allocation.C_out >= 512rather than 256 because at 256 the win is thin onceMgets large (1.03x to 1.19xat
M = 16384and above), which isn't worth the extra pass. Happy to move it if you'd ratherhave the 1.2x.