[Inductor][Triton] Support TMA before strict 3.4 cutoff

njriasan · facebook-github-bot · commit 011cfd4703a4 · 2025-08-04T10:36:26.000-07:00
Summary: Inductor's 3.4 Triton release is the most common used variant of Triton, but if someone is working with an alternative version of Triton this may not match. This moves the version check from 3.4 Triton to any variant that has support for the TMA APIs.

Test Plan:
Relying on CI. Should be a NFC.

Rollback Plan:

Reviewed By: davidberard98

Differential Revision: D79378792
diff --git a/torch/_inductor/codegen/triton.py b/torch/_inductor/codegen/triton.py
@@ -26,7 +26,7 @@
 from torch._prims_common import is_integer_dtype
 from torch.utils._ordered_set import OrderedSet
 from torch.utils._sympy.functions import CeilDiv, FloorDiv, ModularIndexing
-from torch.utils._triton import has_triton_package
+from torch.utils._triton import has_triton_package, has_triton_stable_tma_api
 
 from ...utils._sympy.symbol import free_symbol_is_type, prefix_str, symbol_is_type, SymT
 from ...utils._sympy.value_ranges import ValueRanges
@@ -1674,7 +1674,6 @@ def augment_key(self, cache_key: str) -> Union[str, tuple[str, str]]:
         else:
             return cache_key
 
-
 @dataclasses.dataclass
 class TMACompatibilityChecker:
     """
@@ -1692,14 +1691,13 @@ def __post_init__(self):
     def can_use_tma(
         self,
     ) -> bool:
-        import triton
 
         if not (
             V.graph.get_current_device_or_throw().type == "cuda"
             and torch.cuda.get_device_capability()[0] >= 9
             and config.triton.use_tensor_descriptor
             and config.assume_aligned_inputs
-            and triton.__version__ >= "3.4.0"
+            and has_triton_stable_tma_api()
             # For CUDA The base ptr needs to be aligned
         ):
             log.debug(
