DeepSeek/DeepEP
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepEP synced 2025-06-26 18:28:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add get_comm_stream. (#256)

Browse Source 
* Add `get_comm_stream`.

* Fix style.
This commit is contained in:
Shangyan Zhou
2025-06-25 13:02:13 +08:00
parent 9eb2f84b3e
commit 85adc566e2
3 changed files with 17 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/deep_ep.hpp
View File

@@ -94,6 +94,8 @@ public:
torch::Tensor get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset, bool use_rdma_buffer) const;
torch::Stream get_comm_stream() const;
void sync(const std::vector<int>& device_ids, const std::vector<std::optional<pybind11::bytearray>>& all_gathered_handles, const std::optional<pybind11::bytearray>& root_unique_id_opt);
std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor, torch::Tensor, std::optional<EventHandle>>