Replace has_warp_spec with HAS_AUTO_WS env variable check (#221)

Summary:
Triton release/3.3.x branch only supports AutoWS which makes the check has_warp_spec = hasattr(tl, "async_task") incorrect. Instead this PR adds an environment variable check HAS_AUTO_WS = os.getenv("ENABLE_AUTO_WS") to replace that.

Pull Request resolved: #221

Reviewed By: xuzhao9

Differential Revision: D74487395

Pulled By: ardaunal

fbshipit-source-id: c6b40ac2692df590937bda7d05e668fe3e556f6c

Author: ardaunal

tritonbench/kernels/triton_fused_attention.py

@@ -27,6 +27,7 @@
 WITH_COMPPIPE = os.getenv("ENABLE_COMPPIPE")
 PEEL_LAST = os.getenv("PEEL_LAST_ITER")
 WITH_TMA = os.getenv("WITH_TMA")
+HAS_AUTO_WS = os.getenv("ENABLE_AUTO_WS")
 
 if HAS_TMA_DESC:
     print(
@@ -313,7 +314,6 @@ def _attn_fwd_inner_ws(
 # We don't run auto-tuning every time to keep the tutorial fast. Uncommenting
 # the code below and commenting out the equivalent parameters is convenient for
 # re-tuning.
-EXPLICIT_WARP_SPEC = hasattr(tl, "async_task")
 HAS_NEW_TMA = hasattr(triton, "set_allocator") and hasattr(tl, "make_tensor_descriptor")
 schedList = ["default", "FA_firstDot", "FA_secondDot"] if WITH_COMPPIPE else ["default"]
 # TODO: incorrect result with PEEL_LAST + FA_firstDot + WarpSpec + TMA
@@ -334,7 +334,7 @@ def _attn_fwd_inner_ws(
             num_buffers_warp_spec=0,
             num_consumer_groups=0,
         )
-        if EXPLICIT_WARP_SPEC
+        if HAS_AUTO_WS == "1"
         else triton.Config(
             {
                 "BLOCK_M": BM,
@@ -367,7 +367,7 @@ def _attn_fwd_inner_ws(
             num_buffers_warp_spec=0,
             num_consumer_groups=0,
         )
-        if EXPLICIT_WARP_SPEC
+        if HAS_AUTO_WS == "1"
         else triton.Config(
             {
                 "BLOCK_M": BM,
@@ -397,7 +397,7 @@ def _attn_fwd_inner_ws(
             reg_dec_producer=dec,
             reg_inc_consumer=inc,
         )
-        if EXPLICIT_WARP_SPEC
+        if HAS_AUTO_WS == "1"
         else triton.Config(
             {"BLOCK_M": BM, "BLOCK_N": BN, "ENABLE_TMA": False, "LOOP_SCHEDULE": sched},
             num_stages=2 if sched == "FA_firstDot" or sched == "FA_secondDot" else 0,
@@ -431,7 +431,7 @@ def _attn_fwd_inner_ws(
                 num_buffers_warp_spec=0,
                 num_consumer_groups=0,
             )
-            if EXPLICIT_WARP_SPEC
+            if HAS_AUTO_WS == "1"
             else triton.Config(
                 {
                     "BLOCK_M": BM,
@@ -487,7 +487,7 @@ def _attn_fwd_inner_ws(
             reg_dec_producer=dec,
             reg_inc_consumer=inc,
         )
-        if EXPLICIT_WARP_SPEC
+        if HAS_AUTO_WS == "1"
         else triton.Config(
             {
                 "BLOCK_M": BM,
@@ -528,7 +528,7 @@ def _attn_fwd_inner_ws(
             reg_dec_producer=dec,
             reg_inc_consumer=inc,
         )
-        if EXPLICIT_WARP_SPEC
+        if HAS_AUTO_WS == "1"
         else triton.Config(
             {
                 "BLOCK_M": BM,