[ROCm] Fixup arch checks for ROCM

[dllehr-amd]

The ROCM stack with PyTorch supports a wide set of gfx architectures. This can be displayed by printing PYTORCH_ROCM_ARCH env. In the absence of PYTORCH_ROCM_ARCH pytorch uses theoutput from rocm_agent_enumerator to choose what to compile for.

vllm supports a subset of these, (gfx908, gfx90a,...)

Due to a need to potentially support multiple architectures at once (ex. docker image) it's important to make sure vllm is compiled with them all unless specified otherwise.

We now gather either the PYTORCH_ROCM_ARCH env or rocm_agent_enumerator output and cross reference with ROCM_SUPPORTED_ARCHS from vllm to generate a list of arches to build for.

[WoosukKwon]

Hi @dllehr-amd, thanks for submitting the PR! This is a bit confusing. It seems like there are 3 environment variables:

1. PyTorch uses TORCH_CUDA_ARCH_LIST env var for specifying the target NVIDIA GPUs.
2. PyTorch w/ ROCm backend uses PYTORCH_ROCM_ARCH env var for specifying the target AMD GPUs.
3. The ROCm flash-attn repo uses GPU_ARCHS env var for specifying the target AMD GPUs.

And here we are using PYTORCH_ROCM_ARCH . Is my understanding correct?

[dllehr-amd]

Hi @WoosukKwon. I made a change to vllm arch's to match the ones supported by our Flash Attention(GPU_ARCHS) As we'd be hard pressed to guarantee support on the other gfx's. There are still your three variables, but ROCM_SUPPORTED_ARCHS is reduced to the amount that Flash Attention supports.

The way this PR works now, is it'll take all of the architectures PyTorch built with, and cross reference with what vllm supports, and only build for those.

As of this PR, if you build it will only build for gfx90a and gfx942. Without any additional input from a user.

Does that make things any more clear?

Thanks!

[WoosukKwon]

LGTM! Thanks for the fix! BTW, it'd be nicer if we can have a CI/CD pipeline for AMD GPUs :)

[jamestwhedbee]

I made a change to vllm arch's to match the ones supported by our Flash Attention(GPU_ARCHS) As we'd be hard pressed to guarantee support on the other gfx's.

@WoosukKwon @dllehr-amd just so folks know, since this PR, the ROCm Flash Attention officially supports gfx_908 (it always did, it just wasn't documented until this PR)

[jamestwhedbee]

I've opened #2792