cuda : fix flash_attn kernel to produce same results as CPU

[ggerganov]

No description provided.

[FSSRepo]

@ggerganov Thank you very much for the help; it's evident that you know what you're doing!

[ggerganov]

@FSSRepo This is ready to merge. I've tried to improve the performance on RTX 2060, but I think it is still worse compared to master in most cases. I suspect there is room for improvement in the "online softmax" loop that would give the necessary speedup, but there might be some additional CUDA tricks that I'm not aware of yet

I propose that we merge this and continue to try to improve the performance

[FSSRepo]

@ggerganov

I suspect there is room for improvement in the "online softmax" loop that would give the necessary speedup, but there might be some additional CUDA tricks that I'm not aware of yet.

It seems that there is high pressure on the registers because each 16x16 core tensor fragment uses 8 registers per thread (256 elements, and there are 32 threads per warp). I assume that the distribution of the 65535 registers per block depends on the number of threads launched per block.

There must be a way to reduce the number of fragments loaded at the same time. Additionally, instead of using half types, half2 types should be utilized to maximize the use of registers, as each one is 32 bits and can accommodate a half2.