Fix conv2 gradients in grouped strided case on Metal

[ericphanson]

Route grouped 2D convolutions with input dilation through the group-aware explicit GEMM fallback instead of falling through to dense Metal kernels. This fixes incorrect GPU dL/dx for grouped strided conv2d VJPs.

Proposed changes

I had unexpected exploding gradients when implementing a mobileclip2 model with MLX, which were fine on cpu. Claude Fable produced the following MWE which shows the divergence between gradients of conv2d with cpu vs mlx:

mwe.py

#!/usr/bin/env python3
"""Minimal MLX-vs-PyTorch conv-vjp comparison.

For every conv configuration used by the FastViT-MCI0 stem/downsample path,
compute y = conv(x, w), loss = sum(y * cot) with a fixed random cotangent, and
compare dL/dx and dL/dw against a float64 PyTorch CPU reference, on both the
MLX GPU (Metal) and CPU streams.

Result (mlx==0.31.2, M5 Max): dL/dx from the Metal kernel is numerically wrong
(relative error ~5-56, i.e. 500%-5600%) exactly when groups > 1 AND stride > 1.
The same configs are correct (~1e-8) on the MLX CPU stream, and all other
configs (stride-1 grouped, strided dense) are correct on both streams. These
buggy configs are the stem's second block (3x3 s2 depthwise + 1x1 s2 grouped)
and every stage downsample (7x7 s2 depthwise + 3x3 s2 depthwise), which is why
the backward signal amplified multiplicatively at those exact depths in the
full-model gradient audits (Q1/Q2) until the stem weights were destroyed and
the forward pass went NaN at `backbone.stem`.

The workaround (`models_mlx.StridedGroupedConv2d`) computes the conv at
stride 1 and subsamples the output, which is mathematically identical and only
exercises vjp paths verified correct here.
"""

from __future__ import annotations

import mlx.core as mx
import numpy as np
import torch
import torch.nn.functional as F

CONFIGS = [
    # (name, in_chs, out_chs, kernel, stride, groups, input_hw)
    ("stem0_kxk   3->64  k3 s2 g1", 3, 64, 3, 2, 1, 32),
    ("stem0_scale 3->64  k1 s2 g1", 3, 64, 1, 2, 1, 32),
    ("stem1_kxk   64dw   k3 s2 g64", 64, 64, 3, 2, 64, 32),
    ("stem1_scale 64dw   k1 s2 g64", 64, 64, 1, 2, 64, 32),
    ("ds_large    64dw   k7 s2 g64", 64, 64, 7, 2, 64, 32),
    ("ds_small    64dw   k3 s2 g64", 64, 64, 3, 2, 64, 32),
    ("ctrl        64dw   k3 s1 g64", 64, 64, 3, 1, 64, 32),
    ("ctrl        64->64 k3 s1 g1", 64, 64, 3, 1, 1, 32),
    ("ctrl        64->64 k3 s2 g1", 64, 64, 3, 2, 1, 32),
]

BATCH = 8
BUGGY_DX_REL_ERR_THRESHOLD = 1e-3


def relative_error(a: np.ndarray, b: np.ndarray) -> float:
    a64 = np.asarray(a, dtype=np.float64)
    b64 = np.asarray(b, dtype=np.float64)
    return float(np.linalg.norm(a64 - b64) / max(np.linalg.norm(b64), 1e-30))


def main() -> None:
    rng = np.random.default_rng(0)
    print(f"mlx {mx.__version__}  torch {torch.__version__}")
    header = (
        f"{'config':30s} {'dx gpu':>12s} {'dx cpu':>12s} {'dw gpu':>12s} {'dw cpu':>12s}"
    )
    print(header)

    gpu_dx_errors: dict[str, float] = {}
    for name, cin, cout, k, s, g, hw in CONFIGS:
        pad = k // 2
        x_np = rng.standard_normal((BATCH, cin, hw, hw)).astype(np.float32)
        w_np = (rng.standard_normal((cout, cin // g, k, k)) * 0.1).astype(np.float32)
        h_out = (hw + 2 * pad - k) // s + 1
        cot_np = rng.standard_normal((BATCH, cout, h_out, h_out)).astype(np.float32)

        x_t = torch.tensor(x_np, dtype=torch.float64, requires_grad=True)
        w_t = torch.tensor(w_np, dtype=torch.float64, requires_grad=True)
        y_t = F.conv2d(x_t, w_t, stride=s, padding=pad, groups=g)
        (y_t * torch.tensor(cot_np, dtype=torch.float64)).sum().backward()
        dx_ref = x_t.grad.numpy()
        dw_ref = w_t.grad.numpy()

        x_mx = mx.array(np.transpose(x_np, (0, 2, 3, 1)))
        w_mx = mx.array(np.transpose(w_np, (0, 2, 3, 1)))
        cot_mx = mx.array(np.transpose(cot_np, (0, 2, 3, 1)))

        def loss_fn(x: mx.array, w: mx.array) -> mx.array:
            return mx.sum(mx.conv2d(x, w, stride=s, padding=pad, groups=g) * cot_mx)

        errors: dict[str, tuple[float, float]] = {}
        for dev_name, dev in (("gpu", mx.gpu), ("cpu", mx.cpu)):
            mx.set_default_device(dev)
            dx_mx, dw_mx = mx.grad(loss_fn, argnums=(0, 1))(x_mx, w_mx)
            mx.eval(dx_mx, dw_mx)
            dx = np.transpose(np.array(dx_mx), (0, 3, 1, 2))
            dw = np.transpose(np.array(dw_mx), (0, 3, 1, 2))
            errors[dev_name] = (relative_error(dx, dx_ref), relative_error(dw, dw_ref))
        mx.set_default_device(mx.gpu)
        gpu_dx_errors[name] = errors["gpu"][0]

        print(
            f"{name:30s} {errors['gpu'][0]:12.3e} {errors['cpu'][0]:12.3e} "
            f"{errors['gpu'][1]:12.3e} {errors['cpu'][1]:12.3e}"
        )

    buggy = [name for name, err in gpu_dx_errors.items() if err > BUGGY_DX_REL_ERR_THRESHOLD]
    print()
    if buggy:
        print("BUG PRESENT: GPU-stream dL/dx wrong for configs (groups>1 and stride>1):")
        for name in buggy:
            print(f"  - {name}: rel_err={gpu_dx_errors[name]:.3e}")
    else:
        print("Bug not present in this MLX version: all GPU-stream conv vjps match reference.")


if __name__ == "__main__":
    main()

I used Codex (GPT 5.5) with that MWE to find and fix the dispatch issue PR'd here. It seems there is a bug in the fast path: when input_dilation != 1 it falls back to an ungrouped convolution instead. Here we check this and route to explicit_gemm_conv_group_ND_gpu instead.

I verified that my original mwe (see above) works with this PR, and the new unit tests work on my PR branch and fail on main.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)