Skip to content

GitLab

Unverified Commit fc48a467 authored by Shangyan Zhou's avatar Shangyan Zhou Committed by GitHub
Browse files

Merge pull request #182 from wzc-wuzhicheng/fix-notify-dispatch 

Fix notify_dispatch: using warp 0 to issue send
parents 9fe9021f d0225df2
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 4 deletions
csrc/kernels/internode.cu
View file @ fc48a467
...@@ -254,10 +254,12 @@ notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped, in ...@@ -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 // TODO: overlap EP barrier and NVL cleaning
for (int i = 0; i < kNumRDMARanks; ++i) { for (int i = 0; i < kNumRDMARanks; ++i) {
if (i != rdma_rank) { if (i != rdma_rank) {
nvshmemi_ibgda_put_nbi_warp<true>(reinterpret_cast<uint64_t>(rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank)), if (warp_id == 0) {
reinterpret_cast<uint64_t>(rdma_recv_num_tokens_mixed.send_buffer(i)), nvshmemi_ibgda_put_nbi_warp<true>(reinterpret_cast<uint64_t>(rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank)),
(NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int), reinterpret_cast<uint64_t>(rdma_recv_num_tokens_mixed.send_buffer(i)),
translate_dst_rdma_rank<kLowLatencyMode>(i, nvl_rank), 0, lane_id, 0); (NUM_MAX_NVL_PEERS + num_rdma_experts + 1) * sizeof(int),
translate_dst_rdma_rank<kLowLatencyMode>(i, nvl_rank), 0, lane_id, 0);
}
} else { } else {
UNROLLED_WARP_COPY(1, lane_id, NUM_MAX_NVL_PEERS + num_rdma_experts + 1, UNROLLED_WARP_COPY(1, lane_id, NUM_MAX_NVL_PEERS + num_rdma_experts + 1,
rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank), rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment