add quantization conv and matmul

CMakeLists.txt

@@ -212,9 +212,9 @@ if(BUILD_PYTHON)
 endif(BUILD_PYTHON)
 
 set(CUTLASS_URL
-    https://github.com/Oneflow-Inc/cutlass/archive/e6f548d80bfdf1167d66adbbbcfc2ee3394f4777.zip)
+    https://github.com/Oneflow-Inc/cutlass/archive/d47b8883b5e3661b41cc8a7a6f4c240c5524647f.zip)
 use_mirror(VARIABLE CUTLASS_URL URL ${CUTLASS_URL})
-set(CUTLASS_MD5 425f8cf064ff47c81124e55490135f5c)
+set(CUTLASS_MD5 7b417720240a443276ce4bb9ef169db1)
 
 include(cuda)
 add_subdirectory(external)

cmake/third_party.cmake

@@ -144,6 +144,7 @@ if(BUILD_CUDA)
   endif()
   include(nccl)
   include(cutlass)
+  include(cutlass-extension)
   include(trt_flash_attention)
 
   list(APPEND oneflow_third_party_libs ${NCCL_LIBRARIES})
@@ -161,6 +162,11 @@ if(BUILD_CUDA)
     list(APPEND oneflow_third_party_libs ${CUTLASS_LIBRARIES})
     list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${CUTLASS_INCLUDE_DIR})
   endif()
+  if(WITH_CUTLASS_EXTENSION)
+    list(APPEND oneflow_third_party_dependencies cutlass-extension)
+    list(APPEND oneflow_third_party_libs ${CUTLASS_EXTENSION_LIBRARIES})
+    list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${CUTLASS_EXTENSION_INCLUDE_DIR})
+  endif()
   list(APPEND oneflow_third_party_dependencies trt_flash_attention)
   list(APPEND oneflow_third_party_libs ${TRT_FLASH_ATTENTION_LIBRARIES})
   list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${TRT_FLASH_ATTENTION_INCLUDE_DIR})

cmake/third_party/cutlass-extension.cmake

@@ -0,0 +1,60 @@
+include(ExternalProject)
+
+set(WITH_CUTLASS_EXTENSION OFF CACHE BOOL "")
+
+if(WITH_CUTLASS_EXTENSION)
+
+  add_definitions(-DWITH_CUTLASS_EXTENSION)
+
+  find_package(Threads)
+
+  set(CUTLASS_EXTENSION_PROJECT cutlass-extension)
+
+  set(CUTLASS_EXTENSION_INSTALL_DIR ${THIRD_PARTY_DIR}/cutlass-extension)
+
+  set(CUTLASS_EXTENSION_INCLUDE_DIR ${CUTLASS_EXTENSION_INSTALL_DIR}/include CACHE PATH "" FORCE)
+  set(CUTLASS_EXTENSION_LIBRARY_DIR ${CUTLASS_EXTENSION_INSTALL_DIR}/lib CACHE PATH "" FORCE)
+  set(CUTLASS_EXTENSION_LIBRARIES ${CUTLASS_EXTENSION_LIBRARY_DIR}/libcutlass_extension.so)
+  set(CUTLASS_EXTENSION_SOURCE_DIR
+      ${CMAKE_CURRENT_BINARY_DIR}/cutlass-extension/src/cutlass-extension/)
+  set(CUTLASS_SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/cutlass/src/cutlass)
+
+  foreach(arch ${CUDA_REAL_ARCHS_LIST})
+    if(arch GREATER_EQUAL 70)
+      list(APPEND CUTLASS_REAL_ARCHS ${arch})
+    endif()
+  endforeach()
+
+  if(THIRD_PARTY)
+    ExternalProject_Add(
+      ${CUTLASS_EXTENSION_PROJECT}
+      PREFIX cutlass-extension
+      GIT_REPOSITORY https://github.com/Oneflow-Inc/oneflow-cutlass-extension.git
+      GIT_TAG master
+      UPDATE_COMMAND ""
+      BUILD_BYPRODUCTS ${CUTLASS_EXTENSION_LIBRARIES}
+      CMAKE_ARGS -DCMAKE_BUILD_TYPE:STRING=${CMAKE_BUILD_TYPE}
+                 -DCMAKE_CXX_FLAGS:STRING=${CMAKE_CXX_FLAGS}
+                 -DCMAKE_CXX_FLAGS_DEBUG:STRING=${CMAKE_CXX_FLAGS_DEBUG}
+                 -DCMAKE_CXX_FLAGS_RELEASE:STRING=${CMAKE_CXX_FLAGS_RELEASE}
+      CMAKE_CACHE_ARGS
+        -DCMAKE_CUDA_COMPILER:STRING=${CUDAToolkit_NVCC_EXECUTABLE}
+        -DCMAKE_C_COMPILER_LAUNCHER:STRING=${CMAKE_C_COMPILER_LAUNCHER}
+        -DCMAKE_CXX_COMPILER_LAUNCHER:STRING=${CMAKE_CXX_COMPILER_LAUNCHER}
+        -DCMAKE_INSTALL_PREFIX:PATH=${CUTLASS_EXTENSION_INSTALL_DIR}
+        -DCMAKE_INSTALL_LIBDIR:PATH=${CUTLASS_EXTENSION_LIBRARY_DIR}
+        -DCMAKE_INSTALL_MESSAGE:STRING=${CMAKE_INSTALL_MESSAGE}
+        -DCMAKE_BUILD_TYPE:STRING=${CMAKE_BUILD_TYPE}
+        -DCUTLASS_ENABLE_EXAMPLES:BOOL=OFF
+        -DCUTLASS_ENABLE_PROFILER:BOOL=OFF
+        -DCUTLASS_ENABLE_LIBRARY:BOOL=ON
+        -DCUTLASS_NVCC_ARCHS:STRING=${CUTLASS_REAL_ARCHS}
+        -DCUTLASS_ENABLE_TESTS:BOOL=OFF
+        -DCUTLASS_UNITY_BUILD_ENABLED:BOOL=ON
+        -DCUTLASS_LIBRARY_DEBUG_POSTFIX:STRING=
+        -DCUTLASS_NVCC_EMBED_PTX:BOOL=OFF
+        -DCUTLASS_DIR:STRING=${CUTLASS_SOURCE_DIR}
+      DEPENDS cutlass)
+
+  endif(THIRD_PARTY)
+endif(WITH_CUTLASS_EXTENSION)

cmake/third_party/cutlass.cmake

@@ -67,6 +67,7 @@ if(WITH_CUTLASS)
         "45_dual_gemm/test_run.h"
         "45_dual_gemm/kernel/dual_gemm.h"
         "45_dual_gemm/device/dual_gemm.h"
+        "45_dual_gemm/dual_gemm_common.h"
         "45_dual_gemm/dual_gemm_run.h"
         "45_dual_gemm/thread/left_silu_and_mul.h"
         "45_dual_gemm/threadblock/dual_mma_multistage.h"