[Metal] implicit-GEMM conv2d reads past the weight buffer when the output channels are not a multiple of the tile width #4630

Summary

In the Metal implicit-GEMM conv2d path, the weight loader reads a full BN-row tile without a bounds check unless BN == 8, and BN can be larger than the number of output channels. The extra rows are never stored, so results are correct, but the kernel reads device memory past the end of the weight buffer.

Where (main @ 264c14f, same in v0.25.1 and v0.32.3)

mlx/backend/metal/conv.cpp, implicit_gemm_conv_2D_gpu (l. 343-344): the tile width depends on M and C, not on O.

int bm = implicit_M >= 8192 && C_per_group >= 64 ? 64 : 32;
int bn = (bm == 64 || implicit_N >= 64) ? 64 : 32;

mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h, load_unsafe (l. 408-412): rows beyond O are read whenever BN != 8.

do_read(read_n + n_rows * TROWS <= gemm_params_->N) {}
...
if (BN != 8 || do_read) {   // unchecked copy of BN rows

loader_channel_n.h (l. 278) has the same condition.

Example: O = 32, M >= 8192, C >= 64 gives bm = 64, hence bn = 64, so each tile loads 64 weight rows from a 32-row weight.

To reproduce

# repro.py  M O
import sys
import numpy as np, mlx.core as mx

M, C, O = int(sys.argv[1]), 128, int(sys.argv[2])
rng = np.random.default_rng(0)
x = mx.array(rng.standard_normal((1, 1, M, C)).astype(np.float16))
w = mx.array(rng.standard_normal((O, 1, 7, C)).astype(np.float16))
mx.eval(mx.conv2d(x, w, padding=(0, 3)))

$ MTL_SHADER_VALIDATION=1 MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION_REPORT_TO_STDERR=1 python repro.py 9601 32
Invalid Metal usage: Invalid load in 'device' or 'constant' address space at offset 66304, executing compute function: "implicit_gemm_conv_2d_float16_bm64_bn64_bk16_wm2_wn2_channel_l_filter_s"
MTLBuffer: <no label>, length:65536, ...

Results with mlx 0.25.2 and 0.32.3 (pip), identical on both:

Impact

Silent on the Mac we tested, since the output is unaffected. We found it while investigating [METAL] Command buffer execution failed: Caused GPU Address Fault Error (0000000b:kIOGPUCommandBufferCallbackErrorPageFault) on A14 iOS devices running a conv of this shape. We have not confirmed on a device that this read is the cause of those faults.

Possible fix (not tested)

Use the checked branch whenever the tile is not fully in range, in both loaders:

-    if (BN != 8 || do_read) {
+    if (do_read) {

Workaround

Pad the weight with zero filters so that O is a multiple of 64 (or at most 16) and slice the output. With O padded from 32 to 64 the validation report disappears and the original output channels are bit-identical.

Environment

macOS 27.0, Apple M5. mlx 0.25.2 and 0.32.3 from pip. Also reproduced through mlx-swift 0.25.2.