feat: add simple ROCm GPU support for RDNA3 (gfx1100-1103)

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",

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"""

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.<nvcc major>; 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()

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),

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 <cuda_runtime_api.h> 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 <hip/hip_runtime.h>
+#include <hip/hip_runtime_api.h>
+
+// --- 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 <cuda_runtime_api.h>
+
+#endif  // __HIP__

python/freetoken/kernel/csrc/pinned_tensor.cpp

@@ -1,5 +1,5 @@
 #include <cstdint>
-#include <cuda_runtime_api.h>
+#include <freetoken/hip_compat.h>
 #include <torch/extension.h>
 
 namespace {

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"])
 
 

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,

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",

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