[DCU] high performance LLM train and inference for DCU

csrc/generation/get_padding_offset_v2.cu

@@ -1,3 +1,17 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 #include "paddle/extension.h"
 
 __global__ void RemovePaddingV2(int64_t *output_data,

csrc/generation/helper.h

@@ -15,12 +15,44 @@
 #pragma once
 
 #include "paddle/extension.h"
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_bfloat16.h>
+#include <hipcub/hipcub.hpp>
+#include <hiprand.h>
+#include <hiprand_kernel.h>
+namespace cub = hipcub;
+#else
 #include <cub/cub.cuh>
 #include <curand_kernel.h>
+#endif
 
 constexpr int kBlockSize = 256; 
 constexpr int kNumWaves = 16; 
 
+#ifdef PADDLE_WITH_HIP
+inline hipError_t GetNumBlocks(int64_t n, int* num_blocks) {
+  int dev;
+  {
+    hipError_t err = hipGetDevice(&dev);
+    if (err != hipSuccess) { return err; }
+  }
+  int sm_count;
+  {
+    hipError_t err = hipDeviceGetAttribute(&sm_count, hipDeviceAttributeMultiprocessorCount, dev);
+    if (err != hipSuccess) { return err; }
+  }
+  int tpm;
+  {
+    hipError_t err = hipDeviceGetAttribute(&tpm, hipDeviceAttributeMaxThreadsPerMultiProcessor, dev);
+    if (err != hipSuccess) { return err; }
+  }
+  *num_blocks = std::max<int>(1, std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
+                                                    sm_count * tpm / kBlockSize * kNumWaves));
+  return hipSuccess;
+}
+#else
 inline cudaError_t GetNumBlocks(int64_t n, int* num_blocks) {
   int dev;
   {
@@ -41,6 +73,7 @@ inline cudaError_t GetNumBlocks(int64_t n, int* num_blocks) {
                                                     sm_count * tpm / kBlockSize * kNumWaves));
   return cudaSuccess;
 }
