fix(rocm): complete RDNA3 runtime path

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

@@ -19,6 +19,14 @@
 #include <hip/hip_runtime.h>
 #include <hip/hip_runtime_api.h>
 
+#ifndef CUDART_CB
+#define CUDART_CB
+#endif
+
+#ifndef __grid_constant__
+#define __grid_constant__
+#endif
+
 // --- 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
@@ -52,6 +60,14 @@
 #define cudaHostAlloc hipHostMalloc
 #endif
 
+#ifndef cudaHostAllocPortable
+#define cudaHostAllocPortable hipHostMallocPortable
+#endif
+
+#ifndef cudaHostAllocMapped
+#define cudaHostAllocMapped hipHostMallocMapped
+#endif
+
 #ifndef cudaHostRegister
 #define cudaHostRegister hipHostRegister
 #endif
@@ -122,6 +138,14 @@
 #define cudaStream_t hipStream_t
 #endif
 
+#ifndef cudaStreamSynchronize
+#define cudaStreamSynchronize hipStreamSynchronize
+#endif
+
+#ifndef cudaLaunchHostFunc
+#define cudaLaunchHostFunc hipLaunchHostFunc
+#endif
+
 #ifndef dim3
 // HIP already provides dim3; this is a no-op guard.
 #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;
   }
 

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

@@ -34,40 +34,64 @@ inline constexpr auto get_mem_package() {
 }
 
 __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
