From 4b9971e40183400f10561e8a00cf62d858a40770 Mon Sep 17 00:00:00 2001 From: sogalin <39478626+sogalin@users.noreply.github.com> Date: Fri, 13 Jun 2025 02:07:48 +0800 Subject: [PATCH] Add gfx950 support for sgl-kernel. (#7092) Co-authored-by: HAI Co-authored-by: Yineng Zhang --- sgl-kernel/setup_rocm.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index b230abfa8..e9db969d3 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -17,6 +17,7 @@ import platform import sys from pathlib import Path +import torch from setuptools import find_packages, setup from torch.utils.cpp_extension import BuildExtension, CUDAExtension @@ -49,6 +50,13 @@ cxx_flags = ["-O3"] libraries = ["hiprtc", "amdhip64", "c10", "torch", "torch_python"] extra_link_args = ["-Wl,-rpath,$ORIGIN/../../torch/lib", f"-L/usr/lib/{arch}-linux-gnu"] +amdgpu_target = torch.cuda.get_device_properties("cuda").gcnArchName.split(":")[0] +if amdgpu_target not in ["gfx942", "gfx950"]: + print( + f"Warning: Unsupported GPU architecture detected '{amdgpu_target}'. Expected 'gfx942' or 'gfx950'." + ) + sys.exit(1) + hipcc_flags = [ "-DNDEBUG", f"-DOPERATOR_NAMESPACE={operator_namespace}", @@ -57,7 +65,7 @@ hipcc_flags = [ "-fPIC", "-std=c++17", "-D__HIP_PLATFORM_AMD__=1", - "--amdgpu-target=gfx942", + f"--amdgpu-target={amdgpu_target}", "-DENABLE_BF16", "-DENABLE_FP8", ]