Use the explicit gemm for wide convs on neural accelerators

[erwinzhang7]

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.

[erwinzhang7]

Following up with two real models since the synthetic
table alone doesn't show if this matters.

Short version is this helps convolution-dominated models and is
invisible in a model where convs sit in front of an autoregressive decoder.
ResNet-50 is the case for it; Whisper is the limit of it.

ResNet-50 forward, fp16, M5 Max. Every stage above layer1 crosses the 512
threshold. Control is this PR's parent 3abd0fd6:

batch 1                control     this PR   speedup
  whole model           4.00 ms    3.38 ms     1.18x
  layer1 (cout 256)     0.40 ms    0.38 ms     1.05x   <- below threshold
  layer2 (cout 512)     0.57 ms    0.54 ms     1.06x
  layer3 (cout 1024)    1.10 ms    1.00 ms     1.10x
  layer4 (cout 2048)    1.00 ms    0.59 ms     1.69x

batch 8
  whole model          19.44 ms   16.11 ms     1.21x
  layer1 (cout 256)     1.42 ms    1.39 ms     1.02x   <- below threshold
  layer2 (cout 512)     1.55 ms    1.44 ms     1.08x
  layer3 (cout 1024)    2.35 ms    2.08 ms     1.13x
  layer4 (cout 2048)    1.66 ms    0.90 ms     1.84x

layer1 is the control: both builds run identical code for it, and it measures
1.02x to 1.05x. The gain then rises monotonically with output channels, which is
the mechanism the description argues for, so the effect tracks the axis it should
rather than appearing at an arbitrary place.

Whisper, where the shapes are real but the impact is not. Its encoder
is two Conv1d layers, and conv_1D_gpu routes them through
dispatch_conv_2D_gpu whenever input dilation is 1 and both channel counts are
multiples of 16, which every Whisper size satisfies. The convs themselves do
speed up, small measuring 1.058 ms to 0.535 ms per 30 s window, with tiny at
cout 384 sitting below the threshold as a control at 0.97x.

But the whole encoder is 20.4 ms per window, 41 ms across the two windows of a
2850 ms transcription of 52 s of speech, about 1.4%, and the two convs are a
fraction of that. End to end it is roughly 1 ms in 2850, which is not
measurable.

[zcbenz]

How does the memory usage change?

[erwinzhang7]

Peak goes up, steady state doesn't. The unfold buffer is freed after the conv.

ResNet-50 fp16, same build with and without this pr:

peak	without	with	delta
batch 1	125.4 MB	146.2 MB	+20.8 MB
batch 8	876.8 MB	933.5 MB	+56.7 MB

Per stage at batch 8: l2 +9.2 MB, l3 +6.1 MB, l4 +26.1 MB. l1 is below the threshold and does not move, which is the control.

I added a limiter. The explicit path is taken only while the unfold fits alongside what is already allocated, currentAllocatedSize() + unfold <= recommendedMaxWorkingSetSize, and a conv that exceeds this falls through to the implicit gemm.

It's dynamic, which matters: the same conv is fine on an idle machine and not fine under an already-loaded one, and a fixed fraction gets that wrong in both directions. Measured, 27.6 GB unfold either way:

allocated before	allocated + unfold vs 107.5 GB	path	growth
1.8 GB	29.3 GB, fits	explicit	26.6 GB
94.2 GB	121.7 GB, exceeds	implicit	0.14 GB

ResNet-50 is unaffected, and fp16 output matches an fp32 reference either side of the threshold.

This is a memory bound, not about the speed. The unfold consistently pays for itself in speed, the explicit path is still 1.14x ahead at 45% of the working set. It's just to not mess up the rest of the process.

[zcbenz]

The unfold path really means to be a fallback when all other algorithms do not apply, we can probably enable it if this is actually making a different for LLM inference, otherwise I intend to close this.

Also we don't dispatch kernels based on RAM usage, which is not reliable in real world inference and makes things too un-deterministic.

[erwinzhang7]

That's fair and good to know. Especially on the RAM point.

The speedup is real but minimal except on ResNet-style stacks with wide channels and real spatial extent. I think I'll investigate more outside this repo, upstreaming only if there's a path it can help with LLM inference.

Closing. Thanks for the look, and for letting me make the case first.