Added benchmarking for new torchao low precision attention api

[howardzhang-cv]

Stack from ghstack (oldest at bottom):

• -> Added benchmarking for new torchao low precision attention api #3865
• Added new API for low precision fp8 attention using FA3 #3857

Summary

• Added new benchmark for new low precision attention API (FP8 FA3 backend)
• Compares LPIPS and runtime against high-precision model
• uses flux.1-schnell model, 4 inference steps, DrawBench prompts
• has options to control number of prompts, torch.compile usage, warmup_iters, using debug prompts, number of inference steps
• Following the guidelines of add performance and accuracy eval of flux-1.schnell #3502
• Can configure the recipe of the run through changing config in top of file

Example Run

python benchmarks/attention/eval_flux_model.py

Results

Config	LPIPS	Runtime	Speedup
w/out compile	0.224	1.04s	1.06x
w/ compile	0.226	0.77s	1.12x

• Tested w/ and w/out torch.compile, using FP8 FA3 backend (no fused rope, no hadamard)
• LPIPS and Speedup are compared to high-precision model that is either w/ or w/out torch.compile (to keep it fair)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

As of commit fe36423 with merge base d4c1ba3 ():

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.