Opitmization for AlphaFold2 model (#4)

* add for alpha_fold2

* add some extra setting

* fix some bugs

* fix some changes

* fix some bugs 2nd

* Add another initition of Gmem_tile_qkv and Gmem_tile_o

* add some compensation for try..catch

* fix mistake in flash_attn_fwd

* commit for code style and bug check

* fix some bugs for flash_attn_with_bias-mask

* add more print for pointer debug

* add some bug test cases.

* backward function

* fix bugs

* make some changes for backward

* Fix compiling error.

* quote all printf debug

* quote all printf debug and fix interface error

* quote all printf debug and fix interface error, fix typo

* remove all printf

* split files

* remove useless debug code

* split fwd and bwd execution function

* split fwd and bwd execution function

* remove useless codes

* remove useless codes

* remove useless codes 3rd times

* remove useless codes 4th times

* Fix compiling error.

* Remove const.

csrc/flash_attn/CMakeLists.txt

@@ -22,6 +22,7 @@ add_library(flashattn SHARED
     ${SOURCES_CU}
     ${SOURCES_CPP}
     flash_attn.cpp
+    flash_attn_with_bias_mask.cpp
   )
 
 target_compile_options(flashattn PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:

csrc/flash_attn/flash_attn.cpp

@@ -26,6 +26,7 @@
  *
  ******************************************************************************/
 
+#include "flash_attn.h"
 #include "fmha.h"
 #include "utils.h"
 #include "cuda_utils.h"
@@ -62,7 +63,7 @@ extern "C" {
 
 static thread_local std::unique_ptr<char[]> flash_attn_err_msg;
 
-static void flash_attn_set_error(const char *msg) {
+void flash_attn_set_error(const char *msg) {
   if (msg == nullptr || *msg == '\0') {
     msg = "unknown error";
   }

csrc/flash_attn/flash_attn.h

@@ -70,6 +70,79 @@ bool flash_attn_bwd(
         uint64_t offset
 );
 
+bool flash_attn_fwd_with_bias_and_mask(
+        const void *q,              // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
+        const void *k,              // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        const void *v,              // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        void *out,                  // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        const int32_t *cu_seqlens_q,   // int32, batch_size+1, starting offset of each sequence
+        const int32_t *cu_seqlens_k,   // int32, batch_size+1, starting offset of each sequence
+        const int total_q,
+        const int total_k,
+        const int batch_size,
+        const int num_heads,
+        const int head_size,
+        const int max_seqlen_q_,
+        const int max_seqlen_k_,
+        const float p_dropout,
+        const float softmax_scale,
+        const bool zero_tensors,
+        const bool is_causal,
+        const bool is_bf16,
+        const int num_splits,        // SMs per attention matrix, can be 1
+        void *softmax_lse_ptr,       // softmax log_sum_exp
+        void *softmax_ptr,
+        void *workspace_ptr,
+        uint64_t *workspace_size,
+        cudaStream_t stream,
+        uint64_t seed,
+        uint64_t offset,
+        const void *attn_mask,
+        const void *attn_bias,
+        const int64_t* mask_dims,
+        const int64_t* bias_dims
+);
+
+bool flash_attn_bwd_with_bias_and_mask(
+        const void *q,              // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
+        const void *k,              // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        const void *v,              // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        void *dq,                   // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
+        void *dk,                   // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        void *dv,                   // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        const void *out,            // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        const void *dout,           // total_q x num_heads, x head_size
+        const int32_t *cu_seqlens_q,   // int32, batch_size+1
+        const int32_t *cu_seqlens_k,   // int32, batch_size+1
+        const int total_q,
+        const int total_k,
+        const int batch_size,
+        const int num_heads,
+        const int head_size,
+        const int max_seqlen_q_,
+        const int max_seqlen_k_,
+        const float p_dropout,
+        const float softmax_scale,
+        const bool zero_tensors,
+        const bool is_causal,
+        const bool is_bf16,
+        const int num_splits,
+        const void *softmax_lse_ptr,
+        void *dsoftmax_ptr,
+        void *dbias_ptr,
+        void *workspace_ptr,
+        uint64_t *workspace_size,
+        cudaStream_t stream,
+        uint64_t seed,
+        uint64_t offset,
+        const void* attn_mask,
+        const void* attn_bias,
+        const int64_t* mask_dims,
+        const int64_t* bias_dims
+);
+
+void flash_attn_set_error(const char *msg);
+
 const char *flash_attn_error();
 
 #ifdef __cplusplus