[HGEMM] Add slicked_k&t_8x8_sliced_k_f16x4 (#31)

DefTruth · web-flow · commit b417e20c18e0 · 2024-09-17T11:17:13.000+08:00
* [Refactor][6/N] CUDA Learn Notes refactor Part-6

* [Refactor][6/N] CUDA Learn Notes refactor Part-6

* [HGEMM] Add slicked_k&amp;t_8x8_sliced_k_f16x4
diff --git a/README.md b/README.md
@@ -77,8 +77,8 @@
 | ✔️ [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|❔|⭐️⭐️⭐️|  
-| ❔ [hgemm_t_8x8_sliced_k_f16x2_f32_kernel](./hgemm)|f16|f32|❔|⭐️⭐️⭐️|  
+| ✔️ [hgemm_sliced_k_f16_kernel](./hgemm)|f16|f16|[link](./hgemm/)|⭐️⭐️⭐️|  
+| ✔️ [hgemm_t_8x8_sliced_k_f16x4_kernel](./hgemm)|f16|f16|[link](./hgemm/)|⭐️⭐️⭐️|  
 | ✔️ [sgemv_k32_f32_kernel](./sgemv/sgemv.cu)|f32|f32|[link](./sgemv/)|⭐️⭐️⭐️|  
 | ✔️ [sgemv_k128_f32x4_kernel](./sgemv/sgemv.cu)|f32|f32|[link](./sgemv/)|⭐️⭐️⭐️|  
 | ✔️ [sgemv_k16_f32_kernel](./sgemv/sgemv.cu)|f32|f32|[link](./sgemv/)|⭐️⭐️⭐️|  
diff --git a/hgemm/README.md b/hgemm/README.md
@@ -0,0 +1,27 @@
+# HGEMM 
+
+## 0x00 说明
+
+包含以下内容：
+
+- [X] hgemm_sliced_k_f16_kernel 
+- [X] hgemm_t_8x8_sliced_k_f16x4_kernel
+- [X] PyTorch bindings
+
+## 测试
+
+```bash
+# 只测试Ada架构 不指定默认编译所有架构 耗时较长
+export TORCH_CUDA_ARCH_LIST=Ada 
+python3 hgemm.py
+```
+
+输出:
+
+```bash
+--------------------------------------------------------------------------------
+      out_f16(sk): [-1.08691406, 14.2890625, -0.57226562], time:0.38200140ms
+out_f16x4(t8x8sk): [-1.08691406, 14.2890625, -0.57226562], time:0.06475449ms
+       out_f16_th: [-1.08398438, 14.3046875, -0.56152344], time:0.02875686ms
+--------------------------------------------------------------------------------
+```
diff --git a/hgemm/hgemm.cu b/hgemm/hgemm.cu
@@ -0,0 +1,229 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <float.h>
+#include <vector>
+#include <algorithm>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <torch/types.h>
+#include <torch/extension.h>
+
+#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])
+
+// -------------------------------------- FP16 -------------------------------------- 
+// HGEMM: Block Tile + K Tile, with smem
+// Block Tile (BM, BN) + K Tile (BK=32)
+// grid((N + BN - 1) / BN, (M + BM - 1) / BM), block(BN, BM)
+// a: MxK, b: KxN, c: MxN, compute: c = a * b, all row major  
+template<const int BM=32, const int BN=32, const int BK=32>
+__global__ void hgemm_sliced_k_f16_kernel(half* a, half* b, half* c, int M, int N, int K) {
+  // [1] Block Tile: 32x32的block处理c上一块32x32的元素计算
+  // [2]     K Tile: 使用共享内存，并将K分块为BK大小的块
+  __shared__ half s_a[BM][BK], s_b[BK][BN]; 
+
+  int bx = blockIdx.x;
+  int by = blockIdx.y;
+  int tx = threadIdx.x;
+  int ty = threadIdx.y;
+  int tid = threadIdx.y * blockDim.x + tx; // tid within the block
+  // load values to shared memory, 32x32 threads working together 
+  // to fetch data along the row direction of a and b both for s_a 
+  // and s_b 32x32x4x2=8KB, we use 32x32 threads within block to 
+  // load 32x32 elements from global memory to shared memory, namely, 
+  // each thread will load 1 element.
+  int load_smem_a_m = tid / 32; // 0~31, tid / 32, tid / BM, threadIdx.y
+  int load_smem_a_k = tid % 32; // 0~31, tid % 32, tid % BK, threadIdx.x
+  int load_smem_b_k = tid / 32; // 0~31, tid / 32, tid / BK, threadIdx.y
+  int load_smem_b_n = tid % 32; // 0~31, tid % 32, tid % BN, threadIdx.x
+  int load_gmem_a_m = by * BM + load_smem_a_m; // global row of a and c
+  int load_gmem_b_n = bx * BN + load_smem_b_n; // global col of b and c
+  // if (load_gmem_a_m >= M || load_gmem_b_n >= N) return;
+  
+  half sum = __float2half(0.f);
+  for (int bk = 0; bk < (K + BK - 1) / BK; ++bk) {
+    int load_gmem_a_k = bk * BK + load_smem_a_k;
+    int load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;
+    s_a[load_smem_a_m][load_smem_a_k] = a[load_gmem_a_addr];
+    int load_gmem_b_k = bk * BK + load_smem_b_k;
+    int load_gmem_b_addr = load_gmem_b_k * N + load_gmem_b_n;
+    s_b[load_smem_b_k][load_smem_b_n] = b[load_gmem_b_addr];
+    __syncthreads();
+    #pragma unroll
+    for (int k = 0; k < BK; ++k) {
+      int comp_smem_a_m = load_smem_a_m;
+      int comp_smem_b_n = load_smem_b_n;
+      sum += s_a[comp_smem_a_m][k] * s_b[k][comp_smem_b_n];
+    }
+    __syncthreads();
+  }
+  int store_gmem_c_m = load_gmem_a_m;
+  int store_gmem_c_n = load_gmem_b_n;
+  int store_gmem_c_addr = store_gmem_c_m * N + store_gmem_c_n;
+  c[store_gmem_c_addr] = sum;
+}
+
+// HGEMM: Block Tile + Thread Tile + K Tile + Vec4, with smem
+// BK:TILE_K=8 BM=BN=128
+// TM=TN=8 增加计算密度 BM/TM=16 BN/TN=16
+// dim3 blockDim(BN/TN, BM/TM);
+// dim3 gridDim((N + BN - 1) / BN, (M + BM - 1) / BM)
+template<const int BM=128, const int BN=128, const int BK=8, const int TM=8, const int TN=8>
+__global__ void hgemm_t_8x8_sliced_k_f16x4_kernel(half* a, half* b, half* c, int M, int N, int K) {
+  // [1]  Block Tile: 一个16x16的block处理C上大小为128X128的一个目标块
+  // [2] Thread Tile: 每个thread负责计算TM*TN(8*8)个元素，增加计算密度
+  // [3]      K Tile: 将K分块，每块BK大小，迭代(K+BK-1/BK)次，
+  //                  每次计算TM*TN个元素各自的部分乘累加
+  // [4]   Vectorize: 减少load和store指令，使用half2
+
+  // 线程总数16x16=256，每个线程负责计算8x8的元素
+  int bx = blockIdx.x;
+  int by = blockIdx.y;
+  int tx = threadIdx.x;
+  int ty = threadIdx.y;
+  int tid = threadIdx.y * blockDim.x + tx; // tid within the block
+  __shared__ half s_a[BM][BK], s_b[BK][BN]; // 2*128*8*2=4KB
+  
+  // 0. 先计算shared memory中的索引
+  // tid和需要加载的smem s_a[BM][BK] 之间的索引关系 BM=128 BK=8 按行读取 A行主序
+  // 对于s_a每行8个数据，每个线程读取4个，需要2个线程；总共128行，需要128x2刚好256线程
+  int load_smem_a_m = tid / 2; // tid/2 (128/8)*(128/8)=256 threads per block, tid/2->[0,128), BM=128 0~127
+  int load_smem_a_k = (tid % 2 == 0) ? 0 : 4;  // (tid%2 == 0) ? 0 : 4, col of s_a 0,4
+  // tid和需要加载的smem s_b[BK][BN] 之间的索引关系 BK=8 BN=128 按行读取 B行主序
+  // 对于s_b每行128个数据，每个线程读4个数据，需要32个线程；总共8行，需要32x8=256个线程
+  int load_smem_b_k = tid / 32; // tid/32, row of s_b 256/32=8 行 0~7
+  int load_smem_b_n = (tid % 32) * 4;  // (tid % 32) * 4, col of s_b 0,4,...,124
+  // 1. 再计算全局内存中的索引
+  // 要加载到s_a中的元素对应到A全局内存中的行数 每个block负责出C中大小为BM*BN的块
+  int load_gmem_a_m = by * BM + load_smem_a_m; // global row of a and c
+  int load_gmem_b_n = bx * BN + load_smem_b_n; // global col of b and c
+
+  half r_c[TM][TN] = {__float2half(0.0f)}; // 8x8
+  // 2. 先对K进行分块，每块BK大小
+  for (int bk = 0; bk < (K + BK - 1) / BK; ++bk) {
+    // 加载数据到共享内存smem s_a BM*BK 128*8 vectorize float4
+    int load_gmem_a_k = bk * BK + load_smem_a_k; // global col of a
+    int load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;
+    HALF2(s_a[load_smem_a_m][load_smem_a_k + 0]) = HALF2(a[load_gmem_a_addr + 0]);
+    HALF2(s_a[load_smem_a_m][load_smem_a_k + 2]) = HALF2(a[load_gmem_a_addr + 2]);
+    // 加载数据到共享内存smem s_b BK*BN 8*128 vectorize float4
+    int load_gmem_b_k = bk * BK + load_smem_b_k; // global row of b
+    int load_gmem_b_addr = load_gmem_b_k * N + load_gmem_b_n; 
+    HALF2(s_b[load_smem_b_k][load_smem_b_n + 0]) = HALF2(b[load_gmem_b_addr + 0]);
+    HALF2(s_b[load_smem_b_k][load_smem_b_n + 2]) = HALF2(b[load_gmem_b_addr + 2]); 
+    __syncthreads();
+    #pragma unroll
+    for (int k = 0; k < BK; k++) {
+      // 3. 每个线程负责计算BM*BN(12x128)中的TM*TN(8x8)个元素
+      #pragma unroll
+      for (int m = 0; m < TM; m++) {
+        #pragma unroll
+        for (int n = 0; n < TN; n++) {
+          // k from 0~7，0 ~ BK, ty and tx range from 0 to 15, 16x8=128
+          int comp_smem_a_m = ty * TM + m;  // 128*8 128/TM(8)=16 M方向 16线程
+          int comp_smem_b_n = tx * TN + n;  // 8*128 128/TN(8)=16 N方向 16线程
+          r_c[m][n] += s_a[comp_smem_a_m][k] * s_b[k][comp_smem_b_n];
+        }
+      }
+    }
+    __syncthreads();
+  }
+  
+  #pragma unroll
+  for (int m = 0; m < TM; ++m) {
+    int store_gmem_c_m = by * BM + ty * TM + m;
+    #pragma unroll
+    for (int n = 0; n < TN; n += 2) {
+      int store_gmem_c_n = bx * BN + tx * TN + n;
+      int store_gmem_c_addr = store_gmem_c_m * N + store_gmem_c_n;
+      HALF2(c[store_gmem_c_addr]) = HALF2(r_c[m][n]);
+    }
+  }
+}
+
+// --------------------- PyTorch bindings for custom kernel -----------------------
+#define STRINGFY(str) #str
+#define TORCH_BINDING_COMMON_EXTENSION(func) \
+  m.def(STRINGFY(func), &func, STRINGFY(func));
+
+#define CHECK_TORCH_TENSOR_DTYPE(T, th_type)                 \
+if(((T).options().dtype() != (th_type))) {                   \
+  std::cout << "Tensor Info:" << (T).options() << std::endl; \
+  throw std::runtime_error("values must be "#th_type);       \
+}
+
+#define CHECK_TORCH_TENSOR_SHAPE(T, S0, S1)           \
+if (((T).size(0) != (S0)) || ((T).size(1) != (S1))) { \
+  throw std::runtime_error("Tensor size mismatch!");  \
+}
+
+// HGEMM: Block Tile + K Tile, with smem
+// Block Tile (BM, BN) + K Tile (BK=32)
+// grid((N + BN - 1) / BN, (M + BM - 1) / BM), block(BN, BM)
+// a: MxK, b: KxN, c: MxN, compute: c = a * b, all row major  
+void hgemm_sliced_k_f16(torch::Tensor a, torch::Tensor b, torch::Tensor c) {
+  CHECK_TORCH_TENSOR_DTYPE(a, torch::kHalf)
+  CHECK_TORCH_TENSOR_DTYPE(b, torch::kHalf)
+  CHECK_TORCH_TENSOR_DTYPE(c, torch::kHalf)
+  const int M = a.size(0);
+  const int K = a.size(1);
+  const int N = b.size(1); 
+  CHECK_TORCH_TENSOR_SHAPE(a, M, K)
+  CHECK_TORCH_TENSOR_SHAPE(b, K, N)
+  CHECK_TORCH_TENSOR_SHAPE(c, M, N)
+  constexpr int BM = 32;
+  constexpr int BN = 32;
+  constexpr int BK = 32;
+
+  dim3 block(BN, BM);
+  dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
+
+  hgemm_sliced_k_f16_kernel<BM, BN, BK><<<grid, block>>>(
+    reinterpret_cast<half*>(a.data_ptr()),
+    reinterpret_cast<half*>(b.data_ptr()),
+    reinterpret_cast<half*>(c.data_ptr()),
+    M, N, K
+  );
+}
+
+// HGEMM: Block Tile + Thread Tile + K Tile + half2x2, with smem
+// BK:TILE_K=8 BM=BN=128
+// TM=TN=8 增加计算密度 BM/TM=16 BN/TN=16
+// dim3 blockDim(BN/TN, BM/TM);
+// dim3 gridDim((N + BN - 1) / BN, (M + BM - 1) / BM)
+void hgemm_t_8x8_sliced_k_f16x4(torch::Tensor a, torch::Tensor b, torch::Tensor c) {
+  CHECK_TORCH_TENSOR_DTYPE(a, torch::kHalf)
+  CHECK_TORCH_TENSOR_DTYPE(b, torch::kHalf)
+  CHECK_TORCH_TENSOR_DTYPE(c, torch::kHalf)
+  const int M = a.size(0);
+  const int K = a.size(1);
+  const int N = b.size(1); 
+  CHECK_TORCH_TENSOR_SHAPE(a, M, K)
+  CHECK_TORCH_TENSOR_SHAPE(b, K, N)
+  CHECK_TORCH_TENSOR_SHAPE(c, M, N)
+  constexpr int BM = 128;
+  constexpr int BN = 128;
+  constexpr int BK = 8; 
+  constexpr int TM = 8;
+  constexpr int TN = 8;
+
+  dim3 block(BN/TN, BM/TM);
+  dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
+
+  hgemm_t_8x8_sliced_k_f16x4_kernel<BM, BN, BK, TM, TN><<<grid, block>>>(
+    reinterpret_cast<half*>(a.data_ptr()),
+    reinterpret_cast<half*>(b.data_ptr()),
+    reinterpret_cast<half*>(c.data_ptr()),
+    M, N, K
+  );
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  TORCH_BINDING_COMMON_EXTENSION(hgemm_sliced_k_f16)
+  TORCH_BINDING_COMMON_EXTENSION(hgemm_t_8x8_sliced_k_f16x4)
+}
diff --git a/hgemm/hgemm.py b/hgemm/hgemm.py
@@ -0,0 +1,69 @@
+import torch
+import time 
+from torch.utils.cpp_extension import load
+from functools import partial
+from typing import Optional
+
+torch.set_grad_enabled(False)
+
+# Load the CUDA kernel as a python module
+lib = load(name='hgemm_lib', 
+           sources=['hgemm.cu'], 
+           extra_cuda_cflags=[
+               "-O3",
+                "-U__CUDA_NO_HALF_OPERATORS__",
+                "-U__CUDA_NO_HALF_CONVERSIONS__",
+                "-U__CUDA_NO_HALF2_OPERATORS__",
+                "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+                "--expt-relaxed-constexpr",
+                "--expt-extended-lambda",
+                "--use_fast_math"
+            ], 
+           extra_cflags=['-std=c++17'])
+
+
+def run_benchmark(perf_func: callable, 
+                  a: torch.Tensor, b: torch.Tensor,
+                  tag: str, out: Optional[torch.Tensor] = None, 
+                  warmup: int = 10, iters: int = 200,
+                  show_all: bool = False):
+    if out is not None: 
+        out.fill_(0)      
+    if out is not None:
+        for i in range(warmup):
+            perf_func(a, b, out)
+    else:
+        for i in range(warmup):
+            _ = perf_func(a, b) 
+    
+    torch.cuda.synchronize()
+    start = time.time()
+    # iters
+    if out is not None:
+        for i in range(iters):
+            perf_func(a, b, out)
+    else:
+        for i in range(iters):
+            out = perf_func(a, b) 
+    torch.cuda.synchronize()
+    end = time.time()
+    total_time = (end - start) * 1000 # ms
+    mean_time = total_time / iters
+    out_info = f"out_{tag}"
+    out_val = out.flatten().detach().cpu().numpy().tolist()[:3]
+    out_val = [round(v, 8) for v in out_val]
+    print(f"{out_info:>17}: {out_val}, time:{mean_time:.8f}ms")
+    if show_all: print(out)
+    return out.clone(), mean_time
+
+
+print("-" * 80)
+M, N, K = 2048, 1024, 128
+a = torch.randn((M, K)).cuda().half().contiguous() 
+b = torch.randn((K, N)).cuda().half().contiguous() 
+c = torch.randn((M, N)).cuda().half().contiguous() 
+run_benchmark(lib.hgemm_sliced_k_f16, a, b, "f16(sk)", c)
+run_benchmark(lib.hgemm_t_8x8_sliced_k_f16x4, a, b, "f16x4(t8x8sk)", c)
+run_benchmark(partial(torch.matmul, out=c), a, b, "f16_th")
+print("-" * 80)
+
