feat(rocm): serve on AMD GPUs through the HIP toolchain

csrc-hip-shim/cuda_runtime_api.h

@@ -0,0 +1,32 @@
+// HIP compatibility shim for FreeToken's host-side extensions.
+// Maps every CUDA runtime API symbol used by kernel/csrc onto its HIP
+// equivalent so the same sources build against ROCm without edits.
+#pragma once
+#include <hip/hip_runtime_api.h>
+
+#define cudaSuccess hipSuccess
+#define cudaError_t hipError_t
+#define cudaGetErrorString hipGetErrorString
+#define cudaStream_t hipStream_t
+
+#define cudaMallocHost hipMallocHost
+#define cudaHostAlloc hipHostAlloc
+#define cudaFreeHost hipFreeHost
+#define cudaHostRegister hipHostRegister
+#define cudaHostGetDevicePointer hipHostGetDevicePointer
+
+#define cudaHostAllocPortable hipHostAllocPortable
+#define cudaHostAllocMapped hipHostAllocMapped
+#define cudaHostRegisterPortable hipHostRegisterPortable
+#define cudaHostRegisterMapped hipHostRegisterMapped
+
+#define cudaGetDevice hipGetDevice
+#define cudaDriverGetVersion hipDriverGetVersion
+#define cudaDeviceGetAttribute hipDeviceGetAttribute
+#define cudaStreamSynchronize hipStreamSynchronize
+#define cudaLaunchHostFunc hipLaunchHostFunc
+
+#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing
+#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeCanUseHostPointerForRegisteredMem
+
+#define CUDART_CB

python/freetoken/kernel/csrc/jit/fast_index_copy.cuh

@@ -33,41 +33,70 @@ inline constexpr auto get_mem_package() {
     }
 }
 
+// AMD ROCm port: the PTX streaming-load/store asm below has no HIP
+// equivalent; under HIP fall back to plain device loads and stores.
+// The cache-policy hints (L1 no-allocate, write-through) are dropped;
+// correctness is unchanged.
+
 __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
+#if defined(__HIP_PLATFORM_AMD__)
+    return *src;
+#else
     uint32_t tmp;
     asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
     return uint1{tmp};
+#endif
 }
 
 __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 {
+#if defined(__HIP_PLATFORM_AMD__)
+    return *src;
+#else
     uint32_t tmp0, tmp1;
     asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
     return uint2{tmp0, tmp1};
+#endif
 }
 
 __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 {
+#if defined(__HIP_PLATFORM_AMD__)
+    return *src;
+#else
     uint32_t tmp0, tmp1, tmp2, tmp3;
     asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src));
     return uint4{tmp0, tmp1, tmp2, tmp3};
