support v100 csrc

RichardWooSJTU · Dec 7, 2023 · 6d92e6c · 6d92e6c
1 parent 5eb2eae
commit 6d92e6c
Show file tree

Hide file tree

Showing 11 changed files with 34 additions and 7 deletions.
diff --git a/csrc/generation/dequant_int8.cu b/csrc/generation/dequant_int8.cu
@@ -116,9 +116,11 @@ std::vector<paddle::Tensor> LaunchDequantInt8(const paddle::Tensor& input,
                 "Only bfloat16, float16 and float32 are supported. ");
 
     switch (data_type) {
+#if CUDA_VERSION >= 11000
         case paddle::DataType::BFLOAT16: 
             return DispatchLaunchDequantInt8<paddle::DataType::BFLOAT16>(input, scale);
             break;
+#endif
         case paddle::DataType::FLOAT16: 
             return DispatchLaunchDequantInt8<paddle::DataType::FLOAT16>(input, scale);
             break;
@@ -145,8 +147,10 @@ std::vector<paddle::DataType> DequantInt8Dtype(const paddle::DataType& input_dty
     paddle::DataType data_type;
     if (dtype == "float32")
         data_type = paddle::DataType::FLOAT32;
+#if CUDA_VERSION >= 11000
     else if (dtype == "bfloat16")
         data_type = paddle::DataType::BFLOAT16;
+#endif
     else if (dtype ==  "float16")
         data_type = paddle::DataType::FLOAT16;
     else 
@@ -163,4 +167,4 @@ PD_BUILD_OP(dequant_int8)
     .Attrs({"dtype: std::string"})
     .SetKernelFn(PD_KERNEL(DequantInt8))
     .SetInferShapeFn(PD_INFER_SHAPE(DequantInt8Shape))
-    .SetInferDtypeFn(PD_INFER_DTYPE(DequantInt8Dtype));
+    .SetInferDtypeFn(PD_INFER_DTYPE(DequantInt8Dtype));
diff --git a/csrc/generation/encode_rotary_qk.cu b/csrc/generation/encode_rotary_qk.cu
@@ -193,11 +193,13 @@ void RotaryQK(const paddle::Tensor& q,
               const int32_t rotary_emb_dims, 
               bool use_neox) {
     switch (q.type()) {
+#if CUDA_VERSION >= 11000
         case paddle::DataType::BFLOAT16: {
             return LaunchRotaryQK<paddle::DataType::BFLOAT16>(
                 q, kv, rotary_emb, seq_lens, rotary_emb_dims, use_neox
             );
         }
+#endif 
         case paddle::DataType::FLOAT16: {
             return LaunchRotaryQK<paddle::DataType::FLOAT16>(
                 q, kv, rotary_emb, seq_lens, rotary_emb_dims, use_neox

diff --git a/csrc/generation/helper.h b/csrc/generation/helper.h
@@ -15,7 +15,6 @@
 #pragma once
 
 #include "paddle/extension.h"
-#include <cub/cub.cuh>
 #include <curand_kernel.h>
 
 constexpr int kBlockSize = 256; 
@@ -70,13 +69,14 @@ class PDTraits<paddle::DataType::FLOAT16> {
   typedef half DataType;
   typedef paddle::float16 data_t;
 };
-
+#if CUDA_VERSION >= 11000
 template <>
 class PDTraits<paddle::DataType::BFLOAT16> {
 public:
   typedef __nv_bfloat16 DataType;
   typedef paddle::bfloat16 data_t;
 };
+#endif
 
 template <typename T, int Size>
 struct alignas(sizeof(T) * Size) AlignedVector {
@@ -100,4 +100,4 @@ HOSTDEVICE inline void Store(const AlignedVector<T, Size>& vec, T* addr) {
   *addr_vec = vec;
 }
 
-constexpr int VEC_16B = 16;
+constexpr int VEC_16B = 16;
diff --git a/csrc/generation/qkv_transpose_split.cu b/csrc/generation/qkv_transpose_split.cu
@@ -135,6 +135,7 @@ std::vector<paddle::Tensor> QKVTransposeSplit(const paddle::Tensor& qkv,
                                               int num_head,
                                               int head_size) {
     switch (qkv.type()) {
+#if CUDA_VERSION >= 11000
         case paddle::DataType::BFLOAT16: {
             return qkv_transpose_split<paddle::DataType::BFLOAT16>(
                 qkv,
@@ -145,6 +146,7 @@ std::vector<paddle::Tensor> QKVTransposeSplit(const paddle::Tensor& qkv,
                 head_size
             );
         }
+#endif 
         case paddle::DataType::FLOAT16: {
             return qkv_transpose_split<paddle::DataType::FLOAT16>(
                 qkv,

diff --git a/csrc/generation/quant_int8.cu b/csrc/generation/quant_int8.cu
@@ -23,8 +23,9 @@
 #include<stdio.h>
 #include<algorithm>
 #include<cuda_fp16.h>
+#if CUDA_VERSION >= 11000
 #include<cuda_bf16.h>
-
+#endif
 
 constexpr int DequantKernelVecSize = 4;
 
@@ -52,11 +53,12 @@ __forceinline__ __device__ half add_mul<half>(half a, half b, half c) {
     return __hmul(__hadd(a, b), c);
 }
 
+#if CUDA_VERSION >= 11000
 template<>
 __forceinline__ __device__ __nv_bfloat16 add_mul<__nv_bfloat16>(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) {
     return __hmul(__hadd(a, b), c);
 }
-
+#endif
 
 
 template <typename data_t>
@@ -210,11 +212,13 @@ std::vector<paddle::Tensor> QuantInt8(const paddle::Tensor& input,
                                       float min_bound) {
     // printf("#### quant int8 scale:%f \n",scale);
     switch (input.type()) {
+#if CUDA_VERSION >= 11000
         case paddle::DataType::BFLOAT16: {
             return LaunchQuantInt8<paddle::DataType::BFLOAT16>(
                 input, shift, smooth, scale, round_type, max_bound, min_bound
             );
         }
+#endif
         case paddle::DataType::FLOAT16: {
             return LaunchQuantInt8<paddle::DataType::FLOAT16>(
                 input, shift, smooth, scale, round_type, max_bound, min_bound
@@ -256,4 +260,4 @@ PD_BUILD_OP(quant_int8)
     .Attrs({"scale: float","round_type: int","max_bound: float", "min_bound: float"})
     .SetKernelFn(PD_KERNEL(QuantInt8))
     .SetInferShapeFn(PD_INFER_SHAPE(QuantInt8Shape))
-    .SetInferDtypeFn(PD_INFER_DTYPE(QuantInt8Dtype));
+    .SetInferDtypeFn(PD_INFER_DTYPE(QuantInt8Dtype));
diff --git a/csrc/generation/rebuild_padding.cu b/csrc/generation/rebuild_padding.cu
@@ -109,6 +109,7 @@ std::vector<paddle::Tensor> RebuildPadding(const paddle::Tensor& tmp_out,
                                            const paddle::Tensor& seq_lens,
                                            const paddle::Tensor& input_ids) {
     switch (tmp_out.type()) {
+#if CUDA_VERSION >= 11000       
         case paddle::DataType::BFLOAT16: {
             return rebuild_padding<paddle::DataType::BFLOAT16>(
                 tmp_out,
@@ -117,6 +118,7 @@ std::vector<paddle::Tensor> RebuildPadding(const paddle::Tensor& tmp_out,
                 input_ids
             );
         }
+#endif 
         case paddle::DataType::FLOAT16: {
             return rebuild_padding<paddle::DataType::FLOAT16>(
                 tmp_out,

diff --git a/csrc/generation/set_alibi_mask_value.cu b/csrc/generation/set_alibi_mask_value.cu
@@ -76,6 +76,7 @@ std::vector<paddle::Tensor> SetMaskValue(const paddle::Tensor& input_data,
                                           const paddle::Tensor& alibi_slopes,
                                           const paddle::Tensor& tgt_pos) {
     switch (input_data.type()) {
+#if CUDA_VERSION >= 11000
         case paddle::DataType::BFLOAT16: {
             return set_mask_value<paddle::DataType::BFLOAT16>(
                 input_data,
@@ -85,6 +86,7 @@ std::vector<paddle::Tensor> SetMaskValue(const paddle::Tensor& input_data,
                 tgt_pos
             );
         }
+#endif 
         case paddle::DataType::FLOAT16: {
             return set_mask_value<paddle::DataType::FLOAT16>(
                 input_data,

diff --git a/csrc/generation/set_mask_value.cu b/csrc/generation/set_mask_value.cu
@@ -14,6 +14,7 @@
 
 #include "paddle/extension.h"
 
+
 template <paddle::DataType D>
 class PDTraits;
 
@@ -31,12 +32,14 @@ public:
   typedef paddle::float16 data_t;
 };
 
+#if CUDA_VERSION >= 11000
 template <>
 class PDTraits<paddle::DataType::BFLOAT16> {
 public:
   typedef __nv_bfloat16 DataType;
   typedef paddle::bfloat16 data_t;
 };
+#endif
 
 template <typename T>
 __global__ void set_value_by_id(const int *seq_lens, const bool *stop_flags, T *output_data, int *sequence_lengths, int bs, int length) {
@@ -77,13 +80,15 @@ std::vector<paddle::Tensor> set_mask_value(const paddle::Tensor& input_data, con
 
 std::vector<paddle::Tensor> SetMaskValue(const paddle::Tensor& input_data, const paddle::Tensor& stop_flags, const paddle::Tensor& seq_lens) {
     switch (input_data.type()) {
+#if CUDA_VERSION >= 11000       
         case paddle::DataType::BFLOAT16: {
             return set_mask_value<paddle::DataType::BFLOAT16>(
                 input_data,
                 stop_flags,
                 seq_lens
             );
         }
+#endif 
         case paddle::DataType::FLOAT16: {
             return set_mask_value<paddle::DataType::FLOAT16>(
                 input_data,

diff --git a/csrc/generation/token_penalty_multi_scores.cu b/csrc/generation/token_penalty_multi_scores.cu
@@ -156,6 +156,7 @@ std::vector<paddle::Tensor> TokenPenaltyMultiScores(const paddle::Tensor& pre_id
                                                     const paddle::Tensor& eos_token_id) {
 
     switch (logits.type()) {
+#if CUDA_VERSION >= 11000
         case paddle::DataType::BFLOAT16: {
             return token_penalty_multi_scores_kernel<paddle::DataType::BFLOAT16>(
                 pre_ids,
@@ -168,6 +169,7 @@ std::vector<paddle::Tensor> TokenPenaltyMultiScores(const paddle::Tensor& pre_id
                 eos_token_id
             );
         }
+#endif 
         case paddle::DataType::FLOAT16: {
             return token_penalty_multi_scores_kernel<paddle::DataType::FLOAT16>(
                 pre_ids,

diff --git a/csrc/generation/transpose_removing_padding.cu b/csrc/generation/transpose_removing_padding.cu
@@ -125,13 +125,15 @@ std::vector<paddle::Tensor> ApplyTransposeRemovingPadding(const paddle::Tensor&
                                                           const paddle::Tensor& seq_lens, 
                                                           const paddle::Tensor& padding_offset) {
     switch (input.type()) {
+#if CUDA_VERSION >= 11000
         case paddle::DataType::BFLOAT16: {
             return apply_transpose_remove_padding<paddle::DataType::BFLOAT16>(
                 input,
                 seq_lens,
                 padding_offset
             );
         }
+#endif 
         case paddle::DataType::FLOAT16: {
             return apply_transpose_remove_padding<paddle::DataType::FLOAT16>(
                 input,

diff --git a/csrc/generation/write_cache_kv.cu b/csrc/generation/write_cache_kv.cu
@@ -154,11 +154,13 @@ void WriteCacheKV(const paddle::Tensor& input_k,
                   const paddle::Tensor& cache_kv,
                   const paddle::Tensor& sequence_lengths_shape) {
     switch (cache_kv.type()) {
+#if CUDA_VERSION >= 11000
         case paddle::DataType::BFLOAT16: {
             return LaunchWriteCacheKV<paddle::DataType::BFLOAT16>(
                 input_k, input_v, cache_kv, sequence_lengths_shape
             );
         }
+#endif 
         case paddle::DataType::FLOAT16: {
             return LaunchWriteCacheKV<paddle::DataType::FLOAT16>(
                 input_k, input_v, cache_kv, sequence_lengths_shape