diff --git a/README.md b/README.md index 2a56a086..d2dfaf3e 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ FreeToken is an edge-native Mixture-of-Experts (MoE) serving engine designed for - **Semantic-Aware Caching**: Features semantic anchor checkpoints for recurrent state and KV caches, allowing agentic context edits (e.g., tool calls, thinking blocks) to avoid redundant context recomputation. - **Elastic Memory Management**: Supports dynamic, runtime VRAM re-allocation between expert caches and KV memory without engine restarts or weight reloading. - **Broad MoE & Ecosystem Support**: Supports frontier open-weight MoE models (e.g., DeepSeek-V4-Flash, Qwen3.6-35B-A3B, GLM-5.2) across various parameter scales and quantization formats (e.g., MXFP4, NVFP4, FP8, BF16), with Anthropic/OpenAI-compatible APIs for seamless integration with real-world coding and tool-calling agents (e.g., Codex, Claude Code, OpenCode, OpenClaw, DeepSeek Harness). -- **Diverse Consumer Hardware**: Scales across consumer laptops, gaming desktops, and workstation GPUs, with native support for NVIDIA RTX 30, RTX 40, and RTX 50 series GPUs. +- **Diverse Consumer Hardware**: Scales across consumer laptops, gaming desktops, and workstation GPUs, with native support for NVIDIA RTX 30/40/50 GPUs and experimental ROCm source support for AMD RDNA3/RDNA4 GPUs. ## Getting Started diff --git a/docs/install.md b/docs/install.md index f5205ab3..850558a1 100644 --- a/docs/install.md +++ b/docs/install.md @@ -2,7 +2,9 @@ ## Requirements -- Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13) +- Linux x86_64 with either: + - NVIDIA GPU, driver r580+ (CUDA 13), or + - AMD RDNA3/RDNA4 GPU (`gfx1100`-`gfx1103`, `gfx1200`, or `gfx1201`) with ROCm 7.14 - Python >= 3.10, with [uv](https://docs.astral.sh/uv/) recommended (plain `pip` + `venv` works too) @@ -15,6 +17,40 @@ uv pip install "freetoken[accel]" CUDA kernels are JIT-compiled on first use, need a CUDA 13 toolkit with `nvcc` on PATH. +### AMD ROCm source install (experimental) + +Use an official ROCm PyTorch image whose PyTorch version satisfies the project's +`torch>=2.11,<2.12` constraint. For RDNA4, the matching ROCm 7.14 image is: + +```bash +VIDEO_GID="$(getent group video | cut -d: -f3)" +RENDER_GID="$(getent group render | cut -d: -f3)" +docker run --rm -it \ + --device=/dev/kfd --device=/dev/dri \ + --group-add="$VIDEO_GID" --group-add="$RENDER_GID" --ipc=host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e PYTORCH_ROCM_ARCH=gfx1201 -e FREETOKEN_ROCM_ARCH=gfx1201 \ + -v "$PWD:/workspace/FreeToken" -w /workspace/FreeToken \ + rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 bash +``` + +Inside the container, preserve the ROCm-enabled PyTorch already supplied by the +image and disable build isolation so it is also used to compile the extensions: + +```bash +python -m pip install --no-build-isolation -e . +``` + +Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the +actual target reported by `rocminfo`. + +The optional native GGUF kernels also require Thrust headers. Install the generic +headers before the first GGUF kernel JIT build: + +```bash +apt-get update && apt-get install -y --no-install-recommends libthrust-dev +``` + ## Method 2: Install from source ```bash diff --git a/pyproject.toml b/pyproject.toml index 8bd653f8..d22ae67c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Environment :: GPU :: NVIDIA CUDA", + "Environment :: GPU :: AMD ROCm", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", @@ -57,7 +58,10 @@ dependencies = [ "torch>=2.11,<2.12", "tqdm>=4.66,<5", "transformers>=5.5,<6", - "triton==3.6.0; platform_system == 'Linux'", + # CUDA torch 2.11 resolves Triton 3.6; AMD's ROCm 7.14 image supplies its + # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the + # runtime-specific wheel selected by the PyTorch distribution. + "triton>=3.6,<3.8; platform_system == 'Linux'", "uvicorn>=0.30,<1", ] diff --git a/python/freetoken/kernel/__main__.py b/python/freetoken/kernel/__main__.py index 7be541a6..5b66d484 100644 --- a/python/freetoken/kernel/__main__.py +++ b/python/freetoken/kernel/__main__.py @@ -6,27 +6,32 @@ def generate_clangd(): import subprocess from freetoken.kernel.utils import DEFAULT_INCLUDE - from freetoken.utils import init_logger + from freetoken.utils import get_rocm_gfx_arch, init_logger, is_rocm from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path logger = init_logger(__name__) logger.info("Generating .clangd file...") include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE - status = subprocess.run( - args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], - capture_output=True, - check=True, - ) - compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] - major, minor = compute_cap.split(".") + + # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries. + if is_rocm(): + arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"] + else: + try: + status = subprocess.run( + args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], + capture_output=True, + check=True, + ) + compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] + major, minor = compute_cap.split(".") + except (subprocess.CalledProcessError, FileNotFoundError, ValueError): + import torch + + major, minor = torch.cuda.get_device_capability() + arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"] compile_flags = ",\n ".join( - [ - "-xcuda", - f"--cuda-gpu-arch=sm_{major}{minor}", - "-std=c++20", - "-Wall", - "-Wextra", - ] + arch_flags + ["-std=c++20", "-Wall", "-Wextra"] + [f"-isystem{path}" for path in include_paths] ) clangd_content = f""" diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebbb..93d03221 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -1,4 +1,4 @@ -"""CUDA toolchain/torch consistency checks. +"""CUDA/HIP toolchain/torch consistency checks. Standalone on purpose: setup.py and the kernel-cache build backend load this file by path, so it must not import the freetoken package. @@ -16,6 +16,11 @@ _TRUE_VALUES = {"1", "true", "yes", "on"} +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + def _nvcc_path() -> str | None: from torch.utils.cpp_extension import CUDA_HOME @@ -49,6 +54,8 @@ def check_nvcc_matches_torch() -> None: nvcc-built binaries link libcudart.so.; at runtime only the torch wheel's own CUDA runtime is guaranteed to be loadable. """ + if _is_rocm(): + return # ROCm uses hipcc, not nvcc if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES: return torch_major = torch_cuda_major() diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 3037ad8d..8293e3a6 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -42,6 +42,24 @@ def is_triton_kernels_installed() -> bool: return _importable("triton_kernels") +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU).""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def driver_hip_version() -> int | None: + """ROCm driver version, or None if undetermined.""" + # TODO(ROCm): flashinfer/sgl_kernel have no ROCm builds — Triton fallback is used. + try: + from freetoken.kernel.pinned import _load_pinned_extension + return int(_load_pinned_extension().driver_cuda_version()) or None + except Exception: + return None + + @functools.cache def driver_cuda_version() -> int | None: """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0), 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/gguf/dispatch.h b/python/freetoken/kernel/csrc/gguf/dispatch.h index f42a2163..bb15096b 100644 --- a/python/freetoken/kernel/csrc/gguf/dispatch.h +++ b/python/freetoken/kernel/csrc/gguf/dispatch.h @@ -5,18 +5,28 @@ #pragma once #include +#include #ifndef WARP_SIZE #define WARP_SIZE 32 #endif -// Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h (CUDA variants). +// HIP's synchronized shuffle API requires a 64-bit mask even on wave32 targets +// such as gfx1201, while CUDA uses a 32-bit mask. +#if defined(__HIP_PLATFORM_AMD__) +#define SGLANG_SHUFFLE_MASK(mask) static_cast(mask) +#else +#define SGLANG_SHUFFLE_MASK(mask) (mask) +#endif + +// Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h. #ifndef SGLANG_SHFL_XOR_SYNC -#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) __shfl_xor_sync((mask), (var), (lane_mask)) +#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) \ + __shfl_xor_sync(SGLANG_SHUFFLE_MASK(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)) + __shfl_xor_sync(SGLANG_SHUFFLE_MASK(mask), (var), (lane_mask), (width)) #endif #define DISPATCH_CASE_FLOAT_TYPES(...) \ diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h new file mode 100644 index 00000000..eb7e71e1 --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -0,0 +1,162 @@ +#pragma once + +// HIP compatibility shim: maps CUDA runtime API names to HIP equivalents so +// the same C++ source compiles under both nvcc and hipcc. Include this instead +// of directly when the file needs the runtime API. +// +// On NVIDIA platforms the CUDA headers are included as-is and every macro below +// resolves to the original CUDA symbol, so there is zero overhead. +// +// Supported ROCm targets: +// gfx1100 — RX 7900 XTX / XT +// gfx1101 — RX 7900 GRE +// gfx1102 — RX 7700 / XT +// gfx1103 — RX 7600 / XT +// gfx1200 — RX 9060 family +// gfx1201 — RX 9070 family / Radeon AI PRO R9700 + +#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM) + +#define FREETOKEN_USE_ROCM 1 + +// --- HIP runtime headers --- +#include +#include + +// --- 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 +// here so call-sites stay unchanged. + +#ifndef cudaSuccess +#define cudaSuccess hipSuccess +#endif + +#ifndef cudaError_t +#define cudaError_t hipError_t +#endif + +#ifndef cudaGetErrorString +#define cudaGetErrorString hipGetErrorString +#endif + +#ifndef cudaGetLastError +#define cudaGetLastError hipGetLastError +#endif + +#ifndef cudaMallocHost +#define cudaMallocHost hipMallocHost +#endif + +#ifndef cudaFreeHost +#define cudaFreeHost hipFreeHost +#endif + +#ifndef cudaHostAlloc +#define cudaHostAlloc hipHostMalloc +#endif + +#ifndef cudaHostRegister +#define cudaHostRegister hipHostRegister +#endif + +#ifndef cudaHostRegisterPortable +#define cudaHostRegisterPortable hipHostRegisterPortable +#endif + +#ifndef cudaHostRegisterMapped +#define cudaHostRegisterMapped hipHostRegisterMapped +#endif + +#ifndef cudaHostAllocPortable +#define cudaHostAllocPortable hipHostMallocPortable +#endif + +#ifndef cudaHostAllocMapped +#define cudaHostAllocMapped hipHostMallocMapped +#endif + +#ifndef cudaHostGetDevicePointer +#define cudaHostGetDevicePointer hipHostGetDevicePointer +#endif + +#ifndef cudaGetDevice +#define cudaGetDevice hipGetDevice +#endif + +#ifndef cudaDriverGetVersion +#define cudaDriverGetVersion hipDriverGetVersion +#endif + +#ifndef cudaDeviceGetAttribute +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#endif + +#ifndef cudaDevAttrUnifiedAddressing +#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaDevAttrCanUseHostPointerForRegisteredMem +// HIP does not expose this attribute; assume UVA identity on ROCm (true on Linux). +// TODO(ROCm): re-enable proper UVA query if HIP adds this attribute. +#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaFuncSetAttribute +#define cudaFuncSetAttribute hipFuncSetAttribute +#endif + +#ifndef cudaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize +#endif + +#ifndef cudaLaunchKernelEx +// ROCm 7 exposes the CUDA-compatible extended launch configuration through HIP. +#define cudaLaunchKernelEx hipLaunchKernelEx +#endif + +#ifndef cudaLaunchConfig_t +#define cudaLaunchConfig_t hipLaunchConfig_t +#endif + +#ifndef cudaLaunchAttribute +#define cudaLaunchAttribute hipLaunchAttribute +#endif + +#ifndef cudaLaunchAttributeProgrammaticStreamSerialization +// PDL (Programmatic Dependent Launch) is NVIDIA-specific. +// TODO(ROCm): PDL has no ROCm equivalent — disabled, may affect overlap scheduling latency. +#define cudaLaunchAttributeProgrammaticStreamSerialization 0 +#endif + +#ifndef cudaStream_t +#define cudaStream_t hipStream_t +#endif + +#ifndef cudaStreamSynchronize +#define cudaStreamSynchronize hipStreamSynchronize +#endif + +#ifndef cudaLaunchHostFunc +#define cudaLaunchHostFunc hipLaunchHostFunc +#endif + +#ifndef CUDART_CB +#define CUDART_CB +#endif + +#ifndef __grid_constant__ +#define __grid_constant__ +#endif + +#ifndef dim3 +// HIP already provides dim3; this is a no-op guard. +#endif + +#else // NVIDIA CUDA path + +#define FREETOKEN_USE_ROCM 0 + +#include + +#endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832..f46d2af2 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 @@ -44,13 +45,21 @@ namespace PDL { template __always_inline __device__ void wait() { if constexpr (kUsePDL) { +#if FREETOKEN_USE_ROCM + // Programmatic dependent launch is NVIDIA-specific. +#else asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif } } template __always_inline __device__ void launch() { if constexpr (kUsePDL) { +#if FREETOKEN_USE_ROCM + // Programmatic dependent launch is NVIDIA-specific. +#else asm volatile("griddepcontrol.launch_dependents;" :::); +#endif } } @@ -115,6 +124,10 @@ public: } auto with_attr(bool use_pdl) -> LaunchKernel & { +#if FREETOKEN_USE_ROCM + RuntimeCheck(!use_pdl, "Programmatic dependent launch is unavailable on ROCm"); + m_config.numAttrs = 0; +#else if (use_pdl) { m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization; m_attr_cache.val.programmaticStreamSerializationAllowed = 1; @@ -123,6 +136,7 @@ public: } else { m_config.numAttrs = 0; } +#endif return *this; } @@ -138,7 +152,9 @@ private: return config; } cudaLaunchConfig_t m_config; +#if !FREETOKEN_USE_ROCM 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..bf313c52 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 { +#if FREETOKEN_USE_ROCM + 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 FREETOKEN_USE_ROCM + 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 FREETOKEN_USE_ROCM + 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 FREETOKEN_USE_ROCM + *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 FREETOKEN_USE_ROCM + *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 FREETOKEN_USE_ROCM + *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) { @@ -75,7 +99,7 @@ __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag auto* flag = reinterpret_cast(const_cast(flag_ptr)); uint32_t sleep_ns = 128; while (atomicAdd(flag, 0) > 0) { -#if __CUDA_ARCH__ >= 700 +#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 700 __nanosleep(sleep_ns); #endif sleep_ns = sleep_ns < 2048 ? (sleep_ns << 1) : 2048; @@ -147,7 +171,7 @@ 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 +293,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 +368,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 +387,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 +553,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/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adf..9355f57a 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,5 +1,5 @@ #include -#include +#include #include namespace { diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 04a16560..4d44ae5d 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -13,6 +13,7 @@ from __future__ import annotations import functools +import hashlib import os import pathlib import shutil @@ -22,6 +23,32 @@ _CSRC = pathlib.Path(__file__).parent / "csrc" / "gguf" +def _staged_rocm_sources() -> pathlib.Path: + """Copy CUDA sources out of the checkout before PyTorch HIPifies them. + + ``torch.utils.cpp_extension.load`` writes generated ``*_hip`` sources next to + the input file. Keeping the staging directory under the extension cache makes + the source checkout stay clean while still allowing normal incremental builds. + """ + cache_root = pathlib.Path( + os.environ.get("TORCH_EXTENSIONS_DIR", pathlib.Path.home() / ".cache" / "torch_extensions") + ) + digest = hashlib.sha256() + digest.update(f"torch={torch.__version__};hip={torch.version.hip}".encode()) + for source in sorted(_CSRC.iterdir()): + if source.is_file() and "_hip." not in source.name and source.suffix != ".hip": + digest.update(source.name.encode()) + digest.update(source.read_bytes()) + staged = cache_root / f"freetoken_gguf_sources_{digest.hexdigest()[:16]}" + shutil.copytree( + _CSRC, + staged, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns("*_hip.*", "*.hip", "__pycache__"), + ) + return staged + + def _host_compiler() -> str | None: """A host compiler nvcc + libtorch headers accept. @@ -47,12 +74,29 @@ def _c_compiler_for(cxx: str) -> str: cc = base.replace("g++", "gcc") return shutil.which(cc) or cc + @functools.cache def _module(): from torch.utils.cpp_extension import load - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] - host_cxx = _host_compiler() + is_rocm = getattr(torch.version, "hip", None) is not None + extra_cuda_cflags = ["-O3"] + extra_ldflags: list[str] = [] + if is_rocm: + from freetoken.kernel.utils import _rocm_link_flags + + extra_ldflags = _rocm_link_flags() + # Ubuntu's generic Thrust headers otherwise select the CUDA backend and + # try to include cuda_runtime_api.h. GGUF only reaches Thrust through a + # libtorch complex-number header, so the backend-neutral C++ path is + # sufficient for HIP compilation. + extra_cuda_cflags.append("-DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_CPP") + csrc = _staged_rocm_sources() + else: + extra_cuda_cflags.append("--expt-relaxed-constexpr") + csrc = _CSRC + + host_cxx = None if is_rocm else _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 @@ -66,9 +110,10 @@ def _module(): # plain `load` of the single source compiles + binds the ggml_* ops. return load( name="freetoken_gguf_kernels", - sources=[str(_CSRC / "gguf_kernel.cu")], - extra_include_paths=[str(_CSRC)], + sources=[str(csrc / "gguf_kernel.cu")], + extra_include_paths=[str(csrc)], extra_cuda_cflags=extra_cuda_cflags, + extra_ldflags=extra_ldflags, verbose=True, ) diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea5735..71bde734 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/python/freetoken/kernel/pynccl.py @@ -27,6 +27,7 @@ def get_buffer(self) -> int: ... @functools.cache def _load_nccl_module() -> Module: + # TODO(ROCm): NCCL -> RCCL migration for multi-GPU tensor parallelism on AMD. return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"]) 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 7a0164b5..b59a588f 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -4,6 +4,7 @@ import os import pathlib import re +from functools import cache from typing import TYPE_CHECKING, List, NamedTuple, Tuple, TypeAlias, Union if TYPE_CHECKING: @@ -19,7 +20,14 @@ DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] +DEFAULT_HIP_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_LDFLAGS = [] +DEFAULT_ROCM_ARCHES = ("gfx1100", "gfx1101", "gfx1102", "gfx1103", "gfx1200", "gfx1201") + + +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None def _cuda_cflags(extra: List[str]) -> List[str]: @@ -40,6 +48,69 @@ def _rank(a: str) -> int: cc = max(arch_list, key=_rank).rstrip("a").replace(".", "") flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"] return flags + + +def _hip_cflags(extra: List[str]) -> List[str]: + """HIP flags for a kernel build on ROCm.""" + # TODO(ROCm): Triton autotune configs need RDNA-specific tuning (wave count, LDS size). + flags = DEFAULT_HIP_CFLAGS + extra + raw_arches = os.getenv("FREETOKEN_ROCM_ARCH") or os.getenv("PYTORCH_ROCM_ARCH", "") + arches = list(dict.fromkeys(re.findall(r"gfx\d+[a-z]?", raw_arches.lower()))) + if not arches: + from freetoken.utils.arch import get_rocm_gfx_arch + + detected = get_rocm_gfx_arch() + arches = [detected] if detected else list(DEFAULT_ROCM_ARCHES) + return flags + [f"--offload-arch={arch}" for arch in arches] + + +@cache +def _rocm_link_flags() -> List[str]: + """Make ROCm's runtime library discoverable to JIT link commands. + + Traditional ROCm installs provide ``libamdhip64.so`` under ``$ROCM_HOME/lib``. + ROCm 7.14 Python SDK images only provide the versioned soname, while TVM-FFI + still links with ``-lamdhip64``. Supply a cache-local unversioned symlink via + an explicit linker search path without modifying the Python environment. + """ + candidates: list[pathlib.Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(pathlib.Path(os.environ["ROCM_HOME"])) + try: + from torch.utils.cpp_extension import ROCM_HOME + + if ROCM_HOME: + candidates.append(pathlib.Path(ROCM_HOME)) + except ImportError: + pass + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(pathlib.Path(next(iter(spec.submodule_search_locations)))) + candidates.append(pathlib.Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + library_dir = rocm_home / "lib" + unversioned = library_dir / "libamdhip64.so" + link_dir = library_dir + if not unversioned.exists(): + versioned = sorted(library_dir.glob("libamdhip64.so.*")) + if not versioned: + continue + link_dir = pathlib.Path.home() / ".cache" / "freetoken" / "rocm-lib" + link_dir.mkdir(parents=True, exist_ok=True) + compat_link = link_dir / "libamdhip64.so" + if not compat_link.exists() and not compat_link.is_symlink(): + try: + compat_link.symlink_to(versioned[-1]) + except FileExistsError: + # Multiple tensor-parallel ranks may prepare the same cache. + pass + + return [f"-L{link_dir}", f"-Wl,-rpath,{library_dir}"] + + raise RuntimeError("Unable to locate libamdhip64 for ROCm JIT linking") + + CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -217,13 +288,20 @@ def load_aot( cpp_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cpp_files] cuda_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cuda_files] + if _is_rocm(): + cuda_cflags = _hip_cflags(extra_cuda_cflags) + runtime_ldflags = _rocm_link_flags() + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags) + runtime_ldflags = [] + return load( name, cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_cuda_cflags=cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) @@ -272,13 +350,20 @@ def load_jit( cuda_sources = [f'#include "{path}"' for path in cuda_paths] cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] + if _is_rocm(): + cuda_cflags = _hip_cflags(extra_cuda_cflags) + runtime_ldflags = _rocm_link_flags() + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags) + runtime_ldflags = [] + return load_inline( name, cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_cuda_cflags=cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index 2e4ad15f..c54579a1 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -1,5 +1,9 @@ from .arch import ( is_arch_supported, + is_rocm, + get_rocm_gfx_arch, + is_gfx11xx_family, + is_gfx12xx_family, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -35,6 +39,10 @@ "load_toolcall_anchor_id", "init_logger", "is_arch_supported", + "is_rocm", + "get_rocm_gfx_arch", + "is_gfx11xx_family", + "is_gfx12xx_family", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3d..3bf8fb61 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -1,14 +1,77 @@ from __future__ import annotations import functools +import os +import re from typing import Tuple +_GFX_ARCH_RE = re.compile(r"gfx\d+[a-z]?") + + +def _gfx_arch_from(value: object) -> str | None: + match = _GFX_ARCH_RE.search(str(value).lower()) + return match.group(0) if match else None + + +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU) instead of CUDA.""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def get_rocm_gfx_arch() -> str | None: + """Return the current AMD GPU target (for example ``gfx1201``). + + Prefer the runtime device because build variables may contain multiple + semicolon-separated targets. Environment variables remain useful for + cross-compilation and systems where no GPU is currently visible. + """ + if not is_rocm(): + return None + + import torch + + if torch.cuda.is_available(): + try: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + for attr in ("gcnArchName", "arch"): + arch = _gfx_arch_from(getattr(props, attr, "")) + if arch: + return arch + except (AttributeError, RuntimeError): + pass + + for env_var in ("FREETOKEN_ROCM_ARCH", "PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"): + arch = _gfx_arch_from(os.getenv(env_var, "")) + if arch: + return arch + return None + + +@functools.cache +def is_gfx11xx_family() -> bool: + """True when the current AMD GPU is RDNA3 (gfx110x).""" + arch = get_rocm_gfx_arch() + return arch is not None and arch.startswith("gfx110") + + +@functools.cache +def is_gfx12xx_family() -> bool: + """True when the current AMD GPU is RDNA4 (gfx120x).""" + arch = get_rocm_gfx_arch() + return arch is not None and arch.startswith("gfx120") + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: import torch import torch.version + if is_rocm(): + return None if not torch.cuda.is_available() or not torch.version.cuda: return None return torch.cuda.get_device_capability() diff --git a/setup.py b/setup.py index cfe41b7d..bac07c6f 100644 --- a/setup.py +++ b/setup.py @@ -1,13 +1,15 @@ from __future__ import annotations import importlib.util +import os 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, CppExtension, ROCM_HOME ROOT = Path(__file__).parent +KERNEL_INCLUDE = str(ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include") def _check_toolchain() -> None: @@ -18,6 +20,43 @@ def _check_toolchain() -> None: module.check_nvcc_matches_torch() +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + +def _rocm_paths() -> tuple[list[str], list[str], str]: + candidates: list[Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(Path(os.environ["ROCM_HOME"])) + if ROCM_HOME: + candidates.append(Path(ROCM_HOME)) + + # ROCm 7.14 PyTorch images ship the SDK as a Python package instead of + # installing it at /opt/rocm. + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(Path(next(iter(spec.submodule_search_locations)))) + candidates.append(Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + include_dir = rocm_home / "include" + library_dir = rocm_home / "lib" + if not (include_dir / "hip" / "hip_runtime.h").exists(): + continue + if (library_dir / "libamdhip64.so").exists(): + return [str(include_dir)], [str(library_dir)], "amdhip64" + versioned = sorted(library_dir.glob("libamdhip64.so.*")) + if versioned: + return [str(include_dir)], [str(library_dir)], f":{versioned[-1].name}" + + searched = ", ".join(str(path) for path in dict.fromkeys(candidates)) + raise RuntimeError( + "A ROCm SDK with HIP headers and libamdhip64 is required to build on ROCm; " + f"searched: {searched}. Set ROCM_HOME to override." + ) + + def _cuda_runtime_paths() -> tuple[list[str], list[str]]: if CUDA_HOME is None: raise RuntimeError( @@ -31,7 +70,21 @@ 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() +IS_ROCM = _is_rocm() + +if IS_ROCM: + runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths() + runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"] + # These extensions contain host code only. BuildExtension supplies the ROCm + # platform defines to the C++ compiler; offload architecture flags belong on + # HIP device sources and would be rejected by the host compiler here. + extra_compile = ["-O3", "-std=c++17"] +else: + runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() + runtime_lib = "cudart" + runtime_link_args = [] + extra_compile = ["-O3", "-std=c++17"] + _check_toolchain() @@ -42,12 +95,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=[KERNEL_INCLUDE, *runtime_include_dirs], + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile, + extra_link_args=runtime_link_args, ), - # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the + # CPU-compute MoE executor for --moe-backend cpu. Links cudart/amdhip64 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 @@ -57,10 +111,11 @@ 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=[KERNEL_INCLUDE, *runtime_include_dirs], + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile + ["-pthread"], + extra_link_args=runtime_link_args, ), ], cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)}, diff --git a/tests/kernels/test_gguf_rocm.py b/tests/kernels/test_gguf_rocm.py new file mode 100644 index 00000000..52936d32 --- /dev/null +++ b/tests/kernels/test_gguf_rocm.py @@ -0,0 +1,27 @@ +import pytest +import torch + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is None, + reason="a ROCm GPU is required", +) + + +def test_q4_0_dequant_matches_torch_reference(): + from freetoken.kernel.gguf import ggml_dequantize + from freetoken.models.gguf.dequant import GGML_Q4_0, dequantize + + scale = torch.tensor([0.5], dtype=torch.float16).view(torch.uint8) + quants = torch.tensor( + [0x10, 0x32, 0x54, 0x76, 0x98, 0xBA, 0xDC, 0xFE] * 2, + dtype=torch.uint8, + ) + packed_cpu = torch.cat((scale, quants)).reshape(1, 18) + expected = dequantize(packed_cpu, GGML_Q4_0, torch.float32).reshape(1, 32) + + actual = ggml_dequantize( + packed_cpu.to("cuda"), GGML_Q4_0, m=1, n=32, dtype=torch.float32 + ) + + torch.testing.assert_close(actual.cpu(), expected) diff --git a/tests/kernels/test_jit_index_store.py b/tests/kernels/test_jit_index_store.py new file mode 100644 index 00000000..7590023a --- /dev/null +++ b/tests/kernels/test_jit_index_store.py @@ -0,0 +1,36 @@ +import pytest +import torch + +from freetoken.kernel import indexing, store_cache + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="a CUDA or ROCm GPU is required" +) + + +def test_indexing_jit_matches_torch_on_cold_and_warm_loads(): + weights = torch.arange(8 * 64, dtype=torch.float32, device="cuda").reshape(8, 64) + + for values in ([7, 2, 0], [1, 6, 3]): + indices = torch.tensor(values, dtype=torch.int32, device="cuda") + actual = indexing(weights, indices) + torch.testing.assert_close(actual, weights[indices.long()]) + + +def test_store_jit_matches_torch_on_cold_and_warm_loads(): + k_cache = torch.zeros((8, 64), dtype=torch.float32, device="cuda") + v_cache = torch.zeros_like(k_cache) + indices = torch.tensor([5, 0, 3], dtype=torch.int64, device="cuda") + + for offset in (0.0, 1000.0): + k = torch.arange(3 * 64, dtype=torch.float32, device="cuda").reshape(3, 64) + k = k + offset + v = k + 500.0 + store_cache(k_cache, v_cache, indices, k, v) + torch.testing.assert_close(k_cache[indices], k) + torch.testing.assert_close(v_cache[indices], v) + + untouched = torch.tensor([1, 2, 4, 6, 7], device="cuda") + torch.testing.assert_close(k_cache[untouched], torch.zeros((5, 64), device="cuda")) + torch.testing.assert_close(v_cache[untouched], torch.zeros((5, 64), device="cuda")) diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py index e61108fd..2fee4f25 100644 --- a/tests/kernels/test_pinned_tensor.py +++ b/tests/kernels/test_pinned_tensor.py @@ -122,7 +122,12 @@ def test_host_device_ptr_is_identity_under_uva(): pytest.skip("non-UVA platform: host_device_ptr rejects unregistered memory instead") # Under UVA cudaHostGetDevicePointer degenerates to identity for any host pointer # (no registration validation); rejection of pageable memory only exists on - # non-identity platforms (Windows/WDDM), where the translation is real. + # non-identity CUDA platforms (Windows/WDDM), where the translation is real. + # HIP validates registration even though registered/pinned memory uses the + # identity address on Linux. Calling it with pageable memory also leaves a + # sticky HIP error, so the pinned identity case above is the relevant check. + if torch.version.hip is not None: + return pageable = torch.empty(64, dtype=torch.uint8) ext = _load_pinned_extension() assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr() diff --git a/tests/utils/test_rocm_arch.py b/tests/utils/test_rocm_arch.py new file mode 100644 index 00000000..95529862 --- /dev/null +++ b/tests/utils/test_rocm_arch.py @@ -0,0 +1,89 @@ +import importlib +import pathlib +from types import SimpleNamespace + +import torch + +from freetoken.utils import arch + + +def _clear_arch_caches() -> None: + arch.get_rocm_gfx_arch.cache_clear() + arch.is_gfx11xx_family.cache_clear() + arch.is_gfx12xx_family.cache_clear() + + +def test_rocm_arch_prefers_visible_device_over_multi_arch_build_env(monkeypatch): + monkeypatch.setattr(arch, "is_rocm", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: SimpleNamespace(gcnArchName="gfx1201:sramecc-:xnack-"), + ) + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1200") + _clear_arch_caches() + + assert arch.get_rocm_gfx_arch() == "gfx1201" + assert arch.is_gfx12xx_family() + assert not arch.is_gfx11xx_family() + + _clear_arch_caches() + + +def test_rocm_arch_falls_back_to_cross_compile_env(monkeypatch): + monkeypatch.setattr(arch, "is_rocm", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201") + _clear_arch_caches() + + assert arch.get_rocm_gfx_arch() == "gfx1200" + + _clear_arch_caches() + + +def test_hip_cflags_emit_one_offload_flag_per_arch(monkeypatch): + from freetoken.kernel.utils import _hip_cflags + + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201") + + flags = _hip_cflags(["-Wno-unused-command-line-argument"]) + + assert "--offload-arch=gfx1200" in flags + assert "--offload-arch=gfx1201" in flags + assert not any(";" in flag for flag in flags) + + +def test_rocm_link_flags_support_versioned_modular_sdk(monkeypatch, tmp_path): + import torch.utils.cpp_extension as cpp_extension + + from freetoken.kernel import utils + + sdk = tmp_path / "sdk" + library_dir = sdk / "lib" + library_dir.mkdir(parents=True) + versioned_runtime = library_dir / "libamdhip64.so.7" + versioned_runtime.write_bytes(b"") + real_find_spec = importlib.util.find_spec + + def find_spec(name: str): + if name == "_rocm_sdk_core": + return SimpleNamespace(submodule_search_locations=[str(sdk)]) + return real_find_spec(name) + + monkeypatch.delenv("ROCM_HOME", raising=False) + monkeypatch.setattr(cpp_extension, "ROCM_HOME", None) + monkeypatch.setattr(importlib.util, "find_spec", find_spec) + monkeypatch.setattr(pathlib.Path, "home", lambda: tmp_path) + utils._rocm_link_flags.cache_clear() + + flags = utils._rocm_link_flags() + + compat_dir = tmp_path / ".cache" / "freetoken" / "rocm-lib" + compat_link = compat_dir / "libamdhip64.so" + assert f"-L{compat_dir}" in flags + assert f"-Wl,-rpath,{library_dir}" in flags + assert compat_link.resolve() == versioned_runtime.resolve() + + utils._rocm_link_flags.cache_clear()