add hip_bf16.h

7477e8f3 · zhuwenwen · e00b0a19 · 7477e8f3
Commit 7477e8f3 authored Mar 23, 2024 by zhuwenwen
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 4 deletions

csrc/cache_kernels.cu csrc/cache_kernels.cu +4 -4

No files found.
--- a/csrc/cache_kernels.cu
+++ b/csrc/cache_kernels.cu
@@ -13,10 +13,10 @@
 #include <map>
 #include <vector>

-// #ifdef USE_ROCM
-//   #include <hip/hip_bf16.h>
-//   typedef __hip_bfloat16 __nv_bfloat16;
-// #endif
+#ifdef USE_ROCM
+  #include <hip/hip_bf16.h>
+  typedef __hip_bfloat16 __nv_bfloat16;
+#endif

 void swap_blocks(
  torch::Tensor& src,