llama : fix FA when KV cache is not used (i.e. embeddings) (#12825)

* ggml : FA supports F32 V

* graph : cast KV to F16 when the KV cache is not used

ggml-ci

* server : add test that exercises embeddings with FA enabled

ggml-ci
This commit is contained in:
Georgi Gerganov 2025-04-08 19:54:51 +03:00 committed by GitHub
parent 78a1ba0a4f
commit a19b5cef16
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 59 additions and 6 deletions

View file

@ -1345,6 +1345,11 @@ static bool ggml_metal_supports_op(const struct ggml_backend_metal_device_contex
case GGML_OP_ARANGE:
return true;
case GGML_OP_FLASH_ATTN_EXT:
if (op->src[0]->ne[0] == 32) {
// head size == 32 (e.g. bert-bge-small)
// TODO: not sure if it is worth adding kernels for this size
return false;
}
if (op->src[1]->type != op->src[2]->type) {
return false;
}