zbtrs2
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[EP] fix a regression when IBGDA is disabled

This commit is contained in:
Xun Sun 2026-02-07 14:41:13 +08:00 committed by GitHub
parent c724318217
commit 83384e9c0c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 0 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
mooncake-ep/src/mooncake_ep_buffer.cpp
View File

@ -130,15 +130,6 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x,
EP_HOST_ASSERT(not(async and return_recv_hook));
if (not return_recv_hook) stream_wait(launch_stream, compute_stream);
// NVLink/P2P path uses `rdma_recv_signal_buffer` for synchronization.
// `cudaMalloc` does not guarantee zeroed memory; stale values can cause
// incorrect counts (or deadlocks). Clearing on the launch stream preserves
// ordering with subsequent kernels.
if (ibgda_disabled_) {
CUDA_CHECK(cudaMemsetAsync(buffer.rdma_recv_signal_buffer, 0,
num_experts * sizeof(int), launch_stream));
}
// Allocate packed tensors
auto packed_recv_x = torch::empty(
{num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank,
@ -268,12 +259,6 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx,
EP_HOST_ASSERT(not(async and return_recv_hook));
if (not return_recv_hook) stream_wait(launch_stream, compute_stream);
// Same rationale as dispatch(): clear receive signal buffer for NVLink/P2P.
if (ibgda_disabled_) {
CUDA_CHECK(cudaMemsetAsync(buffer.rdma_recv_signal_buffer, 0,
num_experts * sizeof(int), launch_stream));
}
// Allocate output tensor
torch::Tensor combined_x;
if (out.has_value()) {