diff --git a/setup.py b/setup.py index b552cb67cb363..8fad433f98b09 100644 --- a/setup.py +++ b/setup.py @@ -287,11 +287,15 @@ if _is_cuda(): }, )) elif _is_hip(): - amd_arch = get_amdgpu_offload_arch() - if amd_arch not in ROCM_SUPPORTED_ARCHS: - raise RuntimeError( - f"Only the following arch is supported: {ROCM_SUPPORTED_ARCHS}" - f"amdgpu_arch_found: {amd_arch}") + amd_archs = os.getenv("GPU_ARCHS") + if amd_archs is None: + amd_archs = get_amdgpu_offload_arch() + for arch in amd_archs.split(";"): + if arch not in ROCM_SUPPORTED_ARCHS: + raise RuntimeError( + f"Only the following arch is supported: {ROCM_SUPPORTED_ARCHS}" + f"amdgpu_arch_found: {arch}") + NVCC_FLAGS += [f"--offload-arch={arch}"] elif _is_neuron(): neuronxcc_version = get_neuronxcc_version()