Replace the FA3 attention backend with FlashAttention-4 (Hopper and Blackwell)

[tolleybot]

Issue

Closes #1235.

Motivation and Context

fa3_backend.py is Hopper-only and noted that Blackwell would need FlashAttention-4. FA4 is
now on PyPI (flash-attn-4, CuTeDSL, weekly betas) with kernels for sm_90 and sm_100+, so one
backend can serve both. Per the discussion on #1235, this PR replaces FA3 with FA4 rather than
adding it alongside.

What changes for a user:

• Hopper: FA4 matches or beats FA3 at every measured point (4–10% whole-model from 10k rows
  up, tie below) and, unlike FA3, is pip install "tabpfn[fa4]" rather than a 30–60 minute
  source build.
• Blackwell: attention dispatch to a fused kernel for the first time. 1.12–1.26× on
  predict() at ≥30k rows with ≥100 features.
• No sequence-length threshold. FA3's _FA3_MIN_SEQLEN_FOR_SPEEDUP = 10_000 existed because
  its fixed num_splits=32 made it 3–6× slower than SDPA at 1k rows. FA4 has no such penalty
  on either architecture (within noise from 100 rows, ahead from ~3k), so is_preferred is
  eligibility alone.
• Wider head-dim eligibility. FA4 accepts any multiple of 8 (8–256 on sm_90, 8–128 on sm_100),
  so it also takes the head_dim-16 feature-attention stages FA3 could not. On Blackwell this is
  where all the gain is: torch's SDPA picks cuDNN attention there, which already runs the
  head_dim-64 ICL calls at FA4 speed.
• bf16 stays on SDPA on Blackwell, with a one-time UserWarning. FA4's bf16 forward on sm_100
  is 18–23% slower than SDPA from ~10k rows (fp16 is not; Hopper is fine in both). TabPFN
  defaults to fp16 autocast on CUDA, so this only affects inference_precision=torch.bfloat16.

Measurements (H100 and GB200, torch 2.14.0+cu130, flash-attn-4 4.0.0b30) are in
#1235 and
summarised in fa4_setup.md.

Split-KV on Hopper (asked on the issue): FA4 implements it on sm_100/110 only, so on sm_90 a short-Q call against a long KV at batch 1, the cached-prediction shape, fills only a handful of SMs: 9.9 ms for (1, 16, 8, 64) × (1, 1M, 1, 64) against 0.4 ms for FA3's 32-way split. The backend now splits outside the kernel on such calls (_split_kv_plan, _fa4_split_kv): KV chunks folded into the batch dimension for one launch, partial outputs combined with the LSE FA4 returns. Cached prediction (fit_with_cache, then predict 16–256 rows) on H100: 27–28 ms at 100k and 300k rows, against 23–25 ms with FA3 and 37–92 ms with SDPA or unsplit FA4.

Implementation

• shared/fa4_backend.py (new): same shape as fa3_backend.py. Lazy cached import of
  flash_attn.cute.flash_attn_func so is_available() is one cheap check at registration;
  ROCm guard; per-architecture head-dim table; bf16 gate on cc ≥ 10. Three FA4 4.0.0b30
  specifics it handles: flash_attn_func returns (out, lse) unconditionally; split-KV is
  sm_100/110 only (sm_12x accepts num_splits=1 only); the kernel launches one grid entry per
  batch element, so batches above 65,535 are chunked.
• shared/scaled_dot_product_attention.py: consult order (FA4, torch-MPS, MLX).
• shared/fa3_backend.py, shared/fa3_setup.md: removed.
• shared/fa4_setup.md: install, dispatch conditions, measured speedups on both GPUs, bf16 rule.
• pyproject.toml: fa4 and fa4-cu13 extras (flash-attn-4>=4.0.0b30; the pre-release floor
  lets pip resolve the beta without --pre; [cu13] selects the CUDA 13 CuTeDSL libs). New
  blackwell pytest marker; hopper now means sm_90 + flash-attn-4.
• Tests: FA4 twins of the FA3 dispatch and equivalence tests, plus long-KV cross-attention,
  batch > 65,535, is_preferred == eligibility, and the bf16 gate/warning with a mocked compute
  capability so it runs anywhere. FA3-only tests removed.

torch.compile: FA4's flash_attn_func is not Dynamo-traceable (FA3's was), so with enable_torch_compile=True each FA4 call is a graph break inside flash_attn/cute/. Measured on v3 predict() on H100: within noise at 30k rows, ~4% at 100k, with FA4 still 1.17× faster than SDPA under compile. enable_torch_compile is off by default.

One thing found on the way, unrelated to this change: FA3's split-KV path crashed the process at
batch ≥ 65536 for any num_splits > 1 (CUDA grid limit), which fa3_attn_func used for
seq_q ≤ 1024 since #1168, so test__fa3_batch_heads_above_cuda_max_grid had been failing
since then without CI noticing. Moot once FA3 is gone.

---

Public API Changes

• No Public API changes
• Yes, Public API changes (Details below)

fa3_backend was an internal module. Dispatch stays automatic and falls back to SDPA when
flash-attn-4 is not installed; there is no user-facing switch. Installing FA3 no longer has any
effect.

---

How Has This Been Tested?

• tests/test_architectures/test_attention_backends.py,
  tests/test_architectures/test_shared/test_attention_backends.py,
  tests/test_architectures/test_attention_dispatch.py: 27 pass on an H100 (sm_90) and on a
  GB200 (sm_100), Python 3.11 / 3.12, torch 2.14.0+cu130, flash-attn-4 4.0.0b30.
