diff --git a/src/ext/ep/legacy/low_latency.cu b/src/ext/ep/legacy/low_latency.cu index c2e5967e..67296291 100644 --- a/src/ext/ep/legacy/low_latency.cu +++ b/src/ext/ep/legacy/low_latency.cu @@ -862,28 +862,7 @@ void combine(void* output, const void* input, const float* inputScales, const in const auto numWarps = kNumWarpGroups * kNumWarpsPerGroup; const auto numSmsBase = cell_div(numExperts, kNumWarpGroups); - // combineRecv's per-token weighted-reduction loop strides by smId over - // numCombinedTokens, so extra blocks parallelize it ~linearly. Grow the grid - // from ceil(numExperts / kNumWarpGroups) toward numCombinedTokens; the - // send-side expert work is guarded by responsibleExpertIdx < numExperts, so - // the extra blocks run recv-only and stay correct. (The feature/ep - // restructure capped the grid at numSmsBase, which regressed the LL combine - // on the intranode/IPC path.) - // - // Cap the grid at the device SM count: combine is launched cooperatively - // (cg::this_grid().sync()), so every block must be co-resident on the GPU -- - // requesting more blocks than SMs would fail the cooperative launch. The SM - // count is a device constant, so query it once and cache it (one GPU per - // process in this EP topology). - static const int deviceNumSms = [] { - int curDev = 0, sms = 0; - CUDA_CHECK(cudaGetDevice(&curDev)); - CUDA_CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, curDev)); - return sms; - }(); - int numSmsWanted = numCombinedTokens > numSmsBase ? numCombinedTokens : numSmsBase; - if (numSmsWanted > deviceNumSms) numSmsWanted = deviceNumSms; - const auto numSms = numSmsWanted; + const auto numSms = numSmsBase; auto atomicCleanFlag = reinterpret_cast(workspace); EP_HOST_ASSERT(sizeof(int) <= NUM_WORKSPACE_BYTES); diff --git a/test/python/ep/CMakeLists.txt b/test/python/ep/CMakeLists.txt index ddb63219..57ab3917 100644 --- a/test/python/ep/CMakeLists.txt +++ b/test/python/ep/CMakeLists.txt @@ -1,26 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # -# Standalone CMake build for mscclpp_ep_bench -- the pure-C++/MPI low-latency EP -# benchmark that calls mscclpp::ep::MoERuntime directly. +# Standalone build for: +# - mscclpp_ep_bench: pure-C++/MPI low-latency EP benchmark. +# - cupti_kernel_timer: optional in-process timer used by ep_bench_ll.py. # -# The EP dispatch/combine symbols live only in the nanobind Python module -# (mscclpp_ep_cpp.so), so this recompiles the two LL translation units -# (moe_runtime.cc + kernels/low_latency.cu) into the binary and links the -# installed libmscclpp.so. Flags mirror src/ext/ep/CMakeLists.txt. -# -# Configure/build (from this directory), e.g. on a GB200 node in the torch env: -# cmake -S . -B build \ -# -DMSCCLPP_SRC=/opt/microsoft/mrc/ep/mscclpp \ -# -DMSCCLPP_EP_NUM_MAX_NVL_PEERS=4 -DCMAKE_CUDA_ARCHITECTURES=100 -# cmake --build build -j -# -# Key cache options: -# MSCCLPP_SRC mscclpp source tree (for EP headers + the two LL TUs) -# MSCCLPP_INSTALL_DIR installed mscclpp package dir (has lib/ + include/); -# autodetected under the active conda env if unset -# MSCCLPP_EP_NUM_MAX_NVL_PEERS 4 for GB200 NVL72, 8 for HGX (default 4) -# CMAKE_CUDA_ARCHITECTURES 100 for GB200 sm_100 (default 100) +# Example: +# cmake -S test/python/ep -B test/python/ep/build \ +# -DMSCCLPP_EP_NUM_MAX_NVL_PEERS=8 \ +# -DCMAKE_CUDA_ARCHITECTURES=90 +# cmake --build test/python/ep/build -j 64 cmake_minimum_required(VERSION 3.25) project(mscclpp_ep_bench LANGUAGES CXX CUDA) @@ -32,93 +21,84 @@ if(NOT CMAKE_BUILD_TYPE) set(CMAKE_BUILD_TYPE Release) endif() if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) - set(CMAKE_CUDA_ARCHITECTURES 100) # GB200 sm_100 + set(CMAKE_CUDA_ARCHITECTURES native) endif() -# --- Paths ----------------------------------------------------------------- -set(MSCCLPP_SRC "/opt/microsoft/mrc/ep/mscclpp" CACHE PATH "mscclpp source tree") -set(MSCCLPP_EP_NUM_MAX_NVL_PEERS "4" CACHE STRING - "Compile-time NUM_MAX_NVL_PEERS for the EP kernels (4 for GB200, 8 for HGX)") +get_filename_component(_default_mscclpp_src "${CMAKE_CURRENT_SOURCE_DIR}/../../.." ABSOLUTE) +set(MSCCLPP_SRC "${_default_mscclpp_src}" CACHE PATH "mscclpp source tree") +set(MSCCLPP_EP_NUM_MAX_NVL_PEERS "8" CACHE STRING + "Compile-time NUM_MAX_NVL_PEERS for the EP kernels (8 for HGX, 4 for GB200)") + +find_package(MPI REQUIRED) +find_package(CUDAToolkit REQUIRED) +find_package(Python 3.10 COMPONENTS Interpreter REQUIRED) +find_package(Threads REQUIRED) -# Installed mscclpp (libmscclpp.so + public headers). Default: active conda env. if(NOT DEFINED MSCCLPP_INSTALL_DIR) - set(_conda "$ENV{CONDA_PREFIX}") - if(NOT _conda) - set(_conda "$ENV{HOME}/miniconda3/envs/torch") - endif() - file(GLOB _cands "${_conda}/lib/python*/site-packages/mscclpp") - if(_cands) - list(GET _cands 0 MSCCLPP_INSTALL_DIR) + execute_process( + COMMAND "${Python_EXECUTABLE}" -c + "import os, mscclpp._mscclpp as m; print(os.path.dirname(m.__file__))" + OUTPUT_VARIABLE MSCCLPP_INSTALL_DIR + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _mscclpp_python_result) + if(NOT _mscclpp_python_result EQUAL 0) + unset(MSCCLPP_INSTALL_DIR) endif() endif() if(NOT MSCCLPP_INSTALL_DIR) - message(FATAL_ERROR "MSCCLPP_INSTALL_DIR not found; pass -DMSCCLPP_INSTALL_DIR=<...>/site-packages/mscclpp") + message(FATAL_ERROR + "MSCCLPP_INSTALL_DIR not found; install mscclpp or pass the package directory containing lib/ and include/") endif() -message(STATUS "mscclpp source : ${MSCCLPP_SRC}") -message(STATUS "mscclpp install: ${MSCCLPP_INSTALL_DIR}") set(EP "${MSCCLPP_SRC}/src/ext/ep") - find_library(MSCCLPP_LIBRARY NAMES mscclpp HINTS "${MSCCLPP_INSTALL_DIR}/lib" REQUIRED) - -# --- Dependencies ---------------------------------------------------------- -find_package(MPI REQUIRED) -find_package(CUDAToolkit REQUIRED) - -# CUPTI (kernel timing). Not a CUDAToolkit:: imported target on all versions. find_path(CUPTI_INCLUDE_DIR cupti.h HINTS "${CUDAToolkit_TARGET_DIR}/include" "${CUDAToolkit_INCLUDE_DIRS}" - /usr/local/cuda/targets/sbsa-linux/include /usr/local/cuda/extras/CUPTI/include + "${CUDAToolkit_ROOT}/extras/CUPTI/include" REQUIRED) find_library(CUPTI_LIBRARY cupti - HINTS "${CUDAToolkit_LIBRARY_DIR}" /usr/local/cuda/targets/sbsa-linux/lib - /usr/local/cuda/extras/CUPTI/lib64 + HINTS "${CUDAToolkit_TARGET_DIR}/lib" "${CUDAToolkit_LIBRARY_DIR}" + "${CUDAToolkit_ROOT}/extras/CUPTI/lib64" REQUIRED) -# --- Target ---------------------------------------------------------------- add_executable(mscclpp_ep_bench mscclpp_ep_bench.cu - ${EP}/moe_runtime.cc - ${EP}/kernels/low_latency.cu -) - -# moe_runtime.cc uses device intrinsics via gpu_data_types.hpp -> it is a CUDA -# translation unit (the nanobind build compiles it with nvcc too). -set_source_files_properties(${EP}/moe_runtime.cc PROPERTIES LANGUAGE CUDA) + "${EP}/moe_runtime.cc" + "${EP}/low_latency/dispatch.cu" + "${EP}/low_latency/combine.cu") target_include_directories(mscclpp_ep_bench PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR} - ${EP} - ${EP}/ht - ${MSCCLPP_SRC}/include - ${MSCCLPP_SRC}/src/core/include - ${MSCCLPP_SRC}/src/ext/include - ${MSCCLPP_INSTALL_DIR}/include - ${CUPTI_INCLUDE_DIR} -) + "${EP}" + "${EP}/include" + "${MSCCLPP_SRC}/include" + "${MSCCLPP_SRC}/src/core/include" + "${MSCCLPP_SRC}/src/ext/include" + "${MSCCLPP_INSTALL_DIR}/include" + "${CUPTI_INCLUDE_DIR}") target_compile_definitions(mscclpp_ep_bench PRIVATE MSCCLPP_USE_CUDA EP_DISPATCH_NCCLEP - NUM_MAX_NVL_PEERS=${MSCCLPP_EP_NUM_MAX_NVL_PEERS} -) + NUM_MAX_NVL_PEERS=${MSCCLPP_EP_NUM_MAX_NVL_PEERS}) target_compile_options(mscclpp_ep_bench PRIVATE $<$:--expt-relaxed-constexpr> - $<$:--expt-extended-lambda> -) + $<$:--expt-extended-lambda>) set_target_properties(mscclpp_ep_bench PROPERTIES CUDA_SEPARABLE_COMPILATION OFF POSITION_INDEPENDENT_CODE ON BUILD_RPATH "${MSCCLPP_INSTALL_DIR}/lib" - INSTALL_RPATH "${MSCCLPP_INSTALL_DIR}/lib" -) + INSTALL_RPATH "${MSCCLPP_INSTALL_DIR}/lib") target_link_libraries(mscclpp_ep_bench PRIVATE - ${MSCCLPP_LIBRARY} - ${CUPTI_LIBRARY} + "${MSCCLPP_LIBRARY}" + "${CUPTI_LIBRARY}" MPI::MPI_CXX CUDA::cudart CUDA::cuda_driver -) + Threads::Threads) + +add_library(cupti_kernel_timer SHARED cupti_kernel_timer.cpp) +target_include_directories(cupti_kernel_timer PRIVATE "${CUPTI_INCLUDE_DIR}") +target_link_libraries(cupti_kernel_timer PRIVATE "${CUPTI_LIBRARY}" Threads::Threads) diff --git a/test/python/ep/cupti_kernel_timer.cpp b/test/python/ep/cupti_kernel_timer.cpp index 82054704..67613860 100644 --- a/test/python/ep/cupti_kernel_timer.cpp +++ b/test/python/ep/cupti_kernel_timer.cpp @@ -11,16 +11,9 @@ // ep_bench's methodology exactly: start() after warmup, stop() after the timed // loop, get_avg_us("dispatch"/"combine") buckets by mangled-name substring. // -// COOPERATIVE-LAUNCH NOTE (GB200 / CUDA 13): the mscclpp LL dispatch/combine -// kernels are launched with cudaLaunchCooperativeKernel. Those are NOT reported -// by CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL on this driver, but they ARE reported -// by CUPTI_ACTIVITY_KIND_KERNEL (the serialized-kernel activity), which is what -// we subscribe to below. KIND_KERNEL only serializes *inter*-kernel concurrency; -// in this dispatch->sync->combine->sync paired loop the kernels already run one -// at a time, so the measured per-kernel GPU duration is unaffected. The activity -// record carries the RAW MANGLED name (e.g. ...low_latency8dispatch...), so the -// caller matches the substring "dispatch"/"combine" (present in the mangled form) -// rather than the demangled "low_latency::dispatch". +// CONCURRENT_KERNEL records ordinary LL kernel launches without serializing +// activity. The record carries the raw mangled name; the caller matches +// "dispatch"/"combine" substrings. // // Build (host-only C++, links libcupti): // g++ -O2 -fPIC -shared cupti_kernel_timer.cpp -o libcupti_kernel_timer.so \ @@ -57,10 +50,8 @@ void CUPTIAPI bufferCompleted(CUcontext, uint32_t, uint8_t* buffer, size_t, size CUpti_Activity* record = nullptr; std::lock_guard lock(g_mutex); while (cuptiActivityGetNextRecord(buffer, validSize, &record) == CUPTI_SUCCESS) { - if (record->kind == CUPTI_ACTIVITY_KIND_KERNEL) { - // CUpti_ActivityKernel10 is the record layout for CUDA 13 CUPTI. start/end - // (GPU HW timestamps, ns) and name have been stable across versions. - auto* k = reinterpret_cast(record); + if (record->kind == CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL) { + auto* k = reinterpret_cast(record); if (k->name) { auto& s = g_stats[k->name]; s.total_ns += (k->end - k->start); @@ -75,8 +66,7 @@ void CUPTIAPI bufferCompleted(CUcontext, uint32_t, uint8_t* buffer, size_t, size extern "C" { -// Clear stats, register the buffer callbacks, and enable concurrent-kernel -// activity recording. Call AFTER warmup (like ep_bench's KernelTimer::start()). +// Clear stats, register callbacks, and enable kernel activity recording after warmup. int kt_start() { { std::lock_guard lock(g_mutex); @@ -84,14 +74,14 @@ int kt_start() { } CUptiResult r = cuptiActivityRegisterCallbacks(bufferRequested, bufferCompleted); if (r != CUPTI_SUCCESS) return static_cast(r); - r = cuptiActivityEnable(CUPTI_ACTIVITY_KIND_KERNEL); + r = cuptiActivityEnable(CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL); return static_cast(r); } // Flush pending buffers and disable recording. Returns CUPTI result code (0=ok). int kt_stop() { cuptiActivityFlushAll(0); - CUptiResult r = cuptiActivityDisable(CUPTI_ACTIVITY_KIND_KERNEL); + CUptiResult r = cuptiActivityDisable(CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL); return static_cast(r); } diff --git a/test/python/ep/ep_bench_ll.py b/test/python/ep/ep_bench_ll.py index bd182bc4..f1d37af8 100644 --- a/test/python/ep/ep_bench_ll.py +++ b/test/python/ep/ep_bench_ll.py @@ -97,6 +97,7 @@ def parse_args() -> argparse.Namespace: "--hidden", type=int, default=int(os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "7168")), + choices=(4096, 7168, 8192, 9216), help="hidden dimension", ) p.add_argument( @@ -104,6 +105,7 @@ def parse_args() -> argparse.Namespace: "--num-topk", type=int, default=int(os.environ.get("MSCCLPP_EP_BENCH_TOPK", "8")), + choices=range(1, 10), help="top-k experts per token", ) p.add_argument( @@ -127,6 +129,19 @@ def parse_args() -> argparse.Namespace: default=int(os.environ.get("MSCCLPP_EP_BENCH_ITERS", "50")), help="timed iterations", ) + p.add_argument( + "--dispatch-dtype", + choices=("bf16", "fp8_e4m3"), + default="bf16", + help="low-latency dispatch payload format", + ) + p.add_argument( + "--combine-mode", + choices=("rank_local_reduce", "direct_send"), + default="rank_local_reduce", + help="low-latency combine algorithm", + ) + p.add_argument("--num-blocks", type=int, default=130, help="total low-latency dispatch blocks") p.add_argument( "--no-kernel-timing", dest="kernel_timing", @@ -149,12 +164,22 @@ def parse_args() -> argparse.Namespace: help="use the in-process CUPTI collector (libcupti_kernel_timer.so, a faithful " "port of ep_bench's KernelTimer): CUPTI Activity API records per-kernel GPU " "time over the post-warmup timed loop, near-zero host perturbation, and works " - "multinode without nsys. Uses CUPTI_ACTIVITY_KIND_KERNEL (which -- unlike " - "CONCURRENT_KERNEL -- captures mscclpp's cudaLaunchCooperativeKernel LL kernels); " - "matches the mangled name substring dispatch/combine. Replaces the torch.profiler pass.", + "multinode without nsys. Matches mangled dispatch/combine kernel names and " + "replaces the torch.profiler pass.", ) p.add_argument("--seed", type=int, default=0xB3C4, help="per-rank RNG seed base") - return p.parse_args() + args = p.parse_args() + if args.hidden not in (4096, 7168, 8192, 9216): + p.error("--hidden must be one of 4096, 7168, 8192, 9216") + if not 1 <= args.num_topk <= 9: + p.error("--num-topk must be in [1, 9]") + if args.num_tokens <= 0 or args.num_experts <= 0: + p.error("--num-tokens and --num-experts must be positive") + if args.num_topk > args.num_experts: + p.error("--num-topk must not exceed --num-experts") + if args.num_warmup < 0 or args.num_iters <= 0: + p.error("--num-warmup must be non-negative and --num-iters must be positive") + return args def init_dist(): @@ -257,7 +282,10 @@ class _InProcCupti: import ctypes import os as _os - so = _os.path.join(_os.path.dirname(_os.path.abspath(__file__)), "libcupti_kernel_timer.so") + so = _os.environ.get( + "MSCCLPP_EP_CUPTI_TIMER_LIB", + _os.path.join(_os.path.dirname(_os.path.abspath(__file__)), "build", "libcupti_kernel_timer.so"), + ) self.lib = ctypes.CDLL(so) self.lib.kt_start.restype = ctypes.c_int self.lib.kt_stop.restype = ctypes.c_int @@ -296,6 +324,19 @@ def main() -> None: iters = args.num_iters assert num_experts % num_ranks == 0, "num_experts must be divisible by num_ranks" num_local_experts = num_experts // num_ranks + dispatch_data_type = { + "bf16": ep.DispatchDataType.BF16, + "fp8_e4m3": ep.DispatchDataType.FP8_E4M3, + }[args.dispatch_dtype] + combine_mode = { + "rank_local_reduce": ep.CombineMode.RANK_LOCAL_REDUCE, + "direct_send": ep.CombineMode.DIRECT_SEND, + }[args.combine_mode] + dispatch_quant = ( + None if dispatch_data_type == ep.DispatchDataType.BF16 else ep.QuantConfig(format=dispatch_data_type) + ) + dispatch_dtype = torch.bfloat16 if dispatch_quant is None else torch.float8_e4m3fn + dispatch_label = "BF16" if dispatch_quant is None else "FP8_E4M3" # bf16 precision anchor (same convention as test_low_latency_multirank.py). rank_offset = 128 @@ -314,14 +355,16 @@ def main() -> None: # ep_bench byte accounting: num_valid_selections = count(topk_idx >= 0). We # keep every selection valid (a full LL load), so this equals num_tokens*top_k. num_valid_selections = int((topk_idx >= 0).sum().item()) - disp_bytes = num_valid_selections * hidden * 2 # BF16 + dispatch_bytes_per_token = hidden * 2 if dispatch_quant is None else hidden + hidden // 128 * 4 + disp_bytes = num_valid_selections * dispatch_bytes_per_token comb_bytes = num_valid_selections * hidden * 2 # BF16 (symmetric, per ep_bench) - num_rdma_bytes = get_low_latency_rdma_size_hint(num_tokens, hidden, num_ranks, num_experts) + num_rdma_bytes = get_low_latency_rdma_size_hint(num_tokens, hidden, num_ranks, num_experts, num_topk) if rank == 0: print( f"[cfg] algorithm=LOW_LATENCY num_ranks={num_ranks} tokens/rank={num_tokens} hidden={hidden} " f"num_experts={num_experts} top_k={num_topk} warmup={warmup} iters={iters} " + f"dispatch_dtype={args.dispatch_dtype} combine_mode={args.combine_mode} " f"num_rdma_bytes={num_rdma_bytes}", flush=True, ) @@ -337,6 +380,9 @@ def main() -> None: max_tokens_per_rank=num_tokens, mode=ep.MoEMode.LOW_LATENCY, num_rdma_qps_per_rank=max(1, num_experts // num_ranks), + low_latency_num_blocks=args.num_blocks, + low_latency_combine_mode=combine_mode, + quant=dispatch_quant, ) assert moe_comm.is_available() if rank == 0: @@ -348,7 +394,12 @@ def main() -> None: # owns its src_info/layout_range/count buffers internally; we only supply the # dispatch output buffer and the combine output tensor. output_buffer = torch.empty( - (num_local_experts, num_ranks * num_tokens, hidden), dtype=torch.bfloat16, device="cuda" + (num_local_experts, num_ranks * num_tokens, hidden), dtype=dispatch_dtype, device="cuda" + ) + expert_output = ( + None + if dispatch_quant is None + else torch.zeros((num_local_experts, num_ranks * num_tokens, hidden), dtype=torch.bfloat16, device="cuda") ) out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") @@ -360,7 +411,7 @@ def main() -> None: def combine_fn(dout): dispatch_out, handle = dout - moe_comm.combine(dispatch_out.tokens, handle, out=out) + moe_comm.combine(dispatch_out.tokens if expert_output is None else expert_output, handle, out=out) stream = torch.cuda.current_stream() @@ -385,20 +436,38 @@ def main() -> None: # In-process CUPTI collector (ep_bench KernelTimer analog). start() after # warmup, stop() after the timed loop -- same window as the CUDA events. _inproc = None - if bool(getattr(args, "cupti_inproc", False)): + inproc_requested = bool(getattr(args, "cupti_inproc", False)) + local_inproc_ready = False + if inproc_requested: try: _inproc = _InProcCupti() - torch.cuda.synchronize() - dist.barrier(group=group) - _rc = _inproc.start() - if _rc != 0: - if rank == 0: - print(f"[warn] in-proc CUPTI kt_start rc={_rc}; disabling", flush=True) - _inproc = None + local_inproc_ready = True except Exception as exc: if rank == 0: print(f"[warn] in-proc CUPTI unavailable ({exc}); host-observed only", flush=True) _inproc = None + ready = torch.tensor(int(local_inproc_ready), dtype=torch.int32, device="cuda") + dist.all_reduce(ready, op=dist.ReduceOp.MIN, group=group) + if ready.item() == 0: + if rank == 0: + print("[warn] in-proc CUPTI unavailable on at least one rank; disabling globally", flush=True) + _inproc = None + else: + torch.cuda.synchronize() + dist.barrier(group=group) + try: + _rc = _inproc.start() + except Exception: + _rc = -1 + started = torch.tensor(int(_rc == 0), dtype=torch.int32, device="cuda") + dist.all_reduce(started, op=dist.ReduceOp.MIN, group=group) + if started.item() == 0: + if _rc == 0: + _inproc.stop() + if rank == 0: + print("[warn] in-proc CUPTI failed to start on at least one rank; disabling globally", flush=True) + _inproc = None + dist.barrier(group=group) d_start = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] d_end = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] @@ -521,7 +590,7 @@ def main() -> None: print(f"\n=== Summary (Low Latency, across {num_ranks} ranks) ===") print("\n--- Host-observed performance ---") - print(f"Dispatch (BF16): avg={g_d_avg:.2f} us, min={g_d_min:.2f} us, max={g_d_max:.2f} us") + print(f"Dispatch ({dispatch_label}): avg={g_d_avg:.2f} us, min={g_d_min:.2f} us, max={g_d_max:.2f} us") print( f" throughput: avg={avg_d_tp:.2f} GB/s, " f"min={d_lo:.2f} GB/s (rank {d_lo_r}), max={d_hi:.2f} GB/s (rank {d_hi_r})" @@ -551,13 +620,11 @@ def main() -> None: f"[avg={g_dk_avg:.2f}, max={g_dk_max:.2f} us -- inflated by profiler recv-spin skew]" ) print(f" throughput @min: {(disp_bytes / 1e9) / (g_dk_min * 1e-6):.2f} GB/s") - print(f"Combine: avg={g_ck_avg:.2f} us, min={g_ck_min:.2f} us, max={g_ck_max:.2f} us") print( - f" throughput: avg={(comb_bytes / 1e9) / (g_ck_avg * 1e-6):.2f} GB/s, " - f"min={(comb_bytes / 1e9) / (g_ck_min * 1e-6):.2f} GB/s, " - f"max={(comb_bytes / 1e9) / (g_ck_max * 1e-6):.2f} GB/s" + f"Combine: min={g_ck_min:.2f} us (representative) " + f"[avg={g_ck_avg:.2f}, max={g_ck_max:.2f} us -- inflated by profiler rank skew]" ) - print(f"Total (D+C): {g_dk_min + g_ck_avg:.2f} us (dispatch min + combine avg)") + print(f" throughput @min: {(comb_bytes / 1e9) / (g_ck_min * 1e-6):.2f} GB/s") print( " NOTE: for an authoritative low-perturbation kernel-only number, run under " "nsys (as ep_bench's CUPTI path does); torch.profiler perturbs the LL recv-spin." @@ -579,18 +646,16 @@ def main() -> None: f"[avg={g_ik_d_avg:.2f}, max={g_ik_d_max:.2f} us -- recv-spin skew on lagging ranks]" ) print(f" throughput @min: {(disp_bytes / 1e9) / (g_ik_d_min * 1e-6):.2f} GB/s") - print(f"Combine: avg={g_ik_c_avg:.2f} us, min={g_ik_c_min:.2f} us, max={g_ik_c_max:.2f} us") print( - f" throughput: avg={(comb_bytes / 1e9) / (g_ik_c_avg * 1e-6):.2f} GB/s, " - f"min={(comb_bytes / 1e9) / (g_ik_c_max * 1e-6):.2f} GB/s, " - f"max={(comb_bytes / 1e9) / (g_ik_c_min * 1e-6):.2f} GB/s" + f"Combine: min={g_ik_c_min:.2f} us (representative) " + f"[avg={g_ik_c_avg:.2f}, max={g_ik_c_max:.2f} us -- rank skew on lagging ranks]" ) - print(f"Total (D+C): {g_ik_d_min + g_ik_c_avg:.2f} us (dispatch min + combine avg)") + print(f" throughput @min: {(comb_bytes / 1e9) / (g_ik_c_min * 1e-6):.2f} GB/s") else: print(" NOTE: in-process CUPTI collector unavailable (see [warn] above).") print( - f"\nByte counts: dispatch={disp_bytes / 1e6:.2f} MB (BF16), " + f"\nByte counts: dispatch={disp_bytes / 1e6:.2f} MB ({dispatch_label}), " f"combine={comb_bytes / 1e6:.2f} MB (BF16), selections={num_valid_selections}" ) diff --git a/test/python/ep/mscclpp_ep_bench.cu b/test/python/ep/mscclpp_ep_bench.cu index d3947f7b..46147dad 100644 --- a/test/python/ep/mscclpp_ep_bench.cu +++ b/test/python/ep/mscclpp_ep_bench.cu @@ -5,8 +5,7 @@ // mscclpp::ep::MoERuntime::dispatch / ::combine directly (no Python), so mscclpp // EP can be compared with NVIDIA NCCL-EP's ep_bench on an equal footing -- // C++ host launch, and CUPTI kernel timing via CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL -// (the same activity kind ep_bench uses; unlike the torch/Kineto in-process path -// it does report mscclpp's cooperative-launch LL kernels). +// (the same activity kind ep_bench uses). // // It mirrors ep_bench's LL measurement methodology and emits the identical // "=== Summary (Low Latency, across N ranks) ===" block so the unified driver @@ -16,7 +15,6 @@ // (the bootstrap uses an MPI_Bcast of a TcpBootstrap UniqueId). #include -#include #include #include #include @@ -28,12 +26,14 @@ #include #include #include +#include +#include #include #include -#include "config.hpp" // mscclpp::ep::getLowLatencyRdmaSizeHint -#include "kernels/api.cuh" // mscclpp::ep::MoEMode, DispatchLayout -#include "moe_runtime.hpp" // mscclpp::ep::MoERuntime +#include "api.cuh" +#include "config.hpp" +#include "moe_runtime.hpp" #define CUDA_CHECK(x) \ do { \ @@ -56,7 +56,7 @@ // --------------------------------------------------------------------------- // KernelTimer: per-kernel GPU timing via the CUPTI Activity API, a faithful -// analog of ep_bench's KernelTimer. Uses CONCURRENT_KERNEL (ep_bench's kind). +// analog of ep_bench's KernelTimer. // Records are bucketed by mangled-name substring ("dispatch"/"combine"). // --------------------------------------------------------------------------- namespace { @@ -79,7 +79,7 @@ void CUPTIAPI bufferCompleted(CUcontext, uint32_t, uint8_t* buffer, size_t, size CUpti_Activity* record = nullptr; while (cuptiActivityGetNextRecord(buffer, validSize, &record) == CUPTI_SUCCESS) { if (record->kind == CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL || record->kind == CUPTI_ACTIVITY_KIND_KERNEL) { - auto* k = reinterpret_cast(record); + auto* k = reinterpret_cast(record); if (k->name) { auto& e = g_kernel_stats[k->name]; e.total_ns += (k->end - k->start); @@ -132,6 +132,11 @@ struct Args { int num_experts = 256; int num_warmup = 10; int num_iters = 50; + int num_blocks = mscclpp::ep::low_latency::MaxDispatchBlocks; + int seed = 0xB3C4; + bool kernel_timing = false; + std::string dispatch_dtype = "bf16"; + std::string combine_mode = "rank_local_reduce"; }; Args parse_args(int argc, char** argv) { @@ -153,6 +158,16 @@ Args parse_args(int argc, char** argv) { a.num_warmup = next(); else if (s == "-i" || s == "--num-iters") a.num_iters = next(); + else if (s == "--num-blocks") + a.num_blocks = next(); + else if (s == "--seed") + a.seed = next(); + else if (s == "--kernel-timing") + a.kernel_timing = true; + else if (s == "--dispatch-dtype" && i + 1 < argc) + a.dispatch_dtype = argv[++i]; + else if (s == "--combine-mode" && i + 1 < argc) + a.combine_mode = argv[++i]; } return a; } @@ -185,11 +200,47 @@ int main(int argc, char** argv) { Args args = parse_args(argc, argv); const int T = args.num_tokens, H = args.hidden, K = args.num_topk, E = args.num_experts; const int W = nRanks, warmup = args.num_warmup, iters = args.num_iters; + if (T <= 0 || E <= 0 || warmup < 0 || iters <= 0) { + if (rank == 0) fprintf(stderr, "tokens, experts, and iters must be positive; warmup must be non-negative\n"); + MPI_Abort(MPI_COMM_WORLD, 1); + } + if (H != 4096 && H != 7168 && H != 8192 && H != 9216) { + if (rank == 0) fprintf(stderr, "hidden must be one of 4096, 7168, 8192, 9216\n"); + MPI_Abort(MPI_COMM_WORLD, 1); + } + if (K <= 0 || K > 9) { + if (rank == 0) fprintf(stderr, "num_topk must be in [1, 9]\n"); + MPI_Abort(MPI_COMM_WORLD, 1); + } + if (K > E) { + if (rank == 0) fprintf(stderr, "num_topk must not exceed num_experts\n"); + MPI_Abort(MPI_COMM_WORLD, 1); + } if (E % W != 0) { if (rank == 0) fprintf(stderr, "num_experts (%d) must be divisible by world_size (%d)\n", E, W); MPI_Abort(MPI_COMM_WORLD, 1); } const int Elocal = E / W; + const auto dispatchDataType = args.dispatch_dtype == "fp8_e4m3" ? mscclpp::ep::low_latency::DispatchDataType::FP8_E4M3 + : mscclpp::ep::low_latency::DispatchDataType::BF16; + const auto combineMode = args.combine_mode == "direct_send" + ? mscclpp::ep::low_latency::CombineMode::DIRECT_SEND + : mscclpp::ep::low_latency::CombineMode::RANK_LOCAL_REDUCE; + if (args.dispatch_dtype != "bf16" && args.dispatch_dtype != "fp8_e4m3") { + if (rank == 0) fprintf(stderr, "unsupported --dispatch-dtype=%s\n", args.dispatch_dtype.c_str()); + MPI_Abort(MPI_COMM_WORLD, 1); + } + if (args.combine_mode != "rank_local_reduce" && args.combine_mode != "direct_send") { + if (rank == 0) fprintf(stderr, "unsupported --combine-mode=%s\n", args.combine_mode.c_str()); + MPI_Abort(MPI_COMM_WORLD, 1); + } + if (args.num_blocks < W + mscclpp::ep::low_latency::DispatchControlBlocks || + args.num_blocks > mscclpp::ep::low_latency::MaxDispatchBlocks) { + if (rank == 0) fprintf(stderr, "--num-blocks must be in [world_size + 2, 130]\n"); + MPI_Abort(MPI_COMM_WORLD, 1); + } + const bool fp8Dispatch = dispatchDataType == mscclpp::ep::low_latency::DispatchDataType::FP8_E4M3; + const char* dispatchLabel = fp8Dispatch ? "FP8_E4M3" : "BF16"; // --- Bootstrap mscclpp::Communicator (TcpBootstrap + MPI_Bcast of UniqueId). --- auto bootstrap = std::make_shared(rank, nRanks); @@ -199,7 +250,7 @@ int main(int argc, char** argv) { bootstrap->initialize(uid); mscclpp::Communicator comm(bootstrap); - const int64_t numRdmaBytes = static_cast(mscclpp::ep::getLowLatencyRdmaSizeHint(T, H, W, E)); + const int64_t numRdmaBytes = static_cast(mscclpp::ep::low_latency::getRdmaSizeHint(T, H, W, E, K)); mscclpp::ep::MoERuntime rt(comm, /*numNvlBytes=*/0, numRdmaBytes, mscclpp::ep::MoEMode::LOW_LATENCY); if (!rt.isAvailable()) { if (rank == 0) fprintf(stderr, "MoERuntime not available\n"); @@ -208,50 +259,81 @@ int main(int argc, char** argv) { if (rank == 0) { printf( "[cfg] algorithm=LOW_LATENCY num_ranks=%d tokens/rank=%d hidden=%d num_experts=%d " - "top_k=%d warmup=%d iters=%d num_rdma_bytes=%lld is_internode=%d\n", - W, T, H, E, K, warmup, iters, (long long)numRdmaBytes, (int)rt.isInternodeAvailable()); + "top_k=%d warmup=%d iters=%d dispatch_dtype=%s combine_mode=%s num_rdma_bytes=%lld is_internode=%d\n", + W, T, H, E, K, warmup, iters, args.dispatch_dtype.c_str(), args.combine_mode.c_str(), (long long)numRdmaBytes, + (int)rt.isInternodeAvailable()); fflush(stdout); } // --- Device buffers (hoisted out of the timed loop). --- const size_t slots = (size_t)W * T; // recv slots per local expert - __nv_bfloat16 *d_x = nullptr, *d_out = nullptr, *d_recv = nullptr; + using Bf16 = mscclpp::ep::low_latency::Bf16; + using Fp8E4M3 = mscclpp::ep::low_latency::Fp8E4M3; + Bf16 *d_x = nullptr, *d_out = nullptr, *d_expert_output = nullptr; + void* d_recv = nullptr; + float* d_scales = nullptr; int64_t *d_topk = nullptr, *d_layout = nullptr; float* d_weights = nullptr; int *d_srcinfo = nullptr, *d_count = nullptr; - CUDA_CHECK(cudaMalloc(&d_x, (size_t)T * H * sizeof(__nv_bfloat16))); - CUDA_CHECK(cudaMalloc(&d_out, (size_t)T * H * sizeof(__nv_bfloat16))); - CUDA_CHECK(cudaMalloc(&d_recv, (size_t)Elocal * slots * H * sizeof(__nv_bfloat16))); + CUDA_CHECK(cudaMalloc(&d_x, (size_t)T * H * sizeof(Bf16))); + CUDA_CHECK(cudaMalloc(&d_out, (size_t)T * H * sizeof(Bf16))); + const size_t recvBytes = (size_t)Elocal * slots * H * (fp8Dispatch ? sizeof(Fp8E4M3) : sizeof(Bf16)); + CUDA_CHECK(cudaMalloc(&d_recv, recvBytes)); + if (fp8Dispatch) { + CUDA_CHECK(cudaMalloc(&d_scales, (size_t)Elocal * slots * (H / 128) * sizeof(float))); + CUDA_CHECK(cudaMalloc(&d_expert_output, (size_t)Elocal * slots * H * sizeof(Bf16))); + CUDA_CHECK(cudaMemset(d_expert_output, 0, (size_t)Elocal * slots * H * sizeof(Bf16))); + } else { + d_expert_output = static_cast(d_recv); + } CUDA_CHECK(cudaMalloc(&d_topk, (size_t)T * K * sizeof(int64_t))); CUDA_CHECK(cudaMalloc(&d_weights, (size_t)T * K * sizeof(float))); CUDA_CHECK(cudaMalloc(&d_srcinfo, (size_t)Elocal * slots * sizeof(int))); CUDA_CHECK(cudaMalloc(&d_layout, (size_t)Elocal * W * sizeof(int64_t))); CUDA_CHECK(cudaMalloc(&d_count, (size_t)Elocal * sizeof(int))); - // Inputs (content is immaterial to timing; give every token K distinct experts). - CUDA_CHECK(cudaMemset(d_x, 0, (size_t)T * H * sizeof(__nv_bfloat16))); + // Inputs (content is immaterial to timing). Route each rank independently to + // random distinct experts so top-k selections span destination ranks. + CUDA_CHECK(cudaMemset(d_x, 0, (size_t)T * H * sizeof(Bf16))); std::vector h_topk((size_t)T * K); - std::vector h_weights((size_t)T * K, 1.0f); - for (int t = 0; t < T; ++t) - for (int j = 0; j < K; ++j) h_topk[(size_t)t * K + j] = ((int64_t)t * K + j) % E; + std::vector h_weights((size_t)T * K); + std::mt19937 rng(args.seed + rank); + std::uniform_int_distribution expertDist(0, E - 1); + std::uniform_real_distribution weightDist(0.5f, 1.5f); + for (int t = 0; t < T; ++t) { + for (int j = 0; j < K; ++j) { + int expert; + bool duplicate; + do { + expert = expertDist(rng); + duplicate = false; + for (int previous = 0; previous < j; ++previous) { + duplicate |= h_topk[(size_t)t * K + previous] == expert; + } + } while (duplicate); + h_topk[(size_t)t * K + j] = expert; + h_weights[(size_t)t * K + j] = weightDist(rng); + } + } CUDA_CHECK(cudaMemcpy(d_topk, h_topk.data(), h_topk.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); CUDA_CHECK(cudaMemcpy(d_weights, h_weights.data(), h_weights.size() * sizeof(float), cudaMemcpyHostToDevice)); const long long num_valid_selections = (long long)T * K; - const double disp_bytes = (double)num_valid_selections * H * 2.0; // BF16 - const double comb_bytes = disp_bytes; + const double dispatchBytesPerToken = fp8Dispatch ? H + (H / 128) * sizeof(float) : H * sizeof(Bf16); + const double disp_bytes = (double)num_valid_selections * dispatchBytesPerToken; + const double comb_bytes = (double)num_valid_selections * H * sizeof(Bf16); cudaStream_t stream; CUDA_CHECK(cudaStreamCreate(&stream)); auto dispatch = [&]() { - rt.dispatch(d_recv, /*outputScales=*/nullptr, d_srcinfo, d_layout, d_count, d_x, d_topk, T, H, K, - /*numMaxDispatchTokensPerRank=*/T, E, /*requiresQuantization=*/false, - mscclpp::ep::DispatchLayout::EXPERT_MAJOR, stream); + rt.dispatch(d_recv, d_scales, d_srcinfo, d_layout, d_count, d_x, d_topk, d_weights, T, H, K, + /*maxTokensPerRank=*/T, E, dispatchDataType, args.num_blocks, stream); }; auto combine = [&]() { - rt.combine(d_out, d_recv, /*inputScales=*/nullptr, d_topk, d_weights, d_srcinfo, d_layout, T, H, K, - /*numMaxDispatchTokensPerRank=*/T, E, /*requiresDequantization=*/false, stream); + rt.combine(d_out, d_expert_output, d_topk, d_weights, d_srcinfo, d_layout, T, H, K, + /*maxTokensPerRank=*/T, E, dispatchDataType, combineMode, + args.num_blocks - mscclpp::ep::low_latency::DispatchControlBlocks, stream); }; // --- Warmup (paired), then per-iter timed (paired), matching ep_bench. --- @@ -266,7 +348,15 @@ int main(int argc, char** argv) { KernelTimer ktimer; CUDA_CHECK(cudaDeviceSynchronize()); MPI_Barrier(MPI_COMM_WORLD); - int kt_rc = ktimer.start(); + int kt_rc = args.kernel_timing ? ktimer.start() : -1; + int localTimerStarted = kt_rc == CUPTI_SUCCESS; + int allTimersStarted = 0; + MPI_Allreduce(&localTimerStarted, &allTimersStarted, 1, MPI_INT, MPI_MIN, MPI_COMM_WORLD); + if (!allTimersStarted) { + if (localTimerStarted) ktimer.stop(); + kt_rc = -1; + } + MPI_Barrier(MPI_COMM_WORLD); std::vector ds(iters), de(iters), cs(iters), ce(iters); for (int i = 0; i < iters; ++i) { @@ -287,7 +377,7 @@ int main(int argc, char** argv) { MPI_Barrier(MPI_COMM_WORLD); } CUDA_CHECK(cudaDeviceSynchronize()); - if (kt_rc == CUPTI_SUCCESS) ktimer.stop(); + if (args.kernel_timing && kt_rc == CUPTI_SUCCESS) ktimer.stop(); // --- Collect per-iter host times (ms->us), trim first (warmup outlier). --- std::vector disp_us, comb_us, tot_us; @@ -318,10 +408,21 @@ int main(int argc, char** argv) { // Kernel-only (CUPTI). Per-rank mean, then cross-rank avg/min/max. double kd = (kt_rc == CUPTI_SUCCESS) ? ktimer.get_avg_us("dispatch") : 0.0; double kc = (kt_rc == CUPTI_SUCCESS) ? ktimer.get_avg_us("combine") : 0.0; + int localKernelOk = (kt_rc == CUPTI_SUCCESS) && (kd > 0.0) && (kc > 0.0); + int allKernelsOk = 0; + MPI_Allreduce(&localKernelOk, &allKernelsOk, 1, MPI_INT, MPI_MIN, MPI_COMM_WORLD); double gkda, gkdmn, gkdmx, gkca, gkcmn, gkcmx; - reduce3(kd, kd, kd, gkda, gkdmn, gkdmx); - reduce3(kc, kc, kc, gkca, gkcmn, gkcmx); - bool kernel_ok = (kt_rc == CUPTI_SUCCESS) && (kd > 0.0) && (kc > 0.0); + double kdMinInput = localKernelOk ? kd : 1e30; + double kcMinInput = localKernelOk ? kc : 1e30; + MPI_Reduce(&kd, &gkda, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD); + MPI_Reduce(&kdMinInput, &gkdmn, 1, MPI_DOUBLE, MPI_MIN, 0, MPI_COMM_WORLD); + MPI_Reduce(&kd, &gkdmx, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD); + MPI_Reduce(&kc, &gkca, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD); + MPI_Reduce(&kcMinInput, &gkcmn, 1, MPI_DOUBLE, MPI_MIN, 0, MPI_COMM_WORLD); + MPI_Reduce(&kc, &gkcmx, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD); + gkda /= W; + gkca /= W; + bool kernel_ok = allKernelsOk != 0; if (std::getenv("MSCCLPP_EP_KDEBUG") && rank == 0) { printf("[kdebug] kt_start rc=%d dispatch=%.1fus x%llu combine=%.1fus x%llu (kind=%s)\n", kt_rc, kd, @@ -332,7 +433,7 @@ int main(int argc, char** argv) { if (rank == 0) { printf("\n=== Summary (Low Latency, across %d ranks) ===\n", W); printf("\n--- Host-observed performance ---\n"); - printf("Dispatch (BF16): avg=%.2f us, min=%.2f us, max=%.2f us\n", gda, gdmn, gdmx); + printf("Dispatch (%s): avg=%.2f us, min=%.2f us, max=%.2f us\n", dispatchLabel, gda, gdmn, gdmx); printf(" throughput: avg=%.2f GB/s\n", (disp_bytes / 1e9) / (gda * 1e-6)); printf("Combine (BF16): avg=%.2f us, min=%.2f us, max=%.2f us\n", gca, gcmn, gcmx); printf(" throughput: avg=%.2f GB/s\n", (comb_bytes / 1e9) / (gca * 1e-6)); @@ -341,17 +442,16 @@ int main(int argc, char** argv) { printf("\n--- Kernel-only performance ---\n"); if (kernel_ok) { - printf("Dispatch: avg=%.2f us, min=%.2f us, max=%.2f us\n", gkda, gkdmn, gkdmx); - printf(" throughput: avg=%.2f GB/s\n", (disp_bytes / 1e9) / (gkda * 1e-6)); - printf("Combine: avg=%.2f us, min=%.2f us, max=%.2f us\n", gkca, gkcmn, gkcmx); - printf(" throughput: avg=%.2f GB/s\n", (comb_bytes / 1e9) / (gkca * 1e-6)); - printf("Total (D+C): %.2f us (kernel dispatch avg + combine avg)\n", gkda + gkca); + printf("Dispatch: min=%.2f us (representative) [avg=%.2f, max=%.2f us -- rank skew]\n", gkdmn, gkda, gkdmx); + printf(" throughput @min: %.2f GB/s\n", (disp_bytes / 1e9) / (gkdmn * 1e-6)); + printf("Combine: min=%.2f us (representative) [avg=%.2f, max=%.2f us -- rank skew]\n", gkcmn, gkca, gkcmx); + printf(" throughput @min: %.2f GB/s\n", (comb_bytes / 1e9) / (gkcmn * 1e-6)); } else { printf(" NOTE: CUPTI kernel timing unavailable (rc=%d) or captured 0 LL kernels.\n", kt_rc); } - printf("\nByte counts: dispatch=%.2f MB (BF16), combine=%.2f MB (BF16), selections=%lld\n", disp_bytes / 1e6, - comb_bytes / 1e6, num_valid_selections); + printf("\nByte counts: dispatch=%.2f MB (%s), combine=%.2f MB (BF16), selections=%lld\n", disp_bytes / 1e6, + dispatchLabel, comb_bytes / 1e6, num_valid_selections); fflush(stdout); } @@ -365,6 +465,8 @@ int main(int argc, char** argv) { cudaFree(d_x); cudaFree(d_out); cudaFree(d_recv); + if (fp8Dispatch) cudaFree(d_expert_output); + cudaFree(d_scales); cudaFree(d_topk); cudaFree(d_weights); cudaFree(d_srcinfo); diff --git a/test/python/ep/run_ep_bench.py b/test/python/ep/run_ep_bench.py index 5f94078d..8ca3d676 100644 --- a/test/python/ep/run_ep_bench.py +++ b/test/python/ep/run_ep_bench.py @@ -12,6 +12,7 @@ Backends (``--ep-lib``): * ``mscclpp`` -- this repo's :mod:`ep_bench_ll` (MoECommunicator LL) launched with ``torchrun``. +* ``mscclpp-cpp`` -- the pure-C++ ``MoERuntime`` benchmark launched with MPI. * ``nccl-ep`` -- NVIDIA NCCL-EP's ``contrib/nccl_ep/ep_bench`` binary launched with ``mpirun`` (HPCX). * ``both`` -- run mscclpp then nccl-ep and print them side by side. @@ -49,6 +50,7 @@ Print the commands without running them:: from __future__ import annotations import argparse +import glob import os import re import shlex @@ -57,11 +59,30 @@ import sys from dataclasses import dataclass, field from typing import Optional -CUDA_INC = "/usr/local/cuda/targets/sbsa-linux/include" -CUDA_LIB = "/usr/local/cuda/targets/sbsa-linux/lib" +CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda") _HERE = os.path.dirname(os.path.abspath(__file__)) +def _find_cupti_paths() -> tuple[str, str]: + target_dirs = sorted(glob.glob(os.path.join(CUDA_HOME, "targets", "*"))) + include_candidates = [os.path.join(path, "include") for path in target_dirs] + include_candidates.append(os.path.join(CUDA_HOME, "extras", "CUPTI", "include")) + library_candidates = [os.path.join(path, "lib") for path in target_dirs] + library_candidates.append(os.path.join(CUDA_HOME, "extras", "CUPTI", "lib64")) + + include_dir = next( + (path for path in include_candidates if os.path.isfile(os.path.join(path, "cupti.h"))), + "", + ) + library_dir = next( + (path for path in library_candidates if glob.glob(os.path.join(path, "libcupti.so*"))), + "", + ) + if not include_dir or not library_dir: + raise SystemExit(f"CUPTI was not found under CUDA_HOME={CUDA_HOME}") + return include_dir, library_dir + + def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser( description="Unified EP low-latency benchmark driver (mscclpp EP vs NCCL-EP)", @@ -84,11 +105,31 @@ def parse_args() -> argparse.Namespace: # Shared problem shape -- passed to whichever backend is selected. p.add_argument("-t", "--num-tokens", type=int, default=128, help="tokens per rank") - p.add_argument("-d", "--hidden", type=int, default=7168, help="hidden dimension") - p.add_argument("-k", "--num-topk", type=int, default=8, help="top-k experts per token") + p.add_argument( + "-d", + "--hidden", + type=int, + default=7168, + choices=(4096, 7168, 8192, 9216), + help="hidden dimension", + ) + p.add_argument("-k", "--num-topk", type=int, default=8, choices=range(1, 10), help="top-k experts per token") p.add_argument("-e", "--num-experts", type=int, default=256, help="global number of experts") p.add_argument("-w", "--num-warmup", type=int, default=10, help="warmup iterations") p.add_argument("-i", "--num-iters", type=int, default=50, help="timed iterations") + p.add_argument( + "--dispatch-dtype", + choices=("bf16", "fp8_e4m3"), + default="bf16", + help="MSCCL++ dispatch format; NCCL-EP runs keep their own configured format", + ) + p.add_argument( + "--combine-mode", + choices=("rank_local_reduce", "direct_send"), + default="rank_local_reduce", + help="MSCCL++ low-latency combine mode", + ) + p.add_argument("--num-blocks", type=int, default=130, help="MSCCL++ low-latency dispatch blocks") # Launch / fabric. p.add_argument("--nproc-per-node", type=int, default=4, help="GPUs (ranks) on this node") @@ -99,17 +140,18 @@ def parse_args() -> argparse.Namespace: "local node. Applies to the mpirun backends (nccl-ep, mscclpp-cpp); the Python " "mscclpp backend is single-node only (torchrun --standalone).", ) - p.add_argument("--iface", default="enP22p1s0f1", help="socket interface name (NCCL/GLOO/UCX)") - p.add_argument("--hca", default="mlx5_0,mlx5_1,mlx5_2,mlx5_3", help="mscclpp HCA devices") + p.add_argument("--iface", default="", help="optional socket interface name (NCCL/GLOO/UCX)") + p.add_argument("--hca", default="", help="optional comma-separated mscclpp HCA devices") # mscclpp backend. p.add_argument("--mscclpp-bench", default=os.path.join(_HERE, "ep_bench_ll.py"), help="path to ep_bench_ll.py") + p.add_argument("--python", default=sys.executable, help="Python interpreter used for the MSCCL++ benchmark") p.add_argument( "--conda-prefix", - default=os.path.join(os.path.expanduser("~"), "miniconda3"), - help="conda installation prefix for the mscclpp torch env", + default="", + help="optional conda installation prefix; used only when --conda-env is set", ) - p.add_argument("--conda-env", default="torch", help="conda env name with torch + mscclpp") + p.add_argument("--conda-env", default="", help="optional conda env name with torch + mscclpp") p.add_argument( "--cupti-inproc", action="store_true", help="mscclpp: also collect in-process CUPTI kernel-only timing" ) @@ -150,7 +192,7 @@ def parse_args() -> argparse.Namespace: # mscclpp-cpp backend (pure C++ MoERuntime binary). p.add_argument( "--mscclpp-cpp-bench", - default="/opt/microsoft/mrc/ep/mscclpp/test/python/ep/build/mscclpp_ep_bench", + default=os.path.join(_HERE, "build", "mscclpp_ep_bench"), help="path to the mscclpp_ep_bench C++ binary (built via test/python/ep/CMakeLists.txt)", ) @@ -162,10 +204,27 @@ def parse_args() -> argparse.Namespace: # and to fail fast on values that would break the launch (spaces, quotes, ...). if args.nodes and not re.fullmatch(r"[0-9A-Za-z._:-]+( [0-9A-Za-z._:-]+)*", args.nodes): raise SystemExit("--nodes must be space-separated hostnames/IPs") - if not re.fullmatch(r"[0-9A-Za-z._:-]+", args.iface): + if args.iface and not re.fullmatch(r"[0-9A-Za-z._:-]+", args.iface): raise SystemExit("--iface must be a valid network interface name") - if not re.fullmatch(r"[0-9A-Za-z._,-]+", args.hca): + if args.hca and not re.fullmatch(r"[0-9A-Za-z._,-]+", args.hca): raise SystemExit("--hca must be comma-separated HCA device names") + if args.conda_env and not args.conda_prefix: + raise SystemExit("--conda-prefix is required when --conda-env is set") + if args.dispatch_dtype != "bf16" and args.ep_lib in ("nccl-ep", "both", "all"): + raise SystemExit("FP8 unified comparison is unsupported because the NCCL-EP command is configured for BF16") + if args.num_tokens <= 0 or args.num_experts <= 0 or args.nproc_per_node <= 0: + raise SystemExit("tokens, experts, and nproc-per-node must be positive") + if args.num_topk > args.num_experts: + raise SystemExit("num-topk must not exceed num-experts") + if args.num_warmup < 0 or args.num_iters <= 0: + raise SystemExit("num-warmup must be non-negative and num-iters must be positive") + num_nodes = max(1, len(args.nodes.split())) + num_ranks = num_nodes * args.nproc_per_node + if args.num_experts % num_ranks != 0: + raise SystemExit("num-experts must be divisible by the total number of ranks") + if args.ep_lib in ("mscclpp", "mscclpp-cpp", "both", "all"): + if not num_ranks + 2 <= args.num_blocks <= 130: + raise SystemExit("num-blocks must be in [total ranks + 2, 130]") return args @@ -195,8 +254,8 @@ class LLResult: _HOST_RE = { - "dispatch": re.compile(r"^Dispatch \(BF16\):\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us"), - "combine": re.compile(r"^Combine \(BF16\):\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us"), + "dispatch": re.compile(r"^Dispatch \([^)]+\):\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us"), + "combine": re.compile(r"^Combine \([^)]+\):\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us"), "total": re.compile(r"^Total \(D\+C\):\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us"), } _RANKS_RE = re.compile(r"=== Summary \(Low Latency, across (\d+) ranks\) ===") @@ -205,6 +264,7 @@ _RANKS_RE = re.compile(r"=== Summary \(Low Latency, across (\d+) ranks\) ===") # ep_bench CUPTI: ``Dispatch: avg=A us, min=M us, max=X us`` _KDISP_REP_RE = re.compile(r"^Dispatch:\s+min=([\d.]+)\s*us \(representative\)\s*\[avg=([\d.]+),\s*max=([\d.]+)") _KDISP_AMM_RE = re.compile(r"^Dispatch:\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us") +_KCOMB_REP_RE = re.compile(r"^Combine:\s+min=([\d.]+)\s*us \(representative\)\s*\[avg=([\d.]+),\s*max=([\d.]+)") # Kernel-only Combine line (both backends): ``Combine: avg=A us, min=M us, max=X us`` (no ``(BF16)``). _KCOMB_RE = re.compile(r"^Combine:\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us") @@ -234,6 +294,10 @@ def parse_ll_summary(text: str, ep_lib: str) -> LLResult: res.kdispatch = Phase(avg=float(m.group(1)), min=float(m.group(2)), max=float(m.group(3))) continue if res.kcombine is None and res.kdispatch is not None: + m = _KCOMB_REP_RE.match(line) + if m: + res.kcombine = Phase(avg=float(m.group(2)), min=float(m.group(1)), max=float(m.group(3))) + continue m = _KCOMB_RE.match(line) if m: res.kcombine = Phase(avg=float(m.group(1)), min=float(m.group(2)), max=float(m.group(3))) @@ -245,29 +309,46 @@ def parse_ll_summary(text: str, ep_lib: str) -> LLResult: # Backend command construction. # ---------------------------------------------------------------------------- def build_mscclpp_cmd(args: argparse.Namespace) -> str: - env = ( - f"MSCCLPP_EP_LOCAL_WORLD_SIZE={args.nproc_per_node} " - f"NCCL_SOCKET_IFNAME={args.iface} GLOO_SOCKET_IFNAME={args.iface} MSCCLPP_SOCKET_IFNAME={args.iface} " - f"MSCCLPP_HCA_DEVICES={args.hca} NCCL_IB_DISABLE=1 NCCL_MNNVL_ENABLE=0 MSCCLPP_EP_FABRIC_IPC=1" - ) + env_vars = { + "MSCCLPP_EP_LOCAL_WORLD_SIZE": str(args.nproc_per_node), + "NCCL_IB_DISABLE": "1", + "NCCL_MNNVL_ENABLE": "0", + } + if args.iface: + env_vars.update( + { + "NCCL_SOCKET_IFNAME": args.iface, + "GLOO_SOCKET_IFNAME": args.iface, + "MSCCLPP_SOCKET_IFNAME": args.iface, + } + ) + if args.hca: + env_vars["MSCCLPP_HCA_DEVICES"] = args.hca + bench = args.mscclpp_bench bench_flags = ( f"-a ll -t {args.num_tokens} -d {args.hidden} -k {args.num_topk} " - f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters}" + f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters} " + f"--dispatch-dtype {args.dispatch_dtype} --combine-mode {args.combine_mode} --num-blocks {args.num_blocks}" ) cupti_build = "" + extra_exports = "" if args.cupti_inproc or args.kernel_only: # In-process CUPTI kernel-only timing (near-zero perturbation, matches - # ep_bench's KernelTimer). Builds the collector next to the bench if missing. + # ep_bench's KernelTimer). Build it under the benchmark build directory. bench_flags += " --cupti-inproc" - env += f" LD_LIBRARY_PATH={CUDA_LIB}:$LD_LIBRARY_PATH" - so = os.path.join(os.path.dirname(bench), "libcupti_kernel_timer.so") + cupti_include, cupti_lib = _find_cupti_paths() + build_dir = os.path.join(os.path.dirname(bench), "build") + so = os.path.join(build_dir, "libcupti_kernel_timer.so") src = os.path.join(os.path.dirname(bench), "cupti_kernel_timer.cpp") + env_vars["MSCCLPP_EP_CUPTI_TIMER_LIB"] = so cupti_build = ( - f"if [ ! -f {shlex.quote(so)} ]; then " + f"mkdir -p {shlex.quote(build_dir)} && " + f"if [ ! -f {shlex.quote(so)} ] || [ {shlex.quote(src)} -nt {shlex.quote(so)} ]; then " f"g++ -O2 -fPIC -shared {shlex.quote(src)} -o {shlex.quote(so)} " - f"-I{CUDA_INC} -L{CUDA_LIB} -lcupti; fi && " + f"-I{shlex.quote(cupti_include)} -L{shlex.quote(cupti_lib)} -lcupti; fi && " ) + extra_exports = f"export LD_LIBRARY_PATH={shlex.quote(cupti_lib)}:${{LD_LIBRARY_PATH:-}} && " elif args.torch_profiler: # Opt-in torch.profiler kernel pass (perturbs the LL recv-spin; the # in-process CUPTI path is preferred for kernel numbers). @@ -276,12 +357,22 @@ def build_mscclpp_cmd(args: argparse.Namespace) -> str: # Default: clean host-observed only (skip the torch.profiler pass, which # is slow and inflates the LL dispatch recv-spin). bench_flags += " --no-kernel-timing" + + activation = "" + python = shlex.quote(args.python) + if args.conda_env: + activation = ( + f"source {shlex.quote(args.conda_prefix)}/etc/profile.d/conda.sh && " + f"conda activate {shlex.quote(args.conda_env)} && " + ) + python = "python" + exports = " ".join(f"{name}={shlex.quote(value)}" for name, value in env_vars.items()) return ( - f"source {shlex.quote(args.conda_prefix)}/etc/profile.d/conda.sh && " - f"conda activate {shlex.quote(args.conda_env)} && unset PYTHONPATH && " + f"{activation}" f"{cupti_build}" - f"export {env} && " - f"torchrun --standalone --nnodes=1 --nproc_per_node={args.nproc_per_node} " + f"export {exports} && " + f"{extra_exports}" + f"{python} -m torch.distributed.run --standalone --nnodes=1 --nproc_per_node={args.nproc_per_node} " f"{shlex.quote(bench)} {bench_flags}" ) @@ -307,9 +398,11 @@ def _mpi_launch(args, np_total): hostfile = ( f"--hostfile {hf} " f'-mca plm_rsh_args "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null" ' ) + iface_arg = f"-mca btl_tcp_if_include {shlex.quote(args.iface)} " if args.iface else "" + root_arg = "--allow-run-as-root " if os.geteuid() == 0 else "" return setup, ( - f"mpirun -np {np_total} {hostfile}--map-by ppr:{args.nproc_per_node}:node --bind-to none " - f"-mca pml ob1 -mca btl self,vader,tcp -mca btl_tcp_if_include {args.iface} " + f"mpirun {root_arg}-np {np_total} {hostfile}--map-by ppr:{args.nproc_per_node}:node --bind-to none " + f"-mca pml ob1 -mca btl self,vader,tcp {iface_arg}" f"-mca coll_hcoll_enable 0 -mca coll_ucc_enable 0 " ) @@ -323,8 +416,6 @@ def build_nccl_ep_cmd(args: argparse.Namespace) -> str: bench_dir = os.path.dirname(os.path.abspath(args.nccl_ep_bench)) nccl_lib = os.path.join(os.path.dirname(os.path.dirname(bench_dir)), "lib") hpcx = args.hpcx or _autodetect_hpcx() - if not hpcx: - raise SystemExit("nccl-ep: no HPCX found under /opt; pass --hpcx") nodes = args.nodes.split() nnodes = max(1, len(nodes)) np_total = nnodes * args.nproc_per_node @@ -334,47 +425,56 @@ def build_nccl_ep_cmd(args: argparse.Namespace) -> str: f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters}" ) setup, mpi_prefix = _mpi_launch(args, np_total) + opal = f"-x OPAL_PREFIX={shlex.quote(hpcx)}/ompi " if hpcx else "" + iface_env = ( + f"-x UCX_NET_DEVICES={shlex.quote(args.iface)} " f"-x NCCL_SOCKET_IFNAME={shlex.quote(args.iface)} " + if args.iface + else "" + ) mpi = ( f"{mpi_prefix}" - f"-x LD_LIBRARY_PATH -x PATH -x CUDA_HOME=/usr/local/cuda -x OPAL_PREFIX={shlex.quote(hpcx)}/ompi " - f"-x UCX_NET_DEVICES={args.iface} -x UCX_TLS=tcp,sm,self,cuda_copy -x UCX_HANDLE_ERRORS=none " - f"-x NCCL_SOCKET_IFNAME={args.iface} -x NCCL_NET_PLUGIN=none " + f"-x LD_LIBRARY_PATH -x PATH -x CUDA_HOME={shlex.quote(CUDA_HOME)} {opal}" + f"{iface_env}-x UCX_TLS=tcp,sm,self,cuda_copy -x UCX_HANDLE_ERRORS=none " + f"-x NCCL_NET_PLUGIN=none " f"-x NCCL_IB_DISABLE=1 -x NCCL_MNNVL_ENABLE={mnnvl} " f"{shlex.quote(args.nccl_ep_bench)} {bench_flags}" ) - return ( - f"source {shlex.quote(hpcx)}/hpcx-init.sh && hpcx_load && " - f"export LD_LIBRARY_PATH={shlex.quote(nccl_lib)}:$LD_LIBRARY_PATH && " - f"{setup}{mpi}" - ) + activation = f"source {shlex.quote(hpcx)}/hpcx-init.sh && hpcx_load && " if hpcx else "" + return f"{activation}" f"export LD_LIBRARY_PATH={shlex.quote(nccl_lib)}:$LD_LIBRARY_PATH && " f"{setup}{mpi}" def build_mscclpp_cpp_cmd(args: argparse.Namespace) -> str: """Pure-C++ mscclpp_ep_bench (MoERuntime), launched with mpirun -- no Python.""" hpcx = args.hpcx or _autodetect_hpcx() - if not hpcx: - raise SystemExit("mscclpp-cpp: no HPCX found under /opt; pass --hpcx") nodes = args.nodes.split() nnodes = max(1, len(nodes)) np_total = nnodes * args.nproc_per_node bench_flags = ( f"-a ll -t {args.num_tokens} -d {args.hidden} -k {args.num_topk} " - f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters}" + f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters} " + f"--dispatch-dtype {args.dispatch_dtype} --combine-mode {args.combine_mode} --num-blocks {args.num_blocks}" ) + if args.kernel_only or args.cupti_inproc: + bench_flags += " --kernel-timing" setup, mpi_prefix = _mpi_launch(args, np_total) + env_exports = ( + f"-x MSCCLPP_EP_LOCAL_WORLD_SIZE={args.nproc_per_node} " f"-x NCCL_IB_DISABLE=1 -x NCCL_MNNVL_ENABLE=0 " + ) + if args.hca: + env_exports += f"-x MSCCLPP_HCA_DEVICES={shlex.quote(args.hca)} " + if args.iface: + env_exports += ( + f"-x NCCL_SOCKET_IFNAME={shlex.quote(args.iface)} " f"-x MSCCLPP_SOCKET_IFNAME={shlex.quote(args.iface)} " + ) mpi = ( f"{mpi_prefix}" f"-x LD_LIBRARY_PATH -x PATH " - f"-x MSCCLPP_EP_LOCAL_WORLD_SIZE={args.nproc_per_node} -x MSCCLPP_HCA_DEVICES={args.hca} " - f"-x NCCL_IB_DISABLE=1 -x NCCL_MNNVL_ENABLE=0 -x MSCCLPP_EP_FABRIC_IPC=1 " - f"-x NCCL_SOCKET_IFNAME={args.iface} -x MSCCLPP_SOCKET_IFNAME={args.iface} " + f"{env_exports}" f"{shlex.quote(args.mscclpp_cpp_bench)} {bench_flags}" ) - return ( - f"source {shlex.quote(hpcx)}/hpcx-init.sh && hpcx_load && " - f"export LD_LIBRARY_PATH={CUDA_LIB}:$LD_LIBRARY_PATH && " - f"{setup}{mpi}" - ) + _, cupti_lib = _find_cupti_paths() + activation = f"source {shlex.quote(hpcx)}/hpcx-init.sh && hpcx_load && " if hpcx else "" + return f"{activation}" f"export LD_LIBRARY_PATH={shlex.quote(cupti_lib)}:$LD_LIBRARY_PATH && " f"{setup}{mpi}" # ---------------------------------------------------------------------------- @@ -421,33 +521,24 @@ def print_unified(results: list, kernel_only: bool = False) -> None: row("Host Combine max", lambda r: r.combine.max) row("Host D+C avg", lambda r: r.total.avg) if has_kernel: - # Kernel-only dispatch/combine, full avg/min/max for an apples-to-apples view. - # NOTE: mscclpp's collector (KIND_KERNEL) serializes kernels, inflating the - # cross-rank dispatch avg/max via recv-spin skew; min is the robust floor. - row("Kernel Dispatch avg", lambda r: r.kdispatch.avg) - row("Kernel Dispatch min", lambda r: r.kdispatch.min) - row("Kernel Dispatch max", lambda r: r.kdispatch.max) - row("Kernel Combine avg", lambda r: r.kcombine.avg) - row("Kernel Combine min", lambda r: r.kcombine.min) - row("Kernel Combine max", lambda r: r.kcombine.max) - row("Kernel D+C (avg)", lambda r: r.kdispatch.avg + r.kcombine.avg) - row("Kernel D+C (min)", lambda r: r.kdispatch.min + r.kcombine.min) + if kernel_only: + row("Kernel Dispatch repr", lambda r: r.kdispatch.min) + row("Kernel Combine repr", lambda r: r.kcombine.min) + else: + row("Kernel Dispatch avg", lambda r: r.kdispatch.avg) + row("Kernel Dispatch min", lambda r: r.kdispatch.min) + row("Kernel Dispatch max", lambda r: r.kdispatch.max) + row("Kernel Combine avg", lambda r: r.kcombine.avg) + row("Kernel Combine min", lambda r: r.kcombine.min) + row("Kernel Combine max", lambda r: r.kcombine.max) elif kernel_only: print( " NOTE: kernel-only requested but kernel timing missing for a backend " "(mscclpp needs --cupti-inproc / libcupti; nccl-ep needs CUPTI-enabled ep_bench)." ) - if len(results) == 2: + if len(results) == 2 and not kernel_only: a, b = results - if kernel_only and has_kernel: - ka_avg, kb_avg = a.kdispatch.avg + a.kcombine.avg, b.kdispatch.avg + b.kcombine.avg - ka_min, kb_min = a.kdispatch.min + a.kcombine.min, b.kdispatch.min + b.kcombine.min - if kb_avg: - print( - f"\nKernel D+C ratio {a.ep_lib}/{b.ep_lib}: avg={ka_avg / kb_avg:.2f}x, " - f"min={ka_min / kb_min:.2f}x" - ) - elif a.total.avg == a.total.avg and b.total.avg == b.total.avg and b.total.avg: + if a.total.avg == a.total.avg and b.total.avg == b.total.avg and b.total.avg: print(f"\nHost D+C ratio {a.ep_lib}/{b.ep_lib} = {a.total.avg / b.total.avg:.2f}x")