Skip to content

Commit

Permalink
ROCm: Allow setting compilation target (vllm-project#2581)
Browse files Browse the repository at this point in the history
  • Loading branch information
rlrs authored and NikolaBorisov committed Jan 31, 2024
1 parent 95521ce commit e453178
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -287,11 +287,15 @@ def get_torch_arch_list() -> Set[str]:
},
))
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()
Expand Down

0 comments on commit e453178

Please sign in to comment.