[MarcosAsh]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure:
  I used AI to help me find the issue and with general questions. All code was written by me.

Padding applies to the dilated input, but conv_general cropped negative padding before dilating, so each unit removed a whole input element. Convolution::vjp cropped the result afterwards, which gave wrong gradient shapes.

x = mx.array([1.0, 2.0, 3.0]).reshape(1, 3, 1)
mx.conv_general(x, mx.ones((1, 1, 1)), padding=([-1], [0]), input_dilation=2).flatten()

# before: [2, 0, 3]    after: [0, 2, 0, 3]

conv_general now slices to the elements that survive and adds the leftover gap back as positive padding. If nothing survives it returns zeros with the padded output shape. The separate crop in vjp is removed.

983/1728 gradient cases, 266/370 forward cases and 45/106 conv_transpose1d cases were wrong before and are correct now. New tests fail on main and pass with the fix. Follows up #4662.

The second commit fixes the CUDA side. cuDNN’s forward conv can’t express input dilation, so it rejects these cases, and mlx then fell back to a backward weight conv. That path drops input_dilation in get_conv_settings, so cuDNN correctly computed a different convolution from the one requested, even with no negative padding (x = [1, 2, 3], input_dilation=2, stride=2, padding=(2, 1) gave [0, 1, 3, 0] instead of [0, 1, 2, 3]). The backward data path looks like the same kind of mismatch, since it uses input_dilation as its stride. Those cases now use the gemm_conv fallback, and input_dilation is part of the conv cache key.

I kept the dgrad sentence hedged (“looks like”) since you haven’t traced those cases. If you trace n=1, k=2, stride 2, padding (1,2) and it confirms, you can state it outright and mention that it drops kernel_strides.

Tested on CPU and on an RTX 5070 Ti (CUDA 12.9). The conv suites pass on both, and CPU and GPU agree on a 1165 case sweep except for 5 cases from conv_out_axis_size truncating toward zero, which is a separate issue. Metal wasn't rebuilt.

[zcbenz]

cuDNN's forward conv can't express input dilation and its backward data conv can't express a stride, but both were picked anyway and computed the wrong thing, even with no negative padding (x = [1, 2, 3], input_dilation=2, stride=2, padding=(2, 1) gave [0, 1, 3, 0] instead of [0, 1, 2, 3]).

This feels like a bug in cudnn, they should throw error instead of giving wrong result, can you update cudnn (in your system env) and cudnn-frontend (in mlx/backend/cuda/CMakeLists.txt) to latest, and see if it still silently gives wrong results? And if it does, can you please report a bug to cudnn-frontend?

[MarcosAsh]

Hi, thanks for the review. I retested with cuDNN 9.27.0.42, cudnn-frontend v1.31.0 and mlx main (243bdf0) on the 5070 Ti, and the example still gives [0, 1, 3, 0] on GPU vs [0, 1, 2, 3] on CPU. 835 of 1689 cases were wrong in the non-negative padding sweep. cuDNN correctly rejects the forward conv with CUDNN_STATUS_BAD_PARAM (shape mismatch in is_valid_spatial_dim), same on 9.10 and 9.27. mlx then falls back to CONV_BACKWARD_WEIGHT and get_conv_settings drops input_dilation, so cuDNN correctly computes the wgrad it was given. x padded to [0,0,1,2,3,0,0] and sampled with dilation 2 gives [0,1,3,0]. The flipped path looks like the same pattern (dgrad uses input_dilation as its stride and drops kernel_strides).

So nothing to report to cudnn-frontend. The bug is in mlx’s mapping, which the second commit fixes. I’ve also fixed the sentence in the PR description that blamed cuDNN.