wehub-resource-sync
126 lines
5.6 KiB
C++
126 lines
5.6 KiB
C++
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include <Python.h>
|
|
#include "pybind11/functional.h"
|
|
#include "pybind11/stl.h"
|
|
|
|
#ifdef PADDLE_WITH_DEEP_EP
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
#include "paddle/fluid/distributed/collective/deep_ep/deep_ep.hpp"
|
|
#elif defined(PADDLE_WITH_XPU)
|
|
#include "paddle/fluid/distributed/collective/deep_ep_xpu/deep_ep.hpp"
|
|
#endif
|
|
#endif
|
|
#include "paddle/fluid/pybind/deep_ep_api.h"
|
|
#include "paddle/utils/pybind.h"
|
|
|
|
namespace py = pybind11;
|
|
|
|
namespace paddle::pybind {
|
|
|
|
void BindDeepEPApi(pybind11::module *m) {
|
|
#ifdef PADDLE_WITH_DEEP_EP
|
|
pybind11::class_<deep_ep::Config>(*m, "Config", py::module_local())
|
|
.def(pybind11::init<int, int, int, int, int>(),
|
|
py::arg("num_sms") = 20,
|
|
py::arg("num_max_nvl_chunked_send_tokens") = 6,
|
|
py::arg("num_max_nvl_chunked_recv_tokens") = 256,
|
|
py::arg("num_max_rdma_chunked_send_tokens") = 6,
|
|
py::arg("num_max_rdma_chunked_recv_tokens") = 256)
|
|
.def("get_nvl_buffer_size_hint",
|
|
&deep_ep::Config::get_nvl_buffer_size_hint)
|
|
.def("get_rdma_buffer_size_hint",
|
|
&deep_ep::Config::get_rdma_buffer_size_hint);
|
|
m->def("get_low_latency_rdma_size_hint",
|
|
&deep_ep::get_low_latency_rdma_size_hint);
|
|
m->def("get_low_latency_rdma_size_hint_two_stage",
|
|
&deep_ep::get_low_latency_rdma_size_hint_two_stage);
|
|
m->def("get_low_latency_nvl_size_hint_two_stage",
|
|
&deep_ep::get_low_latency_nvl_size_hint_two_stage);
|
|
|
|
pybind11::class_<deep_ep::EventHandle>(*m, "EventHandle", py::module_local())
|
|
.def(pybind11::init<>())
|
|
.def("current_stream_wait", &deep_ep::EventHandle::current_stream_wait)
|
|
.def("calc_stream_wait", &deep_ep::EventHandle::CalcStreamWait)
|
|
.def("comm_stream_wait", &deep_ep::EventHandle::CommStreamWait);
|
|
|
|
m->def("get_event_handle_from_calc_stream",
|
|
&deep_ep::GetEventHandleFromCalcStream);
|
|
m->def("get_event_handle_from_comm_stream",
|
|
&deep_ep::GetEventHandleFromCommStream);
|
|
|
|
m->def("get_event_handle_from_custom_stream",
|
|
&deep_ep::GetEventHandleFromCustomStream);
|
|
|
|
pybind11::class_<deep_ep::Buffer>(*m, "Buffer", py::module_local())
|
|
.def(pybind11::init<int, int, int64_t, int64_t, bool, int>())
|
|
.def("is_available", &deep_ep::Buffer::is_available)
|
|
.def("get_num_rdma_ranks", &deep_ep::Buffer::get_num_rdma_ranks)
|
|
.def("get_rdma_rank", &deep_ep::Buffer::get_rdma_rank)
|
|
.def("get_root_rdma_rank", &deep_ep::Buffer::get_root_rdma_rank)
|
|
.def("get_local_device_id", &deep_ep::Buffer::get_local_device_id)
|
|
.def("get_comm_stream",
|
|
[](deep_ep::Buffer &self) {
|
|
int device_id = self.get_local_device_id();
|
|
cudaStream_t comm_stream = self.get_comm_stream();
|
|
auto s = phi::Stream(reinterpret_cast<phi::StreamId>(comm_stream));
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
return phi::CUDAStream(GPUPlace(device_id), s);
|
|
#elif defined(PADDLE_WITH_XPU)
|
|
return phi::XPUCUDAStream(phi::XPUPlace(device_id), s);
|
|
#endif
|
|
})
|
|
.def("get_local_ipc_handle", &deep_ep::Buffer::get_local_ipc_handle)
|
|
.def("get_local_nvshmem_unique_id",
|
|
&deep_ep::Buffer::get_local_nvshmem_unique_id)
|
|
.def("sync", &deep_ep::Buffer::sync)
|
|
.def("get_dispatch_layout",
|
|
[](deep_ep::Buffer &self,
|
|
py::handle topk_idx,
|
|
int num_experts,
|
|
std::optional<deep_ep::EventHandle> &previous_event,
|
|
bool async,
|
|
bool allocate_on_comm_stream) {
|
|
auto topk_idx_tensor = CastPyArg2Tensor(topk_idx.ptr(), 0);
|
|
return self.get_dispatch_layout_api(topk_idx_tensor,
|
|
num_experts,
|
|
previous_event,
|
|
async,
|
|
allocate_on_comm_stream);
|
|
})
|
|
.def("intranode_dispatch", &deep_ep::Buffer::intranode_dispatch_api)
|
|
.def("intranode_combine", &deep_ep::Buffer::intranode_combine_api)
|
|
.def("internode_dispatch", &deep_ep::Buffer::internode_dispatch_api)
|
|
.def("internode_combine", &deep_ep::Buffer::internode_combine_api)
|
|
.def("barrier_all", &deep_ep::Buffer::barrier_all)
|
|
.def("clean_low_latency_buffer",
|
|
&deep_ep::Buffer::clean_low_latency_buffer)
|
|
.def("clean_low_latency_two_stage_buffer",
|
|
&deep_ep::Buffer::clean_low_latency_two_stage_buffer)
|
|
.def("low_latency_dispatch", &deep_ep::Buffer::low_latency_dispatch_api)
|
|
.def("low_latency_combine", &deep_ep::Buffer::low_latency_combine_api)
|
|
.def("low_latency_dispatch_two_stage",
|
|
&deep_ep::Buffer::low_latency_dispatch_two_stage_api)
|
|
.def("low_latency_combine_two_stage",
|
|
&deep_ep::Buffer::low_latency_combine_two_stage_api)
|
|
.def("m2n_low_latency_dispatch_two_stage",
|
|
&deep_ep::Buffer::m2n_low_latency_dispatch_two_stage_api)
|
|
.def("m2n_low_latency_combine_two_stage",
|
|
&deep_ep::Buffer::m2n_low_latency_combine_two_stage_api);
|
|
#endif
|
|
}
|
|
|
|
} // namespace paddle::pybind
|