[RMSNorm][Half] support fp16x8 packed RMSNorm (xlite-dev#22)

* Update elementwise.cu

* Update relu.cu

* Update elementwise.cu

* Update rms_norm.cu

* Update rms_norm.py

* Update README.md

* Update README.md

Author: DefTruth

README.md

@@ -66,8 +66,10 @@
 | ✔️ [layer_norm_f16_f32_kernel(per token)](./layer-norm/layer_norm.cu)|f16|f32|[link](./layer-norm/)|⭐️⭐️|  
 | ✔️ [rms_norm_f32_kernel(per token)](./rms-norm/rms_norm.cu)|f32|f32|[link](./rms-norm/)|⭐️⭐️|  
 | ✔️ [rms_norm_f32x4_kernel(per token)](./rms-norm/rms_norm.cu)|f32|f32|[link](./rms-norm/)|⭐️⭐️|  
-| ❔ [rms_norm_f16_kernel(per token)](./rms-norm/rms_norm.cu)|f16|f16|❔|⭐️⭐️|  
-| ❔ [rms_norm_f16x2_kernel(per token)](./rms-norm/rms_norm.cu)|f16|f16|❔|⭐️⭐️|  
+| ✔️ [rms_norm_f16_f16_kernel(per token)](./rms-norm/rms_norm.cu)|f16|f16|[link](./rms-norm/)|⭐️⭐️|  
+| ✔️ [rms_norm_f16x2_f16_kernel(per token)](./rms-norm/rms_norm.cu)|f16|f16|[link](./rms-norm/)|⭐️⭐️|  
+| ✔️ [rms_norm_f16x8_f16_kernel(per token)](./rms-norm/rms_norm.cu)|f16|f16|[link](./rms-norm/)|⭐️⭐️|  
+| ✔️ [rms_norm_f16_f32_kernel(per token)](./rms-norm/rms_norm.cu)|f16|f32|[link](./rms-norm/)|⭐️⭐️| 
 | ✔️ [sgemm_sliced_k_f32_kernel](./sgemm/sgemm.cu)|f32|f32|[link](./sgemm/)|⭐️⭐️⭐️|  
 | ✔️ [sgemm_t_8x8_sliced_k_f32x4_kernel](./sgemm/sgemm.cu)|f32|f32|[link](./sgemm/)|⭐️⭐️⭐️|  
 | ❔ [hgemm_sliced_k_f16_f32_kernel](./hgemm)|f16|f32|❔|⭐️⭐️⭐️|  

elementwise/elementwise.cu

@@ -13,6 +13,8 @@
 #define WARP_SIZE 32
 #define INT4(value) (reinterpret_cast<int4*>(&(value))[0])
 #define FLOAT4(value) (reinterpret_cast<float4*>(&(value))[0])
+#define HALF2(value) (reinterpret_cast<half2*>(&(value))[0])
+#define BFLOAT2(value) (reinterpret_cast<__nv_bfloat162*>(&(value))[0])
 
 // -------------------------------------- FP32 -------------------------------------- 
 // ElementWise Add  
@@ -40,7 +42,6 @@ __global__ void elementwise_add_f32x4_kernel(float* a, float* b, float* c, int N
   }
 }
 
-
 // -------------------------------------- FP16 -------------------------------------- 
 // ElementWise Add  
 // grid(N/256), block(256)
@@ -54,12 +55,12 @@ __global__ void elementwise_add_f16_kernel(half* a, half* b, half* c, int N) {
 __global__ void elementwise_add_f16x2_kernel(half* a, half* b, half* c, int N) {
   int idx = 2 * (blockIdx.x * blockDim.x + threadIdx.x);
   if (idx < N) {
-    half2 reg_a = (reinterpret_cast<half2*>(&(a[idx]))[0]);
-    half2 reg_b = (reinterpret_cast<half2*>(&(b[idx]))[0]);
+    half2 reg_a = HALF2(a[idx]);
+    half2 reg_b = HALF2(b[idx]);
     half2 reg_c;
     reg_c.x = __hadd(reg_a.x, reg_b.x);
     reg_c.y = __hadd(reg_a.y, reg_b.y);
-    (reinterpret_cast<half2*>(&(c[idx]))[0]) = reg_c;
+    HALF2(c[idx]) = reg_c;
   }
 }
 
@@ -137,4 +138,4 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   TORCH_BINDING_COMMON_EXTENSION(elementwise_add_f32x4_v2)
   TORCH_BINDING_COMMON_EXTENSION(elementwise_add_f16_v2)
   TORCH_BINDING_COMMON_EXTENSION(elementwise_add_f16x2_v2)
-}
+}

