Add support for qk hidden dim different from v hidden dim

[smallscientist1]

We add support for

• different hidden dimension between qk and v.
• not equal num_heads_k and num_heads_v, such as (num_heads_q, num_heads_k, num_heads_v) = (32, 4, 16).

For different hidden dimension between qk and v, we have supported:

• FlashAttention-2 with QKHeadDim=32, VHeadDim=64
• FlashAttention-2 with QKHeadDim=64, VHeadDim=128
• FlashAttention-2 with QKHeadDim=96, VHeadDim=192
• FlashAttention-2 with QKHeadDim=128, VHeadDim=256
• FlashAttention-2 with QKHeadDim=192, VHeadDim=128

For headdim not supported, you can use the autotuner to generate the implementation. Details are in autotuner.md.

Performance

We test the performance speedup compare to padding qk&v hidden_dim to the same length.

Test

We add unittest in tests/test_flash_attn_headdim.py tests/test_flash_attn_head.py.

[iqiancheng]

hi~ @smallscientist1
Regarding the combinations of qk and v dimensions you've implemented in FlashAttention-2, which configuration have you found to offer the best balance between performance and model effectiveness? Specifically, among the combinations:

QKHeadDim=32, VHeadDim=64
QKHeadDim=64, VHeadDim=128
QKHeadDim=96, VHeadDim=192
QKHeadDim=128, VHeadDim=256

Which one stands out in terms of computational efficiency and model quality?

[xiayuqing0622]

hi~ @smallscientist1 Regarding the combinations of qk and v dimensions you've implemented in FlashAttention-2, which configuration have you found to offer the best balance between performance and model effectiveness? Specifically, among the combinations:

QKHeadDim=32, VHeadDim=64 QKHeadDim=64, VHeadDim=128 QKHeadDim=96, VHeadDim=192 QKHeadDim=128, VHeadDim=256

Which one stands out in terms of computational efficiency and model quality?

In terms of model quality, it's too early to make a definitive assessment since the work is still in progress. However, several teams we've collaborated with expressed a need for this combination, so we implemented it. Additionally, anticipating that others might find it useful, we created this PR to benefit the broader community.