fix(rocm): route tensor parallel communication through RCCL

docs/install.md

@@ -2,7 +2,9 @@
 
 ## Requirements
 
-- Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13)
+- Linux x86_64 with either:
+  - NVIDIA GPU, driver r580+ (CUDA 13), or
+  - AMD RDNA3/RDNA4 GPU (`gfx1100`-`gfx1103`, `gfx1200`, or `gfx1201`) with ROCm 7.14
 - Python >= 3.10, with [uv](https://docs.astral.sh/uv/) recommended (plain
   `pip` + `venv` works too)
 
@@ -15,6 +17,33 @@ uv pip install "freetoken[accel]"
 
 CUDA kernels are JIT-compiled on first use, need a CUDA 13 toolkit with `nvcc` on PATH.
 
+### AMD ROCm source install (experimental)
+
+Use an official ROCm PyTorch image whose PyTorch version satisfies the project's
+`torch>=2.11,<2.12` constraint. For RDNA4, the matching ROCm 7.14 image is:
+
+```bash
+VIDEO_GID="$(getent group video | cut -d: -f3)"
+RENDER_GID="$(getent group render | cut -d: -f3)"
+docker run --rm -it \
+  --device=/dev/kfd --device=/dev/dri \
+  --group-add="$VIDEO_GID" --group-add="$RENDER_GID" --ipc=host \
+  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
+  -e PYTORCH_ROCM_ARCH=gfx1201 -e FREETOKEN_ROCM_ARCH=gfx1201 \
+  -v "$PWD:/workspace/FreeToken" -w /workspace/FreeToken \
+  rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 bash
+```
+
+Inside the container, preserve the ROCm-enabled PyTorch already supplied by the
+image and disable build isolation so it is also used to compile the extensions:
+
+```bash
+python -m pip install --no-build-isolation -e .
+```
+
+Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the
+actual target reported by `rocminfo`.
+
 ## Method 2: Install from source
 
 ```bash

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",
@@ -57,7 +58,10 @@ dependencies = [
     "torch>=2.11,<2.12",
     "tqdm>=4.66,<5",
     "transformers>=5.5,<6",
-    "triton==3.6.0; platform_system == 'Linux'",
+    # CUDA torch 2.11 resolves Triton 3.6; AMD's ROCm 7.14 image supplies its
+    # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the
+    # runtime-specific wheel selected by the PyTorch distribution.
+    "triton>=3.6,<3.8; platform_system == 'Linux'",
     "uvicorn>=0.30,<1",
 ]
 

python/freetoken/engine/engine.py

@@ -423,7 +423,17 @@ def __init__(self, config: EngineConfig):
             self._warmup_prefill()
 
     def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
-        if config.tp_info.size == 1 or config.use_pynccl:
+        use_pynccl = config.use_pynccl
+        if config.tp_info.size > 1 and use_pynccl:
+            from freetoken.kernel.backend import is_rocm
+
+            if is_rocm():
+                logger.warning_rank0(
+                    "PyNCCL is NVIDIA-only; using PyTorch's ROCm/RCCL process group instead"
+                )
+                use_pynccl = False
+
+        if config.tp_info.size == 1 or use_pynccl:
             torch.distributed.init_process_group(
                 backend="gloo",
                 rank=config.tp_info.rank,

python/freetoken/kernel/__main__.py

@@ -6,27 +6,32 @@ def generate_clangd():
     import subprocess
 
     from freetoken.kernel.utils import DEFAULT_INCLUDE
-    from freetoken.utils import init_logger
+    from freetoken.utils import get_rocm_gfx_arch, init_logger, is_rocm
     from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
 
     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.
+    if is_rocm():
+        arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"]
+    else:
+        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(".")
+        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+            import torch
+
+            major, minor = torch.cuda.get_device_capability()
+        arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"]
     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/cpu_moe/cpu_moe_ext.cpp

@@ -29,7 +29,7 @@
 #include <thread>
 #include <vector>
 
-#include <cuda_runtime_api.h>
+#include <freetoken/hip_compat.h>
 #include <torch/extension.h>
 
 #if defined(__linux__)

python/freetoken/kernel/csrc/include/freetoken/hip_compat.h

@@ -0,0 +1,162 @@
+#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:
+//   gfx1100 — RX 7900 XTX / XT
+//   gfx1101 — RX 7900 GRE
+//   gfx1102 — RX 7700 / XT
+//   gfx1103 — RX 7600 / XT
+//   gfx1200 — RX 9060 family
+//   gfx1201 — RX 9070 family / Radeon AI PRO R9700
+
+#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM)
+
+#define FREETOKEN_USE_ROCM 1
+
+// --- 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 cudaHostAllocPortable
+#define cudaHostAllocPortable hipHostMallocPortable
+#endif
+
+#ifndef cudaHostAllocMapped
+#define cudaHostAllocMapped hipHostMallocMapped
+#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
+// ROCm 7 exposes the CUDA-compatible extended launch configuration through HIP.
+#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 cudaStreamSynchronize
+#define cudaStreamSynchronize hipStreamSynchronize
+#endif
+
+#ifndef cudaLaunchHostFunc
+#define cudaLaunchHostFunc hipLaunchHostFunc
+#endif
+
+#ifndef CUDART_CB
+#define CUDART_CB
+#endif
+
+#ifndef __grid_constant__
+#define __grid_constant__
+#endif
+
+#ifndef dim3
+// HIP already provides dim3; this is a no-op guard.
+#endif
+
+#else  // NVIDIA CUDA path
+
+#define FREETOKEN_USE_ROCM 0
+
+#include <cuda_runtime_api.h>
+
+#endif

python/freetoken/kernel/csrc/include/freetoken/utils.cuh

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <freetoken/hip_compat.h>
 #include <freetoken/utils.h>
 
 #include <dlpack/dlpack.h>
@@ -115,6 +116,10 @@ public:
   }
 
   auto with_attr(bool use_pdl) -> LaunchKernel & {
+#ifdef __HIP__
+    (void)use_pdl;
+    m_config.numAttrs = 0;
+#else
     if (use_pdl) {
       m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization;
       m_attr_cache.val.programmaticStreamSerializationAllowed = 1;
@@ -123,6 +128,7 @@ public:
     } else {
       m_config.numAttrs = 0;
     }
+#endif
     return *this;
   }