feat: add sm90 megamoe phase1 interfaces

2026-06-17 23:54:49 +08:00
parent 062cb160cf
commit 1dafd0f173
12 changed files with 1520 additions and 130 deletions
--- a/csrc/python_api.cpp
+++ b/csrc/python_api.cpp
@@ -7,6 +7,7 @@
 #include "apis/gemm.hpp"
 #include "apis/layout.hpp"
 #include "apis/mega.hpp"
+#include "apis/sm90_mega.hpp"
 #include "apis/runtime.hpp"

 #ifndef TORCH_EXTENSION_NAME
@@ -24,5 +25,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    deep_gemm::gemm::register_apis(m);
    deep_gemm::layout::register_apis(m);
    deep_gemm::mega::register_apis(m);
+    deep_gemm::mega::register_sm90_apis(m);
    deep_gemm::runtime::register_apis(m);
 }