[FA]:Optimize FlashAttention for N_CTX <= 512

[quintinwang5]

This change can make the cache behavior of N_CTX=512 better. We can get 20%+ performance gain from the change, but it may be harmful to large N_CTX. So restrict this change to N_CTX <= 512. CI data

[chengjunlu]

The code change LGTM.
Can you explain the difference of the memory foot print between the old one and the new one?

[quintinwang5]

The code change LGTM. Can you explain the difference of the memory foot print between the old one and the new one?

Actually, I'm still working on figuring it out. It's puzzling.

[whitneywhtsang]

Did you get the idea from XeTLA, i.e., is the implementation closer to XeTLA now?

[quintinwang5]

Did you get the idea from XeTLA, i.e., is the implementation closer to XeTLA now?

It was mainly by many tries with profiling data. XeTLA has a different arange that's suitable for all the shapes.

[quintinwang5]

The code change LGTM. Can you explain the difference of the memory foot print between the old one and the new one?

The reason is: we use BLOCK_M to split Q along Y-axis. For different BLOCK_M, they shares the same K, V. So if these workgroups which process different BLOCK_M are scheduled continuously, K, V's cache can be reused. Taking (32, 32, 512, 64) as an example, the old nd_range is {(4, 32, 32), (128, 1, 1)}. There are 4 x 128(BLOCK_M) = 512 blocks. 4 blocks are not consecutive, the stride is 32x32 workgroups. If we change it to {(32, 32, 4), (128, 1, 1)}, 4 blocks are consecutive now.