From fabd930c39bdda3aa661590cbcab285bc82209aa Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 3 Jul 2026 09:39:43 +0000
Subject: [PATCH] Clear TF32 flags when enable_tf32 is off so the opt-out is
 strict fp32

---
 .../core/training/diffusion_train_common.py   |  9 ++++++-
 .../tests/test_diffusion_train_perf.py        | 27 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 5b035cd57e..234bbdd0c4 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -406,7 +406,8 @@ def _apply_perf_flags(
     cudnn_benchmark: bool = False,
 ) -> dict:
     """Set the run-scoped torch backend knobs: TF32 matmuls + high fp32 matmul precision
-    (under ``cfg.enable_tf32``), plus cudnn autotuning when the caller opts in. Autotune is
+    when ``cfg.enable_tf32`` is on, strict fp32 (all TF32 flags cleared) when it is off,
+    plus cudnn autotuning when the caller opts in. Autotune is
     for the conv-heavy SDXL U-Net only: measured on B200, it DOUBLES peak VRAM (fp32 VAE
     conv workspaces) while the DiT loop -- pure matmuls once the latent cache is built --
     gains nothing from it. Returns a snapshot for ``_restore_perf_flags``. Best-effort:
@@ -424,6 +425,12 @@ def _apply_perf_flags(
             torch.backends.cuda.matmul.allow_tf32 = True
             torch.backends.cudnn.allow_tf32 = True
             torch.set_float32_matmul_precision("high")
+        else:
+            # The opt-out is a strict-fp32 A/B mode, so actively clear the flags rather
+            # than inherit ambient state (cudnn TF32 defaults to ON in torch).
+            torch.backends.cuda.matmul.allow_tf32 = False
+            torch.backends.cudnn.allow_tf32 = False
+            torch.set_float32_matmul_precision("highest")
         if cudnn_benchmark:
             torch.backends.cudnn.benchmark = True
     except Exception:  # noqa: BLE001 -- perf flags are never fatal
diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py
index ecfe966b5f..4339fbb0e7 100644
--- a/studio/backend/tests/test_diffusion_train_perf.py
+++ b/studio/backend/tests/test_diffusion_train_perf.py
@@ -328,3 +328,30 @@ def test_perf_flags_cpu_roundtrip():
     snap = _apply_perf_flags(_cfg(), "cpu")
     assert isinstance(snap, dict)
     _restore_perf_flags(snap)  # no exception
+
+
+def test_perf_flags_tf32_off_clears_flags():
+    # enable_tf32=False is the strict-fp32 A/B mode: it must actively clear the TF32 flags
+    # (cudnn TF32 defaults ON in torch) rather than inherit ambient state, and restore must
+    # put the ambient values back. The flag attributes are plain Python state, present and
+    # settable on CPU-only torch builds, so this runs without a GPU.
+    import torch
+
+    before = (
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.backends.cudnn.allow_tf32,
+        torch.get_float32_matmul_precision(),
+    )
+    snap = _apply_perf_flags(_cfg(enable_tf32 = False), "cuda")
+    try:
+        assert torch.backends.cuda.matmul.allow_tf32 is False
+        assert torch.backends.cudnn.allow_tf32 is False
+        assert torch.get_float32_matmul_precision() == "highest"
+    finally:
+        _restore_perf_flags(snap)
+    after = (
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.backends.cudnn.allow_tf32,
+        torch.get_float32_matmul_precision(),
+    )
+    assert after == before