Decompose small kernel-depth 3D convs into 2D convs

[katlun-lgtm]

Closes #3625.

Problem

mx.conv_general with 5D inputs (3D convolution) is 2–5× slower than decomposing the
same op into per-frame 2D convolutions with a Python loop. This hits video-generation
workloads (VAE decoders with CausalConv3d).

The reason isn't the explicit-gemm fallback — for the common case (mod16 channels,
idil==1, groups==1) dispatch_conv_3D_gpu already routes to
implicit_gemm_conv_3D_gpu. The real cause is that the 2D dispatch has a Winograd
kernel for 3×3 stride-1 convs (dispatch_conv_2D_gpu → winograd_conv_2D_gpu), and the
3D path has no Winograd or 3×3-specialized kernel — so a 3×3×3 conv misses Winograd
entirely. Decomposing a small kernel-depth 3D conv into KD 2D convs lets each 2D conv
hit the tuned 2D path.

Change

small_kd_conv_3D_gpu: for small kernel depth, for each depth tap kd we build a
zero-copy strided view of the input frames ([OD, H, W, C] at depth offset kd) and the
weight depth-slice ([O, KH, KW, C]), run conv_2D_gpu, and accumulate. The accumulator
buffer is repointed into out via copy_shared_buffer (the pad_and_slice pattern).

A guard in dispatch_conv_3D_gpu takes this path only when it is valid and faster:

• idil == 1 (all dims), groups == 1, N == 1;
• KD <= 7 (the KD-1 accumulate adds erode the win for large KD);
• depth stride == 1 and kernel_dilation == 1, and pad[0] == 0;
• mod16 channels (so the per-frame 2D convs hit the fast path).

Everything else falls through to the existing implicit gemm unchanged.

(Skeleton follows @Ved235's sketch in the issue; the depth-tap decomposition idea is
theirs.)

Results (M3 Max, mlx built from this branch)

Correctness — the fast path vs the CPU reference (fp32, exact):

KD	max|gpu − cpu|	max|cpu|
1	3e-5	61.0
2	1e-4	91.6
3	1e-4	106.8
5	2e-4	142.2

Speed (bf16, 3×3×3), native 3D vs the per-frame 2D decomposition:

case (T×H×W, C→O)	before	after	per-frame 2D
41×60×104, 256→256	71 ms	28 ms	27 ms
41×120×208, 512→512	1128 ms	335 ms	336 ms
17×64×64, 256→256	20 ms	9 ms	8 ms

i.e. 2.4–3.4× → ~1.0× (parity) with the per-frame 2D path, and numerically identical
to it.

Tests

python/tests/test_conv.py:

• test_conv_3D_small_kd_decomposition — the fast path vs CPU across 5 shapes
  (Cout != Cin, KD ∈ {1,2,3,5}, 1×1 spatial), fp32 exact.
• test_conv_3D_small_kd_fallback_cases — depth stride > 1, depth padding, and non-mod16
  channels stay correct (must take the fall-back path).

Full test_conv.py suite passes with no regressions.

Open questions for maintainers

1. Accumulation — currently the first tap owns the buffer and later taps do an in-place
   binary_op_gpu_inplace(..., "Add"). Prefer that, or a beta-accumulate added to
   conv_2D_gpu?
2. KD threshold — fixed at 7, or a cost heuristic vs the 3D implicit gemm?
3. Scope — is the decomposition the accepted fix, or do you want a native 3D Winograd
   later? Also: should the fast path be extended to N > 1 and depth padding, or is the
   N == 1 / pad[0] == 0 guard fine for a first PR?

[zcbenz]

Looks good to me, thanks!

[katlun-lgtm]

Gentle ping on this one — it's approved with all CI green. Happy to rebase onto the latest main if you'd like it current; otherwise it's ready to merge whenever you're batching. Thanks!

[zcbenz]

@katlun-lgtm For changes that I'm not 100% sure I keep them opened for a while and ask other maintainers for a second view.

[katlun-lgtm]

Checking in on this one. It's been approved since July 3, CI went green on 237ac5f (21/21), and it's still clean against current main — no conflicts, nothing outstanding on my side.

As I understand it the only thing holding it is the second opinion @zcbenz mentioned on July 7:

For changes that I'm not 100% sure I keep them opened for a while and ask other maintainers for a second view.

That's a reasonable ask for this change — it decomposes small kernel-depth 3D convs into 2D convs, so the decomposition boundaries are the part worth a second pair of eyes. It just hasn't come up since, and I'd rather it not quietly age out.

If another maintainer can give that second view, I'm glad to answer questions, add cases to the test, or rerun the benchmarks on whatever shapes you'd want to see.