From 5eec2582ce89376e3f5016d47ff73ca1b4e8d445 Mon Sep 17 00:00:00 2001 From: Christian Stewart Date: Mon, 24 Aug 2026 08:08:16 +0000 Subject: [PATCH] feat(rocm): serve on AMD GPUs through the HIP toolchain FreeToken built and served only against CUDA: setup.py linked its host-side extensions against cudart, the kernel JIT compiled with nvcc-only flags, and several triton call sites passed NVIDIA-only launch options. On an AMD GPU the engine could not even finish booting. Build _pinned_tensor and _cpu_moe against HIP through a shim header that maps every CUDA runtime symbol they use onto its HIP equivalent. Teach the kernel JIT the same trick: drop --expt-relaxed-constexpr (hipcc rejects it; relaxed constexpr is already its default) and force-include a shim, packaged with the kernel csrc tree, that maps the CUDA launch-config surface onto hipLaunchKernelEx. PDL launch attributes have no equivalent on this runtime; every kernel served on AMD builds with use_pdl=false, so the shim drops attributes instead of setting them. Three smaller fixes complete the port: - norm.py and activation.py omit the triton launch_pdl keyword under HIP; the AMD launcher rejects it outright. NVIDIA keeps the upstream launch_pdl=pdl call for both true and false. - attention.py floors block_h at 16 under HIP because RDNA3 WMMA cannot select an instruction for tl.dot below M=16. Padded head lanes were already masked. - fast_index_copy.cuh guards its PTX streaming-load/store inline asm: CUDA keeps it verbatim; HIP has no equivalent asm, so those builds fall back to plain device loads and stores. The cache-policy hints are dropped on HIP only; correctness is unchanged. setup.py detects the backend from torch: ROCm builds of torch take the HIP branch, CUDA builds keep the original cudart link and nvcc toolchain check unchanged, and anything else fails with a clear error. Every other edit is gated on HIP detection at run time. ROCm installs should resolve torch from the ROCm wheel index first; the pinned PyPI triton conflicts with the ROCm-bundled one. --- csrc-hip-shim/cuda_runtime_api.h | 32 ++++++++ .../kernel/csrc/jit/fast_index_copy.cuh | 29 +++++++ python/freetoken/kernel/csrc/jit_hip_shim.h | 81 +++++++++++++++++++ python/freetoken/kernel/triton/activation.py | 5 +- python/freetoken/kernel/triton/attention.py | 5 ++ python/freetoken/kernel/triton/norm.py | 12 ++- python/freetoken/kernel/utils.py | 10 +++ python/freetoken/utils/arch.py | 7 ++ setup.py | 55 ++++++++++--- 9 files changed, 220 insertions(+), 16 deletions(-) create mode 100644 csrc-hip-shim/cuda_runtime_api.h create mode 100644 python/freetoken/kernel/csrc/jit_hip_shim.h diff --git a/csrc-hip-shim/cuda_runtime_api.h b/csrc-hip-shim/cuda_runtime_api.h new file mode 100644 index 00000000..3e338979 --- /dev/null +++ b/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 + +#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 diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..bcdba878 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/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) { diff --git a/python/freetoken/kernel/csrc/jit_hip_shim.h b/python/freetoken/kernel/csrc/jit_hip_shim.h new file mode 100644 index 00000000..cc39e730 --- /dev/null +++ b/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 +#include +#include + +#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 +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(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 diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533..337060b3 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/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 diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84..459b9e68 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/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) diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f..2cc1c154 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/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, ) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b5..2f4748bb 100644 --- a/python/freetoken/kernel/utils.py +++ b/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: diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3d..ce3e6a5e 100644 --- a/python/freetoken/utils/arch.py +++ b/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 diff --git a/setup.py b/setup.py index cfe41b7d..b231815b 100644 --- a/setup.py +++ b/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)},