Skip to content

Use the explicit gemm for wide convs on neural accelerators - #4214

Open
erwinzhang7 wants to merge 1 commit into
ml-explore:mainfrom
erwinzhang7:nax-conv
Open

Use the explicit gemm for wide convs on neural accelerators#4214
erwinzhang7 wants to merge 1 commit into
ml-explore:mainfrom
erwinzhang7:nax-conv

Conversation

@erwinzhang7

Copy link
Copy Markdown
Contributor

conv.cpp picks between the implicit gemm and unfolding into an explicit gemm. That choice was
tuned when both paths ran on BlockMMA. steel_matmul now takes the NAX path on hardware with
neural 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 that
currently reach the implicit gemm, so nothing winograd takes. n=8 per cell:

C_out fp16 bf16
128 1.10x 0.97x
256 1.08x 0.98x
512 2.03x 1.85x
1024 2.73x 2.46x
2048 3.40x 3.02x
4096 3.75x 3.28x

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*K elements, while the gemm that follows does
2*M*N*K, so the work per materialized element scales with N. Holding C_out fixed and moving
K from 1152 to 9216 leaves the ratio flat (2.62x, 2.47x, 2.58x, 2.58x); holding K fixed and
moving C_out walks it across the whole range in the table above. The same behaviour shows up
for 3x3, 5x5 and 7x7 filters, stride 1 and 2, dilation 1 and 2, aligned and unaligned channels,
and M from 1024 to 65536, all within about 0.1x.

Gating on is_nax_available() because without the accelerators the explicit path is
worse. On an M4 Pro (applegpu_g16s) it loses at every size, reaching 1.01x only at
C_out = 4096:

C_out 128 256 512 1024 2048 4096
M4 Pro 0.39x 0.55x 0.71x 0.87x 0.95x 1.01x

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=0 so both paths
run 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.

dtype M5 Max M4 Pro
float16 7.51e-04 5.88e-04
bfloat16 5.49e-03 5.43e-03
float32 9.07e-06 5.19e-06

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_rows tiles the unfold against
maxBufferLength and reuses one buffer, so this doesn't introduce an unbounded allocation.

C_out >= 512 rather than 256 because at 256 the win is thin once M gets large (1.03x to 1.19x
at M = 16384 and above), which isn't worth the extra pass. Happy to move it if you'd rather
have the 1.2x.

@angeloskath
angeloskath requested a review from jagrit06 August 12, 2026 19:57
@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants