Update mask.h

timt51 · web-flow · commit 88e1cc9c47a6 · 2025-02-08T09:38:28.000-05:00
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) ? 0 : alibi_slope);
+                                        tensor(make_coord(i, mi), make_coord(j, nj)) += ((col_idx == row_idx) ? 0 : alibi_slope);
 
                                     }
                                 }

Original file line number	Diff line number	Diff line change
`@@ -181,7 +181,7 @@ struct Mask {`
`181`	`181`	`tensor(make_coord(i, mi), make_coord(j, nj)) += ((col_idx == (col_idx_limit_right - 1)) ? 0 : alibi_slope);`
`182`	`182`
`183`	`183`	`} else {`
`184`		`- tensor(make_coord(i, mi), make_coord(j, nj)) += (((row_idx + max_seqlen_k - max_seqlen_q - col_idx) == 0) ? 0 : alibi_slope);`
	`184`	`+ tensor(make_coord(i, mi), make_coord(j, nj)) += ((col_idx == row_idx) ? 0 : alibi_slope);`
`185`	`185`
`186`	`186`	`}`
`187`	`187`	`}`