fix rocm build error

Author: jiafei96

csrc/custom_marlin/gptq_marlin/gptq_marlin.cu

@@ -34,7 +34,7 @@ template <typename T> inline std::string str(T x) { return std::to_string(x); }
 
 namespace gptq_marlin {
 
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800) || defined(__HIP_PLATFORM_AMD__)
 
     __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
         int const* __restrict__ perm_int_ptr,

csrc/custom_marlin/gptq_marlin/gptq_marlin.cuh

@@ -38,7 +38,7 @@ using I4 = Vec<int, 4>;
 
 constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
 
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800) || defined(__HIP_PLATFORM_AMD__)
 // No support for async
 #else
 

csrc/custom_marlin/gptq_marlin/gptq_marlin_dtypes.cuh

@@ -8,6 +8,11 @@
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
 
+#ifdef __HIP_PLATFORM_AMD__
+typedef __hip_bfloat16 nv_bfloat16;
+typedef __hip_bfloat162 nv_bfloat162;
+#endif
+
 namespace gptq_marlin {
 
 template <typename scalar_t> class ScalarType {};

csrc/custom_marlin/gptq_marlin/gptq_marlin_repack.cu

@@ -9,7 +9,7 @@ static constexpr int repack_threads = 256;
 static constexpr int tile_k_size = tile_size;
 static constexpr int tile_n_size = tile_k_size * 4;
 
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800) || defined(__HIP_PLATFORM_AMD__)
 
 template <int const num_threads, int const num_bits, bool const has_perm>
 __global__ void marlin_repack_kernel(

ktransformers/operators/linear.py

@@ -16,7 +16,6 @@
 from torch import Tensor, nn
 if not torch.xpu.is_available():
     import KTransformersOps
-    import vLLMMarlin
 from ktransformers.util.custom_loader import GGUFLoader, SafeTensorLoader
 from ktransformers.util.utils import InferenceState
 if not torch.xpu.is_available():
@@ -520,6 +519,7 @@ def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor = None) -> torch.Ten
         # padding x.shape[0] to avoid CUDA illegal memory access error
         x, orig_size_m = self._pad_input(x)
 
+        import vLLMMarlin
         x = vLLMMarlin.gptq_marlin_gemm(
             x,
             self.marlin_q_w,