Paged Attention support for FA3

[kadeng]

Adding support for Paged Attention / block table to the Flash-Attention 3 Kernel.

Limits:

• Not yet integrated with FA3 kvcache + split kv + gqa parallelization #1236 which modifes the same code (and more).
• No fp8 support yet - as that support also seems missing in the non-paged varlen version so far.

Test Plan:

cd hopper
pytest test_flash_attn.py -k "test_flash_attn_varlen_paged"  -s

Update

• Rebased on top of main, which includes FA3 kvcache + split kv + gqa parallelization #1236 now
• Original version of this PR ( not rebased ) still available on this branch

[kadeng]

Hi Alex, I think that makes sense and should be straightforward to implement. The main difference would be that the version in here is based on the varlen variant and flash_attn_with_kvcache would be based off the fixed length variant + sequence length info for KV. So the difference would be mostly in the memory format of the Q Tensor (varlen packed or not). But I would not put it into this PR, it is large enough already. I want to get this merged first. best, Kai Alex Ng ***@***.***> schrieb am Di., 22. Okt. 2024, 09:11:

…

Hi, do you have plans to support block table of flash_attn_with_kvcache()? — Reply to this email directly, view it on GitHub <#1268 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AADVNP6UIGZQPZVGWKVE35TZ4X3DXAVCNFSM6AAAAABPXFJ5JKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDIMRYGQZTMNBVGQ> . You are receiving this because you authored the thread.Message ID: ***@***.***>