Repository navigation
FlashAttention-4 backend for Blackwell #1235
Description
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 Add an attention-backend registry #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.
Activity
alanprior commentedon Sep 4, 2026
Hello @tolleybot, thanks for submitting this issue. We will aim to review it and provide initial feedback on it during the next days.
jmkuebler commentedon Sep 7, 2026
@tolleybot Hi Donald, thanks for opening this issue.
We would indeed welcome such a submission. The best approach would be to actually replace the existing FA3 backend with an FA4 backend. It would then automatically dispatch..
If you ahve time to work on it and open a PR that would be nice. To properly review it, it would be great to have some benchmarking numbers both on Hopper and Blackwell.
We should also have FA4 as an optional install dependency and detect at registration as we do for others.
Hopper:
Blackwell:
tolleybot commentedon Sep 10, 2026
Thanks, that is helpful and I would be glad to work on it. Two questions before I start.
On replacing FA3: since FA4 covers Hopper as well as Blackwell, my plan would be to benchmark FA4 against FA3 on Hopper first, and only remove
fa3_backend.pyif FA4 matches or beats it there. Does that sound like a good plan?One threshold or two?
_FA3_MIN_SEQLEN_FOR_SPEEDUPis a single constant measured on H100. Hopper and Blackwell are unlikely to cross at the same sequence length, so I expect the benchmark to produce two, which would makeis_preferredarchitecture-dependent. Worth agreeing before the numbers arrive.jmkuebler commentedon Sep 14, 2026
@tolleybot thank you very much for working on this.
Yes this sounds good! If we were to loose lots of performance on Hopper, we might want to keep FA3 explicitily.
I would aim for one threshold if they are in the same order. But please let's measure them both and then we can assess whether to make the hardware split or not.
tolleybot commentedon Sep 14, 2026
Both halves are in: Hopper (H100) and Blackwell (GB200). Short version: FA4 has no
short-sequence penalty on either architecture, so one threshold works and it can be 0; on Hopper
FA4 matches or beats FA3 everywhere so nothing is lost by replacing it; and on Blackwell the win
comes entirely from the head_dim-16 feature-attention stages, because cuDNN attention already
runs the head_dim-64 ICL calls at FA4 speed there.
This is the set you listed. Hopper: sdpa / FA3 / FA3-CuTeDSL-via-FA4 (the "FA4" columns on
H100 below). Blackwell: sdpa / FA4.
Setup
07cfc8eplus a draftfa4_backend.pymirroring
fa3_backend.py8d3a3b8TabPFNClassifier(n_estimators=1)predict()wall-clock, fp16 autocast,synthetic data,
n_test = n_train/10, each variant in a fresh process"FA4 hd≥64" is restricted to the call sites FA3 serves
Whole-model, ratio = SDPA time / variant time (>1 is faster than SDPA):
Absolute SDPA times for scale: H100 0.04 s at 1k → 17.6 s at 300k×100; GB200 0.07 s → 10.3 s.
The gain grows with n_train: on H100 FA4 is 1.13–1.19× at 30k, 1.17–1.25× at 100k and
1.27–1.30× at 300k, i.e. 3.7 s off a 17.6 s pass; on GB200 it's 1.12–1.26× at ≥30k for
n_features ≥ 100 (1.2 s off 10.3 s at 300k×100), and ~1.0 at n_features = 10 where cuDNN already
matches it.
On your two questions:
Replace FA3? Yes. On Hopper FA4 ≥ FA3 at every point from 10k up (4–10% whole-model), tie below.
Kernel-level there are two spots FA3 is ahead: ~10% at n_train=30k self-attention, and 2–9% at
1M short-Q cross-attention because FA4 b30 has no split-KV on sm_90. Neither shows
whole-model. FA4 also serves head_dim 16, which FA3 cannot, and that is worth 5–10% on Hopper at
n_features ≥ 100 on its own.
One threshold or two? One, and it can be 0. Neither architecture shows a low-end penalty in
fp16 from n_train=100 (H100) / 300 (GB200) up; the ≤3k rows are ties within noise. The 10k FA3
crossover came from FA3's fixed
num_splits=32, which costs 3–6× at 1k kernel-level; FA4 has noequivalent.
Blackwell specifically. Torch's SDPA picks cuDNN attention on sm_100 and it is at parity with
FA4 for the head_dim-64 ICL calls (the "hd≥64" column is 1.00 ± 0.05 from 3k up). All of the
GB200 gain, 1.12–1.26× at n_features ≥ 100 and n_train ≥ 30k, is the head_dim-16 stages, which
make hundreds of calls with
seq_kv = n_train(1,240 at 100k×100 vs 49 ICL calls). So I'd proposethe FA4 eligibility gate be FA4's real range (8..256 step 8 on sm_90, 8..128 on sm_100) rather
than a copy of FA3's set; without that, FA4 buys nothing on Blackwell.
bf16 caveat. Kernel-level on sm_100, FA4 bf16 is 18–23% slower than SDPA bf16 at every point
from 10k up (and ~20% slower than FA4's own fp16); fp16 has no such dip, and Hopper has none in
either dtype. TabPFN defaults to fp16 autocast so the tables above are unaffected, but a user
passing
inference_precision=torch.bfloat16would regress on Blackwell. I'd gate bf16 out of FA4on compute capability ≥ 10 for now and re-check per beta. Happy to do it the other way if you'd
rather keep the dtype handling uniform.
Numerical equivalence: FA4 matches SDPA within the FA3 test tolerance (
5e-3) on all sixcanonical shapes × fp16/bf16, plus (1, 256, 100k, 8, 1, 64) and batch=70k, on both sm_90 and
sm_100.
One thing found on the way: FA3's split-KV path (
num_splits > 1, whichfa3_attn_funcusesfor
seq_q ≤ 1024since #1168) crashes at batch ≥ 65536, sotest__fa3_batch_heads_above_cuda_max_gridhas been failing since then, unnoticed becausehoppertests don't run in CI. Low impact on v3, but worth a separate issue if FA3 stays.Caveats:
n_estimators=1(ensemble multiplies all columns equally); synthetic data; FA4 is aweekly beta and b30 is what's measured.
If this looks right I'll open the PR with the FA3 removal.
jmkuebler commentedon Sep 18, 2026
Hi @tolleybot thanks for the results, which look quite promising! Yes let's open a PR, and replace FA3:
We can also remove the gating it seems then for the input sizes.
But let's please gate on BF16 inputs and emit a warning once there for blackwell.
jmkuebler commentedon Sep 18, 2026
Ciould you please also check whether FA4 support split KV on Hopper?
tolleybot commentedon Sep 18, 2026
No, FA4 has no split-KV on Hopper (sm_100/110 only, and nothing open upstream for sm_90). What that means for performance is in #1291, which is now open.
tolleybot commentedon Sep 22, 2026
Correcting a number from my results comment above. It said:
Both halves of that second claim were wrong. The 2–9% came from runs at batch ≥ 10 or with Q above FA3's split threshold, not the batch-1, few-rows shape that cached prediction actually uses. At the right shape (batch 1, fp16, MQA 8/1, head_dim 64, H100, ms per call):
And it did show whole-model:
fit_with_cachethen predict 16–256 rows gave FA3 23–25 ms against FA4's 90 ms at 300k rows.Fixed rather than documented. FA4 returns the LSE, so split-KV can be done outside the kernel. After that change, same benchmark:
The remaining gap is FA4's per-call dispatch overhead. The path is inference-only; under autograd FA4 runs unsplit with a backward test against SDPA. Details and the implementation are in #1291.