diff --git a/src/ext/ep/legacy/low_latency.cu b/src/ext/ep/legacy/low_latency.cu index 67296291..c2e5967e 100644 --- a/src/ext/ep/legacy/low_latency.cu +++ b/src/ext/ep/legacy/low_latency.cu @@ -862,7 +862,28 @@ void combine(void* output, const void* input, const float* inputScales, const in const auto numWarps = kNumWarpGroups * kNumWarpsPerGroup; const auto numSmsBase = cell_div(numExperts, kNumWarpGroups); - const auto numSms = numSmsBase; + // 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; 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 new file mode 100644 index 00000000..ddb63219 --- /dev/null +++ b/test/python/ep/CMakeLists.txt @@ -0,0 +1,124 @@ +# 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. +# +# 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) + +cmake_minimum_required(VERSION 3.25) +project(mscclpp_ep_bench LANGUAGES CXX CUDA) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 17) +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 +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)") + +# 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) + endif() +endif() +if(NOT MSCCLPP_INSTALL_DIR) + message(FATAL_ERROR "MSCCLPP_INSTALL_DIR not found; pass -DMSCCLPP_INSTALL_DIR=<...>/site-packages/mscclpp") +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 + REQUIRED) +find_library(CUPTI_LIBRARY cupti + HINTS "${CUDAToolkit_LIBRARY_DIR}" /usr/local/cuda/targets/sbsa-linux/lib + /usr/local/cuda/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) + +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} +) + +target_compile_definitions(mscclpp_ep_bench PRIVATE + MSCCLPP_USE_CUDA + EP_DISPATCH_NCCLEP + NUM_MAX_NVL_PEERS=${MSCCLPP_EP_NUM_MAX_NVL_PEERS} +) + +target_compile_options(mscclpp_ep_bench PRIVATE + $<$:--expt-relaxed-constexpr> + $<$:--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" +) + +target_link_libraries(mscclpp_ep_bench PRIVATE + ${MSCCLPP_LIBRARY} + ${CUPTI_LIBRARY} + MPI::MPI_CXX + CUDA::cudart + CUDA::cuda_driver +) diff --git a/test/python/ep/cupti_kernel_timer.cpp b/test/python/ep/cupti_kernel_timer.cpp new file mode 100644 index 00000000..82054704 --- /dev/null +++ b/test/python/ep/cupti_kernel_timer.cpp @@ -0,0 +1,122 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// In-process CUPTI kernel timer for the mscclpp LL benchmark. +// name, and exposes an extern "C" ABI so it can be driven from Python via ctypes +// (no cuda-python CUPTI bindings exist in this env, but libcupti.so is loadable). +// +// This is near-zero host perturbation (out-of-band buffer callbacks), unlike +// torch.profiler's in-process tracing which serialized the LL dispatch recv-spin +// and inflated one rank's device time into the millisecond range. It matches +// 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". +// +// Build (host-only C++, links libcupti): +// g++ -O2 -fPIC -shared cupti_kernel_timer.cpp -o libcupti_kernel_timer.so \ +// -I/targets/sbsa-linux/include -L/targets/sbsa-linux/lib -lcupti +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +struct KernelStat { + uint64_t total_ns = 0; + uint64_t count = 0; +}; + +std::map g_stats; +std::mutex g_mutex; + +constexpr size_t kBufSize = 8 * 1024 * 1024; // 8 MB, matches ep_bench + +void CUPTIAPI bufferRequested(uint8_t** buffer, size_t* size, size_t* maxNumRecords) { + // 8-byte aligned; aligned_alloc requires size to be a multiple of alignment. + *buffer = static_cast(aligned_alloc(8, kBufSize)); + *size = kBufSize; + *maxNumRecords = 0; +} + +void CUPTIAPI bufferCompleted(CUcontext, uint32_t, uint8_t* buffer, size_t, size_t validSize) { + 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 (k->name) { + auto& s = g_stats[k->name]; + s.total_ns += (k->end - k->start); + s.count += 1; + } + } + } + free(buffer); +} + +} // namespace + +extern "C" { + +// Clear stats, register the buffer callbacks, and enable concurrent-kernel +// activity recording. Call AFTER warmup (like ep_bench's KernelTimer::start()). +int kt_start() { + { + std::lock_guard lock(g_mutex); + g_stats.clear(); + } + CUptiResult r = cuptiActivityRegisterCallbacks(bufferRequested, bufferCompleted); + if (r != CUPTI_SUCCESS) return static_cast(r); + r = cuptiActivityEnable(CUPTI_ACTIVITY_KIND_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); + return static_cast(r); +} + +// Average GPU execution time (microseconds) over every recorded kernel whose +// mangled name contains `substr`. Returns 0 if none matched. +double kt_get_avg_us(const char* substr) { + std::lock_guard lock(g_mutex); + uint64_t total_ns = 0, count = 0; + for (const auto& kv : g_stats) { + if (kv.first.find(substr) != std::string::npos) { + total_ns += kv.second.total_ns; + count += kv.second.count; + } + } + return count ? static_cast(total_ns) / static_cast(count) / 1000.0 : 0.0; +} + +// Number of recorded kernel instances whose name contains `substr`. +long kt_get_count(const char* substr) { + std::lock_guard lock(g_mutex); + uint64_t count = 0; + for (const auto& kv : g_stats) { + if (kv.first.find(substr) != std::string::npos) count += kv.second.count; + } + return static_cast(count); +} + +} // extern "C" diff --git a/test/python/ep/ep_bench_ll.py b/test/python/ep/ep_bench_ll.py new file mode 100644 index 00000000..bd182bc4 --- /dev/null +++ b/test/python/ep/ep_bench_ll.py @@ -0,0 +1,610 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Unified low-latency EP benchmark for MSCCL++ EP — an apples-to-apples port of +NCCL-EP's ``contrib/nccl_ep/ep_bench.cu`` low-latency (LL) flow, with the NCCL-EP +API (``ncclEpDispatch`` / ``ncclEpCombine``) replaced by the MSCCL++ EP high-level +``MoECommunicator.dispatch`` / ``MoECommunicator.combine`` (feature/ep API). + +Why this exists +--------------- +``ep_bench`` is the reference NCCL-EP micro-benchmark. To compare MSCCL++ EP +against it fairly we must measure *the same thing the same way*. This script is a +line-for-line reimplementation of ``ep_bench``'s LL measurement methodology, only +swapping the collective API underneath: + +* **Paired** dispatch→sync→combine→sync→barrier per iteration (``runPairedBenchmark``). +* **Per-iteration CUDA events** recorded on the stream *around each kernel launch*; + the ``cudaStreamSynchronize`` and ``MPI_Barrier`` (here ``dist.barrier``) happen + **outside** the timed region, exactly as in ``ep_bench``. +* **Skip the first timed iteration** (warmup outlier) — matches ``ep_bench``'s + ``calc_stats`` which trims ``times[0]`` when ``num_iters > 1``. +* **Byte accounting** identical to ``calculateLowLatencyBytes``: + ``bytes = num_valid_selections * hidden * 2`` (BF16) for *both* dispatch and + combine, where ``num_valid_selections = count(topk_idx >= 0)``. +* **Cross-rank reduction** identical to ``printLowLatencyResults``: latency + ``avg = mean``, ``min = MIN``, ``max = MAX``; per-rank throughput min/max are + tagged with the owning rank (``MPI_MINLOC`` / ``MPI_MAXLOC`` analog). +* **Output** mirrors ``ep_bench``'s ``=== Summary (Low Latency, across N ranks) ===`` + block so the two runs can be diffed directly. + +CLI mirrors ``ep_bench``'s LL-relevant flags (long + short): + -t/--num-tokens tokens per rank (ep_bench LL default 128) + -d/--hidden hidden dim (7168) + -k/--num-topk top-k experts per token (8) + -e/--num-experts global experts (256) + -w/--num-warmup warmup iterations (10) + -i/--num-iters timed iterations (50) + +Fidelity note +------------- +``ep_bench`` is C++/MPI; MSCCL++ EP's LL API is Python/torch, so this harness is +Python. The *measurement* is identical: both bracket the same dispatch/combine +kernels with CUDA events and report GPU-side host-observed time. The only +difference is host-side launch latency, which sits *outside* the recorded events +for the async kernels and is the same definitional gap ``ep_bench`` has (larger +in Python, but not counted in the kernel elapsed time). For a pure kernel number, +run under ``nsys``/CUPTI as with ``ep_bench``'s ``--- Kernel-only ---`` section. + +Launch +------ +Manual per-rank env (DSM hostnames break torchrun rendezvous on these nodes): + RANK=.. LOCAL_RANK=.. WORLD_SIZE=.. MASTER_ADDR=.. MASTER_PORT=.. \ + python ep_bench_ll.py -t 128 -d 7168 -k 8 -e 256 -w 10 -i 50 +Single node (4/8 GPU): + torchrun --standalone --nproc_per_node=4 ep_bench_ll.py -e 128 +""" + +from __future__ import annotations + +import argparse +import os +import random + +# Quiet ProcessGroupNCCL's heartbeat monitor before importing torch.distributed +# (same rationale as test_low_latency_multirank.py). +os.environ.setdefault("TORCH_NCCL_ENABLE_MONITORING", "0") + +import torch +import torch.distributed as dist + + +# ---------------------------------------------------------------------------- +# CLI — mirrors ep_bench.cu's getopt flags for the LL path. +# ---------------------------------------------------------------------------- +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description="MSCCL++ EP low-latency benchmark (ep_bench parity)", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + # Env fallbacks keep the existing MSCCLPP_EP_BENCH_* launchers working. + p.add_argument( + "-a", + "--algorithm", + default="ll", + choices=["ll", "low-latency"], + help="algorithm mode (only LL is implemented here)", + ) + p.add_argument( + "-t", + "--num-tokens", + type=int, + default=int(os.environ.get("MSCCLPP_EP_BENCH_TOKENS", "128")), + help="tokens per rank (ep_bench LL max_tokens_per_rank)", + ) + p.add_argument( + "-d", + "--hidden", + type=int, + default=int(os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "7168")), + help="hidden dimension", + ) + p.add_argument( + "-k", + "--num-topk", + type=int, + default=int(os.environ.get("MSCCLPP_EP_BENCH_TOPK", "8")), + help="top-k experts per token", + ) + p.add_argument( + "-e", + "--num-experts", + type=int, + default=int(os.environ.get("MSCCLPP_EP_BENCH_EXPERTS", "256")), + help="global number of experts", + ) + p.add_argument( + "-w", + "--num-warmup", + type=int, + default=int(os.environ.get("MSCCLPP_EP_BENCH_WARMUP", "10")), + help="warmup iterations", + ) + p.add_argument( + "-i", + "--num-iters", + type=int, + default=int(os.environ.get("MSCCLPP_EP_BENCH_ITERS", "50")), + help="timed iterations", + ) + p.add_argument( + "--no-kernel-timing", + dest="kernel_timing", + action="store_false", + help="disable the CUPTI/torch.profiler kernel-only measurement pass " + "(on by default, mirrors ep_bench's CUPTI KernelTimer)", + ) + p.add_argument( + "--cupti-region", + action="store_true", + help="bracket ONLY the timed loop with cudaProfilerStart/Stop (for nsys " + "--capture-range=cudaProfilerApi) so an external CUPTI collector times " + "exactly the post-warmup dispatch/combine kernels, like ep_bench's " + "KernelTimer.start()-after-warmup. Skips the in-process torch.profiler " + "pass; kernel numbers come from nsys.", + ) + p.add_argument( + "--cupti-inproc", + action="store_true", + 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.", + ) + p.add_argument("--seed", type=int, default=0xB3C4, help="per-rank RNG seed base") + return p.parse_args() + + +def init_dist(): + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://{os.environ.get('MASTER_ADDR', '127.0.0.1')}:{os.environ.get('MASTER_PORT', '29500')}", + world_size=world_size, + rank=rank, + ) + return rank, world_size, local_rank, dist.new_group(list(range(world_size))) + + +def _reduce_scalar(value: float, op, group) -> float: + t = torch.tensor([value], dtype=torch.float64, device="cuda") + dist.all_reduce(t, op=op, group=group) + return t.item() + + +def _gather_scalars(value: float, num_ranks: int, group) -> list: + t = torch.tensor([value], dtype=torch.float64, device="cuda") + out = [torch.zeros_like(t) for _ in range(num_ranks)] + dist.all_gather(out, t, group=group) + return [float(x.item()) for x in out] + + +def _profile_paired_kernels(dispatch_fn, combine_fn, iters: int, stream, group, rank: int): + """Kernel-only dispatch/combine device time (us/iter) via torch.profiler. + + Mirrors ep_bench's CUPTI ``KernelTimer``: it profiles the SAME paired + ``dispatch -> sync -> combine -> sync -> barrier`` loop used for the + host-observed measurement. Profiling the *paired* loop (rather than isolated + dispatch-only / combine-only loops) is essential: the LL dispatch kernel + ends with a cross-rank receive spin-wait, and without the per-iter barrier + the ranks drift out of lockstep so that spin balloons to milliseconds on the + laggards. The barrier keeps every rank aligned at each iteration boundary, so + the recv-wait stays bounded -- exactly why ep_bench times the paired loop. + + Kernels are bucketed by name substring ``dispatch`` / ``combine`` (the mscclpp + LL kernels demangle to ``mscclpp::ep::low_latency::dispatch<...>`` / + ``::combine<...>``), matching ep_bench's ``get_avg_us("dispatch"/"combine")``. + All other device activity (the pacing barrier's NCCL kernel, memcpy/memset) + is ignored. + """ + from torch.profiler import profile, ProfilerActivity + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + dout = dispatch_fn() + stream.synchronize() + combine_fn(dout) + stream.synchronize() + dist.barrier(group=group) + torch.cuda.synchronize() + + disp_us = 0.0 + comb_us = 0.0 + dbg = [] + for e in prof.key_averages(): + dev_us = getattr(e, "self_device_time_total", None) + if dev_us is None: + dev_us = getattr(e, "self_cuda_time_total", 0.0) + if not dev_us or dev_us <= 0: + continue + low = str(e.key).lower() + if "memcpy" in low or "memset" in low: + continue # CUPTI KernelTimer counts KERNEL activities only + if "dispatch" in low: + disp_us += dev_us + elif "combine" in low: + comb_us += dev_us + dbg.append((dev_us, str(e.key))) + + if os.environ.get("MSCCLPP_EP_KDEBUG", "0") == "1" and rank == 0: + dbg.sort(reverse=True) + print(f"[kdebug] top device activities (self device us/iter over {iters} iters):", flush=True) + for us, name in dbg[:10]: + print(f" {us / iters:8.2f} us/iter {name[:90]}", flush=True) + + return disp_us / iters, comb_us / iters + + +class _InProcCupti: + """In-process CUPTI kernel timer, a faithful analog of ep_bench's KernelTimer. + + Loads ``libcupti_kernel_timer.so`` (built from cupti_kernel_timer.cpp, sitting + next to this file) via ctypes and drives the CUPTI Activity API directly: + ``start()`` after warmup, ``stop()`` after the timed loop, then + ``avg_us("dispatch"/"combine")`` buckets recorded kernels by mangled-name + substring -- exactly ep_bench's methodology, with near-zero host perturbation + (out-of-band buffer callbacks), so the LL dispatch recv-spin is measured + cleanly rather than being serialized by an in-process tracer. + """ + + def __init__(self): + import ctypes + import os as _os + + so = _os.path.join(_os.path.dirname(_os.path.abspath(__file__)), "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 + self.lib.kt_get_avg_us.restype = ctypes.c_double + self.lib.kt_get_avg_us.argtypes = [ctypes.c_char_p] + self.lib.kt_get_count.restype = ctypes.c_long + self.lib.kt_get_count.argtypes = [ctypes.c_char_p] + + def start(self) -> int: + return int(self.lib.kt_start()) + + def stop(self) -> int: + return int(self.lib.kt_stop()) + + def avg_us(self, substr: str) -> float: + return float(self.lib.kt_get_avg_us(substr.encode())) + + def count(self, substr: str) -> int: + return int(self.lib.kt_get_count(substr.encode())) + + +def main() -> None: + args = parse_args() + rank, num_ranks, local_rank, group = init_dist() + from mscclpp import CommGroup + import mscclpp.ep as ep + from mscclpp.ep._cpp import get_low_latency_rdma_size_hint + + ep_group = CommGroup(torch_group=group) + + num_tokens = args.num_tokens + hidden = args.hidden + num_topk = args.num_topk + num_experts = args.num_experts + warmup = args.num_warmup + 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 + + # bf16 precision anchor (same convention as test_low_latency_multirank.py). + rank_offset = 128 + assert num_ranks - rank_offset < 257, "too many ranks for bf16 precision anchor" + + torch.manual_seed(args.seed + rank) + random.seed(args.seed + rank) + + # ---- Inputs (mirror ep_bench setupLowLatencyTensors: BF16 tokens + routing). + x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * (rank - rank_offset) + x[:, -128:] = torch.arange(num_tokens, device="cuda").to(torch.bfloat16).view(-1, 1) + scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs() + 1 + topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=True)[1].to(torch.int64) + topk_weights = torch.randn((num_tokens, num_topk), dtype=torch.float32, device="cuda").abs() + + # 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 + 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) + 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"num_rdma_bytes={num_rdma_bytes}", + flush=True, + ) + + # High-level MoE communicator (feature/ep). LOW_LATENCY mode selects the LL + # backend; dispatch/combine run the full (send+recv) op inline on the stream. + moe_comm = ep.MoECommunicator( + comm=ep_group, + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden, + topk=num_topk, + max_tokens_per_rank=num_tokens, + mode=ep.MoEMode.LOW_LATENCY, + num_rdma_qps_per_rank=max(1, num_experts // num_ranks), + ) + assert moe_comm.is_available() + if rank == 0: + print(f"[cfg] MoECommunicator is_internode={moe_comm.is_internode()}", flush=True) + + # ---- Hoist dispatch/combine output tensors out of the timed loop (ep_bench + # preallocates all EP tensors before benchmarking; matching that keeps the + # timed region kernel-bound rather than allocator-bound). The communicator + # 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" + ) + out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + + def dispatch_fn(): + # MoECommunicator.dispatch runs the full (send+recv) LL dispatch inline on + # the stream and returns (DispatchOutput, DispatchHandle) -- the analog of + # ncclEpDispatch + ncclEpComplete. + return moe_comm.dispatch(x, topk_idx, topk_weights, output_buffer=output_buffer) + + def combine_fn(dout): + dispatch_out, handle = dout + moe_comm.combine(dispatch_out.tokens, handle, out=out) + + stream = torch.cuda.current_stream() + + # ---- runPairedBenchmark: warmup (paired), then per-iter timed (paired). ---- + for _ in range(warmup): + dout = dispatch_fn() + stream.synchronize() + combine_fn(dout) + stream.synchronize() + dist.barrier(group=group) + + # CUPTI/nsys region: capture ONLY the post-warmup timed kernels, matching + # ep_bench's KernelTimer.start() (called after warmup). An external nsys run + # with --capture-range=cudaProfilerApi records exactly the dispatch/combine + # kernels between these two calls. + _cupti = bool(getattr(args, "cupti_region", False)) + if _cupti: + torch.cuda.synchronize() + dist.barrier(group=group) + torch.cuda.cudart().cudaProfilerStart() + + # 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)): + 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 + except Exception as exc: + if rank == 0: + print(f"[warn] in-proc CUPTI unavailable ({exc}); host-observed only", flush=True) + _inproc = None + + d_start = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + d_end = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + c_start = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + c_end = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + + for i in range(iters): + d_start[i].record(stream) + dout = dispatch_fn() + d_end[i].record(stream) # record before sync + stream.synchronize() # sync outside timing + c_start[i].record(stream) # record after sync, before combine + combine_fn(dout) + c_end[i].record(stream) # record before sync + stream.synchronize() # sync outside timing + dist.barrier(group=group) # keep ranks in lockstep, outside timing + + torch.cuda.synchronize() + if _cupti: + torch.cuda.cudart().cudaProfilerStop() + ck_disp_us = ck_comb_us = 0.0 + inproc_ok = False + if _inproc is not None: + _inproc.stop() + dist.barrier(group=group) + ck_disp_us = _inproc.avg_us("dispatch") + ck_comb_us = _inproc.avg_us("combine") + n_disp = _inproc.count("dispatch") + n_comb = _inproc.count("combine") + inproc_ok = ck_disp_us > 0 and ck_comb_us > 0 + if os.environ.get("MSCCLPP_EP_KDEBUG", "0") == "1" and rank == 0: + print( + f"[kdebug inproc] dispatch: {ck_disp_us:.1f}us x{n_disp} " f"combine: {ck_comb_us:.1f}us x{n_comb}", + flush=True, + ) + + # ---- Collect per-iter times (ms->us) and trim the first (warmup outlier). -- + disp_us = [d_start[i].elapsed_time(d_end[i]) * 1e3 for i in range(iters)] + comb_us = [c_start[i].elapsed_time(c_end[i]) * 1e3 for i in range(iters)] + tot_us = [d_start[i].elapsed_time(c_end[i]) * 1e3 for i in range(iters)] + if iters > 1: + disp_us, comb_us, tot_us = disp_us[1:], comb_us[1:], tot_us[1:] + + def stats(times): + return sum(times) / len(times), min(times), max(times) + + d_avg, d_min, d_max = stats(disp_us) + c_avg, c_min, c_max = stats(comb_us) + t_avg, t_min, t_max = stats(tot_us) + + # per-rank throughput (GB/s) uses this rank's own byte count / its avg time. + d_tp = (disp_bytes / 1e9) / (d_avg * 1e-6) + c_tp = (comb_bytes / 1e9) / (c_avg * 1e-6) + t_tp = ((disp_bytes + comb_bytes) / 1e9) / (t_avg * 1e-6) + + # ---- Cross-rank reduction (mirror printLowLatencyResults). ---- + g_d_avg = _reduce_scalar(d_avg, dist.ReduceOp.SUM, group) / num_ranks + g_d_min = _reduce_scalar(d_min, dist.ReduceOp.MIN, group) + g_d_max = _reduce_scalar(d_max, dist.ReduceOp.MAX, group) + g_c_avg = _reduce_scalar(c_avg, dist.ReduceOp.SUM, group) / num_ranks + g_c_min = _reduce_scalar(c_min, dist.ReduceOp.MIN, group) + g_c_max = _reduce_scalar(c_max, dist.ReduceOp.MAX, group) + g_t_avg = _reduce_scalar(t_avg, dist.ReduceOp.SUM, group) / num_ranks + g_t_min = _reduce_scalar(t_min, dist.ReduceOp.MIN, group) + g_t_max = _reduce_scalar(t_max, dist.ReduceOp.MAX, group) + + # ---- Kernel-only pass (torch.profiler / Kineto-CUPTI) — ep_bench parity. ---- + # Measures device-side kernel time (strips host launch latency). Dispatch and + # combine are profiled in isolation so no kernel-name matching is required. + kernel_ok = False + g_dk_avg = g_dk_min = g_dk_max = 0.0 + g_ck_avg = g_ck_min = g_ck_max = 0.0 + if args.kernel_timing and not _cupti and not bool(getattr(args, "cupti_inproc", False)): + try: + dk_us, ck_us = _profile_paired_kernels(dispatch_fn, combine_fn, iters, stream, group, rank) + torch.cuda.synchronize() + dist.barrier(group=group) + g_dk_avg = _reduce_scalar(dk_us, dist.ReduceOp.SUM, group) / num_ranks + g_dk_min = _reduce_scalar(dk_us, dist.ReduceOp.MIN, group) + g_dk_max = _reduce_scalar(dk_us, dist.ReduceOp.MAX, group) + g_ck_avg = _reduce_scalar(ck_us, dist.ReduceOp.SUM, group) / num_ranks + g_ck_min = _reduce_scalar(ck_us, dist.ReduceOp.MIN, group) + g_ck_max = _reduce_scalar(ck_us, dist.ReduceOp.MAX, group) + kernel_ok = g_dk_avg > 0 and g_ck_avg > 0 + except Exception as exc: # profiler unavailable / hiccup: keep host numbers valid + if rank == 0: + print(f"[warn] kernel-only pass failed ({exc}); reporting host-observed only", flush=True) + + # ---- In-process CUPTI reduction (ep_bench KernelTimer analog). ---- + g_ik_d_avg = g_ik_d_min = g_ik_d_max = 0.0 + g_ik_c_avg = g_ik_c_min = g_ik_c_max = 0.0 + g_inproc_ok = 0 + if bool(getattr(args, "cupti_inproc", False)): + g_ik_d_avg = _reduce_scalar(ck_disp_us, dist.ReduceOp.SUM, group) / num_ranks + g_ik_d_min = _reduce_scalar(ck_disp_us if inproc_ok else 1e18, dist.ReduceOp.MIN, group) + g_ik_d_max = _reduce_scalar(ck_disp_us, dist.ReduceOp.MAX, group) + g_ik_c_avg = _reduce_scalar(ck_comb_us, dist.ReduceOp.SUM, group) / num_ranks + g_ik_c_min = _reduce_scalar(ck_comb_us if inproc_ok else 1e18, dist.ReduceOp.MIN, group) + g_ik_c_max = _reduce_scalar(ck_comb_us, dist.ReduceOp.MAX, group) + g_inproc_ok = int(_reduce_scalar(1.0 if inproc_ok else 0.0, dist.ReduceOp.MIN, group)) + + d_tp_all = _gather_scalars(d_tp, num_ranks, group) + c_tp_all = _gather_scalars(c_tp, num_ranks, group) + t_tp_all = _gather_scalars(t_tp, num_ranks, group) + + if rank == 0: + # avg throughput uses rank-0 byte count / global avg time (as ep_bench does). + avg_d_tp = (disp_bytes / 1e9) / (g_d_avg * 1e-6) + avg_c_tp = (comb_bytes / 1e9) / (g_c_avg * 1e-6) + avg_t_tp = ((disp_bytes + comb_bytes) / 1e9) / (g_t_avg * 1e-6) + + def minmax_rank(vals): + lo = min(range(num_ranks), key=lambda r: vals[r]) + hi = max(range(num_ranks), key=lambda r: vals[r]) + return vals[lo], lo, vals[hi], hi + + d_lo, d_lo_r, d_hi, d_hi_r = minmax_rank(d_tp_all) + c_lo, c_lo_r, c_hi, c_hi_r = minmax_rank(c_tp_all) + t_lo, t_lo_r, t_hi, t_hi_r = minmax_rank(t_tp_all) + + 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" 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})" + ) + print(f"Combine (BF16): avg={g_c_avg:.2f} us, min={g_c_min:.2f} us, max={g_c_max:.2f} us") + print( + f" throughput: avg={avg_c_tp:.2f} GB/s, " + f"min={c_lo:.2f} GB/s (rank {c_lo_r}), max={c_hi:.2f} GB/s (rank {c_hi_r})" + ) + print(f"Total (D+C): avg={g_t_avg:.2f} us, min={g_t_min:.2f} us, max={g_t_max:.2f} us") + print( + f" throughput: avg={avg_t_tp:.2f} GB/s, " + f"min={t_lo:.2f} GB/s (rank {t_lo_r}), max={t_hi:.2f} GB/s (rank {t_hi_r})" + ) + + print("\n--- Kernel-only performance (device kernel time via torch.profiler/CUPTI) ---") + if kernel_ok: + # The LL dispatch kernel ends with a cross-rank receive spin-wait, so + # its device time includes wait skew. torch.profiler's host tracing + # overhead makes one rank lag, inflating that rank's dispatch device + # time into the ms range; the cross-rank MIN (the rank that did not + # wait) is the representative kernel floor and matches ep_bench's + # low-perturbation CUPTI number. Combine has little recv-spin and is + # stable across ranks. throughput uses the representative (min) time. + print( + f"Dispatch: min={g_dk_min:.2f} us (representative) " + 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" + ) + print(f"Total (D+C): {g_dk_min + g_ck_avg:.2f} us (dispatch min + combine avg)") + 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." + ) + else: + print(" NOTE: kernel-only pass disabled or unavailable.") + + if bool(getattr(args, "cupti_inproc", False)): + print("\n--- Kernel-only performance (in-process CUPTI Activity API, ep_bench KernelTimer analog) ---") + if g_inproc_ok: + # The LL dispatch kernel ends with a cross-rank receive spin-wait, + # so a lagging rank's device time includes wait skew (same effect + # as nsys's max outlier). The cross-rank MIN (the rank that did not + # wait) is the representative kernel floor; it matches the nsys + # CUPTI number and ep_bench's low-perturbation figure. Combine has + # little recv-spin and is stable across ranks. + print( + f"Dispatch: min={g_ik_d_min:.2f} us (representative) " + 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" + ) + print(f"Total (D+C): {g_ik_d_min + g_ik_c_avg:.2f} us (dispatch min + combine avg)") + else: + print(" NOTE: in-process CUPTI collector unavailable (see [warn] above).") + + print( + f"\nByte counts: dispatch={disp_bytes / 1e6:.2f} MB (BF16), " + f"combine={comb_bytes / 1e6:.2f} MB (BF16), selections={num_valid_selections}" + ) + + +if __name__ == "__main__": + try: + main() + finally: + if dist.is_initialized(): + try: + dist.barrier() + except Exception: + pass + try: + dist.destroy_process_group() + except Exception: + pass diff --git a/test/python/ep/mscclpp_ep_bench.cu b/test/python/ep/mscclpp_ep_bench.cu new file mode 100644 index 00000000..d3947f7b --- /dev/null +++ b/test/python/ep/mscclpp_ep_bench.cu @@ -0,0 +1,377 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +// +// mscclpp_ep_bench: a pure-C++/MPI low-latency EP benchmark that calls +// 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). +// +// It mirrors ep_bench's LL measurement methodology and emits the identical +// "=== Summary (Low Latency, across N ranks) ===" block so the unified driver +// (run_ep_bench.py) parses it with no changes. +// +// Scope: low-latency (LL), BF16, EXPERT_MAJOR layout. Single- or multi-node +// (the bootstrap uses an MPI_Bcast of a TcpBootstrap UniqueId). + +#include +#include +#include +#include +#include + +#include +#include +#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 + +#define CUDA_CHECK(x) \ + do { \ + cudaError_t _e = (x); \ + if (_e != cudaSuccess) { \ + fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(_e), __FILE__, __LINE__); \ + MPI_Abort(MPI_COMM_WORLD, 1); \ + } \ + } while (0) + +#define CUPTI_CHECK(x) \ + do { \ + CUptiResult _e = (x); \ + if (_e != CUPTI_SUCCESS) { \ + const char* _s = nullptr; \ + cuptiGetResultString(_e, &_s); \ + fprintf(stderr, "CUPTI error %s at %s:%d\n", _s ? _s : "?", __FILE__, __LINE__); \ + } \ + } while (0) + +// --------------------------------------------------------------------------- +// 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). +// Records are bucketed by mangled-name substring ("dispatch"/"combine"). +// --------------------------------------------------------------------------- +namespace { + +struct KernStat { + uint64_t total_ns = 0; + uint64_t count = 0; +}; +std::map g_kernel_stats; +int g_activity_kind = CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL; + +void CUPTIAPI bufferRequested(uint8_t** buffer, size_t* size, size_t* maxNumRecords) { + constexpr size_t kBufSize = 8 * 1024 * 1024; + *buffer = static_cast(aligned_alloc(8, kBufSize)); + *size = kBufSize; + *maxNumRecords = 0; +} + +void CUPTIAPI bufferCompleted(CUcontext, uint32_t, uint8_t* buffer, size_t, size_t validSize) { + 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); + if (k->name) { + auto& e = g_kernel_stats[k->name]; + e.total_ns += (k->end - k->start); + e.count += 1; + } + } + } + free(buffer); +} + +class KernelTimer { + public: + KernelTimer() { + if (const char* env = std::getenv("MSCCLPP_EP_BENCH_KERNEL_KIND")) { + if (std::string(env) == "kernel") g_activity_kind = CUPTI_ACTIVITY_KIND_KERNEL; + } + } + int start() { + g_kernel_stats.clear(); + CUPTI_CHECK(cuptiActivityRegisterCallbacks(bufferRequested, bufferCompleted)); + return cuptiActivityEnable(static_cast(g_activity_kind)); + } + void stop() { + CUPTI_CHECK(cuptiActivityFlushAll(1)); + CUPTI_CHECK(cuptiActivityDisable(static_cast(g_activity_kind))); + } + // Mean GPU time (us) over all kernels whose (mangled) name contains substr. + double get_avg_us(const char* substr) const { + uint64_t total_ns = 0, count = 0; + for (const auto& kv : g_kernel_stats) { + if (kv.first.find(substr) != std::string::npos) { + total_ns += kv.second.total_ns; + count += kv.second.count; + } + } + return count ? (static_cast(total_ns) / count) / 1e3 : 0.0; + } + uint64_t get_count(const char* substr) const { + uint64_t count = 0; + for (const auto& kv : g_kernel_stats) + if (kv.first.find(substr) != std::string::npos) count += kv.second.count; + return count; + } +}; + +struct Args { + int num_tokens = 128; + int hidden = 7168; + int num_topk = 8; + int num_experts = 256; + int num_warmup = 10; + int num_iters = 50; +}; + +Args parse_args(int argc, char** argv) { + Args a; + for (int i = 1; i < argc; ++i) { + std::string s = argv[i]; + auto next = [&]() -> int { return (i + 1 < argc) ? std::atoi(argv[++i]) : 0; }; + if (s == "-a" || s == "--algorithm") { + ++i; /* ll only */ + } else if (s == "-t" || s == "--num-tokens") + a.num_tokens = next(); + else if (s == "-d" || s == "--hidden") + a.hidden = next(); + else if (s == "-k" || s == "--num-topk") + a.num_topk = next(); + else if (s == "-e" || s == "--num-experts") + a.num_experts = next(); + else if (s == "-w" || s == "--num-warmup") + a.num_warmup = next(); + else if (s == "-i" || s == "--num-iters") + a.num_iters = next(); + } + return a; +} + +struct Stat { + double avg, mn, mx; +}; +Stat stats(const std::vector& v) { + double s = 0, mn = 1e30, mx = -1e30; + for (double x : v) { + s += x; + mn = std::min(mn, x); + mx = std::max(mx, x); + } + return {v.empty() ? 0.0 : s / v.size(), mn, mx}; +} + +} // namespace + +int main(int argc, char** argv) { + MPI_Init(&argc, &argv); + int rank = 0, nRanks = 1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &nRanks); + + int localRank = 0; + if (const char* env = std::getenv("OMPI_COMM_WORLD_LOCAL_RANK")) localRank = std::atoi(env); + CUDA_CHECK(cudaSetDevice(localRank)); + + 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 (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; + + // --- Bootstrap mscclpp::Communicator (TcpBootstrap + MPI_Bcast of UniqueId). --- + auto bootstrap = std::make_shared(rank, nRanks); + mscclpp::UniqueId uid; + if (rank == 0) uid = bootstrap->createUniqueId(); + MPI_Bcast(&uid, sizeof(uid), MPI_BYTE, 0, MPI_COMM_WORLD); + bootstrap->initialize(uid); + mscclpp::Communicator comm(bootstrap); + + const int64_t numRdmaBytes = static_cast(mscclpp::ep::getLowLatencyRdmaSizeHint(T, H, W, E)); + 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"); + MPI_Abort(MPI_COMM_WORLD, 1); + } + 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()); + 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; + 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_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))); + 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; + 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; + + 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); + }; + 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); + }; + + // --- Warmup (paired), then per-iter timed (paired), matching ep_bench. --- + for (int w = 0; w < warmup; ++w) { + dispatch(); + CUDA_CHECK(cudaStreamSynchronize(stream)); + combine(); + CUDA_CHECK(cudaStreamSynchronize(stream)); + MPI_Barrier(MPI_COMM_WORLD); + } + + KernelTimer ktimer; + CUDA_CHECK(cudaDeviceSynchronize()); + MPI_Barrier(MPI_COMM_WORLD); + int kt_rc = ktimer.start(); + + std::vector ds(iters), de(iters), cs(iters), ce(iters); + for (int i = 0; i < iters; ++i) { + CUDA_CHECK(cudaEventCreate(&ds[i])); + CUDA_CHECK(cudaEventCreate(&de[i])); + CUDA_CHECK(cudaEventCreate(&cs[i])); + CUDA_CHECK(cudaEventCreate(&ce[i])); + } + for (int i = 0; i < iters; ++i) { + CUDA_CHECK(cudaEventRecord(ds[i], stream)); + dispatch(); + CUDA_CHECK(cudaEventRecord(de[i], stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + CUDA_CHECK(cudaEventRecord(cs[i], stream)); + combine(); + CUDA_CHECK(cudaEventRecord(ce[i], stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + MPI_Barrier(MPI_COMM_WORLD); + } + CUDA_CHECK(cudaDeviceSynchronize()); + if (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; + for (int i = 0; i < iters; ++i) { + float d_ms = 0, c_ms = 0, t_ms = 0; + CUDA_CHECK(cudaEventElapsedTime(&d_ms, ds[i], de[i])); + CUDA_CHECK(cudaEventElapsedTime(&c_ms, cs[i], ce[i])); + CUDA_CHECK(cudaEventElapsedTime(&t_ms, ds[i], ce[i])); + if (i == 0 && iters > 1) continue; + disp_us.push_back(d_ms * 1e3); + comb_us.push_back(c_ms * 1e3); + tot_us.push_back(t_ms * 1e3); + } + Stat d = stats(disp_us), c = stats(comb_us), tt = stats(tot_us); + + // --- Cross-rank reduction (MPI), mirroring ep_bench / ep_bench_ll. --- + auto reduce3 = [&](double avg, double mn, double mx, double& g_avg, double& g_min, double& g_max) { + MPI_Reduce(&avg, &g_avg, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD); + MPI_Reduce(&mn, &g_min, 1, MPI_DOUBLE, MPI_MIN, 0, MPI_COMM_WORLD); + MPI_Reduce(&mx, &g_max, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD); + g_avg /= W; + }; + double gda, gdmn, gdmx, gca, gcmn, gcmx, gta, gtmn, gtmx; + reduce3(d.avg, d.mn, d.mx, gda, gdmn, gdmx); + reduce3(c.avg, c.mn, c.mx, gca, gcmn, gcmx); + reduce3(tt.avg, tt.mn, tt.mx, gta, gtmn, gtmx); + + // 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; + 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); + + 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, + (unsigned long long)ktimer.get_count("dispatch"), kc, (unsigned long long)ktimer.get_count("combine"), + g_activity_kind == CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL ? "CONCURRENT_KERNEL" : "KERNEL"); + } + + 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(" 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)); + printf("Total (D+C): avg=%.2f us, min=%.2f us, max=%.2f us\n", gta, gtmn, gtmx); + printf(" throughput: avg=%.2f GB/s\n", ((disp_bytes + comb_bytes) / 1e9) / (gta * 1e-6)); + + 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); + } 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); + fflush(stdout); + } + + for (int i = 0; i < iters; ++i) { + cudaEventDestroy(ds[i]); + cudaEventDestroy(de[i]); + cudaEventDestroy(cs[i]); + cudaEventDestroy(ce[i]); + } + cudaStreamDestroy(stream); + cudaFree(d_x); + cudaFree(d_out); + cudaFree(d_recv); + cudaFree(d_topk); + cudaFree(d_weights); + cudaFree(d_srcinfo); + cudaFree(d_layout); + cudaFree(d_count); + + MPI_Barrier(MPI_COMM_WORLD); + MPI_Finalize(); + return 0; +} diff --git a/test/python/ep/run_ep_bench.py b/test/python/ep/run_ep_bench.py new file mode 100644 index 00000000..5f94078d --- /dev/null +++ b/test/python/ep/run_ep_bench.py @@ -0,0 +1,483 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Unified EP low-latency benchmark driver. + +Runs the *same* low-latency dispatch/combine benchmark -- identical tokens, +experts, hidden size, top-k, warmup and iteration counts -- against a +selectable expert-parallel library, then prints one normalized summary so the +libraries can be compared apples-to-apples. + +Backends (``--ep-lib``): + +* ``mscclpp`` -- this repo's :mod:`ep_bench_ll` (MoECommunicator LL) launched + with ``torchrun``. +* ``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. + +Both backends emit the identical ``=== Summary (Low Latency, across N ranks) ===`` +block (``ep_bench_ll.py`` was written to mirror ``ep_bench``), so a single parser +reads either one. + +NCCL-EP dynamically links its shared libraries (``libnccl.so``, ``libnccl_ep.so``). +Point the driver at the correct build with ``--nccl-lib-path`` (falls back to the +``NCCL_LIB_PATH`` environment variable, else the ``lib`` directory beside the +``--nccl-ep-bench`` build tree); that directory is prepended to ``LD_LIBRARY_PATH`` +for the ``ep_bench`` process so the intended NCCL is loaded. + +Scope: single node (``--nproc-per-node`` GPUs). Multi-node runs use the existing +per-backend launchers (mscclpp: run_ep_bench_ll_multinode.sh; nccl-ep: mpirun with +a hostfile); this driver focuses on the common single-node comparison. + +Examples +-------- +Compare both libraries, 4 GPUs, e128:: + + python run_ep_bench.py --ep-lib both -e 128 -t 128 -d 7168 -k 8 -w 10 -i 50 \ + --nccl-lib-path /opt/microsoft/mrc/ep/nccl/build/lib + +Just mscclpp with in-process CUPTI kernel timing:: + + python run_ep_bench.py --ep-lib mscclpp -e 128 --cupti-inproc + +Print the commands without running them:: + + python run_ep_bench.py --ep-lib both -e 128 --dry-run +""" + +from __future__ import annotations + +import argparse +import os +import re +import shlex +import subprocess +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" +_HERE = os.path.dirname(os.path.abspath(__file__)) + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description="Unified EP low-latency benchmark driver (mscclpp EP vs NCCL-EP)", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument( + "--ep-lib", + required=True, + choices=["mscclpp", "mscclpp-cpp", "nccl-ep", "both", "all"], + help="which expert-parallel library to benchmark. mscclpp=MoECommunicator (Python), " + "mscclpp-cpp=MoERuntime (pure C++), nccl-ep=ep_bench. both=mscclpp+nccl-ep; all=the three.", + ) + p.add_argument( + "-a", + "--algorithm", + default="ll", + choices=["ll", "low-latency"], + help="algorithm mode (only low-latency is wired up here)", + ) + + # 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("-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") + + # Launch / fabric. + p.add_argument("--nproc-per-node", type=int, default=4, help="GPUs (ranks) on this node") + p.add_argument( + "--nodes", + default="", + help="space-separated node IPs for a multi-node run (first = master). Empty = single " + "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") + + # 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( + "--conda-prefix", + default=os.path.join(os.path.expanduser("~"), "miniconda3"), + help="conda installation prefix for the mscclpp torch env", + ) + p.add_argument("--conda-env", default="torch", help="conda env name with torch + mscclpp") + p.add_argument( + "--cupti-inproc", action="store_true", help="mscclpp: also collect in-process CUPTI kernel-only timing" + ) + p.add_argument( + "--torch-profiler", + action="store_true", + help="mscclpp: run the torch.profiler kernel pass (default: host-observed only)", + ) + p.add_argument( + "--kernel-only", + action="store_true", + help="compare KERNEL execution time only, stripping host/Python launch overhead " + "(what ep_bench's CUPTI reports). mscclpp uses in-process CUPTI; nccl-ep uses " + "ep_bench's built-in CUPTI KernelTimer. The unified table then leads with the " + "kernel dispatch/combine times and a kernel D+C ratio.", + ) + + # nccl-ep backend. + p.add_argument( + "--nccl-lib-path", + default=os.environ.get("NCCL_LIB_PATH", ""), + help="directory with libnccl.so / libnccl_ep.so; prepended to LD_LIBRARY_PATH " + "for ep_bench (falls back to $NCCL_LIB_PATH, else derived from --nccl-ep-bench)", + ) + p.add_argument( + "--nccl-ep-bench", + default="/opt/microsoft/mrc/ep/nccl/build/test/nccl_ep/ep_bench", + help="path to the NCCL-EP ep_bench binary", + ) + p.add_argument("--hpcx", default="", help="HPCX install dir (for mpirun); autodetected under /opt if empty") + p.add_argument( + "--layout", + default="em", + choices=["em", "rm", "fl"], + help="nccl-ep dispatch layout (em=expert-major, matches mscclpp LL)", + ) + + # 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", + help="path to the mscclpp_ep_bench C++ binary (built via test/python/ep/CMakeLists.txt)", + ) + + p.add_argument("--dry-run", action="store_true", help="print the backend command(s) and exit") + args = p.parse_args() + + # These free-form values are interpolated into shell command strings that are + # executed via bash; constrain them to safe characters to prevent injection + # 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): + raise SystemExit("--iface must be a valid network interface name") + if not re.fullmatch(r"[0-9A-Za-z._,-]+", args.hca): + raise SystemExit("--hca must be comma-separated HCA device names") + + return args + + +# ---------------------------------------------------------------------------- +# Parsing the common "=== Summary (Low Latency ...) ===" block. +# ---------------------------------------------------------------------------- +@dataclass +class Phase: + avg: float = float("nan") + min: float = float("nan") + max: float = float("nan") + + +@dataclass +class LLResult: + ep_lib: str + num_ranks: int = 0 + dispatch: Phase = field(default_factory=Phase) + combine: Phase = field(default_factory=Phase) + total: Phase = field(default_factory=Phase) + # Kernel-only dispatch/combine (avg/min/max) from mscclpp --cupti-inproc or + # ep_bench's CUPTI KernelTimer, if present. + kdispatch: Optional[Phase] = None + kcombine: Optional[Phase] = None + ok: bool = False + + +_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"), + "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\) ===") +# Kernel-only Dispatch line, two formats (both carry avg/min/max): +# mscclpp in-process CUPTI: ``Dispatch: min=M us (representative) [avg=A, max=X us -- ...]`` +# 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") +# 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") + + +def parse_ll_summary(text: str, ep_lib: str) -> LLResult: + res = LLResult(ep_lib=ep_lib) + for raw in text.splitlines(): + line = raw.strip() + m = _RANKS_RE.search(line) + if m: + res.num_ranks = int(m.group(1)) + continue + for name, rx in _HOST_RE.items(): + m = rx.match(line) + if m: + ph = Phase(float(m.group(1)), float(m.group(2)), float(m.group(3))) + setattr(res, name, ph) + # Kernel-only dispatch, first occurrence only. The host lines carry + # ``(BF16)`` so they never match these bare ``Dispatch:``/``Combine:`` forms. + if res.kdispatch is None: + m = _KDISP_REP_RE.match(line) + if m: # mscclpp: printed order is min, avg, max + res.kdispatch = Phase(avg=float(m.group(2)), min=float(m.group(1)), max=float(m.group(3))) + continue + m = _KDISP_AMM_RE.match(line) + if m: # ep_bench: printed order is avg, min, max + 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_RE.match(line) + if m: + res.kcombine = Phase(avg=float(m.group(1)), min=float(m.group(2)), max=float(m.group(3))) + res.ok = res.dispatch.avg == res.dispatch.avg # not NaN + return res + + +# ---------------------------------------------------------------------------- +# 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" + ) + 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}" + ) + cupti_build = "" + 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. + 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") + src = os.path.join(os.path.dirname(bench), "cupti_kernel_timer.cpp") + cupti_build = ( + f"if [ ! -f {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 && " + ) + 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). + pass + else: + # 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" + return ( + f"source {shlex.quote(args.conda_prefix)}/etc/profile.d/conda.sh && " + f"conda activate {shlex.quote(args.conda_env)} && unset PYTHONPATH && " + f"{cupti_build}" + f"export {env} && " + f"torchrun --standalone --nnodes=1 --nproc_per_node={args.nproc_per_node} " + f"{shlex.quote(bench)} {bench_flags}" + ) + + +def _autodetect_hpcx() -> str: + import glob + + cands = sorted(glob.glob("/opt/hpcx-*")) + return cands[0] if cands else "" + + +def _mpi_launch(args, np_total): + """Common mpirun prefix. Multi-node when --nodes lists >1 IP (writes a + hostfile, adds an SSH launcher); otherwise a plain single-node launch.""" + nodes = args.nodes.split() + setup = "" + hostfile = "" + if len(nodes) > 1: + slots = args.nproc_per_node + lines = "\\n".join(f"{ip} slots={slots}" for ip in nodes) + hf = "/tmp/ep_unified_hostfile" + setup = f"printf '{lines}\\n' > {hf} && " + hostfile = ( + f"--hostfile {hf} " f'-mca plm_rsh_args "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null" ' + ) + 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"-mca coll_hcoll_enable 0 -mca coll_ucc_enable 0 " + ) + + +def build_nccl_ep_cmd(args: argparse.Namespace) -> str: + nccl_lib = args.nccl_lib_path + if not nccl_lib: + # Derive the libnccl / libnccl_ep directory from the ep_bench binary + # instead of hard-coding it: /build/test/nccl_ep/ep_bench -> + # /build/lib. + 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 + mnnvl = 1 if nnodes > 1 else 0 + bench_flags = ( + f"-a ll -L {args.layout} -t {args.num_tokens} -d {args.hidden} -k {args.num_topk} " + f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters}" + ) + setup, mpi_prefix = _mpi_launch(args, np_total) + 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 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}" + ) + + +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}" + ) + setup, mpi_prefix = _mpi_launch(args, np_total) + 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"{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}" + ) + + +# ---------------------------------------------------------------------------- +# Run + report. +# ---------------------------------------------------------------------------- +def run_backend(ep_lib: str, cmd: str, dry_run: bool) -> Optional[LLResult]: + print(f"\n########## ep-lib={ep_lib} ##########", flush=True) + print(f"$ {cmd}\n", flush=True) + if dry_run: + return None + proc = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True) + sys.stdout.write(proc.stdout) + if proc.returncode != 0: + sys.stderr.write(proc.stderr[-4000:]) + print(f"[warn] {ep_lib} exited rc={proc.returncode}", flush=True) + res = parse_ll_summary(proc.stdout, ep_lib) + if not res.ok: + print(f"[warn] could not parse a Low-Latency summary from {ep_lib} output", flush=True) + return None + return res + + +def print_unified(results: list, kernel_only: bool = False) -> None: + results = [r for r in results if r is not None] + if not results: + return + has_kernel = all(r.kdispatch is not None and r.kcombine is not None for r in results) + title = "kernel-only" if (kernel_only and has_kernel) else "host-observed" + print(f"\n=== Unified EP Low-Latency Summary ({title}, us) ===") + hdr = f"{'metric':<24}" + "".join(f"{r.ep_lib:>14}" for r in results) + print(hdr) + print("-" * len(hdr)) + + def row(label, fn): + print(f"{label:<24}" + "".join(f"{fn(r):>14.2f}" for r in results)) + + if not (kernel_only and has_kernel): + # Host-observed dispatch/combine/total, full avg/min/max. + row("Host Dispatch avg", lambda r: r.dispatch.avg) + row("Host Dispatch min", lambda r: r.dispatch.min) + row("Host Dispatch max", lambda r: r.dispatch.max) + row("Host Combine avg", lambda r: r.combine.avg) + row("Host Combine min", lambda r: r.combine.min) + 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) + 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: + 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: + print(f"\nHost D+C ratio {a.ep_lib}/{b.ep_lib} = {a.total.avg / b.total.avg:.2f}x") + + +def main() -> None: + args = parse_args() + if args.ep_lib == "both": + libs = ["mscclpp", "nccl-ep"] + elif args.ep_lib == "all": + libs = ["mscclpp", "mscclpp-cpp", "nccl-ep"] + else: + libs = [args.ep_lib] + + builders = { + "mscclpp": build_mscclpp_cmd, + "mscclpp-cpp": build_mscclpp_cpp_cmd, + "nccl-ep": build_nccl_ep_cmd, + } + if len(args.nodes.split()) > 1 and "mscclpp" in libs: + print( + "[warn] --nodes multi-node ignored for the Python 'mscclpp' backend " + "(torchrun --standalone is single-node); use mscclpp-cpp for multi-node.", + flush=True, + ) + results = [] + for lib in libs: + cmd = builders[lib](args) + results.append(run_backend(lib, cmd, args.dry_run)) + if not args.dry_run: + print_unified(results, kernel_only=args.kernel_only) + + +if __name__ == "__main__": + main()