register bf16 dtype for top_p_sampling kernel (PaddlePaddle#67769)

changeyoung98 · Aug 28, 2024 · 4eb9813 · 4eb9813
1 parent dddda15
commit 4eb9813
Showing 1 changed file with 19 additions and 1 deletion.
diff --git a/paddle/phi/kernels/gpu/top_p_sampling_kernel.cu b/paddle/phi/kernels/gpu/top_p_sampling_kernel.cu
@@ -26,6 +26,11 @@ namespace cub = hipcub;
 #include <cub/cub.cuh>
 #endif
 
+#if defined(__CUDACC__) && CUDA_VERSION >= 11060
+#define CUDA_BFLOAT16_AVALIABLE
+#include <cuda_bf16.h>
+#endif
+
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/backends/gpu/gpu_device_function.h"
 #include "paddle/phi/core/kernel_registry.h"
@@ -55,7 +60,7 @@ struct DataTypeTraits<phi::dtype::float16> {
   using DataType = half;
 };
 
-#ifdef PADDLE_CUDA_BF16
+#ifdef CUDA_BFLOAT16_AVALIABLE
 template <>
 struct DataTypeTraits<phi::dtype::bfloat16> {
   using DataType = __nv_bfloat16;
@@ -1241,6 +1246,18 @@ void TopPSamplingKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
+#ifdef CUDA_BFLOAT16_AVALIABLE
+PD_REGISTER_KERNEL(top_p_sampling,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::TopPSamplingKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+#else
 PD_REGISTER_KERNEL(top_p_sampling,
                    GPU,
                    ALL_LAYOUT,
@@ -1250,3 +1267,4 @@ PD_REGISTER_KERNEL(top_p_sampling,
                    int,
                    int64_t,
                    phi::dtype::float16) {}
+#endif