From 09cc1fc72d2ee097ed1bf81bf99bcfe68f6ce40f Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 23 Jul 2026 23:37:03 +0000 Subject: [PATCH] [ROCm] Fix intranode queue ordering for shared buffers Signed-off-by: Andreas Karatzas --- csrc/kernels/intranode.cu | 247 +++++++++++++++++++----- tests/test_intranode_pipeline.py | 309 +++++++++++++++++++++++++++++++ 2 files changed, 506 insertions(+), 50 deletions(-) create mode 100644 tests/test_intranode_pipeline.py diff --git a/csrc/kernels/intranode.cu b/csrc/kernels/intranode.cu index aa19cb4..4c8f259 100644 --- a/csrc/kernels/intranode.cu +++ b/csrc/kernels/intranode.cu @@ -8,6 +8,58 @@ namespace deep_ep { namespace intranode { +#ifdef USE_ROCM +struct WorkgroupWarpBarrier { + int count; + int phase; +}; + +__device__ __forceinline__ void +init_workgroup_warp_barrier(WorkgroupWarpBarrier* barrier) { + barrier->count = 0; + barrier->phase = 0; +} + +__device__ __forceinline__ void +wait_workgroup_warp_barrier(WorkgroupWarpBarrier* barrier, + int expected_warps, + bool is_warp_leader) { + // Each participating wave publishes its work through its leader. The last + // leader then releases the new phase to every other wave in this group. + syncwarp(); + if (is_warp_leader) { + const int phase = __hip_atomic_load( + &barrier->phase, + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_WORKGROUP); + const int ticket = __hip_atomic_fetch_add( + &barrier->count, + 1, + __ATOMIC_ACQ_REL, + __HIP_MEMORY_SCOPE_WORKGROUP); + if (ticket == expected_warps - 1) { + __hip_atomic_store( + &barrier->count, + 0, + __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_WORKGROUP); + __hip_atomic_store( + &barrier->phase, + phase ^ 1, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_WORKGROUP); + } else { + while (__hip_atomic_load( + &barrier->phase, + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_WORKGROUP) == phase) + __builtin_amdgcn_s_sleep(1); + } + } + syncwarp(); +} +#endif + template __global__ void notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped, @@ -215,6 +267,15 @@ dispatch(int4* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* recv_to const bool is_sender = sm_id % 2 == 0; EP_DEVICE_ASSERT(num_sms % 2 == 0); +#ifdef USE_ROCM + // Rank groups can execute a different number of queue rounds. A full + // workgroup barrier would therefore deadlock or pair unrelated rounds. + __shared__ WorkgroupWarpBarrier rank_barriers[kNumRanks]; + if (thread_id < kNumRanks) + init_workgroup_warp_barrier(&rank_barriers[thread_id]); + __syncthreads(); +#endif + // Several warps are response for a single rank const auto num_threads_per_rank = kNumThreads / kNumRanks; const auto num_channels = num_sms / 2; @@ -292,7 +353,12 @@ dispatch(int4* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* recv_to auto start_time = wall_clock64(); while (send_lane_id == 0) { // NOTES: we only consider the worst case, because counting the real numbers are time-consuming - int num_used_slots = cached_channel_tail_idx - ld_volatile_global(channel_head_idx.buffer()); + int num_used_slots = cached_channel_tail_idx - +#if defined(USE_ROCM) + ld_acquire_sys_global(channel_head_idx.buffer()); +#else + ld_volatile_global(channel_head_idx.buffer()); +#endif if (num_recv_buffer_tokens - num_used_slots >= num_max_send_tokens) break; @@ -369,21 +435,17 @@ dispatch(int4* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* recv_to // Move tail index // NOTES: here all warps should share the same new tail -#if defined(USE_ROCM) - if (num_threads_per_rank > kWarpSize){ - __syncthreads(); - }else{ - syncwarp(); - } +#if defined(USE_ROCM) + wait_workgroup_warp_barrier( + &rank_barriers[responsible_rank], + num_send_warps_per_rank, + send_lane_id == 0); #else asm volatile("bar.sync %0, %1;" :: "r"(responsible_rank), "r"(num_threads_per_rank)); #endif if (send_warp_id_in_rank == 0 and send_lane_id == 0) -#if defined(USE_ROCM) - st_relaxed_sys_global(channel_tail_idx.buffer(), cached_channel_tail_idx); -#else - st_release_sys_global(channel_tail_idx.buffer(), cached_channel_tail_idx); -#endif + st_release_sys_global( + channel_tail_idx.buffer(), cached_channel_tail_idx); } } else { // Workers for receiving and copying into buffer @@ -422,11 +484,7 @@ dispatch(int4* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* recv_to while (num_tokens_to_recv > 0) { // NOTES: unlike the sender, the receiver must ensure that the tail indices hold by different warps are same while (recv_thread_id_in_rank == 0) { -#if defined(USE_ROCM) - cached_channel_tail_idx = ld_relaxed_sys_global(channel_tail_idx.buffer()); -#else cached_channel_tail_idx = ld_acquire_sys_global(channel_tail_idx.buffer()); -#endif // Ready to copy if (cached_channel_head_idx != cached_channel_tail_idx) { @@ -443,12 +501,11 @@ dispatch(int4* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* recv_to } // Synchronize queue tail -#if defined(USE_ROCM) - if (num_threads_per_rank > kWarpSize){ - __syncthreads(); - }else{ - syncwarp(); - } +#if defined(USE_ROCM) + wait_workgroup_warp_barrier( + &rank_barriers[responsible_rank], + num_recv_warps_per_rank, + recv_lane_id == 0); #else asm volatile("bar.sync %0, %1;" :: "r"(responsible_rank), "r"(num_threads_per_rank)); #endif @@ -497,17 +554,22 @@ dispatch(int4* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* recv_to // Move queue cached_channel_head_idx += num_recv_tokens; total_offset += num_recv_tokens; -#if defined(USE_ROCM) - if (num_threads_per_rank > kWarpSize){ - __syncthreads(); - }else{ - syncwarp(); - } +#if defined(USE_ROCM) + wait_workgroup_warp_barrier( + &rank_barriers[responsible_rank], + num_recv_warps_per_rank, + recv_lane_id == 0); #else asm volatile("bar.sync %0, %1;" :: "r"(responsible_rank), "r"(num_threads_per_rank)); #endif if (recv_warp_id_in_rank == num_recv_warps_per_rank - 1 and recv_lane_id == 0) - st_relaxed_sys_global(channel_head_idx.buffer(), cached_channel_head_idx); +#if defined(USE_ROCM) + st_release_sys_global( + channel_head_idx.buffer(), cached_channel_head_idx); +#else + st_relaxed_sys_global( + channel_head_idx.buffer(), cached_channel_head_idx); +#endif // Exit num_tokens_to_recv -= num_recv_tokens; @@ -666,6 +728,13 @@ combine(dtype_t* recv_x, float* recv_topk_weights, const bool is_sender = sm_id % 2 == 0; const int responsible_channel = sm_id / 2; +#ifdef USE_ROCM + __shared__ WorkgroupWarpBarrier rank_barriers[kNumRanks]; + if (thread_id < kNumRanks) + init_workgroup_warp_barrier(&rank_barriers[thread_id]); + __syncthreads(); +#endif + EP_DEVICE_ASSERT(num_topk <= kWarpSize); constexpr int kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t); @@ -716,7 +785,12 @@ combine(dtype_t* recv_x, float* recv_topk_weights, int num_round_tokens = min(num_max_send_tokens, token_end_idx - static_cast(token_idx)); while (lane_id == 0) { // NOTES: we only consider the worst case, because counting the real numbers are time-consuming - int num_used_slots = current_channel_tail_idx - ld_volatile_global(channel_head_idx.buffer()); + int num_used_slots = current_channel_tail_idx - +#if defined(USE_ROCM) + ld_acquire_sys_global(channel_head_idx.buffer()); +#else + ld_volatile_global(channel_head_idx.buffer()); +#endif if (num_recv_buffer_tokens - num_used_slots >= num_round_tokens) break; @@ -755,21 +829,17 @@ combine(dtype_t* recv_x, float* recv_topk_weights, current_channel_tail_idx += num_round_tokens; // Move tail index -#if defined(USE_ROCM) - if (num_threads_per_rank > kWarpSize){ - __syncthreads(); - }else{ - syncwarp(); - } +#if defined(USE_ROCM) + wait_workgroup_warp_barrier( + &rank_barriers[send_rank_id], + num_send_warps_per_rank, + lane_id == 0); #else asm volatile("bar.sync %0, %1;" :: "r"(send_rank_id), "r"(num_threads_per_rank)); #endif if (lane_id == 0 and send_warp_id_in_rank == 0) -#if defined(USE_ROCM) - st_relaxed_sys_global(channel_tail_idx.buffer(), current_channel_tail_idx); -#else - st_release_sys_global(channel_tail_idx.buffer(), current_channel_tail_idx); -#endif + st_release_sys_global( + channel_tail_idx.buffer(), current_channel_tail_idx); } } else { // Workers for receiving @@ -780,11 +850,21 @@ combine(dtype_t* recv_x, float* recv_topk_weights, EP_DEVICE_ASSERT(thread_id >= 0 and kNumThreads % kWarpSize == 0); // Shared head, tail and retired flags for receiver warps +#if defined(USE_ROCM) + __shared__ int warp_channel_head_idx[num_recv_warps][kNumRanks]; + __shared__ int channel_tail_idx[kNumRanks]; + __shared__ int warp_retired[num_recv_warps]; +#else __shared__ volatile int warp_channel_head_idx[num_recv_warps][kNumRanks]; __shared__ volatile int channel_tail_idx[kNumRanks]; __shared__ volatile bool warp_retired[num_recv_warps]; +#endif if (thread_id < num_recv_warps) +#if defined(USE_ROCM) + warp_retired[thread_id] = 0; +#else warp_retired[thread_id] = false; +#endif if (lane_id < kNumRanks) warp_channel_head_idx[recv_warp_id][lane_id] = 0; if (thread_id < kNumRanks) @@ -805,24 +885,49 @@ combine(dtype_t* recv_x, float* recv_topk_weights, bool retired = true; #pragma unroll for (int i = 1; i < num_recv_warps; ++ i) +#if defined(USE_ROCM) + retired = retired and ld_acquire_cta(&warp_retired[i]); +#else retired = retired and warp_retired[i]; +#endif if (retired) break; - // Update queue tail #if defined(USE_ROCM) - channel_tail_idx[lane_id] = ld_relaxed_sys_global(channel_tail_idx_ptr); + // Bridge the producer's system-scoped tail publication to the + // receiver waves through a workgroup-scoped release. + st_release_cta( + &channel_tail_idx[lane_id], + ld_acquire_sys_global(channel_tail_idx_ptr)); #else - channel_tail_idx[lane_id] = ld_acquire_sys_global(channel_tail_idx_ptr); + channel_tail_idx[lane_id] = + ld_acquire_sys_global(channel_tail_idx_ptr); #endif // Update minimum head int min_head = std::numeric_limits::max(); #pragma unroll - for (int i = 1; i < num_recv_warps; ++ i) if (not warp_retired[i]) - min_head = min(min_head, warp_channel_head_idx[i][lane_id]); + for (int i = 1; i < num_recv_warps; ++ i) +#if defined(USE_ROCM) + if (not ld_acquire_cta(&warp_retired[i])) + min_head = min( + min_head, + ld_acquire_cta( + &warp_channel_head_idx[i][lane_id])); +#else + if (not warp_retired[i]) + min_head = min( + min_head, + warp_channel_head_idx[i][lane_id]); +#endif if (min_head != std::numeric_limits::max() and min_head > last_head) - st_relaxed_sys_global(channel_head_idx_ptr, last_head = min_head); +#if defined(USE_ROCM) + st_release_sys_global( + channel_head_idx_ptr, last_head = min_head); +#else + st_relaxed_sys_global( + channel_head_idx_ptr, last_head = min_head); +#endif } } else { // Receivers @@ -861,15 +966,40 @@ combine(dtype_t* recv_x, float* recv_topk_weights, expected_head = ld_nc_global(send_head + token_idx * kNumRanks + lane_id); auto start_time = wall_clock64(); - while (__any_sync(kFullWarpMask, channel_tail_idx[lane_id] <= expected_head and expected_head >= 0)) { - // while (channel_tail_idx[lane_id] <= expected_head and expected_head >= 0) { +#if defined(USE_ROCM) + int observed_tail = lane_id < kNumRanks and expected_head >= 0 + ? ld_acquire_cta( + &channel_tail_idx[lane_id]) + : 0; + while (__any_sync( + kFullWarpMask, + lane_id < kNumRanks and expected_head >= 0 and + observed_tail <= expected_head)) { // Timeout check long long int elapsed_time = wall_clock64() > start_time ? wall_clock64() - start_time : 0; if (elapsed_time > NUM_TIMEOUT_CYCLES) { printf("DeepEP timeout for combine receivers, rank %d, responsible_channel = %d, expect = %d\n", rank, responsible_channel, expected_head); trap(); } + if (lane_id < kNumRanks and expected_head >= 0) + observed_tail = ld_acquire_cta( + &channel_tail_idx[lane_id]); } +#else + while (__any_sync( + kFullWarpMask, + channel_tail_idx[lane_id] <= expected_head and + expected_head >= 0)) { + // Timeout check + long long int elapsed_time = wall_clock64() > start_time + ? wall_clock64() - start_time + : 0; + if (elapsed_time > NUM_TIMEOUT_CYCLES) { + printf("DeepEP timeout for combine receivers, rank %d, responsible_channel = %d, expect = %d\n", rank, responsible_channel, expected_head); + trap(); + } + } +#endif syncwarp(); // Broadcast current heads @@ -933,14 +1063,31 @@ combine(dtype_t* recv_x, float* recv_topk_weights, } // Update head +#if defined(USE_ROCM) + syncwarp(); + if (lane_id < kNumRanks) + st_release_cta( + &warp_channel_head_idx[recv_warp_id][lane_id], + (expected_head < 0) + ? -expected_head - 1 + : expected_head + 1); +#else if (lane_id < kNumRanks) - warp_channel_head_idx[recv_warp_id][lane_id] = (expected_head < 0) ? -expected_head - 1 : expected_head + 1; + warp_channel_head_idx[recv_warp_id][lane_id] = + (expected_head < 0) + ? -expected_head - 1 + : expected_head + 1; +#endif } // Retired syncwarp(); if (lane_id == 0) +#if defined(USE_ROCM) + st_release_cta(&warp_retired[recv_warp_id], 1); +#else warp_retired[recv_warp_id] = true; +#endif } } } diff --git a/tests/test_intranode_pipeline.py b/tests/test_intranode_pipeline.py new file mode 100644 index 0000000..7a7abe1 --- /dev/null +++ b/tests/test_intranode_pipeline.py @@ -0,0 +1,309 @@ +import argparse + +import torch +import torch.distributed as dist + +import deep_ep +import deep_ep_cpp +from utils import init_dist + + +NUM_EXPERTS = 4 +NUM_TOPK = 2 +HIDDEN = 7168 + + +def make_batch(rank: int, iteration: int, microbatch: int): + """Create unequal batches large enough to wrap the 128-slot queues.""" + if rank == 0: + num_tokens = 2048 + (iteration * 17 + microbatch * 31) % 127 + else: + num_tokens = 2176 + (iteration * 29 + microbatch * 43) % 131 + + generator = torch.Generator(device="cuda") + generator.manual_seed(100_000 * iteration + 1_000 * microbatch + rank) + x = torch.randn( + (num_tokens, HIDDEN), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + if iteration % 2 == 0: + # Each source sends to only one destination. The destination alternates + # between the two shared-buffer sequences, leaving one rank group idle + # while the other wraps its queues repeatedly. + destination = rank if microbatch == 0 else 1 - rank + topk_idx = torch.tensor( + [destination * 2, destination * 2 + 1], + dtype=deep_ep.topk_idx_t, + device="cuda", + ).expand(num_tokens, -1).contiguous() + else: + scores = torch.rand( + (num_tokens, NUM_EXPERTS), + dtype=torch.float32, + device="cuda", + generator=generator, + ) + topk_idx = scores.topk( + NUM_TOPK, dim=-1, sorted=False + ).indices.to(deep_ep.topk_idx_t) + topk_weights = torch.rand( + (num_tokens, NUM_TOPK), + dtype=torch.float32, + device="cuda", + generator=generator, + ) + + # DeepEP sends a token once to each distinct destination rank. An identity + # expert therefore returns one copy for every destination rank selected. + rank_idx = topk_idx // (NUM_EXPERTS // dist.get_world_size()) + expected_multiplier = torch.stack( + [ + (rank_idx == destination).any(dim=1) + for destination in range(dist.get_world_size()) + ], + dim=1, + ).sum(dim=1) + return x, topk_idx, topk_weights, expected_multiplier + + +def dispatch(buffer: deep_ep.Buffer, config: deep_ep.Config, batch): + x, topk_idx, topk_weights, expected_multiplier = batch + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + _, + ) = buffer.get_dispatch_layout( + topk_idx, + NUM_EXPERTS, + async_finish=False, + ) + recv_x, recv_topk_idx, recv_topk_weights, _, handle, _ = buffer.dispatch( + x=x, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + topk_idx=topk_idx, + topk_weights=topk_weights, + expert_alignment=1, + config=config, + async_finish=False, + ) + return ( + recv_x, + recv_topk_idx, + recv_topk_weights, + handle, + expected_multiplier, + ) + + +def check_dispatch( + receiver_rank: int, + source_batches, + recv_x: torch.Tensor, + recv_topk_idx: torch.Tensor, + recv_topk_weights: torch.Tensor, + handle, +): + rank_prefix_matrix, _, _, recv_src_idx, _, _ = handle + invalid_index = ( + NUM_EXPERTS // dist.get_world_size() + if deep_ep_cpp.AITER_MOE + else -1 + ) + local_expert_begin = receiver_rank * ( + NUM_EXPERTS // dist.get_world_size() + ) + local_expert_end = local_expert_begin + ( + NUM_EXPERTS // dist.get_world_size() + ) + + expected_x = [] + expected_topk_idx = [] + expected_topk_weights = [] + start = 0 + for source_rank, source_batch in enumerate(source_batches): + end = rank_prefix_matrix[source_rank][receiver_rank].item() + source_indices = recv_src_idx[start:end].long() + source_x, source_topk_idx, source_topk_weights, _ = source_batch + + expected_x.append(source_x[source_indices]) + source_indices_2d = source_indices[:, None].expand(-1, NUM_TOPK) + selected_experts = source_topk_idx.gather(0, source_indices_2d) + selected_weights = source_topk_weights.gather(0, source_indices_2d) + is_local = (selected_experts >= local_expert_begin) & ( + selected_experts < local_expert_end + ) + expected_topk_idx.append( + torch.where( + is_local, + selected_experts - local_expert_begin, + invalid_index, + ) + ) + expected_topk_weights.append( + torch.where(is_local, selected_weights, 0.0) + ) + start = end + + assert start == recv_x.size(0) + assert torch.equal(recv_x, torch.cat(expected_x)) + assert torch.equal(recv_topk_idx, torch.cat(expected_topk_idx)) + assert torch.equal( + recv_topk_weights, + torch.cat(expected_topk_weights), + ) + + +def test_loop(local_rank: int, num_processes: int, iterations: int): + rank, num_ranks, group = init_dist(local_rank, num_processes) + assert num_ranks == 2, "The ordering coverage requires two ranks" + buffer = deep_ep.Buffer( + group, + int(512e6), + 0, + low_latency_mode=False, + num_qps_per_rank=1, + explicitly_destroy=True, + ) + dispatch_config = deep_ep.Buffer.get_dispatch_config(num_ranks) + combine_config = deep_ep.Buffer.get_combine_config(num_ranks) + + try: + for iteration in range(iterations): + batch0 = make_batch(rank, iteration, 0) + batch1 = make_batch(rank, iteration, 1) + + # This is the order used when two vLLM microbatches share one + # DeepEP buffer. It requires each producer's queue-tail publication + # to make the corresponding payload visible to the consumer. + ( + recv0, + recv_topk_idx0, + recv_topk_weights0, + handle0, + multiplier0, + ) = dispatch(buffer, dispatch_config, batch0) + ( + recv1, + recv_topk_idx1, + recv_topk_weights1, + handle1, + multiplier1, + ) = dispatch(buffer, dispatch_config, batch1) + out0, out_topk_weights0, _ = buffer.combine( + x=recv0, + handle=handle0, + topk_weights=recv_topk_weights0, + config=combine_config, + async_finish=False, + ) + out1, out_topk_weights1, _ = buffer.combine( + x=recv1, + handle=handle1, + topk_weights=recv_topk_weights1, + config=combine_config, + async_finish=False, + ) + torch.cuda.synchronize() + + source_batches0 = [ + make_batch(source_rank, iteration, 0) + for source_rank in range(num_ranks) + ] + source_batches1 = [ + make_batch(source_rank, iteration, 1) + for source_rank in range(num_ranks) + ] + check_dispatch( + rank, + source_batches0, + recv0, + recv_topk_idx0, + recv_topk_weights0, + handle0, + ) + check_dispatch( + rank, + source_batches1, + recv1, + recv_topk_idx1, + recv_topk_weights1, + handle1, + ) + + expected0 = batch0[0] * multiplier0[:, None] + expected1 = batch1[0] * multiplier1[:, None] + expected_topk_weights0 = batch0[2] + expected_topk_weights1 = batch1[2] + for microbatch, output, expected, output_weights, expected_weights in ( + ( + 0, + out0, + expected0, + out_topk_weights0, + expected_topk_weights0, + ), + ( + 1, + out1, + expected1, + out_topk_weights1, + expected_topk_weights1, + ), + ): + if not torch.equal(output, expected): + max_abs = ( + (output - expected).abs().float().max().item() + ) + raise AssertionError( + f"rank={rank} iteration={iteration} " + f"microbatch={microbatch} max_abs={max_abs}" + ) + if not torch.equal(output_weights, expected_weights): + max_abs = ( + (output_weights - expected_weights) + .abs() + .max() + .item() + ) + raise AssertionError( + f"rank={rank} iteration={iteration} " + f"microbatch={microbatch} " + f"topk_weights_max_abs={max_abs}" + ) + + if rank == 0: + print( + f"[intranode] passed {iterations} shared-buffer iterations", + flush=True, + ) + dist.barrier(group) + finally: + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Test shared-buffer intranode dispatch/combine ordering" + ) + parser.add_argument( + "--iterations", + type=int, + default=10, + help="Number of two-microbatch sequences to run (default: 10)", + ) + args = parser.parse_args() + + num_processes = 2 + torch.multiprocessing.spawn( + test_loop, + args=(num_processes, args.iterations), + nprocs=num_processes, + )