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..56ab93df 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 #include #if defined(__linux__) diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index 99539d1f..00a1b540 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -19,6 +19,14 @@ #include #include +#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 diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832..f21b9585 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -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; } diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..2f784211 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/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(flag_dtype) - .with_device(device) + .with_device(device) .verify(sync_flag); return static_cast(sync_flag.data_ptr()); } @@ -344,17 +369,17 @@ struct FastIndexCopyKernel { TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(src); TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(dst); TensorMatcher({L}) .with_dtype(indices_dtype) - .with_device(device) + .with_device(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(num_indices_dtype) - .with_device(device) + .with_device(device) .verify(num_indices_tensor); num_indices_data_ptr = static_cast(num_indices_tensor.data_ptr()); @@ -529,14 +554,14 @@ struct MultiIndexCopyKernel { auto indices_dtype = SymbolicDType{}; auto num_indices_dtype = SymbolicDType{}; - TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); - TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) .verify(dst_indices).verify(src_indices); const int64_t* valid_length = nullptr; if (num_indices.has_value()) { - TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) .verify(num_indices.value()); valid_length = static_cast(num_indices.value().data_ptr()); } diff --git a/python/freetoken/kernel/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu index ca0e1db2..aca58383 100644 --- a/python/freetoken/kernel/csrc/jit/index.cu +++ b/python/freetoken/kernel/csrc/jit/index.cu @@ -114,15 +114,15 @@ struct IndexKernel { TensorMatcher({-1, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device(device_) .verify(weights); TensorMatcher({L, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device(device_) .verify(output); TensorMatcher({L}) // .with_dtype(indices_dtype_) - .with_device(device_) + .with_device(device_) .verify(indices); const auto device = device_.unwrap(); diff --git a/python/freetoken/kernel/csrc/jit/store.cu b/python/freetoken/kernel/csrc/jit/store.cu index 8d84d76e..162dfdfe 100644 --- a/python/freetoken/kernel/csrc/jit/store.cu +++ b/python/freetoken/kernel/csrc/jit/store.cu @@ -72,18 +72,18 @@ struct StoreKernel { TensorMatcher({-1, D}) // .with_strides({X, 1}) - .with_device(device_) + .with_device(device_) .with_dtype(dtype_) .verify(k_cache) .verify(v_cache); TensorMatcher({L, D}) // .with_strides({Y, 1}) - .with_device(device_) + .with_device(device_) .with_dtype(dtype_) .verify(k) .verify(v); TensorMatcher({L}) // - .with_device(device_) + .with_device(device_) .with_dtype(indices_dtype_) .verify(indices); diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533..0b7c945c 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/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 diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 61d3a0e7..e52095cc 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/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())} diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f..9071e1df 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_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, ) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index dfdcd4c7..6cb298fc 100644 --- a/python/freetoken/kernel/utils.py +++ b/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] diff --git a/setup.py b/setup.py index 8ba0640a..698ad780 100644 --- a/setup.py +++ b/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() diff --git a/tests/kernels/test_rocm_arch_flags.py b/tests/kernels/test_rocm_arch_flags.py new file mode 100644 index 00000000..6235ffe2 --- /dev/null +++ b/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", + ] diff --git a/tests/kernels/test_rocm_dlpack_devices.py b/tests/kernels/test_rocm_dlpack_devices.py new file mode 100644 index 00000000..35f97dbe --- /dev/null +++ b/tests/kernels/test_rocm_dlpack_devices.py @@ -0,0 +1,38 @@ +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] + + +def _read(relative: str) -> str: + return (ROOT / relative).read_text(encoding="utf-8") + + +def test_rocm_dlpack_devices_are_explicit_not_macro_aliased(): + utils = _read("python/freetoken/kernel/csrc/include/freetoken/utils.cuh") + fast = _read("python/freetoken/kernel/csrc/jit/fast_index_copy.cuh") + index = _read("python/freetoken/kernel/csrc/jit/index.cu") + store = _read("python/freetoken/kernel/csrc/jit/store.cu") + pynccl = _read("python/freetoken/kernel/csrc/src/pynccl.cu") + + # Do not globally rewrite DLPack CUDA tokens under HIP. Besides being hard to + # reason about, that would also silently change the still-CUDA/NCCL-only + # multi-GPU wrapper, which is outside this RDNA3 single-GPU PR's scope. + assert "#define kDLCUDA kDLROCM" not in utils + assert "#define kDLCUDAHost kDLROCMHost" not in utils + + # Single-GPU JIT paths that execute on ROCm explicitly accept ROCm DLPack + # devices while retaining CUDA acceptance for the existing NVIDIA path. + assert ( + ".with_device()" + in fast + ) + assert "dev.device_type == kDLCUDA || dev.device_type == kDLROCM" in fast + assert fast.count(".with_device(device)") >= 6 + assert index.count(".with_device(device_)") == 3 + assert store.count(".with_device(device_)") == 3 + + # Preserve the PR's stated boundary: RCCL/multi-GPU migration is not being + # claimed by a side effect of a preprocessor alias. + assert "kDLROCM" not in pynccl + assert "device_type == kDLCUDA" in pynccl