ggml : require mask when using ALiBi

ggml-ci
This commit is contained in:
Georgi Gerganov 2024-05-10 17:13:11 +03:00
parent 397b1f8f9d
commit 0faf92e74c
No known key found for this signature in database
GPG key ID: 449E073F9DC10735
2 changed files with 10 additions and 1 deletions

9
ggml.c
View file

@ -5657,6 +5657,10 @@ static struct ggml_tensor * ggml_soft_max_impl(
GGML_ASSERT(mask->ne[1] >= a->ne[1]);
}
if (max_bias > 0.0f) {
GGML_ASSERT(mask);
}
bool is_node = false;
if (a->grad) {
@ -6440,6 +6444,7 @@ struct ggml_tensor * ggml_flash_attn_ext(
float max_bias) {
GGML_ASSERT(ggml_can_mul_mat(k, q));
// TODO: check if vT can be multiplied by (k*qT)
if (mask) {
GGML_ASSERT(ggml_is_contiguous(mask));
GGML_ASSERT(mask->ne[2] == 1);
@ -6449,6 +6454,10 @@ struct ggml_tensor * ggml_flash_attn_ext(
//GGML_ASSERT(ggml_can_repeat_rows(mask, qk));
}
if (max_bias > 0.0f) {
GGML_ASSERT(mask);
}
bool is_node = false;
if (q->grad || k->grad || v->grad) {