Gate solve_tril dot precision on whether TF32 is faster, not whether it exists

fla/ops/gated_delta_rule/chunk_fwd.py

@@ -15,9 +15,9 @@
 from fla.ops.utils import prepare_chunk_indices, solve_tril
 from fla.ops.utils.cache import fla_cache_autotune
 from fla.ops.utils.op import exp2
-from fla.utils import IS_INTEL, IS_TF32_SUPPORTED, autotune_cache_kwargs
+from fla.utils import IS_INTEL, IS_TF32_FASTER_THAN_FP32, autotune_cache_kwargs
 
-if IS_TF32_SUPPORTED:
+if IS_TF32_FASTER_THAN_FP32:
     SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32')
 else:
     SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee')

fla/ops/gdn2/chunk_intra.py

@@ -29,9 +29,9 @@
 from fla.ops.utils import prepare_chunk_indices
 from fla.ops.utils.cache import fla_cache_autotune
 from fla.ops.utils.op import exp2, gather
-from fla.utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs
+from fla.utils import IS_GATHER_SUPPORTED, IS_TF32_FASTER_THAN_FP32, autotune_cache_kwargs
 
-if IS_TF32_SUPPORTED:
+if IS_TF32_FASTER_THAN_FP32:
     SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32')
 else:
     SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee')

fla/ops/kda/chunk_intra.py

@@ -15,9 +15,9 @@
 from fla.ops.utils import prepare_chunk_indices
 from fla.ops.utils.cache import fla_cache_autotune
 from fla.ops.utils.op import exp2, gather
-from fla.utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs
+from fla.utils import IS_GATHER_SUPPORTED, IS_TF32_FASTER_THAN_FP32, autotune_cache_kwargs
 
-if IS_TF32_SUPPORTED:
+if IS_TF32_FASTER_THAN_FP32:
     SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32')
 else:
     SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee')

fla/utils/__init__.py

@@ -44,8 +44,10 @@
     IS_NVIDIA_HOPPER,
     IS_NVIDIA_SM100,
     IS_NVIDIA_SM120,
+    IS_TF32_FASTER_THAN_FP32,
     IS_TF32_SUPPORTED,
     IS_TMA_SUPPORTED,
+    TF32_NOT_FASTER_CAPABILITIES,
     Backend,
     autocast_custom_bwd,
     autocast_custom_fwd,

fla/utils/_device.py

@@ -150,6 +150,22 @@ def map_triton_backend_to_torch_device() -> str:
 
 # Nvidia Ampere or newer, haven't check AMD and intel yet.
 IS_TF32_SUPPORTED = (IS_NVIDIA and torch.cuda.get_device_capability(0)[0] >= 8)
+
+# Whether TF32 is AVAILABLE (above) is a different question from whether it is
+# FASTER. On the datacenter dies it is, by a lot: A100 runs fp32 at 19.5 and
+# tf32 at 156 TFLOPS. On the GeForce/Ada dies it is not — an RTX 4090 (sm_89)
+# runs fp32 and tf32 alike at 82.6 TFLOPS, as do the other 8.9 parts (L40S,
+# RTX 6000 Ada) at their own clocks. There, choosing tf32 for a dot whose
+# operands are already fp32 spends 13 mantissa bits and buys no throughput.
+#
+# Only capabilities we can source are listed, so no unlisted part changes
+# behaviour. (RTX 30xx at 8.6 and Blackwell GeForce at 12.0 are believed to be
+# in the same position; left out until someone can confirm them.)
+TF32_NOT_FASTER_CAPABILITIES = frozenset({(8, 9)})
+IS_TF32_FASTER_THAN_FP32 = (
+    IS_TF32_SUPPORTED
+    and torch.cuda.get_device_capability(0) not in TF32_NOT_FASTER_CAPABILITIES
+)
 IS_GATHER_SUPPORTED = hasattr(triton.language, 'gather')
 IS_TMA_SUPPORTED = (
     IS_NVIDIA