Skip to content

ep: gate intranode worst-token notify host-counter writes for CUDA-graph replay safety #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
fergusfinn wants to merge 1 commit into from
Closed

ep: gate intranode worst-token notify host-counter writes for CUDA-graph replay safety #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion ep/include/intranode.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@ void notify_dispatch(int const* num_tokens_per_rank,
int* moe_recv_counter_mapped, int num_ranks,
int const* num_tokens_per_expert,
int* moe_recv_expert_counter_mapped, int num_experts,
int num_tokens, bool const* is_token_in_rank,
int num_tokens, int num_worst_tokens,
bool const* is_token_in_rank,
int* channel_prefix_matrix, int* rank_prefix_matrix_copy,
int num_memset_int, int expert_alignment,
void** buffer_ptrs, int** barrier_signal_ptrs, int rank,
Expand Down
15 changes: 11 additions & 4 deletions ep/src/intranode.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ template <int kNumRanks>
__global__ void notify_dispatch(
int const* num_tokens_per_rank, int* moe_recv_counter_mapped,
int const* num_tokens_per_expert, int* moe_recv_expert_counter_mapped,
int num_experts, int num_tokens, int num_channels,
int num_experts, int num_tokens, int num_worst_tokens, int num_channels,
bool const* is_token_in_rank, int* channel_prefix_matrix,
int* rank_prefix_matrix_copy, int num_memset_int, int expert_alignment,
void** buffer_ptrs, int** barrier_signal_ptrs, int rank) {
Expand Down Expand Up @@ -59,7 +59,11 @@ __global__ void notify_dispatch(
for (int i = 1; i < kNumRanks; ++i)
local_per_rank_buffer[i * kNumRanks + thread_id] +=
local_per_rank_buffer[(i - 1) * kNumRanks + thread_id];
if (thread_id == rank)
// Graph-mode (num_worst_tokens>0) must not write the host-mapped
// counter: replays run with no host poll, so a sum written mid-flight
// by a lagging replay clobbers the -1 reset of a host-synced eager
// dispatch on the same Buffer, which then reads a stale count.
if (thread_id == rank && num_worst_tokens == 0)
*moe_recv_counter_mapped =
local_per_rank_buffer[(kNumRanks - 1) * kNumRanks + rank];
}
Expand All @@ -73,7 +77,8 @@ __global__ void notify_dispatch(
for (int i = 0; i < kNumRanks; ++i)
sum += local_per_expert_buffer[i * num_experts_per_rank + thread_id];
sum = (sum + expert_alignment - 1) / expert_alignment * expert_alignment;
moe_recv_expert_counter_mapped[thread_id] = sum;
if (num_worst_tokens == 0)
moe_recv_expert_counter_mapped[thread_id] = sum;
}
__syncthreads();

Expand Down Expand Up @@ -122,7 +127,8 @@ void notify_dispatch(int const* num_tokens_per_rank,
int* moe_recv_counter_mapped, int num_ranks,
int const* num_tokens_per_expert,
int* moe_recv_expert_counter_mapped, int num_experts,
int num_tokens, bool const* is_token_in_rank,
int num_tokens, int num_worst_tokens,
bool const* is_token_in_rank,
int* channel_prefix_matrix, int* rank_prefix_matrix_copy,
int num_memset_int, int expert_alignment,
void** buffer_ptrs, int** barrier_signal_ptrs, int rank,
Expand All @@ -131,6 +137,7 @@ void notify_dispatch(int const* num_tokens_per_rank,
LAUNCH_KERNEL(&cfg, notify_dispatch<ranks>, num_tokens_per_rank, \
moe_recv_counter_mapped, num_tokens_per_expert, \
moe_recv_expert_counter_mapped, num_experts, num_tokens, \
num_worst_tokens, \
num_channels, is_token_in_rank, channel_prefix_matrix, \
rank_prefix_matrix_copy, num_memset_int, expert_alignment, \
buffer_ptrs, barrier_signal_ptrs, rank); \
Expand Down
4 changes: 2 additions & 2 deletions ep/src/uccl_ep.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -748,8 +748,8 @@ class Buffer {
uccl::intranode::notify_dispatch(
num_tokens_per_rank, moe_recv_counter_mapped, num_ranks,
num_tokens_per_expert, moe_recv_expert_counter_mapped, num_experts,
num_tokens, is_token_in_rank, channel_prefix_matrix, rank_prefix_matrix,
num_memset_int, expert_alignment, buffer_ptrs_gpu,
num_tokens, num_worst_tokens, is_token_in_rank, channel_prefix_matrix,
rank_prefix_matrix, num_memset_int, expert_alignment, buffer_ptrs_gpu,
barrier_signal_ptrs_gpu, rank, comm_stream, num_channels);

int num_recv_tokens = -1;
Expand Down