Add TLX attention (WS pipelined pingpong hopper)

Author: yf225

tritonbench/operators/flash_attention/operator.py

@@ -136,6 +136,18 @@
 except (ImportError, IOError, AttributeError, TypeError):
     HAS_XFORMERS = False
 
+# [Optional] TLX backend
+try:
+    import triton.language.extra.tlx as tlx
+
+    from .tlx_attn_ws_pipelined_pingpong_hopper import (
+        attention as tlx_attn_ws_pipelined_pingpong_hopper,
+    )
+
+    HAS_TLX = True
+except (ImportError, IOError, AttributeError):
+    HAS_TLX = False
+
 from typing import Any, Generator, List
 
 from tritonbench.utils.input import input_filter
@@ -299,6 +311,16 @@ def triton_tutorial_flash_v2_tma(
             q, k, v, self.causal, self.sm_scale, "tma"
         )
 
+    @register_benchmark(enabled=HAS_TLX)
+    def tlx_attn_ws_pipelined_pingpong_hopper(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ) -> Callable:
+        # TLX flash attention with Hopper optimizations
+        return lambda: tlx_attn_ws_pipelined_pingpong_hopper(q, k, v, self.sm_scale)
+
     def xformers_preprocess(
         self,
         q: torch.Tensor,
@@ -341,7 +363,9 @@ def xformers_splitk(
             fhma_input, needs_gradient=need_gradient
         )
 
-    @register_benchmark(enabled=False, label=f"cudnn-{torch.backends.cudnn.version()}")
+    @register_benchmark(
+        enabled=False
+    )  # , label=f"cudnn-{torch.backends.cudnn.version()}")
     def cudnn(self, q, k, v):
         os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1"