metal : FA support F32 K and V

[ggerganov]

target #16528

• Add Metal FA kernels for F32 K and V
• Add Metal FA kernels for head size 32
• Remove K and V casts with cacheless contexts (we should keep the casts for now)

Sample command for testing:

llama-embedding -hf ggml-org/bge-small-en-v1.5-Q8_0-GGUF -e -p "$(printf 'hello %.0s' {1..510})" --pooling cls -c 512 -fa on

[ggerganov]

@JohannesGaessler @jeffbolznv Would it be possible to add support for F32 K and V tensors in the respective backends?

The issue is that these casts on master increase significantly the compute buffer size for embedding models (see #15586 (comment)):

llama.cpp/src/llama-graph.cpp

Lines 1313 to 1326 in 4b2dae3

	// this can happen when KV cache is not used (e.g. an embedding model with non-causal attn)
	if (k->type == GGML_TYPE_F32) {
	k = ggml_cast(ctx0, k, GGML_TYPE_F16);
	}
	if (v->type == GGML_TYPE_F32) {
	v = ggml_cast(ctx0, v, GGML_TYPE_F16);
	}
	cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,
	hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
	cb(cur, LLAMA_TENSOR_NAME_FATTN, il);

If we remove the casts, the memory usage should be significantly reduced for this use case. But to remove them, the FA implementation has to support k->type == GGML_TYPE_F32 && v->type == GGML_TYPE_F32 in the ggml_flash_attn_ext() operator.

[JohannesGaessler]

It's definitely possible but it will require additional considerations w.r.t. SRAM limits. For the tile kernel what would need to be done is to determine FP16 vs. FP32 use via a template parameter rather than the FAST_FP16_AVAILABLE macro and to determine which kernel parameters, if any, can be made to fit in SRAM if you now need twice as much to store KV data vs. FP16.

[JohannesGaessler]

Do the models for which this is relevant use GQA?

[JohannesGaessler]

We could also consider making the operations preceding FA write back their data as FP16 in the first place. In terms of performance that would definitely preferable for all CUDA/ROCm GPUs except for Pascal.

[ggerganov]

Do the models for which this is relevant use GQA?

Generally yes.

If it would make the implementation simpler, maybe we can treat F32 K and V as just another "quantization" type, where the dequantize function is a cast to F16?

[JohannesGaessler]

For CUDA that can definitely be done with comparatively little effort but it would not eliminate the additional memory use, it would just shift it from the compute buffer to the buffer pool in the CUDA backend.

[jeffbolznv]

I think this should be relatively straightforward in the vulkan backend, I'll look into it. This comment is how I'd expect to implement it (we dequantize while loading, so no extra memory usage):

If it would make the implementation simpler, maybe we can treat F32 K and V as just another "quantization" type, where the dequantize function is a cast to F16?

[jeffbolznv]

Done for Vulkan in #16543

[JohannesGaessler]

Basic CUDA support in #16546 .