microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-15 11:44:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update EP benchmarks for current low-latency API

Browse Source 
Refresh the Python and C++ benchmark paths for BF16 and FP8 dispatch, current MoERuntime signatures, active kernel sources, portable CUPTI discovery, realistic routing, and safe unified reporting. Remove the merged change to the inactive legacy implementation.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: efbacae6-f679-430b-bc16-b45ae162fc76
This commit is contained in:
Binyang Li
2026-07-13 21:45:58 +00:00
parent 97e8c5b5be
commit ad0b63d21b
6 changed files with 460 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files

122
test/python/ep/CMakeLists.txt
View File

@@ -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
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
$<$<COMPILE_LANGUAGE:CUDA>:--expt-extended-lambda>
)
$<$<COMPILE_LANGUAGE:CUDA>:--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)

26
test/python/ep/cupti_kernel_timer.cpp
View File

@@ -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<std::mutex> 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<CUpti_ActivityKernel10*>(record);
if (record->kind == CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL) {
auto* k = reinterpret_cast<CUpti_ActivityKernel9*>(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<std::mutex> lock(g_mutex);
@@ -84,14 +74,14 @@ int kt_start() {
}
CUptiResult r = cuptiActivityRegisterCallbacks(bufferRequested, bufferCompleted);
if (r != CUPTI_SUCCESS) return static_cast<int>(r);
r = cuptiActivityEnable(CUPTI_ACTIVITY_KIND_KERNEL);
r = cuptiActivityEnable(CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL);
return static_cast<int>(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<int>(r);
}

123
test/python/ep/ep_bench_ll.py
View File

@@ -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}"
)

182
test/python/ep/mscclpp_ep_bench.cu
View File

@@ -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 <cuda.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cupti.h>
#include <mpi.h>
@@ -28,12 +26,14 @@
#include <cstring>
#include <map>
#include <mscclpp/core.hpp>
#include <mscclpp/gpu_data_types.hpp>
#include <random>
#include <string>
#include <vector>
#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<CUpti_ActivityKernel10*>(record);
auto* k = reinterpret_cast<CUpti_ActivityKernel9*>(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<mscclpp::TcpBootstrap>(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<int64_t>(mscclpp::ep::getLowLatencyRdmaSizeHint(T, H, W, E));
const int64_t numRdmaBytes = static_cast<int64_t>(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<Bf16*>(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<int64_t> h_topk((size_t)T * K);
std::vector<float> 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<float> h_weights((size_t)T * K);
std::mt19937 rng(args.seed + rank);
std::uniform_int_distribution<int> expertDist(0, E - 1);
std::uniform_real_distribution<float> 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<cudaEvent_t> 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<double> 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);

237
test/python/ep/run_ep_bench.py
View File

@@ -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")