Upgrade CUTLASS to 3.1 (facebookincubator#584)

Summary: Pull Request resolved: facebookincubator#584 ATT Reviewed By: muchulee8, chenyang78, alexanderguzhva, tenpercent Differential Revision: D45028066 fbshipit-source-id: 049ffee0f4e8f5f7e6fbd0b3517e2db7e9520c9f
henrylhtsang · Apr 20, 2023 · be47ae5 · be47ae5
1 parent 7b97317
commit be47ae5
Show file tree

Hide file tree

Showing 24 changed files with 11,907 additions and 10 deletions.
diff --git a/.gitmodules b/.gitmodules
@@ -1,6 +1,6 @@
 [submodule "3rdparty/cutlass"]
 	path = 3rdparty/cutlass
-	url = https://github.com/AITemplate/cutlass.git
+	url = https://github.com/facebookincubator/cutlass-fork.git
 [submodule "3rdparty/cub"]
 	path = 3rdparty/cub
 	url = https://github.com/NVIDIA/cub.git

diff --git a/3rdparty/cutlass b/3rdparty/cutlass
diff --git a/python/aitemplate/backend/cuda/attention/mem_eff_attention.py b/python/aitemplate/backend/cuda/attention/mem_eff_attention.py
@@ -43,8 +43,8 @@
 #include <cuda_fp16.h>
 #include "cutlass/cutlass.h"
 #include "cutlass/gemm/device/default_gemm_configuration.h"
-#include "gemm_kernel_utils.h"
 
+#include "mem_eff_attention/gemm_kernel_utils.h"
 #include "mem_eff_attention/kernel_forward.h"
 
 
@@ -215,12 +215,12 @@
 #include "cutlass/util/reference/host/tensor_norm.h"
 
 #include "cutlass/gemm/device/default_gemm_configuration.h"
-#include "gemm_kernel_utils.h"
 #include "cutlass/gemm/device/gemm_grouped.h"
 
 #include "cutlass/fast_math.h"
 
-#include "default_fmha_grouped.h"
+#include "mem_eff_attention/gemm_kernel_utils.h"
+#include "mem_eff_attention/default_fmha_grouped.h"
 
 using namespace gemm_kernel_utils;
 

diff --git a/python/aitemplate/utils/mk_cutlass_lib/extra_enum.py b/python/aitemplate/utils/mk_cutlass_lib/extra_enum.py
@@ -131,9 +131,9 @@ class EpiloguePermuteLayout(enum.Enum):
   NoPermute = enum_auto()
 
 EpiloguePermuteLayoutTag = {
-  EpiloguePermuteLayout.Permute5D_20314: 'cutlass::layout::Tensor5DPermute20314',
-  EpiloguePermuteLayout.Permute4D_0213: 'cutlass::layout::Tensor4DPermute0213',
-  EpiloguePermuteLayout.Permute4DBMM_0213: 'cutlass::layout::Tensor4DPermuteBMM0213',
+  EpiloguePermuteLayout.Permute5D_20314: 'cutlass::layout::Tensor5DPermute20314RowMajor',
+  EpiloguePermuteLayout.Permute4D_0213: 'cutlass::layout::Tensor4DPermute0213RowMajor',
+  EpiloguePermuteLayout.Permute4DBMM_0213: 'cutlass::layout::Tensor4DPermuteBMM0213RowMajor',
   EpiloguePermuteLayout.NoPermute: 'cutlass::layout::NoPermute',
   # EpiloguePermuteLayout.Permute3DBMM_021: 'cutlass::layout::Tensor3DPermute021BMM',
 }