Grouped/depthwise conv reads past the weights: the implicit-GEMM weight loader bounds-checks against total output channels, not per-group #4664

In the Metal implicit-GEMM conv path the weight loader checks each tile row against the total output-channel count params->O, but the loader pointer has already been advanced into a single group. The per-group rows run out before read_n reaches params->O, so the tile keeps reading: earlier groups read into the next group's rows, and the last group reads past the end of the weight buffer. The extra rows do not affect the result (they are masked at store time), so this is silent on the machines we tested, but it is an out-of-bounds device read reachable from the public API.

This is related to #4630 and its PR #4653, but it is a different code path. #4630 is the ungrouped case where a partial output-channel tile skips the bound check via the BN != 8 shortcut. The fix in #4653 removes that shortcut so partial tiles take the checked branch. That does not help the grouped case: there the checked branch itself compares against params->O (all groups) instead of the per-group channel count, so grouped and depthwise convs still overread.

A depthwise conv1d with non-zero padding is a plain way to hit it (for example a FastConformer depthwise conv, kernel 9, padding 4): with padding it misses the separable depthwise kernel and is routed to the grouped implicit GEMM.

To reproduce

# repro1.py
import mlx.core as mx
mx.random.seed(0)
x = mx.random.normal((1, 300, 256))        # [N, L, C]
w = mx.random.normal((256, 9, 1))          # [O, K, C/groups], depthwise
y = mx.conv1d(x, w, padding=4, groups=256)
mx.eval(y)
print("done", y.shape)

$ MTL_SHADER_VALIDATION=1 MTL_SHADER_VALIDATION_REPORT_TO_STDERR=1 \
  MTL_SHADER_VALIDATION_ABORT_ON_FAULT=1 python repro1.py
done (1, 300, 256)
Invalid Metal usage: Invalid load in 'device' or 'constant' address space at offset 9324,
executing compute function:
"implicit_gemm_conv_2d_float32_bm32_bn8_bk16_wm4_wn1_channel_1_filter_l"
MTLBuffer: <no label>, length:9216, ...

The weight is 256 * 9 * 4 = 9216 bytes; the kernel loads at offset 9324, past the end.

The shipped shape that first surfaced this was conv1d with input [1, 300, 1280], weight [1280, 9, 1], groups 1280, padding 4: the invalid load is at offset 49152 of a 49152-byte buffer. With larger weights the extra rows often land in page slack and shader validation stays quiet, so a weight whose buffer ends on or near its last row (as above) shows it most reliably.

The same overread is reachable for grouped (not only depthwise) convs, and for a conv2d whose weight layout is non-standard (a transposed weight fails the separable depthwise_conv_2D_gpu condition wt_strides[1] == wS[1] and falls through to the grouped implicit GEMM).

Expected

No read outside the weight buffer. Results are already correct; only the out-of-bounds read needs removing.

Cause

The kernel advances the weight base into one group:

• B += tid.z * N * K; where N = gemm_params->N is the per-group output-channel count (

  mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h

  Line 134 in 3c40e8f

  B += tid.z * N * K;

  )

The small-channel weight loader then bounds-checks the group-relative row index read_n against the total params->O:

• zero-fill guard read_n >= params->O (

  mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h

  Line 264 in 3c40e8f

  if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) {

  )
• per-row copy guard (read_n + i) < params->O (

  mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h

  Line 293 in 3c40e8f

  if (((read_n + i) < params->O)) {

  )

The large-channel loader has the same params->O guard:

• mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h

  Line 422 in 3c40e8f

  if ((read_n + i) < params->O) {

  (and the 3D weight loader at

  mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h

  Line 925 in 3c40e8f

  if ((read_n + i) < params->O) {

  )

gemm_params->N is already the per-group count and the loader's do_read uses it, so for a group beyond the first the per-group rows run out before read_n reaches params->O, and the loader reads into the next group or past the end of the buffer.

Routing that reaches this path:

• depthwise conv1d with padding (the separable kernel needs padding == 0):

  mlx/mlx/backend/metal/conv.cpp

  Line 1559 in 3c40e8f

  if (is_idil_one && (groups == C) && groups == O && wt_strides[0] == 1 &&

• grouped implicit GEMM dispatch:

  mlx/mlx/backend/metal/conv.cpp

  Line 1438 in 3c40e8f

  return implicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params);

Suggested fix

Bound the row index against the per-group count rather than the total. gemm_params->N already holds the per-group output-channel count and equals params->O when groups == 1, so replacing params->O with gemm_params->N in the loader guards fixes the grouped case without changing the ungrouped one. This is on top of #4653 (both the BN != 8 shortcut and the params->O bound need addressing for the grouped partial-tile case).

(Reported from source inspection plus the shader-validation repro above; I have not written the patch.)

Versions

mlx 0.32.3 and mlx-metal 0.32.3 (pip)
main @ 3c40e8f92d0cce0866fc2f436f9b5096cf3b0ea7
macOS 27.0 (build 26A428), Xcode 27.0 (27A266a)
Apple M4 Max, 128 GB