chore: import upstream snapshot with attribution

2026-07-13 12:24:32 +08:00
commit b5ecf06f65
107 changed files with 24727 additions and 0 deletions
@@ -0,0 +1,39 @@
+#include <pybind11/pybind11.h>
+#include <torch/python.h>
+
+#include <deep_ep/common/compiled.cuh>
+
+#include "jit/api.hpp"
+#include "elastic/buffer.hpp"
+#include "legacy/buffer.hpp"
+
+#ifndef TORCH_EXTENSION_NAME
+#define TORCH_EXTENSION_NAME _C
+#endif
+
+bool is_sm90_compiled() {
+#ifndef DISABLE_SM90_FEATURES
+    return true;
+#else
+    return false;
+#endif
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "DeepEP: an efficient expert-parallel communication library";
+
+    // Whether support FP8 and TMA features
+    m.def("is_sm90_compiled", []() { return deep_ep::kEnableSM90Features; });
+
+    // The integer type of top-k indices
+    m.attr("topk_idx_t") = py::cast(c10::CppTypeToScalarType<deep_ep::topk_idx_t>::value);
+
+    // JIT API
+    deep_ep::jit::register_apis(m);
+
+    // Register legacy buffer APIs
+    deep_ep::legacy::register_apis(m);
+
+    // Register elastic buffer (DeepEP V2) APIs
+    deep_ep::elastic::register_apis(m);
+}