Add ROCm support

Author: ardfork

cuda_ext.py

@@ -11,21 +11,30 @@
 library_dir = os.path.dirname(os.path.abspath(__file__))
 extension_name = "exllama_ext"
 
+if torch.version.hip:
+    # FIXME: To build, I had to comment "flags += ['-fno-gpu-rdc']" in torch/utils/cpp_extension.py.
+    # I am not sure if it's possible to find a way to build without editing that file.
+    # If building without gpu-rdc, build will error with "lld: error: undefined hidden symbol: __llvm_amdgcn_rcp_f16".
+    extra_cuda_cflags= ["-U__HIP_NO_HALF_CONVERSIONS__", "-fgpu-rdc"]
+else:
+    extra_cuda_cflags = []
+
 exllama_ext = load(
     name = extension_name,
     sources = [
         os.path.join(library_dir, "exllama_ext/cuda_buffers.cu"),
         os.path.join(library_dir, "exllama_ext/cpu_func/rep_penalty.cpp"),
-        os.path.join(library_dir, "exllama_ext/cuda_func/column_remap.cu"),
-        os.path.join(library_dir, "exllama_ext/cuda_func/half_matmul.cu"),
-        os.path.join(library_dir, "exllama_ext/cuda_func/q4v2_matmul.cu"),
-        os.path.join(library_dir, "exllama_ext/cuda_func/q4v2_mlp.cu"),
-        os.path.join(library_dir, "exllama_ext/cuda_func/q4v2_recons.cu"),
-        os.path.join(library_dir, "exllama_ext/cuda_func/q4v2_sequential.cu"),
-        os.path.join(library_dir, "exllama_ext/cuda_func/rms_norm.cu"),
-        os.path.join(library_dir, "exllama_ext/cuda_func/rope.cu"),
+        os.path.join(library_dir, "exllama_ext/cu_func/column_remap.cu"),
+        os.path.join(library_dir, "exllama_ext/cu_func/half_matmul.cu"),
+        os.path.join(library_dir, "exllama_ext/cu_func/q4v2_matmul.cu"),
+        os.path.join(library_dir, "exllama_ext/cu_func/q4v2_mlp.cu"),
+        os.path.join(library_dir, "exllama_ext/cu_func/q4v2_recons.cu"),
+        os.path.join(library_dir, "exllama_ext/cu_func/q4v2_sequential.cu"),
+        os.path.join(library_dir, "exllama_ext/cu_func/rms_norm.cu"),
+        os.path.join(library_dir, "exllama_ext/cu_func/rope.cu"),
         os.path.join(library_dir, "exllama_ext/exllama_ext.cpp")
     ],
+    extra_cuda_cflags = extra_cuda_cflags
     # verbose = True,
     # extra_cflags = ["-ftime-report", "-DTORCH_USE_CUDA_DSA"]
 )

exllama_ext/cuda_func/column_remap.cu renamed to exllama_ext/cu_func/column_remap.cu

@@ -1,8 +1,14 @@
 #ifndef _column_remap_cuh
 #define _column_remap_cuh
 
+#if USE_ROCM
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp16.h>
+#define cudaError_t hipError_t
+#else
 #include <cuda_runtime.h>
 #include <cuda_fp16.h>
+#endif
 #include <cstdint>
 
 cudaError_t column_remap_cuda

exllama_ext/cuda_func/column_remap.cuh renamed to exllama_ext/cu_func/column_remap.cuh

@@ -102,7 +102,7 @@ cudaError_t half_matmul_cublas_cuda
     const half alpha = __float2half(1.0f);
     const half beta = __float2half(0.0f);
 
-    cublasHgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, width, height, dim, &alpha, w, width, x, dim, &beta, out, width);
+    cublasHgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, width, height, dim, reinterpret_cast<const rocblas_half*>(&alpha), reinterpret_cast<const rocblas_half*>(w), width, reinterpret_cast<const rocblas_half*>(x), dim, reinterpret_cast<const rocblas_half*>(&beta), reinterpret_cast<rocblas_half*>(out), width);
 
 //     cudaDeviceSynchronize();
 //     _cuda_check(cudaGetLastError());

exllama_ext/cuda_func/half_matmul.cu renamed to exllama_ext/cu_func/half_matmul.cu

@@ -1,10 +1,19 @@
 #ifndef _half_matmul_cuh
 #define _half_matmul_cuh
 
+#if USE_ROCM
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp16.h>
+#include <rocblas/rocblas.h>
+#include <ATen/hip/HIPContext.h>
+#define cudaError_t hipError_t
+#define cublasHandle_t rocblas_handle
+#else
 #include <cuda_runtime.h>
 #include <cuda_fp16.h>
-#include <cstdint>
 #include <ATen/cuda/CUDAContext.h>
+#endif
+#include <cstdint>
 
 cudaError_t half_matmul_cuda
 (

exllama_ext/cuda_func/half_matmul.cuh renamed to exllama_ext/cu_func/half_matmul.cuh

@@ -1,8 +1,14 @@
 #ifndef _q4v2_matmul_cuh
 #define _q4v2_matmul_cuh
 
+#if USE_ROCM
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp16.h>
+#define cudaError_t hipError_t
+#else
 #include <cuda_runtime.h>
 #include <cuda_fp16.h>
+#endif
 #include <cstdint>
 #include <cstdio>
 

exllama_ext/cuda_func/q4v2_matmul.cu renamed to exllama_ext/cu_func/q4v2_matmul.cu

@@ -1,8 +1,14 @@
 #ifndef _q4v2_mlp_cuh
 #define _q4v2_mlp_cuh
 
+#if USE_ROCM
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp16.h>
+#define cudaError_t hipError_t
+#else
 #include <cuda_runtime.h>
 #include <cuda_fp16.h>
+#endif
 #include <cstdint>
 
 cudaError_t q4v2_mlp_cuda