Vulkan with FlashAttention: Extreme performance degredation

[remon-nashid]

Is this expected? Using the latest llama.cpp-vulkan build on an AMD 7900 xtx card. Following are the llama-bench results without and with flash attention.

Note that I've reproduced these results with various models from 3B and up to 32B.

FA disabled

model	size	params	backend	threads	test	t/s
qwen2 3B Q6_K	2.36 GiB	3.09 B	Vulkan,BLAS,RPC	16	pp512	2798.42 ± 107.47
qwen2 3B Q6_K	2.36 GiB	3.09 B	Vulkan,BLAS,RPC	16	tg128	86.11 ± 0.74

FA enabled

model	size	params	backend	threads	fa	test	t/s
qwen2 3B Q6_K	2.36 GiB	3.09 B	Vulkan,BLAS,RPC	16	1	pp512	78.29 ± 4.37
qwen2 3B Q6_K	2.36 GiB	3.09 B	Vulkan,BLAS,RPC	16	1	tg128	2.85 ± 0.56

I couldn't find this reported elsewhere. Usually people complain about FA with ROCm instead of Vulkan. However, please let me know if this has been reported or if there are efforts that I could follow.

Thanks

[daniandtheweb]

Right now flash attention on vulkan is only supported on some NVIDIA drivers with the coopmat2 extension, on any other gpu using flash attention causes the computation to be offloaded to the CPU, that's the reason for the performance loss.

[remon-nashid]

Thank you @daniandtheweb ! Are you aware of any work being done to address that problem?

[daniandtheweb]

I don't really know if there's currently any work being done to add flash attention to vulkan. Currently there seem to be quite a bit of work done on the vulkan backend in general so it may be in the devs roadmap.
If you need flash attention you can use ROCm for now, compiling it with support for ROCWMMA makes the flash attention performance even better than without flash attention at all and for RDNA3 cards (or at least my 7800XT) ROCm performance is still better than vulkan's.