Fix compile error when using rocm (#2648)

2024-02-02 01:35:09 +08:00
parent cd9e60c76c
commit 923797fea4
3 changed files with 9 additions and 1 deletions
--- a/csrc/cache_kernels.cu
+++ b/csrc/cache_kernels.cu
@@ -4,13 +4,20 @@

 #include "cuda_compat.h"
 #include "dispatch_utils.h"
+#ifdef ENABLE_FP8_E5M2
 #include "quantization/fp8_e5m2_kvcache/quant_utils.cuh"
+#endif

 #include <algorithm>
 #include <cassert>
 #include <map>
 #include <vector>

+#ifdef USE_ROCM
+  #include <hip/hip_bf16.h>
+  typedef __hip_bfloat16 __nv_bfloat16;
+#endif
+
 void swap_blocks(
  torch::Tensor& src,
  torch::Tensor& dst,