+#endif
 }
 
 __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) {
+#if defined(__HIP_PLATFORM_AMD__)
+    *dst = value;
+#else
     uint32_t tmp = value.x;
     asm volatile("st.global.wt.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
+#endif
 }
 
 __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) {
+#if defined(__HIP_PLATFORM_AMD__)
+    *dst = value;
+#else
     uint32_t tmp0 = value.x;
     uint32_t tmp1 = value.y;
     asm volatile("st.global.wt.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
+#endif
 }
 
 __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) {
+#if defined(__HIP_PLATFORM_AMD__)
+    *dst = value;
+#else
     uint32_t tmp0 = value.x;
     uint32_t tmp1 = value.y;
     uint32_t tmp2 = value.z;
     uint32_t tmp3 = value.w;
     asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
+#endif
 }
 
 __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) {

python/freetoken/kernel/csrc/jit_hip_shim.h

@@ -0,0 +1,81 @@
+// HIP compatibility shim for FreeToken's JIT-compiled kernels.
+// Force-included (-include) into kernel TUs built with the HIP toolchain.
+// Maps the CUDA launch API surface used by csrc/include/freetoken/utils.cuh
+// onto HIP equivalents. PDL (programmatic dependent launch) attributes have no
+// HIP equivalent on this runtime; every kernel served on AMD builds with
+// use_pdl=false, so the attribute path stays compile-only dead code and the
+// launcher drops attrs instead of setting them.
+#pragma once
+#include <hip/hip_runtime.h>
+#include <cstddef>
+#include <utility>
+
+#ifndef __HIP_PLATFORM_AMD__
+#error "This shim is only meaningful for the AMD HIP toolchain"
+#endif
+
+typedef hipStream_t cudaStream_t;
+typedef hipError_t cudaError_t;
+#define cudaSuccess hipSuccess
+#define cudaGetErrorString hipGetErrorString
+#define cudaDriverGetVersion hipDriverGetVersion
+#define cudaDeviceGetAttribute hipDeviceGetAttribute
+#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing
+#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeCanUseHostPointerForRegisteredMem
+#define cudaMallocHost hipMallocHost
+#define cudaHostAlloc hipHostAlloc
+#define cudaFreeHost hipFreeHost
+#define cudaHostRegister hipHostRegister
+#define cudaHostGetDevicePointer hipHostGetDevicePointer
+#define cudaHostAllocPortable hipHostAllocPortable
+#define cudaHostAllocMapped hipHostAllocMapped
+#define cudaHostRegisterPortable hipHostRegisterPortable
+#define cudaHostRegisterMapped hipHostRegisterMapped
+#define cudaGetDevice hipGetDevice
+#define cudaStreamSynchronize hipStreamSynchronize
+#define cudaLaunchHostFunc hipLaunchHostFunc
+#define CUDART_CB
+
+enum cudaLaunchAttributeID_shim {
+  cudaLaunchAttributeProgrammaticStreamSerialization = 99,
+};
+
+struct cudaLaunchAttribute {
+  int id;
+  union {
+    int programmaticStreamSerializationAllowed;
+  } val;
+};
+
+struct cudaLaunchConfig_t {
+  dim3 gridDim;
+  dim3 blockDim;
+  size_t dynamicSmemBytes;
+  cudaStream_t stream;
+  cudaLaunchAttribute* attrs;
+  unsigned int numAttrs;
+};
+
+template <typename... KernelArgs, typename... Params>
+static inline hipError_t cudaLaunchKernelEx(const cudaLaunchConfig_t* config,
+                                            void (*kernel)(KernelArgs...),
+                                            Params&&... args) {
+  hipLaunchConfig_t hcfg;
+  hcfg.gridDim = config->gridDim;
+  hcfg.blockDim = config->blockDim;
+  hcfg.dynamicSmemBytes = config->dynamicSmemBytes;
+  hcfg.stream = config->stream;
+  hcfg.attrs = nullptr;
+  hcfg.numAttrs = 0;
+  return hipLaunchKernelEx(&hcfg, kernel, std::forward<Params>(args)...);
+}
+
+// ROCm 7.2 HIP does not provide the CUDA 11.4+ __grid_constant__ parameter
+// annotation; an empty define keeps by-value kernel parameters valid.
+#ifndef __grid_constant__
+#define __grid_constant__
+#endif
+
+#define cudaFuncSetAttribute hipFuncSetAttribute
+#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
+#define cudaGetLastError hipGetLastError

python/freetoken/kernel/triton/activation.py

@@ -129,12 +129,15 @@ def _act_and_mul(
     M = x2.shape[0]
     grid = lambda meta: (M, triton.cdiv(d, meta["BLOCK_D"]))
     pdl = _pdl_supported()
+    # The AMD triton launcher rejects the launch_pdl keyword
+    # outright; NVIDIA keeps the upstream launch_pdl=pdl call.
+    launch_kwargs = {} if getattr(torch.version, "hip", None) else {"launch_pdl": pdl}
     # Fixed via H100 sweep (72-config grid; 512/w4/s3 within 11% everywhere,
     # 1024/w4/s2 best at rows>=4096).
     block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512)
     num_stages = 2 if block_d == 1024 else 3
     _act_and_mul_kernel[grid](
-        o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl,
+        o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, **launch_kwargs,
         BLOCK_D=block_d, num_warps=4, num_stages=num_stages,
     )
     return out

python/freetoken/kernel/triton/attention.py

@@ -396,6 +396,11 @@ def decode_paged_attention(
     # (e.g. 6), where block_h rounds up and the kernel masks the extra lanes.
     valid_block_h = min(16, group)
     block_h = triton.next_power_of_2(valid_block_h)
+    if getattr(torch.version, "hip", None):
+        # AMD WMMA requires M >= 16 for tl.dot; the kernel already masks
+        # padded head lanes, so raise the tile floor instead of failing
+        # instruction selection on small GQA groups.
+        block_h = max(block_h, 16)
     block_d = triton.next_power_of_2(head_dim)
     block_dv = triton.next_power_of_2(head_dim)
 

python/freetoken/kernel/triton/norm.py

@@ -30,7 +30,7 @@
 import triton.language as tl
 from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait
 
-from freetoken.utils.arch import is_sm90_supported
+from freetoken.utils.arch import is_hip, is_sm90_supported
 
 _HEUR = {"BLOCK": lambda a: triton.next_power_of_2(a["H"])}
 
@@ -142,9 +142,12 @@ def _rmsnorm(input, weight, eps, out, gemma: bool):
     # PDL only on the contiguous (decode-replay) path: on the strided qk-norm's
     # 32k-CTA prefill grids the per-CTA gdc_wait poll costs more than it hides.
     pdl = contig and is_sm90_supported()
+    # The AMD triton launcher rejects the launch_pdl keyword
+    # outright; NVIDIA keeps the upstream launch_pdl=pdl call.
+    launch_kwargs = {} if is_hip() else {"launch_pdl": pdl}
     _rmsnorm_kernel[(A, B)](
         out, input, weight, eps, H, sxa, sxb, soa, sob,
-        CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma,
+        CONTIG=contig, ENABLE_PDL=pdl, **launch_kwargs, GEMMA=gemma,
         num_warps=_num_warps(A * B), num_stages=1,
     )
     return out
@@ -170,9 +173,12 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool):
     _, _, sra, srb = _leading(residual)
     contig = input.ndim == 2 and input.is_contiguous() and residual.is_contiguous()
     pdl = contig and is_sm90_supported()
+    # The AMD triton launcher rejects the launch_pdl keyword
+    # outright; NVIDIA keeps the upstream launch_pdl=pdl call.
+    launch_kwargs = {} if is_hip() else {"launch_pdl": pdl}
     _fused_add_rmsnorm_kernel[(A, B)](
         input, residual, weight, eps, H, sxa, sxb, sra, srb,
-        CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma,
+        CONTIG=contig, ENABLE_PDL=pdl, **launch_kwargs, GEMMA=gemma,
         num_warps=_num_warps(A * B), num_stages=1,
     )
 

python/freetoken/kernel/utils.py

@@ -12,6 +12,8 @@
 KERNEL_PATH = pathlib.Path(__file__).parent / "csrc"
 KERNEL_CACHE_PACKAGE = "freetoken_kernel_cache"
 KERNEL_CACHE_DIR_ENV = "FREETOKEN_KERNEL_CACHE_DIR"
+HIP_JIT_SHIM = KERNEL_PATH / "jit_hip_shim.h"
+
 DISABLE_KERNEL_CACHE_ENV = "FREETOKEN_DISABLE_KERNEL_CACHE"
 DISABLE_KERNEL_CACHE_VERSION_CHECK_ENV = "FREETOKEN_DISABLE_KERNEL_CACHE_VERSION_CHECK"
 DISABLE_JIT_ENV = "FREETOKEN_DISABLE_JIT"
@@ -30,7 +32,15 @@ def _cuda_cflags(extra: List[str]) -> List[str]:
     PTX→SASS JIT (driver-only, no CUDA toolkit). One top PTX suffices: the loader always
     JIT-forwards from the highest compatible PTX. When the env is unset (runtime JIT), this is a
     no-op and tvm-ffi targets only the local GPU."""
+    import torch
+
     flags = DEFAULT_CUDA_CFLAGS + extra
+    if getattr(torch.version, "hip", None):
+        # AMD ROCm port: hipcc rejects nvcc's --expt-relaxed-constexpr (its
+        # relaxed constexpr is already the default) and needs the HIP
+        # launch-API shim force-included.
+        flags += ["-include", str(HIP_JIT_SHIM)]
+        flags = [f for f in flags if f != "--expt-relaxed-constexpr"]
     arch_list = os.getenv("TVM_FFI_CUDA_ARCH_LIST", "").split()
     if arch_list:
         def _rank(a: str) -> int:

python/freetoken/utils/arch.py

@@ -14,6 +14,13 @@ def _get_torch_cuda_version() -> Tuple[int, int] | None:
     return torch.cuda.get_device_capability()
 
 
+def is_hip() -> bool:
+    """Running on the ROCm/HIP build of torch."""
+    import torch
+
+    return bool(getattr(torch.version, "hip", None))
+
+
 def is_arch_supported(major: int, minor: int = 0) -> bool:
     """capability >= (major, minor). Open-ended: newer archs also pass. Only use this
     for family-portable features (e.g. PDL); arch-specific kernels (sm_90a/sm_100a

setup.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
 import importlib.util
+import os
 from pathlib import Path
 
+import torch
+
 from setuptools import setup
 from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension
 
@@ -31,9 +34,37 @@ 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()
-_check_toolchain()
+def _hip_runtime_paths() -> tuple[list[str], list[str], list[str]]:
+    # AMD ROCm port: build the same sources against HIP instead. The shim
+    # header maps the CUDA runtime symbols csrc uses onto their HIP
+    # equivalents, so the sources stay untouched.
+    rocm = Path(os.environ.get("ROCM_PATH", "/opt/rocm"))
+    if not rocm.exists():
+        raise RuntimeError(f"ROCm not found at {rocm}; set ROCM_PATH")
+    return (
+        [str(ROOT / "csrc-hip-shim"), str(rocm / "include")],
+        [str(rocm / "lib")],
+        ["amdhip64"],
+    )
+
 
+if getattr(torch.version, "hip", None):
+    on_hip = True
+    include_dirs, library_dirs, libraries = _hip_runtime_paths()
+elif getattr(torch.version, "cuda", None):
+    on_hip = False
+    include_dirs, library_dirs = _cuda_runtime_paths()
+    _check_toolchain()
+    libraries = ["cudart"]
+else:
+    raise RuntimeError(
+        "freetoken requires a CUDA or ROCm build of torch; "
+        f"found torch {torch.__version__}"
+    )
+
+common_compile_args = ["-O3", "-std=c++17"] + (
+    ["-D__HIP_PLATFORM_AMD__"] if on_hip else []
+)
 
 setup(
     ext_modules=[
@@ -42,13 +73,13 @@ 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=include_dirs,
+            library_dirs=library_dirs,
+            libraries=libraries,
+            extra_compile_args=common_compile_args,
         ),
-        # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the
-        # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels
+        # CPU-compute MoE executor for --moe-backend cpu. Links the GPU runtime
+        # library 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
         # (scalar fallback) -- no global -march is set.
@@ -57,10 +88,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=include_dirs,
+            library_dirs=library_dirs,
+            libraries=libraries,
+            extra_compile_args=[*common_compile_args, "-pthread"],
         ),
     ],
     cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},