Warn once when float32 ops silently run at TF32 precision

[pierre427]

Addresses the diagnostic half of #3860.

float32 matmul-family ops default to TF32-class reduced precision (MLX_ENABLE_TF32 defaults to 1) on both backends on current hardware — CUDA tensor cores and the M5 Metal NAX route — and the default is undocumented. #3860 has two independent downstream reports of the resulting failure mode: op-level tests stay green while downstream results shift (near-tie argmax flips on sm_120; mlx-lm's test_ssm state comparison on M5), and the bisection points at hardware/environment rather than any commit. This repo's own test harness already pins MLX_ENABLE_TF32=0 (python/tests/mlx_tests.py) for the same reason.

This adds env::tf32_active_for_fp32(): returns exactly env::enable_tf32(), but prints a one-line warning to stderr, once per process, and only when MLX_ENABLE_TF32 is unset — anyone who set =1 or =0 explicitly has already made the choice and never sees it:

[mlx] float32 matmul-family ops are running at reduced (TF32) precision, which is the default when MLX_ENABLE_TF32 is unset. Set MLX_ENABLE_TF32=0 before the first operation for full float32 precision, or set MLX_ENABLE_TF32=1 explicitly to silence this warning.

It is called at exactly the points where a float32 op takes the reduced-precision route, by reordering the existing short-circuit (env::enable_tf32() || dtype != float32) → (dtype != float32 || env::tf32_active_for_fp32()), so processes that never engage TF32 (matvec-only, non-fp32, opted out) never print:

• Metal: the NAX gates in steel matmul, gather_mm_rhs, segmented_mm, qmm, gather_qmm, gather_qmm_rhs, and SDPA full attention
• CUDA: cuBLAS compute type (float32 case), cuDNN conv graph, grouped CUTLASS GEMM

No behavior change beyond the log — every gate computes the same boolean as before. The complex64 cuBLAS case (also FAST_TF32 today) is intentionally left untouched; probably worth folding into the docs follow-up discussed in #3860.

Testing

On M5 Max (macOS 26.5), source build with NAX kernels active:

scenario	warnings
default env, fp32 GEMM ×2 in one process	exactly 1
default env, fp32 matvec only	0
default env, bf16 GEMM	0
MLX_ENABLE_TF32=0, fp32 GEMM	0
explicit MLX_ENABLE_TF32=1, fp32 GEMM	0
default env, fp32 quantized_matmul (NAX qmm route)	exactly 1
default env, fp32 quantized_matmul (split-K route, not reduced)	0

Numerics unchanged: default fp32 GEMM max rel err vs an fp64 reference is 8.3e-4 (TF32 signature), and 2.5e-7 with MLX_ENABLE_TF32=0. test_blas.py (24 tests), test_quantized.py (32), and test_fast_sdpa.py (16) all pass.

The CUDA side compiles the same helper but I have no NVIDIA hardware to run it — @stoyoda0012-cyber offered an sm_120 review in #3860.

🤖 Generated with Claude Code

[stoyoda0012-cyber]

Running the sm_120 pass this week — source build on WSL2, results with numbers to follow.

One design question I'd like to resolve before merge, since it's the item I most want hardware to answer: the Metal gates sit inside the NAX dispatch (shape-aware), so matvec never reaches the helper — consistent with your 0-warning row. The CUDA main gate is dtype_to_compute_type(), which selects the compute type without consulting shape. If sm_120 gemv still passes through it, a matvec-only fp32 workload would warn while staying exact fp32 — a false positive, and the one case where "fires on actual engagement" wouldn't hold on the CUDA half. My scenario 2 is exactly this; I'll report either way, and if it does fire I'll propose a shape guard.

Same question likely applies to the cuDNN conv path: if an fp32-conv-only workload warns, the message text ("matmul-family ops") won't help the reader locate it.

[pierre427]

Both points are well taken.

On the gemv/matvec asymmetry: you've read the structure exactly right — the Metal gates live inside the NAX dispatch, which is shape-aware, so matvec (and any shape that routes to a non-NAX kernel) never reaches the helper; that's where the 0-warning gemv row comes from. The CUDA gate at dtype_to_compute_type() is dtype-only by construction. If your scenario-2 run shows a matvec-only fp32 workload warning while staying exact, that's a real false positive on the CUDA half and I'll take your shape guard (or hoist the helper call below a shape check at the cublas site — whichever you prefer, since you have the hardware to validate it). Your measurement should settle it; no need to guess ahead of the sm_120 numbers.

