Option to opt out of TF32 for fp32 conv: small-kD conv3d is ~5000× less precise on M5 since 0.32

☑️ I understand it is strictly prohibited to use AI to write issues.
What happens

On an M5 Max, fp32 mx.conv3d with batch 1, a small kernel depth and no depth padding is much less
precise on the GPU in 0.32 than in 0.31. The relative L2 error against the CPU result goes from about
1e-6 to about 6.4e-3. With batch 2, or with depth padding, the error stays at about 1e-6.

Repro

import mlx.core as mx

mx.random.seed(7)

def rel(a, b):
    d = a - b
    return (mx.sqrt((d * d).sum()) / mx.sqrt((b * b).sum())).item()

def gpu_vs_cpu(f):
    with mx.stream(mx.cpu):
        c = f()
        mx.eval(c)
    g = f()
    mx.eval(g)
    return rel(g, c)

x = mx.random.normal((1, 5, 64, 64, 128))
w = mx.random.normal((128, 3, 3, 3, 128)) * 0.03

print(mx.__version__)
print("conv3d N=1, pad (0,1,1):", gpu_vs_cpu(lambda: mx.conv3d(x, w, padding=(0, 1, 1))))
print("conv3d N=2, pad (0,1,1):", gpu_vs_cpu(lambda: mx.conv3d(mx.concatenate([x, x]), w, padding=(0, 1, 1))))
print("conv3d N=1, pad (1,1,1):", gpu_vs_cpu(lambda: mx.conv3d(x, w, padding=(1, 1, 1))))

relative L2, GPU vs CPU	0.31.2	0.32.3	0.32.3, `MLX_ENABLE_TF32=0`
conv3d N=1, pad (0,1,1)	1.31e-06	6.40e-03	4.52e-06
conv3d N=2, pad (0,1,1)	1.31e-06	1.31e-06	1.31e-06
conv3d N=1, pad (1,1,1)	1.24e-06	1.24e-06	1.24e-06

Cause

