Add cross attention support for low precision attention API

[howardzhang-cv]

Stack from ghstack (oldest at bottom):

• -> Add cross attention support for low precision attention API #4142

Summary

• Added support for cross attention, different Q and KV sequence length in the low precision attention API

Results

Wan2.1-T2V-1.3B-Diffusers with 480x832 frame size and 81 frames

Config	Median Time (s)	Speedup
bf16 baseline	100.85	1.00x
fp8_attn	67.15	1.50x
bf16 baseline + torch.compile	81.32	1.00x
fp8_attn + compile	46.88	1.73x

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 84473e5 with merge base 34322b5 ():

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/test_low_bit_optim.py::TestFSDP2::test_fsdp2

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

[drisspg]

nit: small test