small improvement in flash attention

[minhthuc2502]

Some benchmarks with FA2:
Tested on GPU NVIDIA H100, model LLama2 7B quantized 8bit:

Method	Speed (tok/s)	Batch size
Standard MHA	85.8	1
Standard MHA	133.6	2
Previous FA2	88.3	1
Previous FA2	136.5	2
Current FA2	90.1	1
Current FA2	139.9	2

In case seq-length ~= 1100 tokens:

Method	Speed (tok/s)	Batch size
Standard MHA	71.2	1
Previous FA2	81.8	1
Current FA2	85	1

In case seq-length ~= 2200 tokens:

Method	Speed (tok/s)	Batch size
Standard MHA	65.6	1
Current FA2	80.9	1