diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index d2049e3f..c67f4457 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -7,7 +7,7 @@ on: jobs: cpplint: - runs-on: ubuntu-22.04 + runs-on: ubuntu-24.04 steps: - name: Check out Git repository @@ -22,7 +22,7 @@ jobs: run: bash ./tools/lint.sh cpp dry pylint: - runs-on: ubuntu-22.04 + runs-on: ubuntu-24.04 steps: - name: Check out Git repository @@ -40,7 +40,7 @@ jobs: run: bash ./tools/lint.sh py dry spelling: - runs-on: ubuntu-22.04 + runs-on: ubuntu-24.04 steps: - name: Check out Git repository diff --git a/.github/workflows/mscclpp-lang.yml b/.github/workflows/mscclpp-lang.yml index a9187e96..7368cfdf 100644 --- a/.github/workflows/mscclpp-lang.yml +++ b/.github/workflows/mscclpp-lang.yml @@ -21,7 +21,7 @@ jobs: - uses: actions/checkout@v4 - name: Set environment variable - run: echo "LD_LIBRARY_PATH=/usr/local/cuda/compat:/usr/local/cuda/lib64" >> $GITHUB_ENV + run: echo "LD_LIBRARY_PATH=/usr/local/cuda/lib64" >> $GITHUB_ENV - name: Install mscclpp run: | diff --git a/src/ext/ep/ht/kernels/internode.cu b/src/ext/ep/ht/kernels/internode.cu index 878005bc..a280f256 100644 --- a/src/ext/ep/ht/kernels/internode.cu +++ b/src/ext/ep/ht/kernels/internode.cu @@ -520,8 +520,7 @@ __global__ void notify_dispatch( sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + num_rdma_experts]; recv_rdma_rank_prefix_sum[i] = sum; } - while (ld_volatile_global(moe_recv_rdma_counter_mapped) != -1) - ; + while (ld_volatile_global(moe_recv_rdma_counter_mapped) != -1); *moe_recv_rdma_counter_mapped = sum; } @@ -552,8 +551,7 @@ __global__ void notify_dispatch( sum += nvl_recv_num_tokens_per_rank.buffer(src_nvl_rank)[src_rdma_rank]; recv_gbl_rank_prefix_sum[i] = sum; } - while (ld_volatile_global(moe_recv_counter_mapped) != -1) - ; + while (ld_volatile_global(moe_recv_counter_mapped) != -1); *moe_recv_counter_mapped = sum; } if (thread_id < num_nvl_experts) { @@ -561,8 +559,7 @@ __global__ void notify_dispatch( #pragma unroll for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) sum += nvl_recv_num_tokens_per_expert.buffer(i)[thread_id]; sum = (sum + expert_alignment - 1) / expert_alignment * expert_alignment; - while (ld_volatile_global(moe_recv_expert_counter_mapped + thread_id) != -1) - ; + while (ld_volatile_global(moe_recv_expert_counter_mapped + thread_id) != -1); moe_recv_expert_counter_mapped[thread_id] = sum; } @@ -933,8 +930,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV *reinterpret_cast(is_token_in_rank + token_idx * num_ranks + lane_id * NUM_MAX_NVL_PEERS); // Acquire sequential lock - while (lane_id == 0 and rdma_send_next_token_idx != token_idx) - ; + while (lane_id == 0 and rdma_send_next_token_idx != token_idx); __syncwarp(); // Acquire next tail @@ -1026,8 +1022,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV // Epilogue // Acquire sequential lock - while (lane_id == 0 and rdma_send_next_token_idx != token_idx) - ; + while (lane_id == 0 and rdma_send_next_token_idx != token_idx); __syncwarp(); // Update last token tail (epilogue). See in-loop note on atomicMax. @@ -2207,7 +2202,7 @@ template 0) ? kNumCombineForwarderWarps / kNumRDMARanks : 1, - int kNumForwarders = kNumRDMARanks* kNumWarpsPerForwarder, + int kNumForwarders = kNumRDMARanks * kNumWarpsPerForwarder, int kNumRDMAReceivers = kNumForwarders + NUM_MAX_NVL_PEERS> __global__ void __launch_bounds__((NUM_MAX_NVL_PEERS + 1 + kNumForwarders) * 32, 1) combine(int4* combined_x, float* combined_topk_weights, const bool* is_combined_token_in_rank, const int4* x, @@ -2949,24 +2944,28 @@ void combine(cudaDataType_t type, void* combined_x, float* combined_topk_weights // latency-bound), NARROW above it (high SM, where the marginal warp costs more scheduling // than it buys). Each branch instantiates its own kernel + sets its own SMEM attribute. const bool cmb_wide = (num_channels <= EP_CMB_TMA_WARPS_MAXCH); -#define COMBINE_FLAT_GATHER_TMA_LAUNCH(num_rdma_ranks, WARPS) \ - { \ - auto tma_func = combine_flat_gather_tma; \ - const size_t cmb_tma_smem = \ - static_cast(WARPS) * kCmbTmaStages * kCmbTmaMaxContrib * kCmbTmaChunkInt4 * sizeof(int4) + \ - static_cast(WARPS) * kCmbTmaStages * sizeof(uint64_t); \ - CUDA_CHECK( \ - cudaFuncSetAttribute(tma_func, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(cmb_tma_smem))); \ - cudaLaunchConfig_t cfg = { \ - static_cast(num_channels * 2), static_cast((WARPS)*32), cmb_tma_smem, stream, nullptr, 0}; \ - cudaLaunchAttribute a[1]; \ - a[0].id = cudaLaunchAttributeCooperative; \ - a[0].val.cooperative = 1; \ - cfg.attrs = a; \ - cfg.numAttrs = 1; \ - LAUNCH_KERNEL(&cfg, tma_func, reinterpret_cast(combined_x), combined_topk_weights, \ - is_combined_token_in_rank, num_combined_tokens, hidden, num_topk, num_ranks, recv_pool_global_ptrs, \ - ep_combine_recv_idx); \ +#define COMBINE_FLAT_GATHER_TMA_LAUNCH(num_rdma_ranks, WARPS) \ + { \ + auto tma_func = combine_flat_gather_tma; \ + const size_t cmb_tma_smem = \ + static_cast(WARPS) * kCmbTmaStages * kCmbTmaMaxContrib * kCmbTmaChunkInt4 * sizeof(int4) + \ + static_cast(WARPS) * kCmbTmaStages * sizeof(uint64_t); \ + CUDA_CHECK( \ + cudaFuncSetAttribute(tma_func, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(cmb_tma_smem))); \ + cudaLaunchConfig_t cfg = {static_cast(num_channels * 2), \ + static_cast((WARPS) * 32), \ + cmb_tma_smem, \ + stream, \ + nullptr, \ + 0}; \ + cudaLaunchAttribute a[1]; \ + a[0].id = cudaLaunchAttributeCooperative; \ + a[0].val.cooperative = 1; \ + cfg.attrs = a; \ + cfg.numAttrs = 1; \ + LAUNCH_KERNEL(&cfg, tma_func, reinterpret_cast(combined_x), combined_topk_weights, \ + is_combined_token_in_rank, num_combined_tokens, hidden, num_topk, num_ranks, recv_pool_global_ptrs, \ + ep_combine_recv_idx); \ } #define COMBINE_FLAT_GATHER_TMA_CASE(num_rdma_ranks) \ if (cmb_wide) \ diff --git a/src/ext/ep/ht/kernels/internode_ncclep.cuh b/src/ext/ep/ht/kernels/internode_ncclep.cuh index 726aa7b5..179d3f2e 100644 --- a/src/ext/ep/ht/kernels/internode_ncclep.cuh +++ b/src/ext/ep/ht/kernels/internode_ncclep.cuh @@ -407,8 +407,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV // consistent ep_combine_recv_idx map. The lock path (high channel count or // non-flat 2-hop, whose coordinator forwards a MONOTONIC rdma_send_channel_tail // to peers) is UNCHANGED -- keeps the spin-wait + plain ++. - while (!ep_use_atomic_route and lane_id == 0 and rdma_send_next_token_idx != token_idx) - ; + while (!ep_use_atomic_route and lane_id == 0 and rdma_send_next_token_idx != token_idx); __syncwarp(); // Acquire next tail (atomic when lock removed: 6 sender warps race; plain ++ @@ -719,8 +718,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV // Epilogue // Acquire sequential lock (lock path only; atomic-route removed the lock) - while (!ep_use_atomic_route and lane_id == 0 and rdma_send_next_token_idx != token_idx) - ; + while (!ep_use_atomic_route and lane_id == 0 and rdma_send_next_token_idx != token_idx); __syncwarp(); // Update last token tail (epilogue). Dead under kEpFlat (coordinator returns diff --git a/src/ext/ep/ht/kernels/intranode_kernel.cu b/src/ext/ep/ht/kernels/intranode_kernel.cu index ef87866a..cf502e2b 100644 --- a/src/ext/ep/ht/kernels/intranode_kernel.cu +++ b/src/ext/ep/ht/kernels/intranode_kernel.cu @@ -379,10 +379,8 @@ __global__ void __launch_bounds__(kNumThreads, 1) // Receive channel offset int total_offset, num_tokens_to_recv; - while (recv_lane_id == 0 and (total_offset = ld_volatile_global(channel_start_offset.buffer())) == 0) - ; - while (recv_lane_id == 0 and (num_tokens_to_recv = ld_volatile_global(channel_end_offset.buffer())) == 0) - ; + while (recv_lane_id == 0 and (total_offset = ld_volatile_global(channel_start_offset.buffer())) == 0); + while (recv_lane_id == 0 and (num_tokens_to_recv = ld_volatile_global(channel_end_offset.buffer())) == 0); if (recv_lane_id == 0) { total_offset = -total_offset - 1, num_tokens_to_recv = -num_tokens_to_recv - 1; if (recv_warp_id_in_rank == 0) @@ -1172,12 +1170,12 @@ bool combine_tma(cudaDataType_t type, void* combined_x, float* combined_topk_wei #define COMBINE_INTRANODE_TMA_LAUNCH(ranks, WARPS) \ { \ auto tma_func = combine_intranode_gather_tma; \ - const size_t tma_smem = static_cast(WARPS) * kStages * (ranks)*kChunkInt4 * sizeof(int4) + \ + const size_t tma_smem = static_cast(WARPS) * kStages * (ranks) * kChunkInt4 * sizeof(int4) + \ static_cast(WARPS) * kStages * sizeof(uint64_t); \ CUDA_CHECK( \ cudaFuncSetAttribute(tma_func, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(tma_smem))); \ cudaLaunchConfig_t cfg = { \ - static_cast(num_blocks), static_cast((WARPS)*32), tma_smem, stream, nullptr, 0}; \ + static_cast(num_blocks), static_cast((WARPS) * 32), tma_smem, stream, nullptr, 0}; \ LAUNCH_KERNEL(&cfg, tma_func, reinterpret_cast(combined_x), combined_topk_weights, send_head, num_tokens, \ hidden, num_topk, num_ranks, recv_pool_ptrs, ep_combine_recv_idx, recv_pool_header_bytes); \ } \ diff --git a/src/ext/ep/kernels/common.cuh b/src/ext/ep/kernels/common.cuh index 179a2545..004758f4 100644 --- a/src/ext/ep/kernels/common.cuh +++ b/src/ext/ep/kernels/common.cuh @@ -56,8 +56,7 @@ MSCCLPP_DEVICE_INLINE void publishSignalDirect(T* slot, T value) { template MSCCLPP_DEVICE_INLINE T waitSignalNonZero(T* slot) { T value; - while ((value = mscclpp::atomicLoad(slot, mscclpp::memoryOrderAcquire)) == T{0}) - ; + while ((value = mscclpp::atomicLoad(slot, mscclpp::memoryOrderAcquire)) == T{0}); return value; } diff --git a/src/ext/ep/kernels/low_latency.cu b/src/ext/ep/kernels/low_latency.cu index c2dcded3..53e6f62f 100644 --- a/src/ext/ep/kernels/low_latency.cu +++ b/src/ext/ep/kernels/low_latency.cu @@ -351,8 +351,7 @@ MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerExpert, int* const auto dstExpertLocalIdx = responsibleExpertIdx % numLocalExperts; const auto numTokensSent = sharedNumTokensSentPerExpert[responsibleExpertIdx - smId * kNumWarpGroups]; - while (ld_acquire_global(atomicFinishCounterPerExpert + responsibleExpertIdx) != FINISHED_SUM_TAG * 2) - ; + while (ld_acquire_global(atomicFinishCounterPerExpert + responsibleExpertIdx) != FINISHED_SUM_TAG * 2); auto* counterPtr = stagedRecvCountBuffer + dstExpertLocalIdx * numRanks + rank; auto* portChannelHandle = dstRank == rank ? nullptr : portChannelHandles + dstExpertLocalIdx * numRanks + dstRank; publishSingleWriterSignal(counterPtr, static_cast(-numTokensSent - 1), rank, dstRank, rdmaBufferPtr, @@ -679,8 +678,7 @@ MSCCLPP_DEVICE_INLINE void combineSend(void* stagedRecv, int64_t* stagedRecvFlag static_assert(kNumWarpsPerGroup > 1, "Requires more than one warp per group"); asm volatile("bar.sync %0, %1;" ::"r"(warpGroupId + 1), "r"(kNumWarpsPerGroup * WARP_SIZE)); if (subWarpId == 1 and laneId == 0) { - while (ld_acquire_global(atomicCleanFlag) == 0) - ; + while (ld_acquire_global(atomicCleanFlag) == 0); auto* flagPtr = stagedRecvFlagBuffer + globalExpertIdx; auto* portChannelHandle = dstRank == rank ? nullptr : portChannelHandles + localExpertIdx * numRanks + dstRank; publishSingleWriterSignal(flagPtr, static_cast(1), rank, dstRank, rdmaBufferPtr, portChannelHandle, diff --git a/test/python/ep/test_low_latency_multirank.py b/test/python/ep/test_low_latency_multirank.py index a00ae1b1..a3d9ed13 100644 --- a/test/python/ep/test_low_latency_multirank.py +++ b/test/python/ep/test_low_latency_multirank.py @@ -5,6 +5,10 @@ Launch with (intra-node, 8 GPUs): torchrun --nproc_per_node=8 test/python/ep/test_low_latency_multirank.py \ --num-tokens 128 --hidden 7168 --num-topk 8 --num-experts 256 + # Optional CUDA graph smoke/benchmark: + torchrun --nproc_per_node=8 test/python/ep/test_low_latency_multirank.py \ + --num-tokens 128 --hidden 7168 --num-topk 8 --num-experts 256 \ + --cuda-graph --bench Launch with (2 nodes, 1 GPU per node -- DeepEP's recommended LL topology): # node 0: @@ -55,6 +59,11 @@ def parse_args(): parser.add_argument("--num-topk", type=int, default=8) parser.add_argument("--num-experts", type=int, default=256) parser.add_argument("--bench", action="store_true", help="Run dispatch/combine benchmark after correctness") + parser.add_argument( + "--cuda-graph", + action="store_true", + help="Capture dispatch/combine in CUDA graphs; correctness captures both in one graph", + ) parser.add_argument("--bench-warmup", type=int, default=5) parser.add_argument("--bench-iters", type=int, default=20) parser.add_argument("--local-rank", "--local_rank", type=int, default=None, help=argparse.SUPPRESS) @@ -215,10 +224,41 @@ def main(): if rank == 0: print("PASS", flush=True) + def _graph_capture(dispatch_buffer, combine_out): + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + dist.barrier(group=group) + with torch.cuda.graph(graph): + graph_dout = moe_comm.dispatch( + x, + topk_idx, + topk_weights, + output_buffer=dispatch_buffer, + ) + graph_combined_x = moe_comm.combine(graph_dout[0].tokens, graph_dout[1], out=combine_out) + return graph, graph_dout, graph_combined_x + + def _run_cuda_graph_correctness(): + graph_dispatch_output_buffer = torch.empty_like(dispatch_output_buffer) + graph_out = torch.empty_like(out) + graph, _, graph_combined_x = _graph_capture(graph_dispatch_output_buffer, graph_out) + graph.replay() + torch.cuda.synchronize() + + graph_diff = (graph_combined_x.float() - expected.float()).abs().max().item() + assert torch.isnan(graph_combined_x).any().item() is False + assert graph_diff < 1e-2, f"rank{rank}: LL CUDA graph combine mismatch diff={graph_diff}" + dist.barrier(group=group) + if rank == 0: + print(f"[cuda graph dispatch+combine] OK max|got-expected|={graph_diff:.4e}", flush=True) + + if args.cuda_graph: + _run_cuda_graph_correctness() + # ------------------------------------------------------------------ - # Optional benchmark. Times dispatch and combine separately, reporting - # per-iter latency (max across ranks) and aggregate effective bandwidth - # (sum across ranks). + # Optional benchmark. In CUDA graph mode, captures dispatch+combine in one + # graph; otherwise times dispatch and combine separately. Reports per-iter + # latency (max across ranks) and aggregate effective bandwidth. # ------------------------------------------------------------------ if not args.bench: return @@ -245,30 +285,48 @@ def main(): dispatch_out_, handle_ = dout moe_comm.combine(dispatch_out_.tokens, handle_, out=out_) - for _ in range(warmup): - _combine(_dispatch(), bench_out) - torch.cuda.synchronize() - dist.barrier(group=group) + if args.cuda_graph: + e2e_graph, e2e_dout, _ = _graph_capture(bench_dispatch_output_buffer, bench_out) + for _ in range(warmup): + e2e_graph.replay() + torch.cuda.synchronize() + assert e2e_dout[0].layout.num_tokens_per_expert is not None + recv_tokens = int(e2e_dout[0].layout.num_tokens_per_expert.sum().item()) - start_ev = torch.cuda.Event(enable_timing=True) - end_ev = torch.cuda.Event(enable_timing=True) - start_ev.record() - dout = None - for _ in range(iters): - dout = _dispatch() - end_ev.record() - torch.cuda.synchronize() - disp_us = start_ev.elapsed_time(end_ev) * 1e3 / iters - assert dout[0].layout.num_tokens_per_expert is not None - recv_tokens = int(dout[0].layout.num_tokens_per_expert.sum().item()) + dist.barrier(group=group) + start_ev = torch.cuda.Event(enable_timing=True) + end_ev = torch.cuda.Event(enable_timing=True) + start_ev.record() + for _ in range(iters): + e2e_graph.replay() + end_ev.record() + torch.cuda.synchronize() + e2e_us = start_ev.elapsed_time(end_ev) * 1e3 / iters + else: + for _ in range(warmup): + _combine(_dispatch(), bench_out) + torch.cuda.synchronize() + dist.barrier(group=group) - dist.barrier(group=group) - start_ev.record() - for _ in range(iters): - _combine(dout, bench_out) - end_ev.record() - torch.cuda.synchronize() - comb_us = start_ev.elapsed_time(end_ev) * 1e3 / iters + start_ev = torch.cuda.Event(enable_timing=True) + end_ev = torch.cuda.Event(enable_timing=True) + start_ev.record() + dout = None + for _ in range(iters): + dout = _dispatch() + end_ev.record() + torch.cuda.synchronize() + disp_us = start_ev.elapsed_time(end_ev) * 1e3 / iters + assert dout[0].layout.num_tokens_per_expert is not None + recv_tokens = int(dout[0].layout.num_tokens_per_expert.sum().item()) + + dist.barrier(group=group) + start_ev.record() + for _ in range(iters): + _combine(dout, bench_out) + end_ev.record() + torch.cuda.synchronize() + comb_us = start_ev.elapsed_time(end_ev) * 1e3 / iters # Dispatch payload: recv_tokens × hidden × bf16 (received on this rank). # Combine payload: recv_tokens × hidden × bf16 as well -- each local expert @@ -278,6 +336,30 @@ def main(): disp_bytes = recv_tokens * hidden * 2 comb_bytes = recv_tokens * hidden * 2 + if args.cuda_graph: + e2e_min_t = torch.tensor([e2e_us], dtype=torch.float64, device="cuda") + e2e_avg_t = torch.tensor([e2e_us], dtype=torch.float64, device="cuda") + e2e_max_t = torch.tensor([e2e_us], dtype=torch.float64, device="cuda") + dist.all_reduce(e2e_min_t, op=dist.ReduceOp.MIN, group=group) + dist.all_reduce(e2e_avg_t, op=dist.ReduceOp.SUM, group=group) + dist.all_reduce(e2e_max_t, op=dist.ReduceOp.MAX, group=group) + e2e_avg_us = e2e_avg_t.item() / num_ranks + e2e_bw_per_rank = (disp_bytes + comb_bytes) / (e2e_avg_us * 1e-6) / 1e9 + if rank == 0: + print( + f"[bench LL cuda_graph] num_ranks={num_ranks} tokens={num_tokens} hidden={hidden} " + f"num_experts={num_experts} num_topk={num_topk} warmup={warmup} iters={iters}", + flush=True, + ) + print( + f" dispatch+combine graph: avg={e2e_avg_us:.1f}us " + f"min={e2e_min_t.item():.1f}us max={e2e_max_t.item():.1f}us " + f"per_rank_bw={e2e_bw_per_rank:.2f} GB/s " + f"agg_bw={e2e_bw_per_rank * num_ranks:.2f} GB/s (BW @ avg time)", + flush=True, + ) + return + # Reduce timings: report min/avg/max and base BW on AVG to match NCCL-EP's # `ep_bench.cu` convention. disp_min_t = torch.tensor([disp_us], dtype=torch.float64, device="cuda")