redo causal and noncausal alibi for diagonal

OpenProteinAI · Feb 7, 2025 · a8d7fcc · a8d7fcc
1 parent c42bf6b
commit a8d7fcc
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/csrc/flash_attn/src/mask.h b/csrc/flash_attn/src/mask.h
@@ -177,7 +177,13 @@ struct Mask {
                             for (int j = 0; j < size<1, 0>(tensor); ++j) {
                                 const int col_idx = col_idx_base + j;
                                 if constexpr (Has_alibi) {
-                                    tensor(make_coord(i, mi), make_coord(j, nj)) += alibi_slope;
+                                    if constexpr (Is_causal) {
+                                        tensor(make_coord(i, mi), make_coord(j, nj)) += ((col_idx == (col_idx_limit_right - 1)) ? 0 : alibi_slope);
+
+                                    } else {
+                                        tensor(make_coord(i, mi), make_coord(j, nj)) += ((col_idx == row_idx) ? 0 : alibi_slope);
+
+                                    }
                                 }
                                 if constexpr (Causal_mask) {
                                     if (col_idx >= col_idx_limit_right) {