relu/relu.cu

@@ -11,6 +11,8 @@
 #define WARP_SIZE 32
 #define INT4(value) (reinterpret_cast<int4*>(&(value))[0])
 #define FLOAT4(value) (reinterpret_cast<float4*>(&(value))[0])
+#define HALF2(value) (reinterpret_cast<half2*>(&(value))[0])
+#define BFLOAT2(value) (reinterpret_cast<__nv_bfloat162*>(&(value))[0])
 
 // -------------------------------------- FP32 -------------------------------------- 
 // Relu x: N, y: N y=max(0,x)
@@ -44,11 +46,11 @@ __global__ void relu_f16_kernel(half* x, half* y, int N) {
 __global__ void relu_f16x2_kernel(half* x, half* y, int N) {
   int idx = 2 * (blockIdx.x * blockDim.x + threadIdx.x);
   if (idx < N) {
-    half2 reg_x = (reinterpret_cast<half2*>(&(x[idx]))[0]);
-    half2 reg_y = (reinterpret_cast<half2*>(&(y[idx]))[0]);
+    half2 reg_x = HALF2(x[idx]);
+    half2 reg_y = HALF2(y[idx]);
     reg_y.x = __hmax(__float2half(0.0f), reg_x.x);
     reg_y.y = __hmax(__float2half(0.0f), reg_x.y);
-    (reinterpret_cast<half2*>(&(y[idx]))[0]) = reg_y;
+    HALF2(y[idx]) = reg_y;
   }
 }
 

rms-norm/README.md

@@ -5,7 +5,11 @@
 包含以下内容：
 
 - [X] rms_norm_f32_kernel 
-- [X] rms_norm_f32x4_kernel(float4向量化版本)
+- [X] rms_norm_f32x4_kernel
+- [X] rms_norm_f16_f16_kernel
+- [X] rms_norm_f16x2_f16_kernel
+- [X] rms_norm_f16x8_f16_kernel
+- [X] rms_norm_f16_f32_kernel
 - [X] PyTorch bindings
 
 ## 测试
@@ -20,8 +24,14 @@ python3 rms_norm.py
 
 ```bash
 --------------------------------------------------------------------------------
-     out_f32: [0.58164781, 0.37353456, 0.12243245], time:0.01189494ms
-   out_f32x4: [0.58164781, 0.37353456, 0.12243245], time:0.00514793ms
-  out_f32_th: [0.58165121, 0.37353677, 0.12243317], time:0.04241681ms
+     out_f32: [0.66361254, 0.69628561, 0.51101440], time:0.01188707ms
+   out_f32x4: [0.66361260, 0.69628561, 0.51101440], time:0.00833464ms
+  out_f32_th: [0.66361588, 0.69628906, 0.51101691], time:0.04334593ms
 --------------------------------------------------------------------------------
-```
+  out_f16f16: [0.66357422, 0.69580078, 0.51074219], time:0.01201081ms
+out_f16x2f16: [0.66357422, 0.69580078, 0.51074219], time:0.00668955ms
+out_f16x8f16: [0.66650391, 0.69921875, 0.51318359], time:0.00398421ms
+  out_f16f32: [0.66357422, 0.69628906, 0.51123047], time:0.01176858ms
+  out_f16_th: [0.66357422, 0.69580078, 0.51074219], time:0.04448509ms
+--------------------------------------------------------------------------------
+```