[HGEMM] Update toy-hgemm library 0.1.0 (#149)

* Update README.md

* Update README.md

* Update hgemm_cublas.cu

* Update hgemm_mma_stage_tn_cute.cu

* Update hgemm.py

* Update makefile

* Update hgemm.cc

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update hgemm.py

* Update setup.py

Author: DefTruth

README.md

@@ -21,12 +21,12 @@
 <div id="hgemm-sgemm"></div>  
 
 <div align='left'>
-  <img src='https://github.com/user-attachments/assets/71927ac9-72b3-4ce9-b0e2-788b5885bc99' height="150px" width="265px">
-  <img src='https://github.com/user-attachments/assets/05ef4f5e-d999-48ea-b58e-782cffb24e85' height="150px" width="265px">
-  <img src='https://github.com/user-attachments/assets/9472e970-c083-4b31-9252-3eeecc761078' height="150px" width="265px">
+  <img src='https://github.com/user-attachments/assets/71927ac9-72b3-4ce9-b0e2-788b5885bc99' height="150px" width="267px">
+  <img src='https://github.com/user-attachments/assets/05ef4f5e-d999-48ea-b58e-782cffb24e85' height="150px" width="267px">
+  <img src='https://github.com/user-attachments/assets/9472e970-c083-4b31-9252-3eeecc761078' height="150px" width="267px">
 </div> 
 
-Currently, on NVIDIA L20, RTX 4090 and RTX 3080 Laptop, compared with cuBLAS's default Tensor Cores math algorithm `CUBLAS_GEMM_DEFAULT_TENSOR_OP`, the `HGEMM (WMMA/MMA)` implemented in this repo (`blue`🔵) can achieve `95%~99%` of its (`orange`🟠) performance. Please check [toy-hgemm library🔥🔥](./kernels/hgemm) for more details.
+Currently, on NVIDIA L20, RTX 4090 and RTX 3080 Laptop, compared with cuBLAS's default Tensor Cores math algorithm `CUBLAS_GEMM_DEFAULT_TENSOR_OP`, the `HGEMM (WMMA/MMA)` implemented in this repo (`blue`🔵) can achieve `99%~100%+` of its (`orange`🟠) performance. Please check [toy-hgemm library🔥🔥](./kernels/hgemm) for more details.
 
 |CUDA Cores|Sliced K(Loop over K)|Tile Block|Tile Thread|
 |:---:|:---:|:---:|:---:|
@@ -35,9 +35,9 @@ Currently, on NVIDIA L20, RTX 4090 and RTX 3080 Laptop, compared with cuBLAS's d
 |✔️|✔️|✔️|✔️|
 |Copy Async|Tile MMA(More Threads)|Tile Warp(More Values)|Multi Stages|  
 |✔️|✔️|✔️|✔️|
-|Reg Double Buffers|Block Swizzle|Warp Swizzle|Collective Store(Warp Shuffle)|
+|Reg Double Buffers|Block Swizzle|Warp Swizzle|SMEM Swizzle(CuTe)|
 |✔️|✔️|✔️|✔️|
-|Row Major(NN)|Col Major(TN)|SGEMM TF32|SMEM Swizzle(CuTe)|
+|Collective Store(Warp Shfl)|Row Major(NN)|Col Major(TN)|SGEMM F32/TF32|
 |✔️|✔️|✔️|✔️|
 
 ## ©️Citations🎉🎉

kernels/hgemm/README.md

@@ -1,20 +1,21 @@
-## 🔥🔥Toy-HGEMM Library: Achieve the performance of cuBLAS
+# 🔥🔥Toy-HGEMM Library: Achieve the performance of cuBLAS
 
-|CUDA Cores|Sliced K(Loop over K)|Tile Block|Tile Thread|
+|CUDA Cores|Sliced K(Loop over K)|Tile Block(BMxBN)|Tile Thread(t 8x8)|
 |:---:|:---:|:---:|:---:|
 |✔️|✔️|✔️|✔️|
 |WMMA(m16n16k16)|MMA(m16n8k16)|Pack LDST(128 bits)|SMEM Padding|
 |✔️|✔️|✔️|✔️|
-|Copy Async|Tile MMA(More Threads)|Tile Warp(More Values)|Multi Stages|  
+|Copy Async(cp.async.cg/ca)|Tile MMA(More Threads)|Tile Warp(More Values)|Multi Stages(2/3/4/5)|  
 |✔️|✔️|✔️|✔️|
-|Reg Double Buffers|Block Swizzle|Warp Swizzle|Collective Store(Warp Shfl)|
+|Register Double Buffers|Block Swizzle(Zigzag N)|Warp Swizzle(Zigzag N)|SMEM Swizzle(CUTLASS/CuTe)|
 |✔️|✔️|✔️|✔️|
-|Row Major(NN)|Col Major(TN)|SGEMM TF32|SMEM Swizzle(CuTe)|
+|Collective Store(Warp Shuffle & Reg Reuse)|Row Major(NN)|Col Major(TN)|SGEMM FP32/TF32|
 |✔️|✔️|✔️|✔️|
 
-
 ## 📖 HGEMM CUDA Kernels in Toy-HGEMM Library 🎉🎉 
 
+<div id="kernels"></div>  
+
 ```C++  
 void hgemm_naive_f16(torch::Tensor a, torch::Tensor b, torch::Tensor c);
 void hgemm_sliced_k_f16(torch::Tensor a, torch::Tensor b, torch::Tensor c);
@@ -49,10 +50,23 @@ void hgemm_mma_m16n8k16_mma2x4_warp4x4x2_stages_dsmem(torch::Tensor a, torch::Te
 void hgemm_mma_m16n8k16_mma2x4_warp4x4x2_stages_dsmem_x4(torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle, int swizzle_stride);
 void hgemm_mma_m16n8k16_mma2x4_warp4x4x2_stages_dsmem_rr(torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle, int swizzle_stride);
 void hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn(torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle, int swizzle_stride);
-void hgemm_mma_stages_tn_cute(torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle, int swizzle_stride);
+void hgemm_mma_stages_block_swizzle_tn_cute(torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle, int swizzle_stride);
 ```
 
-## 📖 安装
+## 📖 目录
+
+- [📖 安装](#install)
+- [📖 测试](#test)
+- [📖 NVIDIA L20 性能数据](#perf-l20)
+- [📖 NVIDIA RTX 4090 性能数据](#perf-4090)
+- [📖 NVIDIA RTX 3080 Laptop 性能数据](#perf-3080)
+- [📖 性能优化笔记](#opt-docs)
+- [📖 参考文献](#ref)
+
+## 📖 安装  
+
+<div id="install"></div>  
+
 本仓库实现的HGEMM可以作为一个python库使用（可选）
 ```bash
 git submodule update --init --recursive --force # 更新cutlass, 必须
@@ -61,6 +75,8 @@ python3 setup.py bdist_wheel && cd dist && python3 -m pip install *.whl # pip un
 
 ## 📖 测试
 
+<div id="test"></div>  
+
 **CUTLASS**: 更新CUTLASS依赖库
 ```bash
 git submodule update --init --recursive --force
@@ -125,9 +141,11 @@ M N K =  16384  16384  16384, Time =   0.07668429   0.07669371   0.07670784 s, A
 
 ## 📖 目前性能  
 
+<div id="perf-l20"></div>  
+
 ### NVIDIA L20  
 
-目前最优的实现，在L20上（理论Tensor Cores FP16算力为 119.5 TFLOPS），整体上能达到cuBLAS大概99%左右的性能。使用WMMA API能达到cuBLAS大概95%~98%左右的性能(105-113 TFLOPS vs 105-115 TFLOPS)，使用MMA API能达到115 TFLOPS，部分case会超越cuBLAS。CuTe版本的HGEMM性能基本持平cuBLAS，部分case会超越cuBLAS，能达到 116-117 TFLOPS。目前通过 SMEM Padding 和 SMEM swizzle的方式缓解bank conflicts。对于 NN layout，使用 SMEM Padding 缓解 bank conflicts；对于 TN layout，通过cutlass cute的 SMEM Swizzle 消除 bank conflicts。
+目前最优的实现，在L20上（理论Tensor Cores FP16算力为 119.5 TFLOPS），整体上能达到cuBLAS大概`99~100+%`左右的性能。使用WMMA API能达到cuBLAS大概`95%~98%`左右的性能(105-113 TFLOPS vs 105-115 TFLOPS)，使用MMA API能达到115 TFLOPS，部分 case 会超越 cuBLAS。CuTe 版本的 HGEMM 实现了 Block Swizzle（L2 Cache friendly）和 SMEM Swizzle（bank conflicts free），性能最优，大规模矩阵乘能达到 116-117 TFLOPS，是 cuBLAS 大概`98%~100%+`左右的性能，很多case会超越cuBLAS。目前通过 SMEM Padding 和 SMEM Swizzle 的方式缓解 bank conflicts。对于 NN layout，使用 SMEM Padding 缓解 bank conflicts；对于 TN layout，通过 CUTLASS/CuTe 的 SMEM Swizzle 消除 bank conflicts。
 
 <div id="NV-L20"></div>
 
@@ -148,6 +166,9 @@ python3 hgemm.py --cute-tn --mma --plot
 ```
 
 ### NVIDIA GeForce RTX 4090
+
+<div id="perf-4090"></div>  
+
 在NVIDIA RTX 4090上(FP16 Tensor Cores算力为330 TFLOPS)，WMMA(m16n16k16)性能表现比MMA(m16n8k16)要更好，大分部MNK下，本仓库的实现能达到cuBLAS 95%~99%的性能，某些case能超过cuBLAS。就本仓库的实现而言，在RTX 4090上，大规模矩阵乘(MNK>=8192)，WMMA表现更优，小规模矩阵乘，MMA表现更优。
 
 <!---
@@ -164,6 +185,8 @@ python3 hgemm.py --cute-tn --mma --wmma-all --plot
 
 ### NVIDIA GeForce RTX 3080 Laptop   
 
+<div id="perf-3080"></div>  
+
 在NVIDIA GeForce RTX 3080 Laptop上测试，使用mma4x4_warp4x4（16 WMMA m16n16k16 ops, warp tile 64x64）以及Thread block swizzle，大部分case能持平甚至超过cuBLAS，使用Windows WSL2 + RTX 3080 Laptop进行测试。
 
 <!--
@@ -179,6 +202,9 @@ python3 hgemm.py --wmma-all --plot
 
 ## 📖 性能优化笔记
 
+<div id="opt-docs"></div>  
+
+
 ### PyTorch HGEMM Profile
 
 在Ada架构下，PyTorch 2.4对FP16使用matmul时，会调用:
@@ -282,7 +308,9 @@ TODO
 
 </details>
 
-## 参考文献 
+## 📖 参考文献 
+
+<div id="ref"></div>  
 
 - [CUDA编程概念】一、什么是bank conflict？](https://zhuanlan.zhihu.com/p/659142274)
 - [解决 bank conflict](https://github.com/PaddleJitLab/CUDATutorial/blob/develop/docs/09_optimize_reduce/02_bank_conflict/README.md)

kernels/hgemm/cublas/hgemm_cublas.cu

@@ -43,6 +43,10 @@ void cublas_tensor_op_nn(half *A, half *B, half *C,  size_t M, size_t N, size_t
   static half alpha = 1.0;
   static half beta = 0.0;
 
+  if (g_handle == nullptr) {
+    init_cublas_handle();
+  }
+
   cublasGemmEx(g_handle, 
                CUBLAS_OP_N, 
                CUBLAS_OP_N, 
@@ -62,6 +66,10 @@ void cublas_tensor_op_tn(half *A, half *B, half *C,  size_t M, size_t N, size_t
   static half alpha = 1.0;
   static half beta = 0.0;
 
+  if (g_handle == nullptr) {
+    init_cublas_handle();
+  }
+
   cublasGemmEx(g_handle, 
                CUBLAS_OP_T, 
                CUBLAS_OP_N, 

kernels/hgemm/cutlass/hgemm_mma_stage_tn_cute.cu

@@ -461,8 +461,8 @@ if (((T).size(0) != (S0)) || ((T).size(1) != (S1))) { \
   );
 
 
-// Multi stages CuTe HGEMM with smem and block swizzle.
-void hgemm_mma_stages_tn_cute(
+// Multi stages CuTe HGEMM with SMEM Swizzle and Block Swizzle.
+void hgemm_mma_stages_block_swizzle_tn_cute(
   torch::Tensor a, torch::Tensor b, torch::Tensor c,
   int stages, bool swizzle, int swizzle_stride) {
   CHECK_TORCH_TENSOR_DTYPE(a, torch::kHalf)

kernels/hgemm/hgemm.py

@@ -58,8 +58,7 @@ def get_args():
 pretty_print_line()
 
 
-hgemm = try_load_hgemm_library(force_build=args.force_build, 
-                               verbose=args.verbose)
+hgemm = try_load_hgemm_library(force_build=args.force_build, verbose=args.verbose)
 
 MAX_TFLOPS = -1
 STATIS_INFO: dict[str, list[float]] = {}
@@ -69,14 +68,12 @@ def get_args():
 CUBLAS_TN_TOTAL_TFLOPS = 0
 
 
-def make_block_swizzle_stride(N: int, K: int):
+def make_block_swizzle_stride(N: int, K: int, swizzle_factor: float = None):
     # make swizzle stride as N/8,N/4,N/2 and multiples of 256
-    if args.swizzle_factor is None:
+    if swizzle_factor is None:
         swizzle_factor = 0.5 if N <= 4096 else 0.25
         if all((N >= 14848, K > 8192, N % 8 == 0)):
             swizzle_factor = 0.125
-    else:
-        swizzle_factor = args.swizzle_factor
 
     swizzle_stride = int(N * swizzle_factor)
     swizzle_stride = swizzle_stride if swizzle_stride >= 256 else 1
@@ -100,7 +97,7 @@ def run_benchmark(perf_func: callable,
     K = a.size(1)
     N = b.size(1) # TN still has shape [K,N]
     if swizzle:
-        swizzle_stride = make_block_swizzle_stride(N, K)
+        swizzle_stride = make_block_swizzle_stride(N, K, args.swizzle_factor)
         swizzle = swizzle if swizzle_stride >= 256 else False
     else:
         swizzle_stride = 1 # means no thread block swizzle
@@ -110,6 +107,10 @@ def run_benchmark(perf_func: callable,
 
     if out is not None: 
         out.fill_(0)      
+    
+    if "cublas" in tag:
+        hgemm.init_cublas_handle()
+
     if out is not None:
         for i in range(warmup):
             if stages > 1:
@@ -177,6 +178,9 @@ def run_benchmark(perf_func: callable,
                 CUBLAS_TOTAL_TFLOPS += TFLOPS
 
     torch.cuda.synchronize()
+    if "cublas" in tag:
+        hgemm.destroy_cublas_handle()
+
     del out_flat
     out_flat = None
     gc.collect()
@@ -262,6 +266,7 @@ def skip_it(tag: str) -> bool:
         save_path = f"{args.save_dir}/{device_name}_{args.save_tag}.png"
     else:
         save_path = f"{args.save_dir}/{device_name}.png"
+    os.makedirs(args.save_dir, exist_ok=True)
     plt.savefig(save_path, dpi=300)
     pretty_print_line(f"plot hgemm TFLOPS done, saved as {save_path}")
 
@@ -383,24 +388,20 @@ def get_mnk(sep: int = args.SEP):
         run_benchmark(hgemm.hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn, a, b_col_major, "tn(mma2x4+warp4x4+stage3+dsmem+swizzle<block>)", c, stages=3, swizzle=True)
         run_benchmark(hgemm.hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn, a, b_col_major, "tn(mma2x4+warp4x4+stage2+dsmem+swizzle<block>)", c, stages=2, swizzle=True)
     if args.enable_cute_tn:
-        run_benchmark(hgemm.hgemm_mma_stages_tn_cute, a, b_col_major, "tn(cute+stage4+swizzle<smem>)", c, stages=4)
-        run_benchmark(hgemm.hgemm_mma_stages_tn_cute, a, b_col_major, "tn(cute+stage3+swizzle<smem>)", c, stages=3)
-        run_benchmark(hgemm.hgemm_mma_stages_tn_cute, a, b_col_major, "tn(cute+stage2+swizzle<smem>)", c, stages=2)
-        run_benchmark(hgemm.hgemm_mma_stages_tn_cute, a, b_col_major, "tn(cute+stage4+swizzle<smem+block>)", c, stages=4, swizzle=True)
-        run_benchmark(hgemm.hgemm_mma_stages_tn_cute, a, b_col_major, "tn(cute+stage3+swizzle<smem+block>)", c, stages=3, swizzle=True)
-        run_benchmark(hgemm.hgemm_mma_stages_tn_cute, a, b_col_major, "tn(cute+stage2+swizzle<smem+block>)", c, stages=2, swizzle=True)
+        run_benchmark(hgemm.hgemm_mma_stages_block_swizzle_tn_cute, a, b_col_major, "tn(cute+stage4+swizzle<smem>)", c, stages=4)
+        run_benchmark(hgemm.hgemm_mma_stages_block_swizzle_tn_cute, a, b_col_major, "tn(cute+stage3+swizzle<smem>)", c, stages=3)
+        run_benchmark(hgemm.hgemm_mma_stages_block_swizzle_tn_cute, a, b_col_major, "tn(cute+stage2+swizzle<smem>)", c, stages=2)
+        run_benchmark(hgemm.hgemm_mma_stages_block_swizzle_tn_cute, a, b_col_major, "tn(cute+stage4+swizzle<smem+block>)", c, stages=4, swizzle=True)
+        run_benchmark(hgemm.hgemm_mma_stages_block_swizzle_tn_cute, a, b_col_major, "tn(cute+stage3+swizzle<smem+block>)", c, stages=3, swizzle=True)
+        run_benchmark(hgemm.hgemm_mma_stages_block_swizzle_tn_cute, a, b_col_major, "tn(cute+stage2+swizzle<smem+block>)", c, stages=2, swizzle=True)
     # TN layout: cublas
     if not args.disable_cublas_tn and any((args.enable_mma_tn, args.enable_cute_tn)):
-        hgemm.init_cublas_handle()
         run_benchmark(hgemm.hgemm_cublas_tensor_op_tn, a, b_col_major, "tn(cublas)", c)
-        hgemm.destroy_cublas_handle()
     # NN layout: cublas/torch
     if (not args.disable_cublas) and any((
         args.enable_mma, args.enable_mma_all, args.enable_wmma, args.enable_wmma_all, 
         args.enable_cuda, args.enable_cuda_all, args.enable_torch)):
-        hgemm.init_cublas_handle()
         run_benchmark(hgemm.hgemm_cublas_tensor_op_nn, a, b, "(cublas)", c)
-        hgemm.destroy_cublas_handle()
     if args.enable_torch:
         run_benchmark(partial(torch.matmul, out=c), a, b, "(torch)")
     torch.cuda.synchronize()

kernels/hgemm/makefile

@@ -1,6 +1,7 @@
 INCLUDE_DIRS=-I ./utils -I ../../third-party/cutlass/include -I ../../third-party/cutlass/tools/util/include
+ARCHS=-gencode arch=compute_80,code=sm_80 -gencode arch=compute_89,code=sm_89
 default:
-	nvcc cutlass/hgemm_mma_stage_tn_cute.cu -o hgemm_cute.bin -O2 -arch=sm_89 -std=c++17 $(INCLUDE_DIRS) --expt-relaxed-constexpr -lcublas
-	nvcc cublas/hgemm_cublas.cu -o hgemm_cublas.bin -O2 -arch=sm_89 -std=c++17 $(INCLUDE_DIRS) --expt-relaxed-constexpr -lcublas
-	nvcc mma/hgemm_mma_stage.cu -o hgemm_mma_stage.bin -O2 -arch=sm_89 -std=c++17 $(INCLUDE_DIRS) --expt-relaxed-constexpr -lcublas
+	nvcc cutlass/hgemm_mma_stage_tn_cute.cu -o hgemm_cute.bin -O2 $(ARCHS) -std=c++17 $(INCLUDE_DIRS) --expt-relaxed-constexpr -lcublas
+	nvcc cublas/hgemm_cublas.cu -o hgemm_cublas.bin -O2 $(ARCHS) -std=c++17 $(INCLUDE_DIRS) --expt-relaxed-constexpr -lcublas
+	nvcc mma/hgemm_mma_stage.cu -o hgemm_mma_stage.bin -O2 $(ARCHS) -std=c++17 $(INCLUDE_DIRS) --expt-relaxed-constexpr -lcublas
 

kernels/hgemm/pybind/hgemm.cc

@@ -48,7 +48,7 @@ void hgemm_mma_m16n8k16_mma2x4_warp4x4x2_stages_dsmem_rr(torch::Tensor a, torch:
 // from hgemm_mma_stage_tn.cu
 void hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn(torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle, int swizzle_stride);
 // from hgemm_mma_stage_tn_cute.cu
-void hgemm_mma_stages_tn_cute(torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle, int swizzle_stride);
+void hgemm_mma_stages_block_swizzle_tn_cute(torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle, int swizzle_stride);
 
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
@@ -96,6 +96,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   // TN: A row major MxK, B col major NxK, C row major MxN
   TORCH_BINDING_COMMON_EXTENSION(hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn)
   // TN: cute hgemm with smem & block swizzle
-  TORCH_BINDING_COMMON_EXTENSION(hgemm_mma_stages_tn_cute)
+  TORCH_BINDING_COMMON_EXTENSION(hgemm_mma_stages_block_swizzle_tn_cute)
 }
 

kernels/hgemm/setup.py

@@ -17,6 +17,8 @@
 generator_flag = []
 cc_flag = []
 cc_flag.append("-gencode")
+cc_flag.append("arch=compute_80,code=sm_80")
+cc_flag.append("-gencode")
 cc_flag.append("arch=compute_89,code=sm_89")