Merge pull request #5469 from Courtesy-Xs/add_vec_traits

Refactor vector utils

extensions/csrc/cuda/decode_kv_cache_memcpy_kernel.cu

@@ -1,7 +1,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <torch/extension.h>
 
-#include "../common/vector_copy_utils.h"
+#include "utils/vector_copy_utils.h"
 #include "../common/micros.h"
 
 template<typename scalar_t, int VecSize>

extensions/csrc/cuda/fused_rotary_emb_and_cache_kernel.cu

@@ -2,7 +2,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <torch/extension.h>
 
-#include "../common/vector_copy_utils.h"
+#include "utils/vector_copy_utils.h"
 #include "../common/micros.h"
 
 template <typename scalar_t, int VecSize>

extensions/csrc/cuda/rms_layernorm_kernel.cu

@@ -10,7 +10,7 @@
 
 #include "block_reduce.h"
 #include "../common/micros.h"
-#include "../common/cuda_type_utils.h"
+#include "utils/cuda_type_utils.h"
 
 #define DISPATCH_RMSNORM_FLOAT_HALF_AND_BFLOAT(DATA_SIZE, TYPE, NAME, ...)  \
   if (DATA_SIZE == 2) {                                                     \

extensions/csrc/cuda/scaled_masked_softmax.h

@@ -6,52 +6,14 @@
 #include <assert.h>
 #include <c10/macros/Macros.h>
 #include <cuda_fp16.h>
-#include <stdint.h>
 
 #include <cfloat>
 #include <limits>
 
+#include "utils/vector_copy_utils.h"
+
 namespace {
 
-template <typename Datatype, int ELEMENTS_PER_LDG>
-__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
-
-template <>
-__device__ __inline__ void copy_vector<c10::BFloat16, 1>(
-    c10::BFloat16 *dst, const c10::BFloat16 *src) {
-  *dst = *src;
-}
-
-template <>
-__device__ __inline__ void copy_vector<c10::BFloat16, 4>(
-    c10::BFloat16 *dst, const c10::BFloat16 *src) {
-  *((float2 *)dst) = *((float2 *)src);
-}
-
-template <>
-__device__ __inline__ void copy_vector<c10::Half, 1>(c10::Half *dst,
-                                                     const c10::Half *src) {
-  *dst = *src;
-}
-
-template <>
-__device__ __inline__ void copy_vector<c10::Half, 4>(c10::Half *dst,
-                                                     const c10::Half *src) {
-  *((float2 *)dst) = *((float2 *)src);
-}
-
-template <>
-__device__ __inline__ void copy_vector<uint8_t, 1>(uint8_t *dst,
-                                                   const uint8_t *src) {
-  *dst = *src;
-}
-
-template <>
-__device__ __inline__ void copy_vector<uint8_t, 4>(uint8_t *dst,
-                                                   const uint8_t *src) {
-  *((half2 *)dst) = *((half2 *)src);
-}
-
 int log2_ceil(int value) {
   int log2_value = 0;
   while ((1 << log2_value) < value) ++log2_value;

extensions/csrc/cuda/scaled_upper_triang_masked_softmax.h

@@ -13,70 +13,6 @@
 
 namespace {
 
-template <typename Datatype, int ELEMENTS_PER_LDG>
-__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
-
-template <>
-__device__ __inline__ void copy_vector<c10::BFloat16, 1>(
-    c10::BFloat16 *dst, const c10::BFloat16 *src) {
-  *dst = *src;
-}
-
-template <>
-__device__ __inline__ void copy_vector<c10::BFloat16, 4>(
-    c10::BFloat16 *dst, const c10::BFloat16 *src) {
-  *((float2 *)dst) = *((float2 *)src);
-}
-
-template <>
-__device__ __inline__ void copy_vector<c10::Half, 1>(c10::Half *dst,
-                                                     const c10::Half *src) {
-  *dst = *src;
-}
-
-template <>
-__device__ __inline__ void copy_vector<c10::Half, 4>(c10::Half *dst,
-                                                     const c10::Half *src) {
-  *((float2 *)dst) = *((float2 *)src);
-}
-
-template <>
-__device__ __inline__ void copy_vector<uint8_t, 1>(uint8_t *dst,
-                                                   const uint8_t *src) {
-  *dst = *src;
-}
-
-template <>
-__device__ __inline__ void copy_vector<uint8_t, 4>(uint8_t *dst,
-                                                   const uint8_t *src) {
-  *((half2 *)dst) = *((half2 *)src);
-}
-
-template <typename Datatype, int ELEMENTS_PER_LDG>
-__device__ __inline__ void copy_zero_vector(Datatype *dst);
-
-template <>
-__device__ __inline__ void copy_zero_vector<c10::BFloat16, 1>(
-    c10::BFloat16 *dst) {
-  *dst = 0.0;
-}
-
-template <>
-__device__ __inline__ void copy_zero_vector<c10::BFloat16, 4>(
-    c10::BFloat16 *dst) {
-  *((float2 *)dst) = make_float2(0.0f, 0.0f);
-}
-
-template <>
-__device__ __inline__ void copy_zero_vector<c10::Half, 1>(c10::Half *dst) {
-  *dst = 0.0;
-}
-
-template <>
-__device__ __inline__ void copy_zero_vector<c10::Half, 4>(c10::Half *dst) {
-  *((float2 *)dst) = make_float2(0.0f, 0.0f);
-}
-
 int log2_ceil(int value) {
   int log2_value = 0;
   while ((1 << log2_value) < value) ++log2_value;

extensions/csrc/cuda/utils/vec_type_traits.h

@@ -0,0 +1,12 @@
+#pragma once
+
+namespace colossalAI {
+namespace cuda {
+namespace utils {
+
+template <typename T, size_t VecSize>
+class VecTypeTraits {};
+
+}  // namespace utils
+}  // namespace cuda
+}  // namespace colossalAI

extensions/csrc/cuda/utils/vector_copy_utils.h

@@ -1,11 +1,12 @@
 
+#pragma once
+
 #include <c10/macros/Macros.h>
 #include <cuda_fp16.h>
+#include <stdint.h>
 
 #include <cfloat>
 
-#include "string"
-
 template <typename Datatype, int ELEMENTS_PER_LDG>
 __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
 
@@ -57,6 +58,18 @@ __device__ __inline__ void copy_vector<c10::Half, 8>(c10::Half *dst,
   *((float4 *)dst) = *((float4 *)src);
 }
 
+template <>
+__device__ __inline__ void copy_vector<uint8_t, 1>(uint8_t *dst,
+                                                   const uint8_t *src) {
+  *dst = *src;
+}
+
+template <>
+__device__ __inline__ void copy_vector<uint8_t, 4>(uint8_t *dst,
+                                                   const uint8_t *src) {
+  *((half2 *)dst) = *((half2 *)src);
+}
+
 template <>
 __device__ __inline__ void copy_vector<float, 1>(float *dst, const float *src) {
   *dst = *src;
@@ -80,6 +93,31 @@ __device__ __inline__ void copy_vector<float, 8>(float *dst, const float *src) {
   *((float4 *)(dst + 4)) = *((float4 *)(src + 4));
 }
 
+template <typename Datatype, int ELEMENTS_PER_LDG>
+__device__ __inline__ void copy_zero_vector(Datatype *dst);
+
+template <>
+__device__ __inline__ void copy_zero_vector<c10::BFloat16, 1>(
+    c10::BFloat16 *dst) {
+  *dst = 0.0;
+}
+
+template <>
+__device__ __inline__ void copy_zero_vector<c10::BFloat16, 4>(
+    c10::BFloat16 *dst) {
+  *((float2 *)dst) = make_float2(0.0f, 0.0f);
+}
+
+template <>
+__device__ __inline__ void copy_zero_vector<c10::Half, 1>(c10::Half *dst) {
+  *dst = 0.0;
+}
+
+template <>
+__device__ __inline__ void copy_zero_vector<c10::Half, 4>(c10::Half *dst) {
+  *((float2 *)dst) = make_float2(0.0f, 0.0f);
+}
+
 template <typename T>
 int get_vec_size(const torch::Tensor &tensor) {
   uint64_t address = reinterpret_cast<uint64_t>(tensor.data_ptr<T>());