Binyli/ep revise (#828)

This pull request makes significant improvements to the MoE (Mixture of
Experts) Python API and documentation, focusing on clarifying and
expanding the Expert Parallel (EP) interface, especially around
quantization, dispatch/combine handles, and overlap configuration. The
changes introduce new data structures, update function signatures, and
improve documentation to better reflect the current and planned
capabilities of the system. Additionally, the base development container
is updated to CUDA 13.0, and minor corrections are made to extension
naming.
This commit is contained in:
Binyang Li
2026-07-06 21:14:29 -07:00
committed by GitHub
parent c398af350d
commit 8e34326d7a
19 changed files with 1585 additions and 1358 deletions

View File

@@ -62,20 +62,20 @@ void register_core(nb::module_& m) {
void* data = reinterpret_cast<void*>(ptr);
self->send(data, size, peer, tag);
},
nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"))
nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>())
.def(
"recv",
[](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) {
void* data = reinterpret_cast<void*>(ptr);
self->recv(data, size, peer, tag);
},
nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"))
nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>())
.def("all_gather", &Bootstrap::allGather, nb::arg("allData"), nb::arg("size"))
.def("barrier", &Bootstrap::barrier)
.def("send", static_cast<void (Bootstrap::*)(const std::vector<char>&, int, int)>(&Bootstrap::send),
nb::arg("data"), nb::arg("peer"), nb::arg("tag"))
nb::arg("data"), nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>())
.def("recv", static_cast<void (Bootstrap::*)(std::vector<char>&, int, int)>(&Bootstrap::recv), nb::arg("data"),
nb::arg("peer"), nb::arg("tag"));
nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>());
nb::class_<UniqueId>(m, "CppUniqueId")
.def(nb::init<>())