+#endif
 
 template<typename T>
 __device__ T max_func(const T a, const T b) {
@@ -74,7 +107,11 @@ class PDTraits<paddle::DataType::FLOAT16> {
 template <>
 class PDTraits<paddle::DataType::BFLOAT16> {
 public:
+#ifdef PADDLE_WITH_HIP
+  typedef hip_bfloat16 DataType;
+#else
   typedef __nv_bfloat16 DataType;
+#endif
   typedef paddle::bfloat16 data_t;
 };
 

csrc/generation/quant_int8.cu

@@ -22,8 +22,13 @@
 #include<sys/mman.h>
 #include<stdio.h>
 #include<algorithm>
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_fp16.h>
+#include <hip/hip_bfloat16.h>
+#else
 #include<cuda_fp16.h>
 #include<cuda_bf16.h>
+#endif
 
 
 constexpr int DequantKernelVecSize = 4;
@@ -52,11 +57,17 @@ __forceinline__ __device__ half add_mul<half>(half a, half b, half c) {
     return __hmul(__hadd(a, b), c);
 }
 
+#ifdef PADDLE_WITH_HIP
+template<>
+__forceinline__ __device__ hip_bfloat16 add_mul<hip_bfloat16>(hip_bfloat16 a, hip_bfloat16 b, hip_bfloat16 c) {
+    return (a + b) * c;
+}
+#else
 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>
@@ -173,8 +184,13 @@ std::vector<paddle::Tensor> LaunchQuantInt8(const paddle::Tensor& input,
     auto output=paddle::full(input_shape, -1, paddle::DataType::INT8, input.place());
     int m = input_shape[0];
     int n = input_shape[1];
+#ifdef PADDLE_WITH_HIP
+    dim3 grid(((n >> 2) + 63) / 64, (m + 7) / 8);
+    dim3 block(64, 8);
+#else
     dim3 grid((n >> 2 + 31) / 32, (m + 31) / 32);
     dim3 block(32, 32);
+#endif
     auto stream = input.stream();
     if (shift && smooth) {
         QuantKernel<DataType_><<<grid, block, 0, stream>>>(reinterpret_cast<const DataType_*>(input.data<data_t>()),

csrc/generation/rebuild_padding.cu

@@ -58,7 +58,12 @@ void InvokeRebuildPadding(T *output_data,
                           const int *padding_offset,
                           const int token_num,
                           const int dim_embed,
-                          cudaStream_t stream) {
+#ifdef PADDLE_WITH_HIP
+                          hipStream_t stream
+#else
+                          cudaStream_t stream
+#endif
+                          ) {
   // src: [token_num, dim_embed]
   // dst: [batch_size * max_seq_len, dim_embed]
   RebuildPaddingKernel<<<token_num, 256, 0, stream>>>(

csrc/generation/rebuild_padding_v2.cu

@@ -1,3 +1,17 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 #include "helper.h"
 
 template <typename T, int VecSize>

csrc/generation/set_value_by_flags_v2.cu

@@ -1,3 +1,17 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 #include "paddle/extension.h"
 
 __global__ void set_value_by_flag_and_id_v2(const bool *stop_flags, 

csrc/generation/step.cu

@@ -1,3 +1,17 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 #include "helper.h"
 
 // #define DEBUG_STEP
@@ -255,7 +269,11 @@ void StepPaddle(const paddle::Tensor& stop_flags,
         max_decoder_block_num
     );
 #ifdef DEBUG_STEP
+#ifdef PADDLE_WITH_HIP
+    hipDeviceSynchronize();
+#else
     cudaDeviceSynchronize();
+#endif
 #endif
     auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
     const int grid_size = cpu_recover_lens.data<int>()[0];
@@ -287,7 +305,11 @@ void StepPaddle(const paddle::Tensor& stop_flags,
             first_token_id
         );
 #ifdef DEBUG_STEP
+#ifdef PADDLE_WITH_HIP
+        hipDeviceSynchronize();
+#else
         cudaDeviceSynchronize();
+#endif
 #endif
     }
 }

csrc/generation/stop_generation_multi_ends_v2.cu

@@ -1,3 +1,17 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 #include "paddle/extension.h"
 #include<stdlib.h>
 #include<string.h>

csrc/generation/token_penalty_multi_scores_v2.cu

@@ -1,3 +1,17 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 #include "helper.h"
 
 

csrc/generation/transpose_removing_padding.cu

@@ -65,7 +65,12 @@ void InvokeTransposeRemovePadding(const T* input_data,
                                   const int head_dim,
                                   const int token_num,
                                   const int* padding_offset,
-                                  cudaStream_t cu_stream) {
+#ifdef PADDLE_WITH_HIP
+                                  hipStream_t cu_stream
+#else
+                                  cudaStream_t cu_stream
+#endif
+                                  ) {
   // [batch_size, num_head, max_len_this_time, head_dim] -> [token_num, num_head,
   // head_dim]
   constexpr int VEC_16B = 16;

csrc/generation/write_int8_cache_kv.cu

@@ -14,8 +14,13 @@
 
 #include "helper.h"
 
+#ifdef PADDLE_WITH_HIP
+constexpr int32_t WARP_SIZE = 64; 
+constexpr int32_t HALF_WARP = 32; 
+#else
 constexpr int32_t WARP_SIZE = 32; 
 constexpr int32_t HALF_WARP = 16; 
+#endif
 constexpr float QUANT_MAX_BOUND = 127.0;
 constexpr float QUANT_MIN_BOUND = -127.0;
 
@@ -47,14 +52,22 @@ struct MaxFunc{
 template<>
 struct MaxFunc<half>{
   __device__ half operator()(half a, half b){
-#if __CUDA_ARCH__ >= 800
+#if (__CUDA_ARCH__ >= 800) || defined(PADDLE_WITH_HIP)
     return __hmax(a, b); 
 #else
     return max(static_cast<float>(a), static_cast<float>(b));
 #endif
   }
 }; 
 
+#ifdef PADDLE_WITH_HIP
+template<>
+struct MaxFunc<hip_bfloat16>{
+  __device__ hip_bfloat16 operator()(hip_bfloat16 a, hip_bfloat16 b){
+    return static_cast<hip_bfloat16>(max(static_cast<float>(a), static_cast<float>(b)));
+  }
+}; 
+#else
 template<>
 struct MaxFunc<__nv_bfloat16>{
   __device__ __nv_bfloat16 operator()(__nv_bfloat16 a, __nv_bfloat16 b){
@@ -65,6 +78,7 @@ struct MaxFunc<__nv_bfloat16>{
 #endif
   }
 }; 
+#endif
 
 template<typename T>
 struct AbsFunc{
@@ -76,14 +90,22 @@ struct AbsFunc{
 template<>
 struct AbsFunc<half>{
   __device__ half operator()(half x){
-  #if __CUDA_ARCH__ >= 800
+  #if (__CUDA_ARCH__ >= 800) || defined(PADDLE_WITH_HIP)
     return __habs(x); 
   #else
     return abs(static_cast<float>(x));
   #endif
   }
 }; 
 
+#ifdef PADDLE_WITH_HIP
+template<>
+struct AbsFunc<hip_bfloat16>{
+  __device__ hip_bfloat16 operator()(hip_bfloat16 x) {
+    return static_cast<hip_bfloat16>(abs(static_cast<float>(x)));
+  }
+}; 
+#else
 template<>
 struct AbsFunc<__nv_bfloat16>{
   __device__ __nv_bfloat16 operator()(__nv_bfloat16 x){
@@ -94,6 +116,7 @@ struct AbsFunc<__nv_bfloat16>{
   #endif
   }
 }; 
+#endif
 
 template <typename T, typename Vec, int VecSize>
 __inline__ __device__ T LocalReduceMax(Vec& vec) {
@@ -109,7 +132,11 @@ template <typename T>
 __inline__ __device__ T WarpReduceAbsMax(T val, unsigned lane_mask) {
   #pragma unroll
   for (int mask = HALF_WARP; mask > 0; mask >>= 1){
+#ifdef PADDLE_WITH_HIP
+    val = MaxFunc<T>()(val, static_cast<T>(__shfl_xor(static_cast<float>(val), mask, WARP_SIZE)));
+#else
     val = MaxFunc<T>()(val, __shfl_xor_sync(lane_mask, val, mask, WARP_SIZE));
+#endif
   }
   return val;
 }
@@ -147,7 +174,7 @@ __global__ void write_cache_k_int8_kernel(const T* k, const int64_t num_head, co
     InVec abs_max_vec;
 #pragma unroll
     for (int i = 0; i < VecSize; ++i) {
-      abs_max_vec[i] = 0.0f;
+      abs_max_vec[i] = static_cast<T>(0.0f);
     }
 
     T local_abs_max;
@@ -205,7 +232,7 @@ __global__ void write_cache_v_int8_kernel(const T* v, const int64_t num_head, co
     InVec abs_max_vec;
   #pragma unroll
     for (int i = 0; i < VecSize; ++i) {
-      abs_max_vec[i] = 0.0f;
+      abs_max_vec[i] = static_cast<T>(0.0f);
     }
 
     T local_abs_max;