follow the original expression more closely

OpenProteinAI · Feb 8, 2025 · f5ce6ee · f5ce6ee
1 parent a645edd
commit f5ce6ee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/csrc/flash_attn/src/mask.h b/csrc/flash_attn/src/mask.h
@@ -181,7 +181,7 @@ struct Mask {
                                         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)) += ((row_idx + max_seqlen_k == max_seqlen_q + col_idx) ? 0 : alibi_slope);
+                                        tensor(make_coord(i, mi), make_coord(j, nj)) += (((row_idx + max_seqlen_k - max_seqlen_q - col_idx) == 0) ? 0 : alibi_slope);
 
                                     }
                                 }