[Core] Support all head sizes up to 256 with FlashAttention backend

[njhill]

We were previously restricting to specific sizes, but the native FA kernels pad and support arbitrary sizes up to 256.

[tlrmchlsmth]

Could you add some unit tests? Looks like we may be able to just extend this list here🤞

vllm/tests/kernels/test_attention.py

Lines 32 to 34 in c2ec430

	# FlashAttention forward only supports head dimension at most 128
	# https://github.com/ROCmSoftwarePlatform/flash-attention/blob/3d2b6f5d037782cc2c906909a46fb7e2e1b48b25/csrc/flash_attn_rocm/flash_api.cpp#L62
	HEAD_SIZES = [64, 80, 96, 112, 120, 128, 192, 256]

[njhill]

Looks like we need to build flash without the FLASHATTENTION_DISABLE_UNEVEN_K flag, have opened vllm-project/flash-attention#21 ... @WoosukKwon wdyt?

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @njhill.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[lgeiger]

@njhill Has there been any progress on this? This would be very useful so we could bring back support for FA in ViTs and multi modal models #12445