Add FA4 fp8 backend to low precision attention api

[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 RoPE fusion compile path for FA4 FP8 low-precision attention (fuse_rope=True), mirroring the FA3 RoPE fusion design
• New elementary block: fp8_fa4_rope_sdpa — fused RoPE + FP8 quantization + low-precision SDPA using the FA4 backend
• FA4-specific custom op registration and compile_with_fp8_fusion entry point, reusing the shared FX graph fusion infrastructure (fusion_utils.py, custom_ops.py)
• Reuses shared Triton quantization kernels and RoPE fusion pass — no new kernels needed, only FA4-specific wiring
• FA4 supports both Hopper (SM 9.x) and Blackwell (SM 10.x) hardware
• Added RoPE SDPA numerical accuracy tests and fuse_rope parametrization for the FA4 backend

New Files

• fp8_fa4/fusion_pass.py: FA4-specific custom op registration, rope_sdpa_fusion_pass, and compile_with_fp8_fusion entry point

Modified Files

• fp8_fa4/attention.py: Added fp8_fa4_rope_sdpa elementary block
• fp8_fa4/init.py: Added fp8_fa4_rope_sdpa export
• fp8_fa4/setup.py: Replaced compile placeholder with real compile_with_fp8_fusion
• test_fp8_attention.py: Wired up rope_sdpa_fn=fp8_fa4_rope_sdpa in FA4 backend config

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()

  # Compile path with RoPE fusion using FA4
  config = LowPrecisionAttentionConfig(
      backend=AttentionBackend.FP8_FA4,
      fuse_rope=True,
  )
  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. Uses RoPE fusion.
Perplexity: 6.19 -> 6.24

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Unrelated Failure

As of commit 97eb634 with merge base 5ebd10d ():

NEW FAILURES - The following jobs have failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t bc268fed67f4d901941336721edfa97449f8e5010cd2253bad12dec5fe7b879b /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_fusion.py::TestDynamicPatternMatcher::test_q_attention_block

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