[Metal][Performance] Add implicit matmul pathway for mx.conv3d by belkakari · Pull Request #3147 · ml-explore/mlx

belkakari · 2026-02-19T17:09:56Z

Proposed changes

The existing conv3d GPU implementation uses an explicit GEMM approach that materializes the full im2col matrix, leading to high memory usage and poor performance compared to PyTorch MPS (25-53% slower and unnecessarily memory-intensive).

This PR introduces an implicit GEMM convolution kernel for 3D (implicit_gemm_conv_3d) on the Metal backend, matching the existing 2D implicit GEMM pattern. The kernel avoids materializing the im2col matrix by computing input positions on-the-fly within the GEMM loop, significantly reducing memory overhead and improving throughput. It dispatches automatically when input channels are 16-aligned and there is no input dilation, falling back to the explicit GEMM kernel otherwise.

What changed

Added Conv3DInputBlockLoaderLargeFilter and Conv3DWeightBlockLoader in loader_channel_l.h, with branch-free bounds checking in load_unsafe() using pre-computed signed dilation factors for correct flip
Added ImplicitGemmConv3DParams struct in params.h.
Added steel_conv_3d.h kernel template and steel_conv_3d.metal instantiations for float32, float16, and bfloat16.
Added implicit_gemm_conv_3D_gpu and dispatch_conv_3D_gpu in conv.cpp with automatic routing: implicit GEMM when C_per_group % 16 == 0 and no input dilation, explicit GEMM otherwise.
Wired up JIT/no-JIT kernel registration in jit_kernels.cpp, nojit_kernels.cpp, kernels.h, and jit/includes.h.
Un-skipped forward 3D conv tests in test_conv.py that were disabled by a stale continue statement, adding forward coverage for the implicit GEMM path (C=16, O=32).
Added GPU 3D conv benchmark (benchmarks/python/conv3d_bench.py) with memory logging and both implicit/explicit GEMM shapes.
Verified performance gains on Apple M4 Pro (float32, float16)
- Implicit GEMM path: 2-3x faster than PyTorch MPS across all tested shapes
- Memory usage comparable to PyTorch for implicit GEMM, dramatically lower than explicit GEMM
- No regressions for explicit GEMM fallback path
Verified correctness: Full conv test suite passes (18 tests, 218 subtests on GPU; 43 subtests for 3D on CPU).

Benchmark script output

If C % 16 implicit GEMM is being used, explicit is used otherwise, note how implicit GEMM is 250% faster then PyTorch with a similar memory footprint and explicit GEMM is up to 50% slower and requires much more memory

========================================================================================================================
  dtype: float16
========================================================================================================================
(N,   D,   H,   W,   C)    (  O, kD, kH, kW,   C)   stride       pads         groups   diff%   MLX peak  MLX act   PT cur   PT drv
(4,  16,  16,  16,  32), ( 32,  3,  3,  3,  32), (1, 1, 1), (1, 1, 1),      1, +269.4%       2.1      2.1      2.1     19.5
(4,  16,  16,  16,  64), ( 64,  3,  3,  3,  64), (1, 1, 1), (1, 1, 1),      1, +255.3%       4.2      4.2      5.0     67.7
(4,  16,  16,  16, 128), (128,  3,  3,  3, 128), (1, 1, 1), (1, 1, 1),      1, +268.3%       8.8      8.8      9.0    104.6
(4,  32,  32,  32,  64), ( 64,  3,  3,  3,  64), (1, 1, 1), (1, 1, 1),      1, +272.5%      32.2     32.2     33.0   1240.6
(4,  32,  32,  32, 128), (128,  3,  3,  3, 128), (1, 1, 1), (1, 1, 1),      1, +238.8%      64.8     64.8     65.0   1432.6
(2,  64,  64,  64,  32), ( 64,  3,  3,  3,  32), (1, 1, 1), (1, 1, 1),      1, +209.5%      96.1     96.1     97.0   1752.7
(1,  64,  64,  64,  64), (128,  3,  3,  3,  64), (1, 1, 1), (1, 1, 1),      1, +246.8%      96.4     96.4     97.0   1881.1
(4,  32,  32,  32,  64), (128,  3,  3,  3,  64), (2, 2, 2), (1, 1, 1),      1, +285.9%      20.4     20.4     37.0   1881.1
(4,  32,  32,  32,  64), (128,  3,  1,  1,  64), (1, 1, 1), (1, 0, 0),      1, +271.5%      48.1     48.1     49.0   1881.1
(4,  32,  32,  32,  64), (128,  1,  3,  3,  64), (1, 1, 1), (0, 1, 1),      1, +273.4%      48.1     48.1     49.0   1881.2
(4,  16,  16,  16,  21), ( 21,  3,  3,  3,  21), (1, 1, 1), (1, 1, 1),      1,  -48.4%      19.1     19.1      1.7    968.4
(4,  16,  16,  16,  65), ( 65,  3,  3,  3,  65), (1, 1, 1), (1, 1, 1),      1,   -3.0%      59.1     59.1      4.7   1288.9
(4,  32,  32,  32,  65), ( 65,  3,  3,  3,  65), (1, 1, 1), (1, 1, 1),      1,  -16.6%     471.5     32.7     33.2   4620.4
(4,  16,  16,  16,   3), ( 32,  3,  3,  3,   3), (1, 1, 1), (1, 1, 1),      1,  +14.3%      21.3      3.6      1.7   3582.2

========================================================================================================================
  dtype: float32
========================================================================================================================
(N,   D,   H,   W,   C)    (  O, kD, kH, kW,   C)   stride       pads         groups   diff%   MLX peak  MLX act   PT cur   PT drv
(4,  16,  16,  16,  32), ( 32,  3,  3,  3,  32), (1, 1, 1), (1, 1, 1),      1, +192.6%       4.1      4.1      4.7   3618.2
(4,  16,  16,  16,  64), ( 64,  3,  3,  3,  64), (1, 1, 1), (1, 1, 1),      1, +235.4%       8.4      8.4      8.7   3622.2
(4,  16,  16,  16, 128), (128,  3,  3,  3, 128), (1, 1, 1), (1, 1, 1),      1, +237.3%      17.7     17.7     20.0   3671.9
(4,  32,  32,  32,  64), ( 64,  3,  3,  3,  64), (1, 1, 1), (1, 1, 1),      1, +204.0%      64.4     64.4     64.4   4671.9
(4,  32,  32,  32, 128), (128,  3,  3,  3, 128), (1, 1, 1), (1, 1, 1),      1, +213.1%     129.7    129.7    160.0   4727.9
(2,  64,  64,  64,  32), ( 64,  3,  3,  3,  32), (1, 1, 1), (1, 1, 1),      1, +212.1%     192.2    192.2    192.2   5631.9
(1,  64,  64,  64,  64), (128,  3,  3,  3,  64), (1, 1, 1), (1, 1, 1),      1, +205.8%     192.8    192.8    192.8   5631.9
(4,  32,  32,  32,  64), (128,  3,  3,  3,  64), (2, 2, 2), (1, 1, 1),      1, +223.5%      40.8     40.8     72.8   5639.9
(4,  32,  32,  32,  64), (128,  3,  1,  1,  64), (1, 1, 1), (1, 0, 0),      1, +213.4%      96.1     96.1     96.8   5639.9
(4,  32,  32,  32,  64), (128,  1,  3,  3,  64), (1, 1, 1), (0, 1, 1),      1, +213.2%      96.3     96.3     96.8   5640.2
(4,  16,  16,  16,  21), ( 21,  3,  3,  3,  21), (1, 1, 1), (1, 1, 1),      1,  -53.8%     221.9     38.1     34.2   5862.0
(4,  16,  16,  16,  65), ( 65,  3,  3,  3,  65), (1, 1, 1), (1, 1, 1),      1,  -12.6%     118.3    118.3     36.9   6439.3
(4,  32,  32,  32,  65), ( 65,  3,  3,  3,  65), (1, 1, 1), (1, 1, 1),      1,  -16.9%     942.9     65.4     65.8  11054.3
(4,  16,  16,  16,   3), ( 32,  3,  3,  3,   3), (1, 1, 1), (1, 1, 1),      1,  -10.7%       7.3      7.3      2.9  10113.1

Checklist

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)

angeloskath · 2026-02-24T01:19:31Z

All right, added the small filter loader for filters <=16^3, added a pad and slice option for unaligned channels. A small channels loader would give us even better performance but I think it is already pretty good.

I also edited the benchmark to measure sequential performance rather than parallel. The following are on an M3 Ultra. TL;DR we are now on average 2.5x faster with highs of 3.7x for larger inputs while before we were 15% slower (The diff is time_torch / time_mlx - 1)

Before this PR

========================================================================================================================
  dtype: float16
========================================================================================================================
(N,   D,   H,   W,   C)    (  O, kD, kH, kW,   C)   stride       pads         groups   diff%   MLX peak  MLX act   PT cur   PT drv
(4,  16,  16,  16,  32),   ( 32,  3,  3,  3,  32),  (1, 1, 1),   (1, 1, 1),   1,       -56.7%      29.1     29.1      2.1     39.6
(4,  16,  16,  16,  64),   ( 64,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,       -36.4%      58.4     58.4      6.0     99.9
(4,  16,  16,  16, 128),   (128,  3,  3,  3, 128),  (1, 1, 1),   (1, 1, 1),   1,        +7.3%     117.7    117.7     10.0    159.1
(4,  32,  32,  32,  64),   ( 64,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,       -35.5%     464.4     32.4     34.0   1529.9
(4,  32,  32,  32, 128),   (128,  3,  3,  3, 128),  (1, 1, 1),   (1, 1, 1),   1,        +5.7%     929.7     65.7     66.0   1963.1
(2,  64,  64,  64,  32),   ( 64,  3,  3,  3,  32),  (1, 1, 1),   (1, 1, 1),   1,       -56.7%     960.2     96.2     98.0   2025.7
(1,  64,  64,  64,  64),   (128,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,       -20.3%     960.8     96.8     98.0   2026.3
(4,  32,  32,  32,  64),   (128,  3,  3,  3,  64),  (2, 2, 2),   (1, 1, 1),   1,      +178.0%      74.8     74.8     38.0   1126.0
(4,  32,  32,  32,  64),   (128,  3,  1,  1,  64),  (1, 1, 1),   (1, 0, 0),   1,       -17.8%      96.1     96.1     50.0   1129.8
(4,  32,  32,  32,  64),   (128,  1,  3,  3,  64),  (1, 1, 1),   (0, 1, 1),   1,       -19.4%     192.3     48.3     50.0   1242.0
(4,  16,  16,  16,  21),   ( 21,  3,  3,  3,  21),  (1, 1, 1),   (1, 1, 1),   1,       -47.8%      38.1     19.1      2.7     65.5
(4,  16,  16,  16,  55),   ( 55,  3,  3,  3,  55),  (1, 1, 1),   (1, 1, 1),   1,       -29.9%      50.2     50.2      4.8    141.7
(4,  32,  32,  32,  55),   ( 55,  3,  3,  3,  55),  (1, 1, 1),   (1, 1, 1),   1,       -34.8%     399.1     27.8     28.8   1464.8
(4,  16,  16,  16,   3),   ( 32,  3,  3,  3,   3),  (1, 1, 1),   (1, 1, 1),   1,       -50.2%       3.6      3.6      2.3     72.3

========================================================================================================================
  dtype: float32
========================================================================================================================
(N,   D,   H,   W,   C)    (  O, kD, kH, kW,   C)   stride       pads         groups   diff%   MLX peak  MLX act   PT cur   PT drv
(4,  16,  16,  16,  32),   ( 32,  3,  3,  3,  32),  (1, 1, 1),   (1, 1, 1),   1,       -60.7%      58.2     58.2      5.3    101.9
(4,  16,  16,  16,  64),   ( 64,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,       -35.5%     116.8    116.8      9.3    162.5
(4,  16,  16,  16, 128),   (128,  3,  3,  3, 128),  (1, 1, 1),   (1, 1, 1),   1,        +6.2%     235.4     19.4     24.0    309.1
(4,  32,  32,  32,  64),   ( 64,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,       -36.0%     928.8     64.8     64.8   1962.5
(4,  32,  32,  32, 128),   (128,  3,  3,  3, 128),  (1, 1, 1),   (1, 1, 1),   1,        +4.3%    1859.4   1859.4    192.0   2885.1
(2,  64,  64,  64,  32),   ( 64,  3,  3,  3,  32),  (1, 1, 1),   (1, 1, 1),   1,       -57.1%    1920.4    192.4    192.4   3018.1
(1,  64,  64,  64,  64),   (128,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,       -21.2%    1921.7    193.7    193.7   3019.4
(4,  32,  32,  32,  64),   (128,  3,  3,  3,  64),  (2, 2, 2),   (1, 1, 1),   1,      +150.3%     149.7    149.7     73.7   1183.4
(4,  32,  32,  32,  64),   (128,  3,  1,  1,  64),  (1, 1, 1),   (1, 0, 0),   1,       -20.4%     192.2     96.2     97.7   1257.9
(4,  32,  32,  32,  64),   (128,  1,  3,  3,  64),  (1, 1, 1),   (0, 1, 1),   1,       -21.6%     384.6     96.6     97.7   1450.3
(4,  16,  16,  16,  21),   ( 21,  3,  3,  3,  21),  (1, 1, 1),   (1, 1, 1),   1,       -56.6%      38.2     38.2     35.0   1071.9
(4,  16,  16,  16,  55),   ( 55,  3,  3,  3,  55),  (1, 1, 1),   (1, 1, 1),   1,       -35.2%     100.3    100.3     37.1   1226.8
(4,  32,  32,  32,  55),   ( 55,  3,  3,  3,  55),  (1, 1, 1),   (1, 1, 1),   1,       -37.6%     798.1     55.6     61.2   1831.8
(4,  16,  16,  16,   3),   ( 32,  3,  3,  3,   3),  (1, 1, 1),   (1, 1, 1),   1,       -57.0%       7.3      7.3      3.7     49.1

With this PR

========================================================================================================================
  dtype: float16
========================================================================================================================
(N,   D,   H,   W,   C)    (  O, kD, kH, kW,   C)   stride       pads         groups   diff%   MLX peak  MLX act   PT cur   PT drv
(4,  16,  16,  16,  32),   ( 32,  3,  3,  3,  32),  (1, 1, 1),   (1, 1, 1),   1,      +151.9%       2.1      2.1      2.1     30.6
(4,  16,  16,  16,  64),   ( 64,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,      +159.8%      30.4      4.4      6.0     83.9
(4,  16,  16,  16, 128),   (128,  3,  3,  3, 128),  (1, 1, 1),   (1, 1, 1),   1,      +212.5%      85.7      9.7     10.0    127.1
(4,  32,  32,  32,  64),   ( 64,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,      +257.5%      48.4     32.4     34.0   1129.9
(4,  32,  32,  32, 128),   (128,  3,  3,  3, 128),  (1, 1, 1),   (1, 1, 1),   1,      +269.0%     129.7     65.7     66.0   1163.1
(2,  64,  64,  64,  32),   ( 64,  3,  3,  3,  32),  (1, 1, 1),   (1, 1, 1),   1,      +148.6%      96.2     96.2     98.0   1161.7
(1,  64,  64,  64,  64),   (128,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,      +270.8%     128.8     96.8     98.0   1162.3
(4,  32,  32,  32,  64),   (128,  3,  3,  3,  64),  (2, 2, 2),   (1, 1, 1),   1,      +112.0%      20.8     20.8     38.0   1054.7
(4,  32,  32,  32,  64),   (128,  3,  1,  1,  64),  (1, 1, 1),   (1, 0, 0),   1,      +226.0%      64.1     48.1     50.0   1097.8
(4,  32,  32,  32,  64),   (128,  1,  3,  3,  64),  (1, 1, 1),   (0, 1, 1),   1,      +252.4%      64.3     48.3     50.0   1098.0
(4,  16,  16,  16,  21),   ( 21,  3,  3,  3,  21),  (1, 1, 1),   (1, 1, 1),   1,       +83.1%      19.3      2.8      2.7     30.0
(4,  16,  16,  16,  55),   ( 55,  3,  3,  3,  55),  (1, 1, 1),   (1, 1, 1),   1,      +119.9%      40.0      6.3      4.8     83.7
(4,  32,  32,  32,  55),   ( 55,  3,  3,  3,  55),  (1, 1, 1),   (1, 1, 1),   1,      +203.4%      94.5     46.3     28.8   1160.2
(4,  16,  16,  16,   3),   ( 32,  3,  3,  3,   3),  (1, 1, 1),   (1, 1, 1),   1,      +129.8%       7.4      1.1      2.3     17.3

========================================================================================================================
  dtype: float32
========================================================================================================================
(N,   D,   H,   W,   C)    (  O, kD, kH, kW,   C)   stride       pads         groups   diff%   MLX peak  MLX act   PT cur   PT drv
(4,  16,  16,  16,  32),   ( 32,  3,  3,  3,  32),  (1, 1, 1),   (1, 1, 1),   1,      +142.4%       4.2      4.2      5.3     45.9
(4,  16,  16,  16,  64),   ( 64,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,      +147.0%      24.8      8.8      9.3     78.5
(4,  16,  16,  16, 128),   (128,  3,  3,  3, 128),  (1, 1, 1),   (1, 1, 1),   1,      +187.9%     163.4     19.4     24.0    237.1
(4,  32,  32,  32,  64),   ( 64,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,      +224.7%      64.8     64.8     64.8   1162.5
(4,  32,  32,  32, 128),   (128,  3,  3,  3, 128),  (1, 1, 1),   (1, 1, 1),   1,      +235.9%     131.4    131.4    192.0   1285.1
(2,  64,  64,  64,  32),   ( 64,  3,  3,  3,  32),  (1, 1, 1),   (1, 1, 1),   1,      +122.4%     256.4    192.4    192.4   1290.1
(1,  64,  64,  64,  64),   (128,  3,  3,  3,  64),  (1, 1, 1),   (1, 1, 1),   1,      +235.6%     193.7    193.7    193.7   1291.4
(4,  32,  32,  32,  64),   (128,  3,  3,  3,  64),  (2, 2, 2),   (1, 1, 1),   1,       +75.6%      41.7     41.7     73.7   1076.1
(4,  32,  32,  32,  64),   (128,  3,  1,  1,  64),  (1, 1, 1),   (1, 0, 0),   1,      +162.3%     128.2     96.2     97.7   1161.9
(4,  32,  32,  32,  64),   (128,  1,  3,  3,  64),  (1, 1, 1),   (0, 1, 1),   1,      +209.0%      96.6     96.6     97.7   1162.3
(4,  16,  16,  16,  21),   ( 21,  3,  3,  3,  21),  (1, 1, 1),   (1, 1, 1),   1,       +55.6%      40.4      5.5     35.0   1074.1
(4,  16,  16,  16,  55),   ( 55,  3,  3,  3,  55),  (1, 1, 1),   (1, 1, 1),   1,       +98.1%      83.9     12.5     37.1   1117.6
(4,  32,  32,  32,  55),   ( 55,  3,  3,  3,  55),  (1, 1, 1),   (1, 1, 1),   1,      +161.8%      92.5     92.5     61.2   1254.7
(4,  16,  16,  16,   3),   ( 32,  3,  3,  3,   3),  (1, 1, 1),   (1, 1, 1),   1,      +102.7%      14.9      3.3      3.7     56.7

Co-authored-by: Gleb Sterkin <g_sterkin@apple.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

[METAL] Add implicit matmul pathway for mx.conv3d

df239c6

belkakari changed the title ~~[METAL] Add implicit matmul pathway for mx.conv3d~~ [Metal][Performance] Add implicit matmul pathway for mx.conv3d Feb 19, 2026

angeloskath added 4 commits February 20, 2026 01:51

Add small filter blockloader

c4d89cd

Add pad and slice for odd channels

846ef29

Fix for non contiguous inputs

829f5c7

Make the benchmark sequential instead of parallel

52c3c69

angeloskath approved these changes Feb 24, 2026

View reviewed changes

angeloskath merged commit 1d8d693 into ml-explore:main Feb 24, 2026
16 checks passed

cocl-pm pushed a commit to cocl-pm/mlx that referenced this pull request Mar 22, 2026

[Metal] Add implicit matmul pathway for mx.conv3d (ml-explore#3147)

062acf2

Co-authored-by: Gleb Sterkin <g_sterkin@apple.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

jzdziarski pushed a commit to jzdziarski/mlx that referenced this pull request Jul 26, 2026

[Metal] Add implicit matmul pathway for mx.conv3d (ml-explore#3147)

b783997

Co-authored-by: Gleb Sterkin <g_sterkin@apple.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

Navigation Menu

Uh oh!

[Metal][Performance] Add implicit matmul pathway for mx.conv3d - #3147

[Metal][Performance] Add implicit matmul pathway for mx.conv3d#3147
angeloskath merged 5 commits into
ml-explore:mainfrom
belkakari:conv3d-implicit-gemm

belkakari commented Feb 19, 2026 •

edited

Loading

Uh oh!

angeloskath commented Feb 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

belkakari commented Feb 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Proposed changes

What changed

Benchmark script output

Checklist

Uh oh!

angeloskath commented Feb 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

belkakari commented Feb 19, 2026 •

edited

Loading