diff --git a/.gitignore b/.gitignore index bf804e07..0fb695f6 100644 --- a/.gitignore +++ b/.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/ diff --git a/pyproject.toml b/pyproject.toml index 8bd653f8..d7de9c15 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e8637..d90e7ab2 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -29,7 +29,7 @@ #include #include -#include +#include "../hip_compat.h" #include #if defined(__linux__) diff --git a/python/freetoken/kernel/csrc/gguf/dispatch.h b/python/freetoken/kernel/csrc/gguf/dispatch.h index f42a2163..17d2a0db 100644 --- a/python/freetoken/kernel/csrc/gguf/dispatch.h +++ b/python/freetoken/kernel/csrc/gguf/dispatch.h @@ -11,6 +11,20 @@ #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 @@ -18,6 +32,7 @@ #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__) \ diff --git a/python/freetoken/kernel/csrc/hip_compat.h b/python/freetoken/kernel/csrc/hip_compat.h new file mode 100644 index 00000000..1acd5b29 --- /dev/null +++ b/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 + +// 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 +#endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832..ea472b63 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -10,6 +10,50 @@ #include #include +// 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 + +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 +#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 __always_inline __device__ void wait() { +#if !(defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)) if constexpr (kUsePDL) { asm volatile("griddepcontrol.wait;" ::: "memory"); } +#endif } template __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 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( + ::TVMFFIEnvGetStream(device.device_type, device.device_id)); + } + + LaunchKernel(const LaunchKernel &) = delete; + LaunchKernel &operator=(const LaunchKernel &) = delete; + + template + auto operator()(T &&kernel, Args &&...args) const -> void { + kernel<<>>( + std::forward(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 diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..649c12dd 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/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(const_cast(flag_ptr)); diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adf..4cb983f2 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,5 +1,5 @@ #include -#include +#include "hip_compat.h" #include namespace { diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 04a16560..dbe392c5 100644 --- a/python/freetoken/kernel/gguf.py +++ b/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. diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533..c26354ec 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -23,6 +23,13 @@ from freetoken.utils.arch import is_sm90_supported +# _fast_tanh/_fast_ex2 below inline raw PTX text (tanh.approx.f32, ex2.approx.f32) via +# tl.inline_asm_elementwise. HIP's inline-asm path doesn't reject PTX outright -- it +# fails much later, deep in register allocation ("couldn't allocate output register +# for constraint 'f'"), since the constraint syntax is generic LLVM inline-asm but the +# instruction text is NVIDIA-ISA-only. Route ROCm through portable tl/libdevice ops. +_IS_HIP = tl.constexpr(torch.version.hip is not None) + SILU = 0 GELU = 1 GELU_TANH = 2 @@ -48,6 +55,8 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): + if _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 +66,8 @@ def _fast_tanh(x): @triton.jit def _fast_ex2(x): + if _IS_HIP: + return tl.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], @@ -129,12 +140,16 @@ def _act_and_mul( M = x2.shape[0] grid = lambda meta: (M, triton.cdiv(d, meta["BLOCK_D"])) pdl = _pdl_supported() + # launch_pdl is a CUDA-Hopper-only Triton launch kwarg; the AMD backend's + # arg-packer rejects it outright (KeyError) even when passed as False, so it + # is only included on the one backend/arch combination that ever sets pdl=True. + pdl_kwargs = {"launch_pdl": pdl} if pdl else {} # 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, **pdl_kwargs, BLOCK_D=block_d, num_warps=4, num_stages=num_stages, ) return out diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84..bc1fd11a 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -396,6 +396,13 @@ 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 torch.version.hip is not None: + # RDNA WMMA has no matrix-core instruction below a 16x16 tile, so a decode + # GQA group smaller than 16 (e.g. 4 here) leaves tl.dot's M dim too small to + # lower on this backend. The kernel already masks lanes >= VALID_BLOCK_H + # (it does this for non-power-of-two groups too), so padding BLOCK_H up to + # 16 is safe -- it only adds masked-out, discarded head lanes. + block_h = max(block_h, 16) block_d = triton.next_power_of_2(head_dim) block_dv = triton.next_power_of_2(head_dim) diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 1d9f744c..d5c923ad 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -59,6 +59,14 @@ def e4m3_native() -> bool: if _native is None: if FORCE_EMU: _native = False + elif torch.version.hip is not None: + # torch.cuda.get_device_capability() on a HIP build returns the gfx/RDNA + # generation number (e.g. (11, 5) for gfx1150), not a CUDA compute + # capability -- comparing it against (8, 9) below is a tuple comparison + # over two unrelated numbering schemes and can false-positive (11 > 8). + # No AMD GPU has this fp8e4nv unit; e4m3_native_cx() (Triton's own, + # backend-aware check) already agrees this must be False. + _native = False else: from freetoken.gpu_select import assigned_visible_gpu diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f..62bd6382 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -142,9 +142,13 @@ 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() + # launch_pdl is a CUDA-Hopper-only Triton launch kwarg; the AMD backend's + # arg-packer rejects it outright (KeyError) even when passed as False, so it + # is only included on the one backend/arch combination that ever sets pdl=True. + pdl_kwargs = {"launch_pdl": pdl} if pdl else {} _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, **pdl_kwargs, num_warps=_num_warps(A * B), num_stages=1, ) return out @@ -170,9 +174,10 @@ 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() + pdl_kwargs = {"launch_pdl": pdl} if pdl else {} _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, **pdl_kwargs, num_warps=_num_warps(A * B), num_stages=1, ) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b5..42f15a5b 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -30,7 +30,13 @@ 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.""" - flags = DEFAULT_CUDA_CFLAGS + extra + import torch + + flags = list(DEFAULT_CUDA_CFLAGS) + if torch.version.hip is not None: + # nvcc-only: hipcc/clang rejects it outright. + flags = [f for f in flags if f != "--expt-relaxed-constexpr"] + flags = flags + extra arch_list = os.getenv("TVM_FFI_CUDA_ARCH_LIST", "").split() if arch_list: def _rank(a: str) -> int: diff --git a/setup.py b/setup.py index cfe41b7d..faf9fc3f 100644 --- a/setup.py +++ b/setup.py @@ -4,13 +4,18 @@ from pathlib import Path from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension +from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, ROCM_HOME, CppExtension ROOT = Path(__file__).parent +IS_ROCM = CUDA_HOME is None and ROCM_HOME is not None def _check_toolchain() -> None: + if IS_ROCM: + # nvcc/CUDA-major checks below are meaningless on a ROCm torch build + # (torch.version.cuda is None there), so _toolchain.py's check is a no-op. + return path = ROOT / "python" / "freetoken" / "kernel" / "_toolchain.py" spec = importlib.util.spec_from_file_location("_freetoken_toolchain", path) module = importlib.util.module_from_spec(spec) @@ -18,20 +23,39 @@ def _check_toolchain() -> None: module.check_nvcc_matches_torch() -def _cuda_runtime_paths() -> tuple[list[str], list[str]]: +def _gpu_runtime_paths() -> tuple[list[str], list[str], list[str], list[str]]: + """Returns (include_dirs, library_dirs, libraries, extra_link_args).""" + if IS_ROCM: + rocm_home = Path(ROCM_HOME) + library_dirs = [d for d in (rocm_home / "lib64", rocm_home / "lib") if d.exists()] + # The pip-vendored rocm-sdk-core ships versioned sonames (libamdhip64.so.7) + # without the bare .so dev symlink `-lamdhip64` needs, so link the exact + # file. At runtime the dynamic linker dedupes on SONAME, so this resolves + # to whichever libamdhip64 torch itself already loaded into the process. + hip_lib = next( + (f for d in library_dirs for f in sorted(d.glob("libamdhip64.so*"))), None + ) + if hip_lib is None: + raise RuntimeError(f"libamdhip64.so* not found under {library_dirs}") + return ( + [str(rocm_home / "include")], + [str(d) for d in library_dirs], + [], + [f"-l:{hip_lib.name}"], + ) if CUDA_HOME is None: raise RuntimeError( - "CUDA_HOME is required to build freetoken.kernel._pinned_tensor " - "because it links against the CUDA runtime API." + "CUDA_HOME (or ROCM_HOME) is required to build freetoken.kernel._pinned_tensor " + "because it links against the CUDA/HIP runtime API." ) cuda_home = Path(CUDA_HOME) library_dirs = [str(cuda_home / "lib64")] if (cuda_home / "lib").exists(): library_dirs.append(str(cuda_home / "lib")) - return [str(cuda_home / "include")], library_dirs + return [str(cuda_home / "include")], library_dirs, ["cudart"], [] -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() +cuda_include_dirs, cuda_library_dirs, cuda_libraries, cuda_extra_link_args = _gpu_runtime_paths() _check_toolchain() @@ -44,7 +68,8 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: ], include_dirs=cuda_include_dirs, library_dirs=cuda_library_dirs, - libraries=["cudart"], + libraries=cuda_libraries, + extra_link_args=cuda_extra_link_args, extra_compile_args=["-O3", "-std=c++17"], ), # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the @@ -59,7 +84,8 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: ], include_dirs=cuda_include_dirs, library_dirs=cuda_library_dirs, - libraries=["cudart"], + libraries=cuda_libraries, + extra_link_args=cuda_extra_link_args, extra_compile_args=["-O3", "-std=c++17", "-pthread"], ), ],