FlashAttention-4 backend for Blackwell
fa3_backend.py notes that FA3 is sm_90-only and that Blackwell will need FA4. FA4 is now available (pip install flash-attn-4, CuTeDSL, covering Hopper and Blackwell), and we'd like to contribute a backend for it. That means one new module in the shape of fa3_backend.py, implementing the AttentionBackend protocol from #1165, with no changes to the architectures themselves.
The shapes line up: v3's ICL head_dim of 64 and its fp16/bf16 dtypes are inside what FA4 supports on SM100 (8 ≤ head_dim ≤ 128, per _validate_head_dims in flash_attn/cute/interface.py as of fa4-v4.0.0.beta28), and flash_attn.cute.flash_attn_func takes q/k/v in the same (B, S, H, D) layout as FA3, so run() is close to identical. One caveat: that 64 comes from the v3 config defaults rather than a released checkpoint, so if the shipped models differ, worth knowing.
Questions
Before we write anything:
- One backend or two? FA4 covers Hopper too (head_dim up to 256 on sm_90). Separate
fa4_backend.pybeside FA3, or a single backend spanning both with FA3 eventually retired? - Does a pre-release dependency belong in-tree? FA4 is still on a weekly beta cadence. Registration is gated on importability, so a missing wheel just degrades to SDPA, but carrying a beta in-tree is your maintenance burden rather than ours. If that's the sticking point, we're equally happy to ship it out-of-tree as a separate package that registers itself, which is what #1165 makes possible.
Benchmarking
_FA3_MIN_SEQLEN_FOR_SPEEDUP was measured on H100 with the FA3 kernels, and we'd rather not assume it carries over to a different kernel on different hardware. We're happy to benchmark FA4 against SDPA across a sequence-length sweep on v3's ICL shapes and post the crossover, including if it shows FA4 doesn't help. That needs no changes to the codebase, just the shapes.
Source: PriorLabs/TabPFN