rocm: fix native extensions, kernel JIT builds, and Triton PTX fallbacks for gfx1150

.gitignore

@@ -6,6 +6,13 @@ __pycache__/
 # C extensions
 *.so
 
+# torch.utils.cpp_extension's ROCm auto-hipify writes translated copies next to the
+# CUDA sources it translates (kernel/csrc/gguf/*.cu -> *.hip, *.cuh -> *_hip.cuh);
+# regenerated on every build, never hand-edited.
+*.hip
+*_hip.cuh
+*_hip.h
+
 # Distribution / packaging
 .Python
 build/

pyproject.toml

@@ -3,7 +3,7 @@
 # own, and a mismatch links the C++ extensions against the wrong libtorch.
 # setuptools floor: 77 is the first release that understands the PEP 639 `license`
 # SPDX string and `license-files` below.
-requires = ["setuptools>=77", "torch>=2.11,<2.12", "wheel"]
+requires = ["setuptools>=77", "torch>=2.11,<2.14", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -54,10 +54,13 @@ dependencies = [
     # floor+ceiling: sglang-kernel 0.4.5 links libtorch symbols only 2.11 has.
     # PyPI's torch 2.11.0 wheel is itself the cu130 build, so plain pip resolves
     # correctly from PyPI alone; uv additionally pins the index below.
-    "torch>=2.11,<2.12",
+    # ROCm note: on AMD (rocm.nightlies.amd.com builds) this range is intentionally
+    # loosened -- those wheels report their own local version segment
+    # (2.13.0a0+rocm...) which the sglang-kernel/cu130 constraint above doesn't apply to.
+    "torch>=2.11,<2.14",
     "tqdm>=4.66,<5",
     "transformers>=5.5,<6",
-    "triton==3.6.0; platform_system == 'Linux'",
+    "triton>=3.6,<3.8; platform_system == 'Linux'",
     "uvicorn>=0.30,<1",
 ]
 

python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp

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

python/freetoken/kernel/csrc/gguf/dispatch.h

@@ -11,13 +11,28 @@
 #endif
 
 // Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h (CUDA variants).
+// HIP's __shfl_xor_sync requires a 64-bit mask unconditionally (amd_warp_sync_functions.h
+// static_asserts sizeof(mask) == 8) regardless of actual wavefront width; the donor's
+// CUDA-style callers pass a 32-bit `unsigned int` mask (e.g. 0xffffffff), so widen it here
+// rather than touching every call site.
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+#ifndef SGLANG_SHFL_XOR_SYNC
+#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) \
+  __shfl_xor_sync((unsigned long long)(mask), (var), (lane_mask))
+#endif
+#ifndef SGLANG_SHFL_XOR_SYNC_WIDTH
+#define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \
+  __shfl_xor_sync((unsigned long long)(mask), (var), (lane_mask), (width))
+#endif
+#else
 #ifndef SGLANG_SHFL_XOR_SYNC
 #define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) __shfl_xor_sync((mask), (var), (lane_mask))
 #endif
 #ifndef SGLANG_SHFL_XOR_SYNC_WIDTH
 #define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \
   __shfl_xor_sync((mask), (var), (lane_mask), (width))
 #endif
+#endif
 
 #define DISPATCH_CASE_FLOAT_TYPES(...)                 \
   AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \

python/freetoken/kernel/csrc/hip_compat.h

@@ -0,0 +1,55 @@
+#pragma once
+
+// Lets pinned_tensor.cpp and cpu_moe_ext.cpp call the CUDA Runtime API names they
+// were written against while actually linking HIP on ROCm builds. Only the calls
+// those two files use are covered -- this is not a general CUDA/HIP compat layer.
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+#include <hip/hip_runtime_api.h>
+
+// CUDA's host-callback calling-convention annotation; empty on POSIX (matches
+// cuda_runtime_api.h's own definition there). hipHostFn_t has no such annotation.
+#define CUDART_CB
+
+using cudaError_t = hipError_t;
+using cudaStream_t = hipStream_t;
+constexpr hipError_t cudaSuccess = hipSuccess;
+constexpr unsigned int cudaHostAllocPortable = hipHostMallocPortable;
+constexpr unsigned int cudaHostAllocMapped = hipHostMallocMapped;
+constexpr unsigned int cudaHostRegisterPortable = hipHostRegisterPortable;
+constexpr unsigned int cudaHostRegisterMapped = hipHostRegisterMapped;
+constexpr hipDeviceAttribute_t cudaDevAttrUnifiedAddressing =
+    hipDeviceAttributeUnifiedAddressing;
+constexpr hipDeviceAttribute_t cudaDevAttrCanUseHostPointerForRegisteredMem =
+    hipDeviceAttributeCanUseHostPointerForRegisteredMem;
+
+inline hipError_t cudaMallocHost(void **ptr, size_t size) {
+  return hipHostMalloc(ptr, size, hipHostMallocDefault);
+}
+inline hipError_t cudaFreeHost(void *ptr) { return hipHostFree(ptr); }
+inline hipError_t cudaHostAlloc(void **ptr, size_t size, unsigned int flags) {
+  return hipHostMalloc(ptr, size, flags);
+}
+inline hipError_t cudaGetDevice(int *device) { return hipGetDevice(device); }
+inline hipError_t cudaDeviceGetAttribute(int *value, hipDeviceAttribute_t attr,
+                                         int device) {
+  return hipDeviceGetAttribute(value, attr, device);
+}
+inline hipError_t cudaHostGetDevicePointer(void **devPtr, void *hostPtr,
+                                           unsigned int flags) {
+  return hipHostGetDevicePointer(devPtr, hostPtr, flags);
+}
+inline hipError_t cudaHostRegister(void *ptr, size_t size, unsigned int flags) {
+  return hipHostRegister(ptr, size, flags);
+}
+inline hipError_t cudaDriverGetVersion(int *v) { return hipDriverGetVersion(v); }
+inline const char *cudaGetErrorString(hipError_t e) { return hipGetErrorString(e); }
+inline hipError_t cudaStreamSynchronize(hipStream_t s) {
+  return hipStreamSynchronize(s);
+}
+inline hipError_t cudaLaunchHostFunc(hipStream_t s, hipHostFn_t fn, void *data) {
+  return hipLaunchHostFunc(s, fn, data);
+}
+
+#else
+#include <cuda_runtime_api.h>
+#endif

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

@@ -10,6 +10,50 @@
 #include <source_location>
 #include <type_traits>
 
+// nvcc implicitly pulls in the CUDA runtime for .cu translation units; hipcc does
+// not do the equivalent for HIP, so it must be included explicitly here. On the
+// HIP path there is no cudaLaunchKernelEx/cudaLaunchConfig_t equivalent (that API
+// is Hopper PDL-specific), so LaunchKernel gets its own HIP-side definition below
+// instead of a name-aliasing shim -- see PDL below for why that also means
+// with_attr(true) is a no-op on this path.
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+#include <hip/hip_runtime.h>
+
+using cudaError_t = hipError_t;
+constexpr hipError_t cudaSuccess = hipSuccess;
+using cudaStream_t = hipStream_t;
+
+inline const char *cudaGetErrorString(hipError_t e) { return hipGetErrorString(e); }
+inline hipError_t cudaGetLastError() { return hipGetLastError(); }
+inline hipError_t cudaFuncSetAttribute(const void *func, hipFuncAttribute attr,
+                                       int value) {
+  return hipFuncSetAttribute(func, attr, value);
+}
+constexpr hipFuncAttribute cudaFuncAttributeMaxDynamicSharedMemorySize =
+    hipFuncAttributeMaxDynamicSharedMemorySize;
+
+inline hipError_t cudaGetDevice(int *device) { return hipGetDevice(device); }
+inline hipError_t cudaDeviceGetAttribute(int *value, hipDeviceAttribute_t attr,
+                                         int device) {
+  return hipDeviceGetAttribute(value, attr, device);
+}
+inline hipError_t cudaHostGetDevicePointer(void **devPtr, void *hostPtr,
+                                           unsigned int flags) {
+  return hipHostGetDevicePointer(devPtr, hostPtr, flags);
+}
+constexpr hipDeviceAttribute_t cudaDevAttrUnifiedAddressing =
+    hipDeviceAttributeUnifiedAddressing;
+constexpr hipDeviceAttribute_t cudaDevAttrCanUseHostPointerForRegisteredMem =
+    hipDeviceAttributeCanUseHostPointerForRegisteredMem;
+
+// CUDA-only kernel-parameter annotation (passes large by-value params via constant
+// memory instead of copying them into local/generic memory first); HIP has no
+// equivalent attribute, so this just falls back to an ordinary by-value parameter.
+#define __grid_constant__
+#else
+#include <cuda_runtime.h>
+#endif
+
 namespace device {
 
 inline constexpr auto kWarpThreads = 32u;
@@ -42,16 +86,24 @@ __always_inline __device__ auto offset(const T *ptr, U... offset) -> const
 
 namespace PDL {
 
+// Programmatic Dependent Launch is a Hopper-only CUDA hardware feature; the PTX
+// below has no HIP/ROCm equivalent. Callers gate kUsePDL off for non-Hopper CUDA
+// targets already, and LaunchKernel::with_attr is a no-op on HIP (see below), so
+// this stays unconditionally a no-op there rather than a compile failure.
 template <bool kUsePDL> __always_inline __device__ void wait() {
+#if !(defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__))
   if constexpr (kUsePDL) {
     asm volatile("griddepcontrol.wait;" ::: "memory");
   }
+#endif
 }
 
 template <bool kUsePDL> __always_inline __device__ void launch() {
+#if !(defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__))
   if constexpr (kUsePDL) {
     asm volatile("griddepcontrol.launch_dependents;" :::);
   }
+#endif
 }
 
 } // namespace PDL
@@ -88,6 +140,50 @@ template <auto F> inline void set_smem_once(std::size_t smem_size) {
       last_smem_size, " bytes");
 }
 
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+
+// HIP has no cudaLaunchKernelEx/cudaLaunchConfig_t analog (that API only exists to
+// carry Hopper PDL attributes, which ROCm hardware has no equivalent for), so this
+// launches via the plain triple-chevron form instead. with_attr(true) is therefore
+// a no-op here -- there is no attribute to carry.
+struct LaunchKernel {
+public:
+  explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device,
+                        std::size_t dynamic_shared_mem_bytes = 0) noexcept
+      : m_grid_dim(grid_dim), m_block_dim(block_dim),
+        m_smem(dynamic_shared_mem_bytes), m_stream(resolve_device(device)) {}
+
+  explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, cudaStream_t stream,
+                        std::size_t dynamic_shared_mem_bytes = 0) noexcept
+      : m_grid_dim(grid_dim), m_block_dim(block_dim),
+        m_smem(dynamic_shared_mem_bytes), m_stream(stream) {}
+
+  static auto resolve_device(DLDevice device) -> cudaStream_t {
+    return static_cast<cudaStream_t>(
+        ::TVMFFIEnvGetStream(device.device_type, device.device_id));
+  }
+
+  LaunchKernel(const LaunchKernel &) = delete;
+  LaunchKernel &operator=(const LaunchKernel &) = delete;
+
+  template <typename T, typename... Args>
+  auto operator()(T &&kernel, Args &&...args) const -> void {
+    kernel<<<m_grid_dim, m_block_dim, m_smem, m_stream>>>(
+        std::forward<Args>(args)...);
+    CUDA_CHECK(::cudaGetLastError());
+  }
+
+  auto with_attr(bool /*use_pdl*/) -> LaunchKernel & { return *this; }
+
+private:
+  dim3 m_grid_dim;
+  dim3 m_block_dim;
+  std::size_t m_smem;
+  cudaStream_t m_stream;
+};
+
+#else
+
 struct LaunchKernel {
 public:
   explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device,
@@ -141,4 +237,6 @@ private:
   cudaLaunchAttribute m_attr_cache;
 };
 
+#endif
+
 } // namespace host

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

@@ -33,6 +33,38 @@ inline constexpr auto get_mem_package() {
     }
 }
 
+// The ld.global.L1::no_allocate / st.global.wt PTX below are cache-policy hints
+// (skip L1 allocate on read, write-through on store) with no HIP equivalent -- AMD
+// ROCm builds fall back to plain loads/stores. Correctness is unchanged; only the
+// cache-policy hint is lost.
+#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+
+__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
+    return *src;
+}
+
+__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 {
+    return *src;
+}
+
+__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 {
+    return *src;
+}
+
+__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) {
+    *dst = value;
+}
+
+__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) {
+    *dst = value;
+}
+
+__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) {
+    *dst = value;
+}
+
+#else
+
 __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
     uint32_t tmp;
     asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
@@ -70,6 +102,8 @@ __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& v
     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) {
     // Exponential backoff to avoid hammering a global atomic in a tight loop.
     auto* flag = reinterpret_cast<int*>(const_cast<int32_t*>(flag_ptr));

python/freetoken/kernel/csrc/pinned_tensor.cpp

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

python/freetoken/kernel/gguf.py

@@ -51,16 +51,22 @@ def _c_compiler_for(cxx: str) -> str:
 def _module():
     from torch.utils.cpp_extension import load
 
-    extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"]
-    host_cxx = _host_compiler()
-    if host_cxx is not None:
-        # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a
-        # libtorch/nvcc-compatible compiler. Force (not setdefault): the system
-        # default (CXX unset -> g++) can be a gcc too new for the torch headers.
-        cxx_path = shutil.which(host_cxx) or host_cxx
-        extra_cuda_cflags += ["-ccbin", cxx_path]
-        os.environ["CXX"] = cxx_path
-        os.environ["CC"] = _c_compiler_for(cxx_path)
+    if torch.version.hip is not None:
+        # Neither issue -ccbin works around applies under hipcc: it has no separate
+        # nvcc-style host pass (its own bundled clang IS the host compiler), and
+        # --expt-relaxed-constexpr is an nvcc-only flag hipcc/clang rejects outright.
+        extra_cuda_cflags = ["-O3"]
+    else:
+        extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"]
+        host_cxx = _host_compiler()
+        if host_cxx is not None:
+            # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a
+            # libtorch/nvcc-compatible compiler. Force (not setdefault): the system
+            # default (CXX unset -> g++) can be a gcc too new for the torch headers.
+            cxx_path = shutil.which(host_cxx) or host_cxx
+            extra_cuda_cflags += ["-ccbin", cxx_path]
+            os.environ["CXX"] = cxx_path
+            os.environ["CC"] = _c_compiler_for(cxx_path)
 
     # gguf_kernel.cu carries its own PYBIND11_MODULE (appended at the end), so a
     # plain `load` of the single source compiles + binds the ggml_* ops.