Enable FA4 Flex attention backend

[fynnsu]

Purpose

Pytorch added support for an Flash Attention 4 backend to Flex Attention which has much better performance on Hopper and Blackwell gpus: https://pytorch.org/blog/flexattention-flashattention-4-fast-and-flexible/

Requires manually installing fa4 first:

uv pip install flash-attn-4 --prerelease allow

or

uv pip install "flash-attn-4[cu13]" --prerelease allow

Note: this is experimental.

By default, if the flash attention package is available and we're using a Hopper or newer gpu, this backend will get used. You can also explicitly enable/disable it with --fa4 on or --fa4 off

Tests

I've added some explicit tests for FA4 on the Hopper runners. I've also updated the buildkite pipeline to install FA4 so that the tests will run.

Checklist

I have filled in:

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan/results, such as providing test command and pasting the results.
• (Optional) The necessary documentation update.
• I (a human) have written or reviewed the code in this pr to the best of my ability.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 4aec2945-b128-4d4a-9837-d5a5d80f9b57

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

FA4 flex-attention backend support

Layer / File(s)	Summary
FA4 setup and runtime control .buildkite/gpu-tests/scripts/run-tests.sh, scripts/train.py	GPU setup conditionally installs flash-attn for Hopper-class GPUs, and training adds the --fa4 mode with startup configuration.
FA4 availability and flex-attention selection src/speculators/models/attention.py	Cached availability detection and configurable auto, on, and off policies select the FLASH backend for flex attention.
CUDA FA4 validation tests/unit/models/test_attention.py	CUDA tests cover basic output, GQA, finite results, and numerical agreement with the Triton backend.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: enabling the FA4 flex attention backend.
Description check	✅ Passed	The description is directly related to the FA4 backend changes and testing updates.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fa4_flex_attention

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[WindChimeRan]

crashed on GB10 Blackwell (maybe it's an upstream bug)

fa4_is_available() returns True on GB10 (sm 12.1 ≥ 9.0 + package present), so --fa4 auto would enable FA4 and then crash at the first masked flex_attention, i.e., it breaks DFlash/Eagle-3 training on GB10 out of the box.

because:

case	result on GB10
FA4, dense (no mask)	runs fwd+bwd
FA4, causal mask	AssertionError: Block sparsity not supported on SM 12.0
FA4, DFlash anchor mask	AssertionError: Block sparsity not supported on SM 12.0

[WindChimeRan]

Could you please run some quick e2e benchmark on, e.g., 100 step speedup on a/b?

[speculatorsbot]

Clean, well-scoped addition. The implementation is minimal — extends the existing flex_attention_forward with an optional kernel_options parameter, adds the CLI plumbing, CI setup, and hardware-gated tests. No design concerns.

Agree with @WindChimeRan's finding that fa4_is_available() would return True on Blackwell GPUs (e.g. GB10, SM 12.x >= 9.0) where FA4's block-sparsity support is incomplete, causing masked flex_attention (causal, anchor) to crash under --fa4 auto. Recommend resolving that before merge.

🤖 Generated with Claude Code using the /pr-review skill

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

• approved-reviews-by = dsikka
• approved-reviews-by = fynnsu
• approved-reviews-by = orestis-z
• approved-reviews-by = rahul-tuli
• approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @fynnsu.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork