Skip to content

Commit

Permalink
rdy for merge
Browse files Browse the repository at this point in the history
  • Loading branch information
Green-Sky committed Sep 7, 2024
1 parent e904b86 commit f69fbfd
Showing 1 changed file with 6 additions and 7 deletions.
13 changes: 6 additions & 7 deletions ggml_extend.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -710,11 +710,10 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*

float scale = (1.0f / sqrt((float)d_head));

if (flash_attn) {
// TODO: remove before merge
LOG_DEBUG("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
}
// is there anything oddly shaped??
//if (flash_attn) {
// LOG_DEBUG("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
//}
// is there anything oddly shaped?? ping Green-Sky if you can trip this assert
GGML_ASSERT(((L_k % 256 == 0) && L_q == L_k) || !(L_k % 256 == 0));

bool can_use_flash_attn = true;
Expand All @@ -725,12 +724,12 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*
can_use_flash_attn = can_use_flash_attn && d_head <= 256; // double check

if (mask != nullptr) {
// TODO: figure out if we can bend t5 to work too
// TODO(Green-Sky): figure out if we can bend t5 to work too
can_use_flash_attn = can_use_flash_attn && mask->ne[2] == 1;
can_use_flash_attn = can_use_flash_attn && mask->ne[3] == 1;
}

// TODO: more pad or disable for funny tensor shapes
// TODO(Green-Sky): more pad or disable for funny tensor shapes

ggml_tensor* kqv = nullptr;
//GGML_ASSERT((flash_attn && can_use_flash_attn) || !flash_attn);
Expand Down

0 comments on commit f69fbfd

Please sign in to comment.