pytorch - 💡(How to fix) Fix FP8 Per-Head Quantized Scaled Dot Product Attention [1 participants]

New Feature for Release

A new prototype op for fp8 per-head quantized scaled dot product attention using the flash attention 3 backend. Prototype op can be called directly from torch.nn.attention.experimental._scaled_dot_product_attention_quantized. However, more detailed e2e usage guides for how to build off of these primitives are in torchao (check torchao 0.17.0 release)

Point(s) of contact

howardzhang-cv

Release Mode (pytorch/pytorch features only)

In-tree

Out-Of-Tree Repo

No response

Description and value to the user

Provide prototype ops for calling the fp8 per-head quantized flash attention 3 backend. Initial tests have shown good speedup for minimal quality degradation for image generation, video generation, and LLMs.

Link to design doc, GitHub issues, past submissions, etc

No response

What feedback adopters have provided

No response

Plan for documentations / tutorials

Tutorial exists

Additional context for tutorials

Pytorch op can be used by itself, in which case tutorial is not needed, since it's a bare op. Users are encouraged to look at the torchao wrapper for this op, which includes e2e model wrapping and direct SDPA replacements (with triton per-head quantization and rope-fused quantization variants). Tutorial can be found in the 0.17.0 torchao release notes.

Marketing/Blog Coverage

Yes

Are you requesting other marketing assistance with this feature?

No response

Release Version

No response

OS / Platform / Compute Coverage

For CUDA SM90 Hopper architectures

Testing Support (CI, test cases, etc..)

Integration tests located in test/nn/attention/test_fa3.py. Additional E2E tests have been done. Results below: <img width="1284" height="735" alt="Image" src="https://github.com/user-attachments/assets/4432dc4d-18e8-4bea-bd79-e80c739000c4" /> Single layer attention tests show fp8 is faster than bf16 past 8192 sequence length (including the time for per-head quantization).

<img width="1266" height="560" alt="Image" src="https://github.com/user-attachments/assets/ac395ef9-87dd-4792-bc99-0f02c127d374" /> Tests with llama3 prefill show fp8 is faster than bf16 past 4096 sequence length. These tests are shown with fusion of the RoPE and per-head quantization computations, which improves performance in e2e model inference (see torchao documentation for more information). Results show minimal quality degradation (perplexity 7.54 -> 7.62).

Additional tests with Wan2.1-14B show 1.18x speedup in video generation (1080p, 81 frames). Naive application leads to harsh quality drops. Suggested to skip quantization on the first and last few layers for good runtime/quality tradeoff.

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel