Add TLX attention (WS pipelined pingpong hopper)

Author: yf225

tritonbench/operators/flash_attention/operator.py

@@ -61,6 +61,10 @@
 from tritonbench.utils.path_utils import add_ld_library_path
 from tritonbench.utils.triton_op import is_fbcode
 
+from .tlx_attn_ws_pipelined_pingpong_hopper import (
+    attention as tlx_attn_ws_pipelined_pingpong_hopper,
+)
+
 
 # [Optional] flash_attn v2
 try:
@@ -136,6 +140,14 @@
 except (ImportError, IOError, AttributeError, TypeError):
     HAS_XFORMERS = False
 
+# [Optional] TLX backend
+try:
+    import triton.language.extra.tlx as tlx
+
+    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,