Add support for AITER attention backend

[lauri9]

What does this PR do?

AITER is AMD’s centralized repository to support high performance AI operators such as attention kernels for AMD ROCm enabled accelerators. This PR adds support for FlashAttention through AITER by introducing a new attention backend.

Test code for Flux inference below. Requires installation of aiter>=0.15.0 and a supported ROCm enabled accelerator.

import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, attention_backend

model_id = "black-forest-labs/FLUX.1-dev"
transformer = FluxTransformer2DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16, device_map="cuda")
transformer.set_attention_backend("aiter")
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16)
pipe.text_encoder.to("cuda")
pipe.text_encoder_2.to("cuda")
pipe.vae.to("cuda")

prompt = "A cat holding a sign that says 'hello world'"

image = pipe(prompt, num_inference_steps=28, guidance_scale=4.0).images[0]
image.save("output.png")

We are interested in following up this PR by eventually also enabling support for context parallelism across multiple devices.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[avjves]

Should we have _check_device_rocm? Not sure if one can accidentally install aiter on NV, but AFAIK nothing here would then prevent from trying to run it.

[lauri9]

Hmm, perhaps we should change this to _check_device_cuda instead to guarantee that the tensors are on an accelerator device. CPU is not meaningful I suppose.

I tried installing aiter in a nvcr container, but it gives an error:

No ROCm runtime is found, using ROCM_HOME='None'
Traceback (most recent call last):
  File "/diffusers_aiter_backend/aiter/aiter/jit/utils/cpp_extension.py", line 90, in get_hip_version
    hipconfig = executable_path("hipconfig")
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/diffusers_aiter_backend/aiter/aiter/jit/utils/cpp_extension.py", line 83, in executable_path
    path is not None
AssertionError: Could not find hipconfig in PATH or ROCM_HOME(None)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/diffusers_aiter_backend/aiter/setup.py", line 16, in <module>
    from jit import core
  File "/diffusers_aiter_backend/aiter/aiter/jit/core.py", line 23, in <module>
    from chip_info import get_gfx
  File "/diffusers_aiter_backend/aiter/aiter/jit/utils/chip_info.py", line 9, in <module>
    from cpp_extension import executable_path
  File "/diffusers_aiter_backend/aiter/aiter/jit/utils/cpp_extension.py", line 173, in <module>
    HIP_VERSION = get_hip_version()
                  ^^^^^^^^^^^^^^^^^
  File "/diffusers_aiter_backend/aiter/aiter/jit/utils/cpp_extension.py", line 94, in get_hip_version
    raise RuntimeError("ROCm version file not found")
RuntimeError: ROCm version file not found

IMO, this is a corner case that is probably not worth adding a separate function to check for - a user that has managed to install ROCm and CUDA on the same platform but only has one type of accelerator has already taken many wrong turns 😁

[avjves]

Yeah, makes sense 😆

[avjves]

Was there some way to not run these on NV devices? I don't see any flags here to not run them

[lauri9]

Running without aiter installed will lead to xfail on these tests. Trying on MI355X env without aiter skips all tests except native (which fails probably due to the brittleness of the tests; comparing numerical values) and cudnn (which fails because kernel is not available).

=============================================================================================================================== short test summary info ================================================================================================================================
XFAIL tests/others/test_attention_backends.py::test_forward[flash_hub] - Backend 'flash_hub' not supported in this environment.
XFAIL tests/others/test_attention_backends.py::test_forward[_flash_3_hub] - Backend '_flash_3_hub' not supported in this environment.
XFAIL tests/others/test_attention_backends.py::test_forward[aiter] - Backend 'aiter' not supported in this environment.
XFAIL tests/others/test_attention_backends.py::test_forward_with_compile[flash_hub] - Backend 'flash_hub' not supported in this environment.
XFAIL tests/others/test_attention_backends.py::test_forward_with_compile[_flash_3_hub] - Backend '_flash_3_hub' not supported in this environment.
XFAIL tests/others/test_attention_backends.py::test_forward_with_compile[aiter] - Backend 'aiter' not supported in this environment.
FAILED tests/others/test_attention_backends.py::test_forward[native] - assert False
FAILED tests/others/test_attention_backends.py::test_forward[_native_cudnn] - RuntimeError: No available kernel. Aborting execution.
FAILED tests/others/test_attention_backends.py::test_forward_with_compile[native] - assert False
FAILED tests/others/test_attention_backends.py::test_forward_with_compile[_native_cudnn] - torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function <built-in function scaled_dot_product_attention>(*(), **{'query': FakeTensor(..., device='cuda:0', size=(1, 24, 384, 128), dtype=torch.bfloat16), 'key': FakeTensor(..., device=...
====================================================================================================================== 4 failed, 6 xfailed, 17 warnings in 45.12s ======================================================================================================================

[avjves]

Ah, I see, then that's not an issue :)

[avjves]

Left a couple of questions, but overall LGTM!