Rocm fix

[Aman-Dwivedi]

Added amd gpu support.
Updated requirements for rocm and added functions in setup.py to detect amd gpus. An example script has also been added by yiakwy-xpu-ml-framework-team.

[avjves]

Hi! This makes a custom requirements file for ROCm specifically - is there a reason for it? Also, that yunchang branch / version it installs is a year old with some changes. Yunchang already supports AMD GPUs in the upstream repo via flash_attn or AITER (the latest way to call FA with AMD GPUs). This seems like a regression.

Also, this breaks the changes made by PR #559 due to the duplicate imports in xfuser/core/long_ctx_attention/ring/ring_flash_attn.py. Currently it's gated like this:

xDiT/xfuser/core/long_ctx_attention/ring/ring_flash_attn.py

Line 18 in cd06115

try:

[feifeibear]

LGTM

[feifeibear]

@Aman-Dwivedi could you please check the duplicate imports problem mentioned before?

[avjves]

This line in the requirements:

yunchang @ git+https://github.com/yiakwy-xpu-ml-framework-team/xDiT-long-context-attention-fork.git@add_amd_gpu_suppport

is also a big problem in general for AMD GPUs. Could it also be removed? 😄

[eppaneamd]

@feifeibear kindly note that this PR should be revisited and its merits re-evaluated.

@Aman-Dwivedi could you elaborate why this PR is needed, how xDiT and yunchang is not working for AMD GPUs currently? Why gfx942 is the only allowed gpu arch? Have you tried newer images than rocm/pytorch:rocm6.2_ubuntu22.04_py3.10_pytorch_release_2.3.0, such as rocm7.0_ubuntu22.04_py3.10_pytorch_release_2.8.0?