[FEATURE] Enhance platform compatibility for ARM (#5746)

sgl-kernel/setup_rocm.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import platform
 import sys
 from pathlib import Path
 
@@ -20,6 +21,7 @@ from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
 root = Path(__file__).parent.resolve()
+arch = platform.machine().lower()
 
 
 def _get_version():
@@ -45,7 +47,7 @@ sources = [
 
 cxx_flags = ["-O3"]
 libraries = ["hiprtc", "amdhip64", "c10", "torch", "torch_python"]
-extra_link_args = ["-Wl,-rpath,$ORIGIN/../../torch/lib", "-L/usr/lib/x86_64-linux-gnu"]
+extra_link_args = ["-Wl,-rpath,$ORIGIN/../../torch/lib", f"-L/usr/lib/{arch}-linux-gnu"]
 
 hipcc_flags = [
     "-DNDEBUG",