metal : add quantized FA support

[ggerganov]

supersed #9735

Extend the FA kernels to support quantized KV cache.

TODO:

• try to extend the non-vec FA kernel as well
• separate code movement in different PR to reduce diff

[slaren]

test-backend-ops fails on M3 Max. On master this test was skipped.

FLASH_ATTN_EXT(hs=256,nh=32,kv=1024,nb=1,mask=1,max_bias=0.000000,logit_softcap=0.000000,type_KV=f16): ggml/src/ggml-metal.m:3193: GGML_ASSERT(smem <= device.maxThreadgroupMemoryLength) failed

Performance looks good, a bit slower than f16 on prompt processing:

model	size	params	backend	ngl	type_k	type_v	fa	test	t/s
llama 7B Q4_0	3.56 GiB	6.74 B	Metal,BLAS	99	f16	f16	1	pp512	832.59 ± 1.26
llama 7B Q4_0	3.56 GiB	6.74 B	Metal,BLAS	99	f16	f16	1	tg128	69.68 ± 0.05
llama 7B Q4_0	3.56 GiB	6.74 B	Metal,BLAS	99	q4_0	q4_0	1	pp512	775.79 ± 0.77
llama 7B Q4_0	3.56 GiB	6.74 B	Metal,BLAS	99	q4_0	q4_0	1	tg128	69.50 ± 0.08
llama 7B Q4_0	3.56 GiB	6.74 B	Metal,BLAS	99	q8_0	q8_0	1	pp512	776.22 ± 0.75
llama 7B Q4_0	3.56 GiB	6.74 B	Metal,BLAS	99	q8_0	q8_0	1	tg128	69.56 ± 0.05

[ggerganov]

Thanks for checking. The shared memory assert should be good now - let me know if it still fails.

[slaren]

test-backend-ops passes now, and perplexity also looks good.

[ddh0]

@slaren could you please clarify which K/V quantization combinations are supported on Metal? and is there an equivalent of the GGML_CUDA_FA_ALL_QUANTS compile option? Thanks :)

[ggerganov]

The Metal backend supports KV quantizations where the K and V are the same time:

llama.cpp/ggml/src/ggml-metal/ggml-metal.m

Lines 1099 to 1103 in 82bca22

	case GGML_OP_FLASH_ATTN_EXT:
	if (op->src[1]->type != op->src[2]->type) {
	return false;
	}
	return has_simdgroup_mm; // TODO: over-restricted for vec-kernels

It is currently implemented for F16, BF16, Q8_0, Q4_0, Q4_1, Q5_0 and Q5_1.