diff --git a/csrc/kernels/internode.cu b/csrc/kernels/internode.cu index 809e2a2..6593f96 100644 --- a/csrc/kernels/internode.cu +++ b/csrc/kernels/internode.cu @@ -254,10 +254,12 @@ notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped, in // TODO: overlap EP barrier and NVL cleaning for (int i = 0; i < kNumRDMARanks; ++i) { if (i != rdma_rank) { - nvshmemi_ibgda_put_nbi_warp(reinterpret_cast(rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank)), - reinterpret_cast(rdma_recv_num_tokens_mixed.send_buffer(i)), - (NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int), - translate_dst_rdma_rank(i, nvl_rank), 0, lane_id, 0); + if (warp_id == 0) { + nvshmemi_ibgda_put_nbi_warp(reinterpret_cast(rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank)), + reinterpret_cast(rdma_recv_num_tokens_mixed.send_buffer(i)), + (NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int), + translate_dst_rdma_rank(i, nvl_rank), 0, lane_id, 0); + } } else { UNROLLED_WARP_COPY(1, lane_id, NUM_MAX_NVL_PEERS + num_rdma_experts + 1, rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank),