Skip to content

Keep the winograd gemm at full float32 precision #4242

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
erwinzhang7 wants to merge 1 commit into
base: main
Choose a base branch
from
+50 −6
Open

Keep the winograd gemm at full float32 precision #4242

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion mlx/backend/metal/conv.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -948,7 +948,13 @@ void winograd_conv_2D_gpu(
/*b_cols = */ conv_params.O,
/*a_transposed = */ false,
/*b_transposed = */ false,
/*copies = */ empty_copies);
/*copies = */ empty_copies,
/*batch_shape = */ {},
/*A_batch_stride = */ {},
/*B_batch_stride = */ {},
// The winograd transforms amplify whatever the gemm rounds away, so
// this one stays at full float32 even when tf32 is allowed elsewhere.
/*allow_tf32 = */ false);
}

// Do output transform
Expand Down
5 changes: 3 additions & 2 deletions mlx/backend/metal/matmul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -862,7 +862,8 @@ void steel_matmul_axpby(
Strides B_batch_stride /* = {} */,
Strides C_batch_stride /* = {} */,
float alpha /* = 1.0f */,
float beta /* = 0.0f */) {
float beta /* = 0.0f */,
bool allow_tf32 /* = true */) {
if (batch_shape.empty()) {
/////////////////////////////////////////////////////////////////////////////
// Check and collapse batch dimensions
Expand Down Expand Up @@ -916,7 +917,7 @@ void steel_matmul_axpby(
int64_t matrix_size = static_cast<int64_t>(M) * N;
bool use_nax = metal::is_nax_available() &&
!issubdtype(a.dtype(), complexfloating) &&
(env::enable_tf32() || a.dtype() != float32);
((allow_tf32 && env::enable_tf32()) || a.dtype() != float32);
char devc = d.get_architecture().back();
int min_tmn_threshold = (devc == 's' || devc == 'd') ? 2048 : 1024;

Expand Down
12 changes: 9 additions & 3 deletions mlx/backend/metal/matmul.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,8 @@ void steel_matmul_axpby(
Strides B_batch_stride = {},
Strides C_batch_stride = {},
float alpha = 1.0f,
float beta = 0.0f);
float beta = 0.0f,
bool allow_tf32 = true);

inline void steel_matmul(
const Stream& s,
Expand All @@ -119,7 +120,8 @@ inline void steel_matmul(
std::vector<array>& copies,
Shape batch_shape = {},
Strides A_batch_stride = {},
Strides B_batch_stride = {}) {
Strides B_batch_stride = {},
bool allow_tf32 = true) {
return steel_matmul_axpby<false>(
/* const Stream& s = */ s,
/* metal::Device& d = */ d,
Expand All @@ -138,7 +140,11 @@ inline void steel_matmul(
/* std::vector<array>& copies = */ copies,
/* Shape batch_shape = */ batch_shape,
/* Strides A_batch_stride = */ A_batch_stride,
/* Strides B_batch_stride = */ B_batch_stride);
/* Strides B_batch_stride = */ B_batch_stride,
/* Strides C_batch_stride = */ {},
/* float alpha = */ 1.0f,
/* float beta = */ 0.0f,
/* bool allow_tf32 = */ allow_tf32);
}

} // namespace mlx::core
31 changes: 31 additions & 0 deletions python/tests/test_conv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@

import math
import os
import subprocess
import sys
import textwrap
import unittest
from itertools import permutations

Expand Down Expand Up @@ -48,6 +51,34 @@ def test_numpy_conv(self):
self.assertEqual(c_mx.shape, c_np.shape)
self.assertTrue(np.allclose(c_mx, c_np, atol=atol))

def test_conv_2d_winograd_float32_precision(self):
# The aligned conv takes winograd and the unaligned one does not, so

@JasonHonKL JasonHonKL Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

curious why need to use sub process for this test ?

@erwinzhang7 erwinzhang7 Aug 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mlx_tests.py sets MLX_ENABLE_TF32=0 at import, before mlx.core is imported, for the whole
suite:

# Use regular fp32 precision for tests
os.environ["MLX_ENABLE_TF32"] = "0"

and enable_tf32() reads the environment once into a static:

inline bool enable_tf32() {
  static bool enable_tf32_ = get_var("MLX_ENABLE_TF32", 1);
  return enable_tf32_;
}

So by the time any test runs, tf32 is off and the value is fixed for the life of the process.
Setting os.environ inside the test does nothing, and this change is a no-op with tf32 off, so an
in-process test would vacuously pass.

The subprocess runs the conv with MLX_ENABLE_TF32=1, which is the default a user gets. It
reports 2.38 before the change and 0.0023 after.

The same two facts mean the suite currently has no coverage of the default
tf32 configuration at all, which is why a 4859x accuracy difference on float32 conv went
unnoticed.

# the two have to agree to float32 precision.
script = textwrap.dedent("""
import mlx.core as mx

# Winograd also wants N * iH * iW >= 4096 and C + O >= 256.
N, S, K, C = 4, 32, 3, 128
x = mx.random.normal((N, S, S, C))
w = mx.random.normal((C, K, K, C))
aligned = mx.conv2d(x, w, padding=1)

# The same conv with an input channel count winograd declines.
xu = mx.pad(x, [(0, 0), (0, 0), (0, 0), (0, 16)])
wu = mx.pad(w, [(0, 0), (0, 0), (0, 0), (0, 16)])
unaligned = mx.conv2d(xu, wu, padding=1)

mx.eval(aligned, unaligned)
print(mx.abs(aligned - unaligned).max().item())
""")
# tf32 is off for the whole suite and enable_tf32() caches the read.
env = dict(os.environ, MLX_ENABLE_TF32="1")
out = subprocess.check_output(
[sys.executable, "-c", script], env=env, text=True
)
gap = float(out.strip())
self.assertLess(gap, 1e-2, f"winograd and gemm disagree by {gap}")

def test_conv_1d_groups_flipped(self):
x = mx.broadcast_to(mx.arange(5).astype(mx.float32), (2, 5)).T
w = mx.broadcast_to(mx.arange(4).astype(mx.float32), (2, 4))
Expand Down