#3785 sends small kernel-depth 3D convs to per-frame 2D convs, so 3×3 kernels use Winograd.
Winograd does its product with steel_matmul. On NAX, steel_matmul runs fp32 as TF32 unless
MLX_ENABLE_TF32=0. The implicit GEMM conv path (used before #3785, and still used with batch 2
or depth padding) does not use TF32. I understand the change is intended for speed
(see ml-explore/mlx-swift#492).

Request

A way to turn off the NAX/TF32 path for one call or one scope, for example a parameter on the conv
functions (and possibly matmul / SDPA). Today the only way is MLX_ENABLE_TF32=0. It applies to the
whole process, is read once at startup, and makes other fp32 work much slower:

op (fp32, M5 Max)	TF32 on	TF32 off
matmul 4096×4096×4096	3.5 ms	9.7 ms
matmul (7800×1536)×(1536×8960)	4.7 ms	14.3 ms
scaled_dot_product_attention B1 H12 L7800 D128	18.4 ms	35.0 ms
conv3d kD=3, 384 ch, input (5, 64, 64)	3.2 ms	3.7 ms

So a model that needs exact conv results (for example a video VAE) cannot get them without slowing
down all other fp32 work in the same process (for example an fp32 DiT).

Environment

macOS 27.2, Apple M5 Max. Precision: mlx 0.31.2 / 0.32.2 / 0.32.3 from pip. Timings: same machine,
measured through mlx-swift 0.32.2 (same core).

NOTE: This draft was submitted by a human in the loop but I did have AI assistance in verification of the actual issue and the wording and contents of the message.

Adding a data point to #4595: this appears to be a different cause (Winograd rather than TF32) producing the same symptom. I'm seeing it on an M3 Pro, which has no TF32 hardware support.

Tested on MLX 0.32.3 and also reproduced on current main (f8aaf49), using an M3 Pro (18-core GPU, 36 GB) on macOS 27.0. setup is float32, 40×40×256 input → 256 output channels, 3×3 convolution, stride 1, padding 1.

I noticed that the selected convolution algorithm depends on batch size. MLX appears to switch to Winograd when batch * h * w >= 4096, so for a 40×40 input, batches 1–2 use one algorithm while batches 3+ use another.

As a result, the exact same image can produce slightly different outputs depending on what other images are present in the batch. Compared against a float64 reference, the maximum absolute error increases from roughly 2.19e-5 to 2.47e-4 (~11×) once Winograd is selected.

This surfaced in RF-DETR, which ranks proposals and keeps only the top-scoring ones. The numerical differences are small, but they can reorder near-tied proposals, causing detections from batch 3+ to differ from the same image run in batch 1. My workaround was to run convs separately for each image.

As far as i can tell, the only current way to disable Winograd is setting MLX_CONV_WINOGRAD_WORKING_SET=1, which is an undocumented side effect.

Minimal reproduction

"""
conv2d float32 error vs a float64 reference, by batch size, with Winograd on and off.

The same image is repeated across the batch, so any change in error comes from the
kernel MLX picks for that batch size, not from the data.

usage:
    python conv_batch_precision.py
    MLX_CONV_WINOGRAD_WORKING_SET=1 python conv_batch_precision.py
"""

import os
import numpy as np
import mlx.core as mx

H = W = 40
C = O = 256

rng = np.random.default_rng(0)
img = rng.normal(size=(1, H, W, C)).astype(np.float32)
wt = (rng.normal(size=(O, 3, 3, C)) * 0.05).astype(np.float32)

# float64 reference: direct 3x3 convolution with padding=1
xp = np.pad(img.astype(np.float64), ((0, 0), (1, 1), (1, 1), (0, 0)))
ref = sum(
    np.einsum(
        "nhwc,oc->nhwo",
        xp[:, i : i + H, j : j + W],
        wt[:, i, j].astype(np.float64),
    )
    for i in range(3)
    for j in range(3)
)[0]

print(
    f"mlx {mx.__version__}, "
    f"MLX_CONV_WINOGRAD_WORKING_SET="
    f"{os.environ.get('MLX_CONV_WINOGRAD_WORKING_SET', 'unset')}"
)

print(
    f"input {H}x{W}x{C} -> {O}, 3x3 stride 1, float32; "
    f"N*H*W >= 4096 from N=3"
)

first = None
for n in (1, 2, 3, 4, 8):
    y = mx.conv2d(
        mx.array(np.repeat(img, n, axis=0)),
        mx.array(wt),
        padding=1,
    )
    y = np.array(y)

    err = max(np.abs(y[k] - ref).max() for k in range(n))

    same = ""
    if first is not None:
        same = (
            "  identical to N=1"
            if np.array_equal(y[0], first)
            else "  differs from N=1"
        )

    if first is None:
        first = y[0]

    print(f"N={n}: max |err| vs float64 = {err:.2e}{same}")

Output:

$ python conv_batch_precision.py

mlx 0.32.3, MLX_CONV_WINOGRAD_WORKING_SET=unset
input 40x40x256 -> 256, 3x3 stride 1, float32; N*H*W >= 4096 from N=3

N=1: max |err| vs float64 = 2.19e-05
N=2: max |err| vs float64 = 2.19e-05  identical to N=1
N=3: max |err| vs float64 = 2.47e-04  differs from N=1
N=4: max |err| vs float64 = 2.47e-04  differs from N=1
N=8: max |err| vs float64 = 2.47e-04  differs from N=1

$ MLX_CONV_WINOGRAD_WORKING_SET=1 python conv_batch_precision.py

mlx 0.32.3, MLX_CONV_WINOGRAD_WORKING_SET=1
input 40x40x256 -> 256, 3x3 stride 1, float32; N*H*W >= 4096 from N=3

N=1: max |err| vs float64 = 2.19e-05
N=2: max |err| vs float64 = 2.19e-05  identical to N=1
N=3: max |err| vs float64 = 2.19e-05  identical to N=1
N=4: max |err| vs float64 = 2.19e-05  identical to N=1
N=8: max |err| vs float64 = 2.19e-05  identical to N=1

Would you be open to a small PR that adds a documented MLX_CONV_WINOGRAD=0 switch? It would provide an explicit way for users to opt out of the Winograd path when they need results that are independent of batch size, rather than relying on the undocumented MLX_CONV_WINOGRAD_WORKING_SET=1 workaround.

I'm good adding an option to opt out TF32 for conv only, maybe a new MLX_ENABLE_CONV_TF32 env?

Would you be open to a small PR that adds a documented MLX_CONV_WINOGRAD=0 switch?

Sounds good to me to have a MLX_CONV_WINOGRAD switch.

@zcbenz

@zcbenz . A conv-only TF32 opt-out would fix the regression I reported. On M5 the 0.32 loss comes almost entirely from TF32 inside the Winograd matmul.

With MLX_ENABLE_TF32=0 our conv results return to ~1e-6 relL2, for example GFPGAN's GPU test. A conv-only switch avoids the 2–3× fp32 matmul slowdown that disabling TF32 everywhere costs.

MLX_CONV_WINOGRAD=0 complements it rather than replacing it. @ishaanzee batch-dependence persists without TF32, and both switches together give fully exact convs.

One note is that both proposals are process-wide environment variables. As far as I know MLX reads them once, like MLX_ENABLE_TF32. A library can't reliably set them for its own parity-critical convs, because the app may already have run MLX work, and Swift apps can't set them per call at all. Maybe a per-call or scoped option would serve libraries rather than workarounds like the zero-channel trick in mlx-exact-conv-swift.

Thanks for looking into this.

One note is that both proposals are process-wide environment variables. As far as I know MLX reads them once, like MLX_ENABLE_TF32. A library can't reliably set them for its own parity-critical convs, because the app may already have run MLX work, and Swift apps can't set them per call at all. Maybe a per-call or scoped option would serve libraries rather than workarounds like the zero-channel trick in mlx-exact-conv-swift.

With recent releases setting the env before invoking any MLX ops in Swift should work. But in the meanwhile I agree we should have a better way to dynamic set it, I will write a proposal later for generalized API that sets flags.

With #4647 it will be possible to turn off TF32 for partial graphs.

Navigation Menu

Repository navigation

Option to opt out of TF32 for fp32 conv: small-kD conv3d is ~5000× less precise on M5 since 0.32 #4595

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Option to opt out of TF32 for fp32 conv: small-kD conv3d is ~5000× less precise on M5 since 0.32 #4595

Description

Activity

ishaanzee commented on Oct 5, 2026

zcbenz commented on Oct 7, 2026

xocialize commented on Oct 7, 2026

zcbenz commented on Oct 8, 2026

zcbenz commented on Oct 9, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions