Skip to content

GitLab

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

Fix warp synchronization. (#215) 

* Fix warp synchronization.

* Another fix.
parent 8aaddf76
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
csrc/kernels/internode.cu
View file @ dd133d39
...@@ -567,6 +567,7 @@ dispatch(int4* recv_x, float* recv_x_scales, int64_t* recv_topk_idx, float* recv ...@@ -567,6 +567,7 @@ dispatch(int4* recv_x, float* recv_x_scales, int64_t* recv_topk_idx, float* recv
// Update last token tail // Update last token tail
if (last_rdma_tail_idx >= 0) if (last_rdma_tail_idx >= 0)
st_release_cta(const_cast<const int*>(rdma_send_channel_tail + lane_id), last_rdma_tail_idx + 1); st_release_cta(const_cast<const int*>(rdma_send_channel_tail + lane_id), last_rdma_tail_idx + 1);
__syncwarp();
// Release sequential lock // Release sequential lock
lane_id == 0 ? (rdma_send_next_token_idx += 1) : 0; lane_id == 0 ? (rdma_send_next_token_idx += 1) : 0;
...@@ -633,6 +634,7 @@ dispatch(int4* recv_x, float* recv_x_scales, int64_t* recv_topk_idx, float* recv ...@@ -633,6 +634,7 @@ dispatch(int4* recv_x, float* recv_x_scales, int64_t* recv_topk_idx, float* recv
nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_tail.buffer(rdma_rank), num_tokens_to_issue, nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_tail.buffer(rdma_rank), num_tokens_to_issue,
translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank), channel_id, dst_rdma_rank == rdma_rank); translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank), channel_id, dst_rdma_rank == rdma_rank);
} }
__syncwarp();
} }
} }
} else if (warp_role == WarpRole::kRDMAAndNVLForwarder) { } else if (warp_role == WarpRole::kRDMAAndNVLForwarder) {
......
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