Merge pull request #22 from Fridge003/refactor-draft

[Draft] Add tvm-ffi support

Author: Fridge003

.gitignore

@@ -21,4 +21,5 @@ deep_gemm/include/cutlass
 stubs/
 
 # Symlinks to compiled extensions
-deep_gemm/*.so
+deep_gemm/*.so
+deep_gemm/_C_build

CMakeLists.txt

@@ -6,6 +6,7 @@ set(CMAKE_VERBOSE_MAKEFILE ON)
 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -O3 -fPIC -Wno-psabi")
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -Wno-psabi")
 set(CUDA_SEPARABLE_COMPILATION ON)
+
 list(APPEND CUDA_NVCC_FLAGS "-DENABLE_FAST_DEBUG")
 list(APPEND CUDA_NVCC_FLAGS "-O3")
 list(APPEND CUDA_NVCC_FLAGS "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage")
@@ -17,13 +18,14 @@ set(TORCH_CUDA_ARCH_LIST "${CUDA_ARCH_LIST}")
 find_package(CUDAToolkit REQUIRED)
 find_package(pybind11 REQUIRED)
 find_package(Torch REQUIRED)
+find_package(tvm_ffi REQUIRED)
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CUDA_STANDARD 17)
 
 include_directories(deep_gemm/include third-party/cutlass/include third-party/cutlass/tools/util/include third-party/fmt/include)
-include_directories(${CUDA_TOOLKIT_ROOT_DIR}/targets/x86_64-linux/include ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS})
-link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs)
+include_directories(${CUDA_TOOLKIT_ROOT_DIR}/targets/x86_64-linux/include ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS} ${tvm_ffi_INCLUDE_DIR} ${tvm_ffi_DLPACK_INCLUDE_DIR})
+link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs ${tvm_ffi_ROOT_DIR}/lib)
 
 # The main Python API entrance
 pybind11_add_module(_C csrc/python_api.cpp)

build.sh

@@ -3,6 +3,10 @@ original_dir=$(pwd)
 script_dir=$(realpath "$(dirname "$0")")
 cd "$script_dir"
 
+# Link CUTLASS includes
+ln -sf $script_dir/third-party/cutlass/include/cutlass deep_gemm/include
+ln -sf $script_dir/third-party/cutlass/include/cute deep_gemm/include
+
 # Remove old dist file, build files, and install
 rm -rf build dist
 rm -rf *.egg-info

csrc/apis/attention.hpp

@@ -255,6 +255,8 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
 
 #endif
 
+#if 0
+
 static void register_apis(pybind11::module_& m) {
 #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
     m.def("fp8_gemm_nt_skip_head_mid", &fp8_gemm_nt_skip_head_mid,
@@ -276,4 +278,6 @@ static void register_apis(pybind11::module_& m) {
 #endif
 }
 
+#endif
+
 } // namespace deep_gemm::attention

csrc/apis/einsum.hpp

@@ -1,8 +1,5 @@
 #pragma once
 
-#include <pybind11/pybind11.h>
-#include <torch/python.h>
-
 #include "../utils/exception.hpp"
 #include "../utils/format.hpp"
 #include "../utils/layout.hpp"
@@ -214,6 +211,8 @@ static void fp8_einsum(const std::string& expr,
 }
 #endif
 
+#if 0
+
 static void register_apis(pybind11::module_& m) {
 #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
     m.def("einsum", &einsum,
@@ -227,4 +226,6 @@ static void register_apis(pybind11::module_& m) {
 #endif
 }
 
+#endif
+
 } // namespace deep_gemm::einsum

csrc/apis/gemm.hpp

@@ -608,6 +608,8 @@ static void cublaslt_gemm_tt(const torch::Tensor& a, const torch::Tensor& b,
     cublaslt_gemm_nt(a.transpose(0, 1), b, d, c);
 }
 
+#if 0
+
 static void register_apis(pybind11::module_& m) {
 
 #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
@@ -725,4 +727,6 @@ static void register_apis(pybind11::module_& m) {
           py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt);
 }
 
+#endif
+
 } // namespace deep_gemm::gemm

csrc/apis/hyperconnection.hpp

@@ -59,6 +59,8 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a,
 
 #endif
 
+#if 0
+
 static void register_apis(pybind11::module_& m) {
 #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
     m.def("tf32_hc_prenorm_gemm", &tf32_hc_prenorm_gemm,
@@ -67,4 +69,6 @@ static void register_apis(pybind11::module_& m) {
 #endif
 }
 
+#endif
+
 } // namespace deep_gemm::hyperconnection

csrc/apis/layout.hpp

@@ -99,6 +99,8 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te
 
 #endif
 
+#if 0
+
 static void register_apis(pybind11::module_& m) {
 
 #if DG_TENSORMAP_COMPATIBLE
@@ -117,4 +119,6 @@ static void register_apis(pybind11::module_& m) {
     m.def("get_mk_alignment_for_contiguous_layout", &get_mk_alignment_for_contiguous_layout);
 }
 
+#endif
+
 } // namespace deep_gemm::layout

csrc/apis/runtime.hpp

@@ -7,6 +7,8 @@
 
 namespace deep_gemm::runtime {
 
+#if 0
+
 static void register_apis(pybind11::module_& m) {
     m.def("set_num_sms", [&](const int& new_num_sms) {
         device_runtime->set_num_sms(new_num_sms);
@@ -34,4 +36,6 @@ static void register_apis(pybind11::module_& m) {
     });
 }
 
+#endif
+
 } // namespace deep_gemm::runtime

csrc/jit_kernels/impls/runtime_utils.hpp

@@ -1,7 +1,7 @@
 #pragma once
 
 #include <cuda.h>
-#include <torch/python.h>
+#include <torch/torch.h>
 
 #include "../heuristics/sm90.hpp"
 #include "../../jit/handle.hpp"