unroll 2 loops, int64_t -> int, 309 µs

[JohannesGaessler]

This PR has three very simple performance optimizations:

1. Change loops in such a way that all threads have the same loop variable. This reduces the number of comparisons that each warp has to do by a factor of 32 so it's faster. I only applied this optimization to two loops but it can be applied to essentially all loops in the kernel.
2. Explicitly tell the compiler to unroll loops with #pragma unroll. Unrolling loops (especially the inner ones) is generally faster but the compiler tends to be conservative with it because loop unrolling increases register pressure and therefore reduces occupancy. This ensures that even in the worst case scenario the performance will not be terrible but in my experience it is better to explicitly tell the compiler to unroll more loops. Refactor the loops as I did for two of them and add more #pragma unroll. Start with the inner loops and work your way outwards while monitoring performance and register pressure.
3. Change loop variables from int64_t to int. CUDA has 32 bit registers so data types with that size are the fastest to work with. The maximum values for the loop variables is much smaller than the maximum value of int so there is no benefit to using int64_t. It only makes the code slower.

On my system (1x RTX 3090) the runtime for the flash_attn_ext_f16 kernel has decreased from 383 µs to 309 µs (1.23x faster).

[ggerganov]

Oh nice! The int counters improve significantly the speed

Will look into the other opts now

Edit: having the same loop value also brings noticeable improvement. Will check the unroll now

[ggerganov]

I found another nice speedup: ggml-org@c51f27c