• FA4 matches SDPA within the FA3 tolerance (atol=rtol=5e-3) on the six canonical shapes ×
  fp16/bf16, on (1, 256, 100k, 8, 1, 64), and on batch=70,000, on both architectures.
• On the GB200 the bf16 gate was checked against the real device: bf16 declined with one
  warning, fp16 preferred at every sequence length.
• Whole-model and kernel-level benchmarks on both GPUs: see FlashAttention-4 backend for Blackwell #1235. Cached prediction with few test rows measured separately (table above).
• ruff check / ruff format clean at the pinned 0.15.12; mypy clean on the changed files.

---

Checklist

• The changes have been tested locally.
• Documentation has been updated (if the public API or usage changes).
• A changelog entry has been added (see changelog/README.md), or "no changelog needed" label requested.
• The code follows the project's style guidelines.
• I have considered the impact of these changes on the public API.

[CLAassistant]

All committers have signed the CLA.

[cursor]

Cursor Bugbot has reviewed your changes using high effort and found 3 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want fixes drafted automatically? Bugbot Autofix can create code changes for findings. A team admin can enable Autofix in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 1e4ce28. Configure here.}

[tolleybot]

Thanks. All three addressed in fd3bb0d: the warning is gated on torch.compiler.is_compiling() (confirmed it broke the graph under fullgraph=True; a test now pins that), changelog fragments cut to one or two sentences, module docstring rewritten to describe FA4 as it is.

[jmkuebler]

Thank you for starting this contribution! I will also run a few test on my end.

Split-KV on Hopper (asked on the issue): not supported by FA4. The sm_90 forward asserts
not is_split_kv in b30, in b31 (current), and on main
(interface.py:1319);
only sm_100/110 implement it, and there is no open issue or PR for sm_90. In the sweep this costs
2–9% against FA3 on the short-Q / long-KV shape at n_train=1M only; FA4 is ahead below that
(1.10× at 100k) and it does not show whole-model. _num_splits_for already asks FA4's heuristic
on sm_100/110, so if sm_90 support lands it is a one-line change.

could you please share the numbers for this experiment? in particular which "short-Q" ranges did you check?

[tolleybot]

Good question, I see an issue in what I measured. The 2–9% figure was wrong for this case: it came from runs at batch ≥ 10 or with Q above FA3's split threshold, not the batch-1, few-rows shape of cached prediction. Measured properly, batch 1, fp16, MQA 8/1, head_dim 64, H100, ms per call:

q	KV	SDPA	FA4	FA3 split-32
16	100k	1.07	1.08	0.09
256	100k	1.07	1.06	0.18
1024	100k	1.07	1.06	0.51
16	1M	9.76	9.90	0.38
256	1M	9.31	9.93	1.10
1024	1M	9.92	9.95	3.95
4096	1M	20.3	20.7	16.1

End-to-end, fit_with_cache then predict 16–256 rows: FA3 23–25 ms, FA4 38 ms at 100k rows and 90 ms at 300k, same as SDPA.

FA4 returns the LSE, so split-KV can be done outside the kernel; a prototype recovers most of the gap (0.43 ms vs FA3's 0.39 at q=16, KV=1M). I'll add that to the backend and update the PR.

[tolleybot]

Done: split-KV outside the kernel for short-Q / long-KV calls on sm_90 (bbb5860, refined in 54d99f3). Cached prediction, fit_with_cache then predict 16–256 rows, H100:

n_train	SDPA	FA4 before	FA4 now	FA3
100k	37–39 ms	38–39	27–28	23
300k	89–92 ms	90–91	28	23–25

The remaining gap to FA3 is FA4's per-call dispatch overhead. Matches SDPA to 6e-5; the path is inference-only, and under autograd FA4 runs unsplit with a backward test against SDPA.

One other change since: the backend now dispatches only on sm_90 and sm_100, the architectures it was measured on. sm_110 and consumer Blackwell (sm_12x) stay on SDPA until someone measures them.

[jmkuebler]

Hi @tolleybot thanks foir your changes. I will try to land #1312 first. Then I would rebenchmark and potentially adjust the splitting. So it will be a bit of time, but hope that is ok.

[tolleybot]

Thanks, and no rush. Since #1312 landed I rebased onto it and re-ran the cached-prediction case at the default n_estimators=8. The batching exposed a problem in my splitting: the long-KV calls now arrive at batch 4, and at batch > 1 I only split when n_train had a divisor ≤ 32, so for an n_train like 99,991 or 300,007 FA4 ran unsplit, at SDPA speed. I've pushed a fix: uneven chunks at batch > 1 go through FA4's varlen path, and splits under 3 are skipped.

fit_with_cache, predict, n_estimators=8, H100, ms:

n_train	n_test	SDPA	FA3	FA4 before	FA4 now
100,000	16	96.9	57.1	67.5	68.5
99,991	16	100.5	59.1	98.1	70.3
99,991	256	102.5	78.7	100.6	78.7
300,007	16	228.8	81.7	223.3	84.4
300,007	256	233.6	139.7	226.6	147.1

Happy for you to adjust the splitting from here as you planned.

[tolleybot]

Quick update since v9.1/v9.2 landed: I rebased locally onto current main (ffd5371, including #1365) and re-ran the cached-prediction benchmark at n_estimators=8. FA4 is within ±3% of the 09-29 numbers above, the split plan is unchanged at every shape I swept, and the attention tests pass. The PR still merges cleanly as is, so I haven't pushed the rebase, but I can if you'd prefer to review on current main.