[torch] Flip --enable-pytorch-flash-attention-windows for releases.

[ScottTodd]

Motivation

Fixes #1040, enabling aotriton for flash attention in pytorch (if it works). This is expected to improve performance in workloads like ComfyUI image generation by upwards of 60% (e.g. 12.6 it/s to 20.0 it/s).

Technical Details

Follow-up to #1432 and depends on pytorch/pytorch#162330.

Note that support is experimental for some GPUs like gfx1100, so the TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 environment variable may be needed to try aotriton on those systems.

Test Plan

Trigger either https://github.com/ROCm/TheRock/actions/workflows/build_windows_pytorch_wheels.yml or https://github.com/ROCm/TheRock/actions/workflows/release_windows_pytorch_wheels.yml across the matrix of GPU families once that PyTorch PR is merged.

We're still going to need automated tests and documentation for this. I'd like numerics tests running somewhere and documentation that shows how to check which pytorch features are enabled in the wheels that a user installs.

Test Result

Test runs:

• https://github.com/ROCm/TheRock/actions/runs/17660396787 using this branch and 7.0.0rc20250908 for gfx110X-dgpu
• https://github.com/ROCm/TheRock/actions/runs/17660456285 using the branch and 7.0.0rc20250908 for gfx1151
• https://github.com/ROCm/TheRock/actions/runs/17662170140 using the branch and 7.0.0rc20250908 for gfx1151
  • Tests not running should be fixed with Use cloudfront_staging_url for Windows pytorch testing workflow. #1469

(may need to retrigger to pick up fixes for flaky checkouts)

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Nem404]

Eyy "jammm self-requested a review" that's what I'm talkin' 'bout 🔥

Now we just wait for Jeff in pytorch/pytorch#162330 👀

[Nem404]

pytorch/pytorch#162330 (comment) 🎉

[ScottTodd]

pytorch/pytorch@62843c1 is merged. Triggering some test runs on this.

[ScottTodd]

Test runs built successfully. Sanity check tests did not run for unrelated reasons. I'd say this is ready for review/merge now.

[Nem404]

Waiting for @jammm's approval - it looks like tomorrow's wheels can include aotriton 💯

[jammm]

Already approved :)

[Nem404]

Already approved :)

Oh, then for Scott to merge :D

[0xDELUXA]

Just wanted to share that I got my hands on a PyTorch wheel built with AOTriton.

On gfx1200, we also need to set the TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 env var - otherwise it only shows a warning and doesn’t use SDPA.

But that’s not really an issue, because once we set it, performance basically doubles. It's really awesome.