Optimize block-sparse forward by up to 15% with pipelining (facebookr…

…esearch#529)
bertmaher · Mar 29, 2023 · 01cffe2 · 01cffe2
1 parent 5b7c501
commit 01cffe2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xformers/ops/triton_fairinternal/block_sparse_mem_eff_kernels.py b/xformers/ops/triton_fairinternal/block_sparse_mem_eff_kernels.py
@@ -197,7 +197,7 @@ def block_sparse_attention(
         BLOCK_DMODEL=Lk,
         IS_CAUSAL=causal,
         num_warps=num_warps,
-        num_stages=1,  # TODO: needs to have num_stages=1 otherwise IMA
+        num_stages=2,
     )
 
     return o, lse