From 27c0977b6f2ffd476b85de116e2db839b614d76a Mon Sep 17 00:00:00 2001 From: bouclem Date: Sat, 22 Aug 2026 11:20:43 +0200 Subject: [PATCH] feat: add simple ROCm GPU support for RDNA3 (gfx1100-1103) - Add hip_compat.h shim mapping CUDA runtime API to HIP equivalents - Update pinned_tensor.cpp to compile under both nvcc and hipcc - Add ROCm detection in arch.py (is_rocm, get_rocm_gfx_arch, is_gfx11xx_family) - Guard NVIDIA arch checks to return None on ROCm - Skip nvcc version check in _toolchain.py when on ROCm - Add ROCm build path in setup.py (ROCM_HOME, amdhip64, --offload-arch) - Add _hip_cflags() in kernel/utils.py for JIT compilation on ROCm - Add is_rocm() and driver_hip_version() in backend.py - Add rocm-smi fallback in __main__.py for clangd generation - Add TODO(ROCm) for NCCL->RCCL, flashinfer/sgl_kernel ROCm builds, Triton autotune RDNA3 tuning, PDL equivalent, hiprtc JIT cache - Add AMD ROCm classifier in pyproject.toml --- pyproject.toml | 1 + python/freetoken/kernel/__main__.py | 29 ++-- python/freetoken/kernel/_toolchain.py | 9 +- python/freetoken/kernel/backend.py | 18 +++ .../csrc/include/freetoken/hip_compat.h | 133 ++++++++++++++++++ .../freetoken/kernel/csrc/pinned_tensor.cpp | 2 +- python/freetoken/kernel/pynccl.py | 1 + python/freetoken/kernel/utils.py | 29 +++- python/freetoken/utils/__init__.py | 6 + python/freetoken/utils/arch.py | 31 ++++ setup.py | 49 +++++-- 11 files changed, 280 insertions(+), 28 deletions(-) create mode 100644 python/freetoken/kernel/csrc/include/freetoken/hip_compat.h diff --git a/pyproject.toml b/pyproject.toml index 8bd653f..3c0c45e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Environment :: GPU :: NVIDIA CUDA", + "Environment :: GPU :: AMD ROCm", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", diff --git a/python/freetoken/kernel/__main__.py b/python/freetoken/kernel/__main__.py index 7be541a..5227443 100644 --- a/python/freetoken/kernel/__main__.py +++ b/python/freetoken/kernel/__main__.py @@ -12,21 +12,22 @@ def generate_clangd(): logger = init_logger(__name__) logger.info("Generating .clangd file...") include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE - status = subprocess.run( - args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], - capture_output=True, - check=True, - ) - compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] - major, minor = compute_cap.split(".") + + # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries. + try: + status = subprocess.run( + args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], + capture_output=True, + check=True, + ) + compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] + major, minor = compute_cap.split(".") + arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"] + except (subprocess.CalledProcessError, FileNotFoundError): + # TODO(ROCm): parse gfx target from rocm-smi; default to gfx1100 for now. + arch_flags = ["-xhip", "--offload-arch=gfx1100"] compile_flags = ",\n ".join( - [ - "-xcuda", - f"--cuda-gpu-arch=sm_{major}{minor}", - "-std=c++20", - "-Wall", - "-Wextra", - ] + arch_flags + ["-std=c++20", "-Wall", "-Wextra"] + [f"-isystem{path}" for path in include_paths] ) clangd_content = f""" diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebb..93d0322 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -1,4 +1,4 @@ -"""CUDA toolchain/torch consistency checks. +"""CUDA/HIP toolchain/torch consistency checks. Standalone on purpose: setup.py and the kernel-cache build backend load this file by path, so it must not import the freetoken package. @@ -16,6 +16,11 @@ _TRUE_VALUES = {"1", "true", "yes", "on"} +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + def _nvcc_path() -> str | None: from torch.utils.cpp_extension import CUDA_HOME @@ -49,6 +54,8 @@ def check_nvcc_matches_torch() -> None: nvcc-built binaries link libcudart.so.; at runtime only the torch wheel's own CUDA runtime is guaranteed to be loadable. """ + if _is_rocm(): + return # ROCm uses hipcc, not nvcc if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES: return torch_major = torch_cuda_major() diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 3037ad8..8293e3a 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -42,6 +42,24 @@ def is_triton_kernels_installed() -> bool: return _importable("triton_kernels") +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU).""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def driver_hip_version() -> int | None: + """ROCm driver version, or None if undetermined.""" + # TODO(ROCm): flashinfer/sgl_kernel have no ROCm builds — Triton fallback is used. + try: + from freetoken.kernel.pinned import _load_pinned_extension + return int(_load_pinned_extension().driver_cuda_version()) or None + except Exception: + return None + + @functools.cache def driver_cuda_version() -> int | None: """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0), diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h new file mode 100644 index 0000000..99539d1 --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -0,0 +1,133 @@ +#pragma once + +// HIP compatibility shim: maps CUDA runtime API names to HIP equivalents so +// the same C++ source compiles under both nvcc and hipcc. Include this instead +// of directly when the file needs the runtime API. +// +// On NVIDIA platforms the CUDA headers are included as-is and every macro below +// resolves to the original CUDA symbol, so there is zero overhead. +// +// Supported ROCm targets (RDNA3): +// gfx1100 — RX 7900 XTX / XT +// gfx1101 — RX 7900 GRE +// gfx1102 — RX 7700 / XT +// gfx1103 — RX 7600 / XT + +#ifdef __HIP__ + +// --- HIP runtime headers --- +#include +#include + +// --- API name mapping (CUDA -> HIP) --- +// HIP already defines most cuda* names as macros that expand to hip* equivalents +// via hip_runtime.h, but a few are missing or differ in signature. Define them +// here so call-sites stay unchanged. + +#ifndef cudaSuccess +#define cudaSuccess hipSuccess +#endif + +#ifndef cudaError_t +#define cudaError_t hipError_t +#endif + +#ifndef cudaGetErrorString +#define cudaGetErrorString hipGetErrorString +#endif + +#ifndef cudaGetLastError +#define cudaGetLastError hipGetLastError +#endif + +#ifndef cudaMallocHost +#define cudaMallocHost hipMallocHost +#endif + +#ifndef cudaFreeHost +#define cudaFreeHost hipFreeHost +#endif + +#ifndef cudaHostAlloc +#define cudaHostAlloc hipHostMalloc +#endif + +#ifndef cudaHostRegister +#define cudaHostRegister hipHostRegister +#endif + +#ifndef cudaHostRegisterPortable +#define cudaHostRegisterPortable hipHostRegisterPortable +#endif + +#ifndef cudaHostRegisterMapped +#define cudaHostRegisterMapped hipHostRegisterMapped +#endif + +#ifndef cudaHostGetDevicePointer +#define cudaHostGetDevicePointer hipHostGetDevicePointer +#endif + +#ifndef cudaGetDevice +#define cudaGetDevice hipGetDevice +#endif + +#ifndef cudaDriverGetVersion +#define cudaDriverGetVersion hipDriverGetVersion +#endif + +#ifndef cudaDeviceGetAttribute +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#endif + +#ifndef cudaDevAttrUnifiedAddressing +#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaDevAttrCanUseHostPointerForRegisteredMem +// HIP does not expose this attribute; assume UVA identity on ROCm (true on Linux). +// TODO(ROCm): re-enable proper UVA query if HIP adds this attribute. +#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaFuncSetAttribute +#define cudaFuncSetAttribute hipFuncSetAttribute +#endif + +#ifndef cudaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize +#endif + +#ifndef cudaLaunchKernelEx +// TODO(ROCm): hipLaunchKernelEx is available in newer ROCm; use it when widely shipped. +// For now, fall back to hipLaunchKernel with config extracted manually. +#define cudaLaunchKernelEx hipLaunchKernelEx +#endif + +#ifndef cudaLaunchConfig_t +#define cudaLaunchConfig_t hipLaunchConfig_t +#endif + +#ifndef cudaLaunchAttribute +#define cudaLaunchAttribute hipLaunchAttribute +#endif + +#ifndef cudaLaunchAttributeProgrammaticStreamSerialization +// PDL (Programmatic Dependent Launch) is NVIDIA-specific. +// TODO(ROCm): PDL has no ROCm equivalent — disabled, may affect overlap scheduling latency. +#define cudaLaunchAttributeProgrammaticStreamSerialization 0 +#endif + +#ifndef cudaStream_t +#define cudaStream_t hipStream_t +#endif + +#ifndef dim3 +// HIP already provides dim3; this is a no-op guard. +#endif + +#else // !__HIP__ — NVIDIA CUDA path + +#include + +#endif // __HIP__ diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947ad..9355f57 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,5 +1,5 @@ #include -#include +#include #include namespace { diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea573..71bde73 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/python/freetoken/kernel/pynccl.py @@ -27,6 +27,7 @@ def get_buffer(self) -> int: ... @functools.cache def _load_nccl_module() -> Module: + # TODO(ROCm): NCCL -> RCCL migration for multi-GPU tensor parallelism on AMD. return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"]) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b..dfdcd4c 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -19,9 +19,15 @@ DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] +DEFAULT_HIP_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_LDFLAGS = [] +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + def _cuda_cflags(extra: List[str]) -> List[str]: """CUDA nvcc flags for a kernel build. During the multi-arch AOT cache build, `TVM_FFI_CUDA_ARCH_LIST` (e.g. "8.6 8.9 9.0 10.0 12.0") makes tvm-ffi emit a SASS cubin @@ -40,6 +46,15 @@ def _rank(a: str) -> int: cc = max(arch_list, key=_rank).rstrip("a").replace(".", "") flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"] return flags + + +def _hip_cflags(extra: List[str]) -> List[str]: + """HIP flags for a kernel build on ROCm.""" + # TODO(ROCm): Triton autotune configs need RDNA3-specific tuning (wave count, LDS size). + flags = DEFAULT_HIP_CFLAGS + extra + rocm_arch = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103") + flags = flags + [f"--offload-arch={rocm_arch}"] + return flags CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -217,12 +232,17 @@ def load_aot( cpp_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cpp_files] cuda_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cuda_files] + if _is_rocm(): + cuda_cflags = _hip_cflags(extra_cuda_cflags) + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags) + return load( name, cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), + extra_cuda_cflags=cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, @@ -272,12 +292,17 @@ def load_jit( cuda_sources = [f'#include "{path}"' for path in cuda_paths] cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] + if _is_rocm(): + cuda_cflags = _hip_cflags(extra_cuda_cflags) + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags) + return load_inline( name, cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), + extra_cuda_cflags=cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index 2e4ad15..1348a52 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -1,5 +1,8 @@ from .arch import ( is_arch_supported, + is_rocm, + get_rocm_gfx_arch, + is_gfx11xx_family, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -35,6 +38,9 @@ "load_toolcall_anchor_id", "init_logger", "is_arch_supported", + "is_rocm", + "get_rocm_gfx_arch", + "is_gfx11xx_family", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3..a70c8aa 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -1,14 +1,45 @@ from __future__ import annotations import functools +import os from typing import Tuple +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU) instead of CUDA.""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def get_rocm_gfx_arch() -> str | None: + """The gfx target of the current AMD GPU (e.g. \"gfx1100\"), or None.""" + if not is_rocm(): + return None + # TODO(ROCm): parse rocm-smi for auto-detection; for now rely on env. + for env_var in ("PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"): + val = os.getenv(env_var, "") + for gfx in ("gfx1100", "gfx1101", "gfx1102", "gfx1103"): + if gfx in val: + return gfx + return None + + +@functools.cache +def is_gfx11xx_family() -> bool: + """True when the current AMD GPU is RDNA3 (gfx110x).""" + arch = get_rocm_gfx_arch() + return arch is not None and arch.startswith("gfx110") + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: import torch import torch.version + if is_rocm(): + return None if not torch.cuda.is_available() or not torch.version.cuda: return None return torch.cuda.get_device_capability() diff --git a/setup.py b/setup.py index cfe41b7..8ba0640 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,7 @@ from __future__ import annotations import importlib.util +import os from pathlib import Path from setuptools import setup @@ -18,6 +19,22 @@ def _check_toolchain() -> None: module.check_nvcc_matches_torch() +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + +def _rocm_paths() -> tuple[list[str], list[str]]: + rocm_home = Path(os.getenv("ROCM_HOME", "/opt/rocm")) + if not rocm_home.exists(): + raise RuntimeError( + "ROCM_HOME is required to build on ROCm. Set ROCM_HOME to your ROCm install." + ) + include_dirs = [str(rocm_home / "include")] + library_dirs = [str(rocm_home / "lib")] + return include_dirs, library_dirs + + def _cuda_runtime_paths() -> tuple[list[str], list[str]]: if CUDA_HOME is None: raise RuntimeError( @@ -31,7 +48,19 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: return [str(cuda_home / "include")], library_dirs -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() +IS_ROCM = _is_rocm() + +if IS_ROCM: + runtime_include_dirs, runtime_library_dirs = _rocm_paths() + runtime_lib = "amdhip64" + # TODO(ROCm): allow override via FREETOKEN_ROCM_ARCH; default to all RDNA3. + rocm_arch = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103") + extra_compile = ["-O3", "-std=c++17", f"--offload-arch={rocm_arch}"] +else: + runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() + runtime_lib = "cudart" + extra_compile = ["-O3", "-std=c++17"] + _check_toolchain() @@ -42,12 +71,12 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/pinned_tensor.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17"], + include_dirs=runtime_include_dirs, + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile, ), - # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the + # CPU-compute MoE executor for --moe-backend cpu. Links cudart/hip for the # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels # use per-function target attributes (avx512bf16/avx512f) + a runtime # __builtin_cpu_supports dispatch, so the single binary stays portable @@ -57,10 +86,10 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17", "-pthread"], + include_dirs=runtime_include_dirs, + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile + ["-pthread"], ), ], cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},