20231088/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

ROCm: Allow setting compilation target (#2581)

Browse Source
This commit is contained in:
Rasmus Larsen 2024-01-29 19:52:31 +01:00 committed by GitHub
parent 1b20639a43
commit ea8489fce2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
setup.py
View File

@ -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()