#1235·TabPFN

FlashAttention-4 backend for Blackwell

Author: tolleybotCreated Sep 2, 2026Updated Sep 18, 2026
Labelsenhancement ✨

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:

  1. One backend or two? FA4 covers Hopper too (head_dim up to 256 on sm_90). Separate fa4_backend.py beside FA3, or a single backend spanning both with FA3 eventually retired?
  2. 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.