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

Allow passing output tensor in low_latency_combine

Browse Source
This commit is contained in:
Dmytro Dzhulgakov
2025-03-10 22:19:21 +00:00
parent ed7487c15e
commit b3b61ef5ef
4 changed files with 17 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/deep_ep.hpp
View File

@@ -143,7 +143,7 @@ public:
low_latency_combine(const torch::Tensor& x, const torch::Tensor& topk_idx, const torch::Tensor& topk_weights,
const torch::Tensor& src_info, const torch::Tensor& layout_range,
int num_max_dispatch_tokens_per_rank, int num_experts,
bool async, bool return_recv_hook);
bool async, bool return_recv_hook, std::optional<torch::Tensor> out = std::nullopt);
};
} // namespace deep_ep