Update for build.toml changes

build.toml

@@ -1,20 +1,17 @@
 [general]
-version = "0.0.1"
+name = "moe"
 
 [torch]
-name = "moe"
 src = [
   "core/scalar_type.hpp",
-  "torch-ext/registration.h",
   "torch-ext/torch_binding.cpp",
   "torch-ext/torch_binding.h",
 ]
 include = ["."]
-pyroot = "torch-ext"
 pyext = ["py", "json"]
 
 [kernel.fp8]
-capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
+cuda-capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
 src = [
   "cuda_compat.h",
   "dispatch_utils.h",
@@ -29,7 +26,7 @@ depends = ["torch"]
 
 
 [kernel.moe]
-capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
+cuda-capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
 src = [
   "cuda_compat.h",
   "dispatch_utils.h",
@@ -39,7 +36,7 @@ src = [
 depends = ["torch"]
 
 [kernel.moe-marlin]
-capabilities = ["8.0", "8.6", "8.7", "8.9", "9.0"]
+cuda-capabilities = ["8.0", "8.6", "8.7", "8.9", "9.0"]
 src = [
   "core/exception.hpp",
   "core/scalar_type.hpp",
@@ -56,7 +53,7 @@ include = ["."]
 depends = ["torch"]
 
 [kernel.activation]
-capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
+cuda-capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
 src = [
   "activation/activation_kernels.cu",
   "activation/cuda_compat.h",

torch-ext/registration.h

@@ -1,27 +0,0 @@
-#pragma once
-
-#include <Python.h>
-
-#define _CONCAT(A, B) A##B
-#define CONCAT(A, B) _CONCAT(A, B)
-
-#define _STRINGIFY(A) #A
-#define STRINGIFY(A) _STRINGIFY(A)
-
-// A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME
-// could be a macro instead of a literal token.
-#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
-
-// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME
-// could be a macro instead of a literal token.
-#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \
-  TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE)
-
-// REGISTER_EXTENSION allows the shared library to be loaded and initialized
-// via python's import statement.
-#define REGISTER_EXTENSION(NAME)                                               \
-  PyMODINIT_FUNC CONCAT(PyInit_, NAME)() {                                     \
-    static struct PyModuleDef module = {PyModuleDef_HEAD_INIT,                 \
-                                        STRINGIFY(NAME), nullptr, 0, nullptr}; \
-    return PyModule_Create(&module);                                           \
-  }