fix

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: manman-ren

tritonbench/kernels/blackwell_triton_fused_attention.py

@@ -239,7 +239,7 @@ def make_tile_config(BM, BN, occ, subtile, vectmul, add2reduce):
     ]
 else:
     # Helper to build config with optional minRegAutoWS/maxRegAutoWS
-    def make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce):
+    def make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce, maxreg):
         config_kwargs = {
             "BLOCK_M": BM,
             "BLOCK_N": BN,
@@ -256,20 +256,21 @@ def make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce):
         # Only add minRegAutoWS/maxRegAutoWS if supported (triton/tree/ws-3.5)
         if HAS_REG_AUTO_WS:
             extra_kwargs["minRegAutoWS"] = 24
-            extra_kwargs["maxRegAutoWS"] = 152
+            extra_kwargs["maxRegAutoWS"] = maxreg
             extra_kwargs["data_partition_factor"] = 2
 
         return triton.Config(config_kwargs, **extra_kwargs)
 
     configs = [
-        make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce)
+        make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce, maxreg)
         for BM in [256]
         for BN in [64, 128]
         for s in NUM_STAGES_OPTIONS
         for w in [4]
         for subtile in [True]
         for vectmul in [1]
         for add2reduce in [False]
+        for maxreg in [152, 192]
     ]
 
 
@@ -384,7 +385,6 @@ def _attn_fwd_tma_dp(
     VECT_MUL: tl.constexpr,
     FADD2_REDUCE: tl.constexpr,
 ):
-    tl.static_assert(BLOCK_N <= HEAD_DIM)
     start_m = pid  # tl.program_id(0)
     # off_hz = tl.program_id(1)
     off_z = off_hz // H
@@ -687,7 +687,7 @@ def grid_debug(META):
             ):
                 extra_kern_args["maxnreg"] = 128
             else:
-                extra_kern_args["maxnreg"] = 80
+                extra_kern_args["maxnreg"] = 128
         if persistent:
             _attn_fwd_persist[grid_persist](
                 sm_scale,

tritonbench/kernels/blackwell_triton_fused_attention_dp.py

@@ -72,7 +72,7 @@ def _attn_fwd_subtile(
         qk -= m_ij[:, None]
     else:
         m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
-        if VECT_MUL:
+        if VECT_MUL == 2 or VECT_MUL == 3:
             qk = _fma_f32x2(qk, qk_scale, -m_ij[:, None])
         else:
             qk = qk * qk_scale - m_ij[:, None]
@@ -88,7 +88,7 @@ def _attn_fwd_subtile(
 
     if SUBTILING:
         acc0, acc1 = acc.reshape([BM, 2, BN // 2]).permute(0, 2, 1).split()
-        if VECT_MUL:
+        if VECT_MUL == 1 or VECT_MUL == 3:
             acc0 = _mul_f32x2(acc0, alpha[:, None])
             acc1 = _mul_f32x2(acc1, alpha[:, None])
         else:
@@ -262,12 +262,12 @@ def make_tile_config(BM, BN, occ, subtile, vectmul, add2reduce):
         for BN in [64, 128]
         for occ in [1, 2]
         for subtile in [True]
-        for vectmul in [False]
+        for vectmul in [0]
         for add2reduce in [False]
     ]
 else:
     # Helper to build config with optional minRegAutoWS/maxRegAutoWS
-    def make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce):
+    def make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce, maxreg):
         config_kwargs = {
             "BLOCK_M": BM,
             "BLOCK_N": BN,
@@ -284,19 +284,20 @@ def make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce):
         # Only add minRegAutoWS/maxRegAutoWS if supported (triton/tree/ws-3.5)
         if HAS_REG_AUTO_WS:
             extra_kwargs["minRegAutoWS"] = 24
-            extra_kwargs["maxRegAutoWS"] = 152
+            extra_kwargs["maxRegAutoWS"] = maxreg
 
         return triton.Config(config_kwargs, **extra_kwargs)
 
     configs = [
-        make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce)
+        make_standard_config(BM, BN, s, w, subtile, vectmul, add2reduce, maxreg)
         for BM in [256]
         for BN in [64, 128]
         for s in NUM_STAGES_OPTIONS
         for w in [4]
         for subtile in [True]
-        for vectmul in [False]
+        for vectmul in [1]
         for add2reduce in [False]
+        for maxreg in [152, 192]
     ]