Add causal flag to fa2 (#1976)

Author: sssshhhhhh

include/ctranslate2/ops/flash_attention.h

@@ -6,7 +6,7 @@ namespace ctranslate2 {
   namespace ops {
     class FlashAttention : public Op {
     public:
-      FlashAttention(float queries_scale, dim_t sliding_window);
+      FlashAttention(float queries_scale, dim_t sliding_window, bool is_causal = true);
 
       void operator()(StorageView& queries,
                       StorageView& keys,
@@ -25,6 +25,7 @@ namespace ctranslate2 {
     private:
       const float _queries_scale;
       const dim_t _sliding_window;
+      const bool _is_causal;
       template <Device D>
       void compute(StorageView& queries,
                    StorageView& keys,

src/ops/flash_attention.cc

@@ -4,9 +4,10 @@
 
 namespace ctranslate2 {
   namespace ops {
-    FlashAttention::FlashAttention(float queries_scale, dim_t sliding_window)
+    FlashAttention::FlashAttention(float queries_scale, dim_t sliding_window, bool is_causal)
     : _queries_scale(queries_scale)
-    ,_sliding_window(sliding_window)
+    , _sliding_window(sliding_window)
+    , _is_causal(is_causal)
     {
     }
 

src/ops/flash_attention_gpu.cu

@@ -232,8 +232,8 @@ namespace ctranslate2 {
         num_heads_k = cached_keys->dim(2);
       }
 
+      bool is_causal = _is_causal;
       // causal=true is the same as causal=false in this case
-      bool is_causal = true;
       if (seqlen_q == 1 && !alibi) { is_causal = false; }
       if (is_causal) { window_size_right = 0; }