alibi.h use the non max seqlen formula which seems more correct actually?

timt51 · web-flow · commit 4bda2db97ba7 · 2025-02-08T09:38:56.000-05:00
diff --git a/csrc/flash_attn/src/alibi.h b/csrc/flash_attn/src/alibi.h
@@ -67,7 +67,7 @@ struct Alibi {
                         #pragma unroll
                         for (int j = 0; j < size<1, 0>(tensor); ++j) {
                             const int col_idx = col_idx_base + j;
-                            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
`@@ -67,7 +67,7 @@ struct Alibi {`
`67`	`67`	`#pragma unroll`
`68`	`68`	`for (int j = 0; j < size<1, 0>(tensor); ++j) {`
`69`	`69`	`const int col_idx = col_idx_base + j;`
`70`		`- tensor(make_coord(i, mi), make_coord(j, nj)) += (((row_idx + max_seqlen_k - max_seqlen_q - col_idx) == 0) ? 0 : alibi_slope);`
	`70`	`+ tensor(make_coord(i, mi), make_coord(j, nj)) += ((col_idx == row_idx) ? 0 : alibi_slope);`
`71`	`71`	`}`
`72`	`72`	`}`
`73`	`73`	`}`