+#ifdef __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 {
+#ifdef __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 {
+#ifdef __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) {
+#ifdef __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) {
+#ifdef __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) {
+#ifdef __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) {
@@ -147,7 +171,8 @@ inline bool host_ptr_identity() {
 }
 
 inline void* device_alias(void* ptr, DLDevice dev) {
-    if (dev.device_type == kDLCUDA || host_ptr_identity()) {
+    if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM ||
+        host_ptr_identity()) {
         return ptr;
     }
     void* mapped = nullptr;
@@ -269,7 +294,7 @@ inline auto get_sync_flag_ptr(
     auto flag_dtype = host::SymbolicDType{};
     host::TensorMatcher({1})
         .with_dtype<int32_t>(flag_dtype)
-        .with_device<kDLCUDA>(device)
+        .with_device<kDLCUDA, kDLROCM>(device)
         .verify(sync_flag);
     return static_cast<int32_t*>(sync_flag.data_ptr());
 }
@@ -344,17 +369,17 @@ struct FastIndexCopyKernel {
 
         TensorMatcher({-1, D})
         .with_dtype(data_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLCUDA, kDLCUDAHost, kDLROCM, kDLROCMHost, kDLCPU>()
         .verify(src);
 
         TensorMatcher({-1, D})
         .with_dtype(data_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLCUDA, kDLCUDAHost, kDLROCM, kDLROCMHost, kDLCPU>()
         .verify(dst);
 
         TensorMatcher({L})
         .with_dtype<int32_t, int64_t>(indices_dtype)
-        .with_device<kDLCUDA>(device)
+        .with_device<kDLCUDA, kDLROCM>(device)
         .verify(src_indices)
         .verify(dst_indices);
 
@@ -363,7 +388,7 @@ struct FastIndexCopyKernel {
             const auto num_indices_tensor = num_indices.value();
             TensorMatcher({1})
                 .with_dtype<int64_t>(num_indices_dtype)
-                .with_device<kDLCUDA>(device)
+                .with_device<kDLCUDA, kDLROCM>(device)
                 .verify(num_indices_tensor);
 
             num_indices_data_ptr = static_cast<const int64_t*>(num_indices_tensor.data_ptr());
@@ -529,14 +554,14 @@ struct MultiIndexCopyKernel {
         auto indices_dtype = SymbolicDType{};
         auto num_indices_dtype = SymbolicDType{};
 
-        TensorMatcher({B}).with_dtype<int64_t>(ptr_dtype).with_device<kDLCUDA>(device)
+        TensorMatcher({B}).with_dtype<int64_t>(ptr_dtype).with_device<kDLCUDA, kDLROCM>(device)
             .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes);
-        TensorMatcher({L}).with_dtype<int32_t, int64_t>(indices_dtype).with_device<kDLCUDA>(device)
+        TensorMatcher({L}).with_dtype<int32_t, int64_t>(indices_dtype).with_device<kDLCUDA, kDLROCM>(device)
             .verify(dst_indices).verify(src_indices);
 
         const int64_t* valid_length = nullptr;
         if (num_indices.has_value()) {
-            TensorMatcher({1}).with_dtype<int64_t>(num_indices_dtype).with_device<kDLCUDA>(device)
+            TensorMatcher({1}).with_dtype<int64_t>(num_indices_dtype).with_device<kDLCUDA, kDLROCM>(device)
                 .verify(num_indices.value());
             valid_length = static_cast<const int64_t*>(num_indices.value().data_ptr());
         }

python/freetoken/kernel/csrc/jit/index.cu

@@ -114,15 +114,15 @@ struct IndexKernel {
 
     TensorMatcher({-1, D}) //
         .with_dtype(weights_dtype_)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .verify(weights);
     TensorMatcher({L, D}) //
         .with_dtype(weights_dtype_)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .verify(output);
     TensorMatcher({L}) //
         .with_dtype<int32_t, int64_t>(indices_dtype_)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .verify(indices);
 
     const auto device = device_.unwrap();

python/freetoken/kernel/csrc/jit/store.cu

@@ -72,18 +72,18 @@ struct StoreKernel {
 
     TensorMatcher({-1, D}) //
         .with_strides({X, 1})
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .with_dtype(dtype_)
         .verify(k_cache)
         .verify(v_cache);
     TensorMatcher({L, D}) //
         .with_strides({Y, 1})
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .with_dtype(dtype_)
         .verify(k)
         .verify(v);
     TensorMatcher({L}) //
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLCUDA, kDLROCM>(device_)
         .with_dtype<int32_t, int64_t>(indices_dtype_)
         .verify(indices);
 

python/freetoken/kernel/triton/activation.py

@@ -20,8 +20,9 @@
 import triton.language as tl
 from triton.language.extra import libdevice
 from triton.language.extra.cuda import gdc_wait, gdc_launch_dependents
+from triton.language import target_info
 
-from freetoken.utils.arch import is_sm90_supported
+from freetoken.utils.arch import is_rocm, is_sm90_supported
 
 SILU = 0
 GELU = 1
@@ -48,6 +49,8 @@ def _pdl_supported() -> bool:
 
 @triton.jit
 def _fast_tanh(x):
+    if target_info.is_hip():
+        return libdevice.tanh(x)
     # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh.
     return tl.inline_asm_elementwise(
         "tanh.approx.f32 $0, $1;", "=f,f", [x],
@@ -57,6 +60,8 @@ def _fast_tanh(x):
 
 @triton.jit
 def _fast_ex2(x):
+    if target_info.is_hip():
+        return libdevice.exp2(x)
     # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu.
     return tl.inline_asm_elementwise(
         "ex2.approx.f32 $0, $1;", "=f,f", [x],
@@ -134,8 +139,9 @@ def _act_and_mul(
     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,
         BLOCK_D=block_d, num_warps=4, num_stages=num_stages,
+        **({} if is_rocm() else {"launch_pdl": pdl}),
     )
     return out
 

python/freetoken/kernel/triton/e4m3_compat.py

@@ -59,6 +59,10 @@ def e4m3_native() -> bool:
     if _native is None:
         if FORCE_EMU:
             _native = False
+        elif torch.version.hip is not None:
+            # ROCm reports gfx1101 as capability (11, 0), which is not a CUDA
+            # compute capability and must not select the native fp8e4nv path.
+            _native = False
         else:
             native = {torch.cuda.get_device_capability(i) >= (8, 9)
                       for i in range(torch.cuda.device_count())}

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_rocm, is_sm90_supported
 
 _HEUR = {"BLOCK": lambda a: triton.next_power_of_2(a["H"])}
 
@@ -144,7 +144,8 @@ def _rmsnorm(input, weight, eps, out, gemma: bool):
     pdl = contig and is_sm90_supported()
     _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, GEMMA=gemma,
+        **({} if is_rocm() else {"launch_pdl": pdl}),
         num_warps=_num_warps(A * B), num_stages=1,
     )
     return out
@@ -172,7 +173,8 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool):
     pdl = contig and is_sm90_supported()
     _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, GEMMA=gemma,
+        **({} if is_rocm() else {"launch_pdl": pdl}),
         num_warps=_num_warps(A * B), num_stages=1,
     )
 

python/freetoken/kernel/utils.py

@@ -52,8 +52,9 @@ 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}"]
+    raw_arches = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103")
+    arches = [arch for arch in re.split(r"[;,\s]+", raw_arches.strip()) if arch]
+    flags = flags + [f"--offload-arch={arch}" for arch in arches]
     return flags
 CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool]
 

setup.py

@@ -9,6 +9,7 @@
 
 
 ROOT = Path(__file__).parent
+KERNEL_INCLUDE = ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include"
 
 
 def _check_toolchain() -> None:
@@ -61,6 +62,8 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]:
     runtime_lib = "cudart"
     extra_compile = ["-O3", "-std=c++17"]
 
+runtime_include_dirs.append(str(KERNEL_INCLUDE))
+
 _check_toolchain()
 
 

tests/kernels/test_rocm_arch_flags.py

@@ -0,0 +1,34 @@
+from freetoken.kernel import utils
+
+
+def test_hip_cflags_expand_default_arches(monkeypatch):
+    monkeypatch.delenv("FREETOKEN_ROCM_ARCH", raising=False)
+
+    flags = utils._hip_cflags([])
+
+    assert flags[-4:] == [
+        "--offload-arch=gfx1100",
+        "--offload-arch=gfx1101",
+        "--offload-arch=gfx1102",
+        "--offload-arch=gfx1103",
+    ]
+    assert all(";" not in flag for flag in flags)
+
+
+def test_hip_cflags_accept_common_arch_separators(monkeypatch):
+    monkeypatch.setenv(
+        "FREETOKEN_ROCM_ARCH",
+        "gfx1100; gfx1101,gfx1102 gfx1103",
+    )
+
+    flags = utils._hip_cflags(["-DFOO=1"])
+
+    assert flags == [
+        "-std=c++20",
+        "-O3",
+        "-DFOO=1",
+        "--offload-arch=gfx1100",
+        "--offload-arch=gfx1101",
+        "--offload-arch=gfx1102",
+        "--offload-arch=gfx1103",
+    ]