Flash Attention 4 support

[jp-gr]

Hello,

Apologies if I am asking a silly question but we are finding this repo useful for FA3 so thought I would ask - are there plans to publish Flash Attention 4 kernels - in particular, targeting the GB200?

Thanks!

[danieldk]

Yes, the kernel itself has been ready for a few months. However, we are waiting for the CUTLASS/CuTeDSL 4.4 release, since it has a fix that is needed for JIT-compiling relative imports. After it's released I'll update the builder and make a PR for the FA4 branch.

[jp-gr]

great news thank you!

[Trgtuan10]

Hi @danieldk, are there any Flash Attention kernels that support the B200 GPU now? I tried FlashAttention-3, but it doesn’t seem to work.
flux_pipe.transformer.set_attention_backend("_flash_3_hub")

[jp-gr]

@danieldk looks like CUTLASS 4.4 has been released, does it contain the expected fix?

[danieldk]

Was merged in #463.