Add FA4 monkey-patch path for low-precision attention

[howardzhang-cv]

Stack from ghstack (oldest at bottom):

• Add FA4 fp8 backend to low precision attention api #3947
• -> Add FA4 monkey-patch path for low-precision attention #3960
• Added prototype low precision attention API to the docs #4056
• soften version guard check for low precision attention API #4058
• remove rope fusion option, do automatically on torch compile #4055

Summary

• Added FA4 FP8 low-precision attention with simple SDPA replacement path, mirroring the FA3 design
• New elementary block: fp8_fa4_sdpa — a direct drop-in replacement for F.scaled_dot_product_attention using the FA4 backend. Reuses the shared FP8 quantization kernels.
• Simple wrapper support via apply_low_precision_attention with AttentionBackend.FP8_FA4 — no torch.compile required.
• FA4 supports both Hopper (SM 9.x) and Blackwell (SM 10.x) hardware
• Added _is_blackwell() and _is_fa4_available() hardware detection utilities
• Added FA4 backend config and numerical accuracy tests (eager SDPA and model-level API)

New Files

• fp8_fa4/init.py: Exports fp8_fa4_sdpa
• fp8_fa4/attention.py: fp8_fa4_sdpa elementary block
• fp8_fa4/setup.py: Thin wrapper calling setup_fp8_backend with FA4 parameters

Modified Files

• config.py: Added FP8_FA4 to AttentionBackend enum
• utils.py: Added _is_blackwell(), _is_fa4_available(), FA4 support in _get_available_backend() and _check_backend_available()
• api.py: Added FA4 dispatch path
• test_fp8_attention.py: Added FA4 backend config, numerical accuracy tests for FA4

Test Plan

python -m pytest test/prototype/attention/test_fp8_attention.py -v

Example Usage

  from torchao.prototype.attention import (
      AttentionBackend,
      LowPrecisionAttentionConfig,
      apply_low_precision_attention,
  )

  model = MyModel()

  # Simple SDPA replacement using FA4 — no torch.compile needed
  config = LowPrecisionAttentionConfig(backend=AttentionBackend.FP8_FA4)
  model = apply_low_precision_attention(model, config)

  # Flash activation is handled internally by the wrapper
  output = model(inputs)

Results

Single-Layer Results

Results directly comparing FA4 SDPA versus FA4 fp8 SDPA (including quantization time):

Llama3 Model Results

Results comparing Llama3 model with FA4 SDPA versus Llama3 using the FA4 fp8 wrapper. Does not use RoPE fusion.
Perplexity: 6.19 -> 6.25

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3960

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 8 Pending

As of commit 6611347 with merge base 5ebd10d ():

NEW FAILURE - The following job has failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[howardzhang-cv]

So this FA4 fp8 low precision backend currently has some issues randomly with the Llama 3 model. For some reason, we get NaNs in the quantized QKV tensors (specifically in torchao/prototype/attention/shared_utils/attention.py in the _fp8_sdpa function, the q_fp8 tensor seems to be the first to become NaN). This issue goes away when we replace the triton kernel with simple PyTorch ops. It also goes away with --compile, and only happens on Blackwell for some reason. Not sure why this happens, I spent a bunch of time trying to fix this, but couldn't find the heart of the issue.