Adding META['enable_tma']

facebookresearch · Sep 4, 2024 · 306492f · 306492f
1 parent ffbfb01
commit 306492f
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/ops/triton/triton_ragged_hstu_attention.py b/ops/triton/triton_ragged_hstu_attention.py
@@ -774,7 +774,8 @@ def grid_tma(META):
             META['BLOCK_N'],
             BLOCK_D_Q,
             k.element_size(),
-            k_buf.numpy()
+            k_buf.numpy(),
+            META['enable_tma']
         )
         triton.runtime.driver.active.utils.fill_2d_tma_descriptor(
             v.data_ptr(),
@@ -783,7 +784,8 @@ def grid_tma(META):
             META['BLOCK_N'],
             BLOCK_D_V,
             v.element_size(),
-            v_buf.numpy()
+            v_buf.numpy(),
+            META['enable_tma']
         )
         desc_k.copy_(k_buf, non_blocking=True)
         desc_v.copy_(v_buf, non_blocking=True)
@@ -905,7 +907,8 @@ def grid_tma(META):
             META['BLOCK_N'],
             BLOCK_D_Q,
             k.element_size(),
-            k_buf.numpy()
+            k_buf.numpy(),
+            META['enable_tma']
         )
         triton.runtime.driver.active.utils.fill_2d_tma_descriptor(
             v.data_ptr(),
@@ -914,7 +917,8 @@ def grid_tma(META):
             META['BLOCK_N'],
             BLOCK_D_V,
             v.element_size(),
-            v_buf.numpy()
+            v_buf.numpy(),
+            META['enable_tma']
         )
         desc_k.copy_(k_buf, non_blocking=True)
         desc_v.copy_(v_buf, non_blocking=True)