On the message text: agreed, and fixed now rather than waiting — d3d6c38a threads the op family from the call site into the helper (default "matmul"; the conv site passes "convolution", SDPA "attention", quantized "quantized matmul", grouped CUTLASS "grouped matmul"), so the once-per-process message names whichever family actually engaged first. No change to when the warning fires. M5-verified: GEMM-first warns "float32 matmul ops", NAX-shape quantized-matmul-first warns "float32 quantized matmul ops", gemv and non-NAX qmm shapes stay silent, explicit =0/=1 stay silent; test_blas / test_quantized / test_fast_sdpa green. The CUDA sites are compile-only from me as before — worth an eye during your pass.

[stoyoda0012-cyber]

Ran the CUDA half on real sm_120 hardware (RTX 5070 Laptop, WSL2
Ubuntu 24.04, CUDA 13.3, branch @ d3d6c38): warning behaves as
designed for GEMM (fires once per process, silent for bf16 and for
either explicit setting), the cuDNN conv site fires with family
"convolution", and numerics match the measurements in #3860
(2.9e-04 default / 2.1e-07 with MLX_ENABLE_TF32=0). Full 774-test
parity suite: exactly 1 warning, failure set identical to the
0.32.0 wheel — no behavioral change.

One asymmetry vs the M5 measurement: fp32 matvec fires the warning
on CUDA (your Metal run had 0), while the matvec result itself
stays fp32-exact (rel err 9.7e-08 vs fp64, same process). The CUDA
gate in dtype_to_compute_type() is shape-blind, so matvec-only
workloads get warned about a reduction that doesn't happen for that
shape. Suggest either shape-gating that call site or documenting the
requested-mode semantics.

Minor observation for the discussion: under pytest's default capture
the once-per-process warning is invisible (swallowed with whichever
test ran the first fp32 GEMM); --capture=no shows it. May be worth
a line in the docs.

[pierre427]

Confirmed — thanks for the hardware run. The 9.7e-08 vs 2.9e-04 split is consistent with cuBLASLt not selecting a TF32 tensor-core kernel for that rank-1 shape while the compute descriptor still requests FAST_TF32; since the CUDA gate at dtype_to_compute_type() is dtype-only, it warns on the request regardless of the kernel cuBLAS actually runs. That's the false positive.

(The Metal half stays silent on matvec for a different reason than I first implied: vector shapes are routed to a dedicated gemv kernel before steel_matmul_axpby, where the fp32/TF32 gate lives — so they never reach the warn, rather than being filtered by a shape-aware use_nax.)

Fix pushed (57466a7): dtype_to_compute_type() goes back to silent enable_tf32() (pure compute-type selection, unchanged), and the warning moves into the CublasGemm constructor, which already holds the output dims (M_ = a_rows, N_ = b_cols, set by init_base), gated on M_ > 1 && N_ > 1. I kept it at the constructor rather than a dispatch site on purpose: AddMM has a second direct CublasGemm construction that bypasses the shared gemm_and_bias helper, and the GEMM-conv fallbacks build CublasGemm directly too — the constructor is the one point every cuBLAS fp32 entry funnels through, so warning coverage matches the old dtype-only gate exactly.

To be precise about what the guard buys: M_ > 1 && N_ > 1 suppresses the rank-1 shapes your run showed staying fp32-exact — it is not a claim to detect actual tensor-core engagement, since cuBLASLt picks the algorithm by runtime heuristic. The warning stays deliberately conservative: it fires when reduced precision is requested on a genuine matrix-matrix shape.

CUDA is still compile-only from me, so would you mind re-checking on the 5070: (1) your scenario-2 matvec now warns 0×, (2) a GEMM-first workload still warns exactly once (family "matmul"), (3) an AddMM / beta != 0 fp32 workload still warns, and (4) the 774-test parity set is unchanged. I'll also add your --capture=no observation to the PR notes — good catch.

[stoyoda0012-cyber]

Re-ran the matrix on real sm_120 hardware against 57466a7
(RTX 5070 Laptop, WSL2, CUDA 13.3, incremental rebuild):

scenario	result
fp32 matvec → in-process control GEMM	(0, 1) — matvec silent, control still warns
fp32 GEMM ×2	1, family "matmul"
fp32 addmm	1, family "matmul"
774-test parity suite (env unset / MLX_ENABLE_TF32=0)	1 warning, failures identical to wheel / 0 warnings, failures identical to wheel
fp32 conv2d	1, family "convolution"

The shape gate resolves the matvec false positive from my earlier
comment — matvec-only workloads are now silent while the in-process
control GEMM proves the machinery still engages, and the addmm
constructor path warns as intended. No behavioral change in either
suite configuration. LGTM from the hardware side.