mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-15 11:44:56 +00:00
Merge branch 'feature/ep' into binyli/ep
This commit is contained in:
Binyang Li
124
test/python/ep/CMakeLists.txt
Normal file
124
test/python/ep/CMakeLists.txt
Normal file
@@ -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
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
|
||||
$<$<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"
|
||||
)
|
||||
|
||||
target_link_libraries(mscclpp_ep_bench PRIVATE
|
||||
${MSCCLPP_LIBRARY}
|
||||
${CUPTI_LIBRARY}
|
||||
MPI::MPI_CXX
|
||||
CUDA::cudart
|
||||
CUDA::cuda_driver
|
||||
)
|
||||
122
test/python/ep/cupti_kernel_timer.cpp
Normal file
122
test/python/ep/cupti_kernel_timer.cpp
Normal file
@@ -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<cuda>/targets/sbsa-linux/include -L<cuda>/targets/sbsa-linux/lib -lcupti
|
||||
#include <cupti.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <map>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
|
||||
namespace {
|
||||
|
||||
struct KernelStat {
|
||||
uint64_t total_ns = 0;
|
||||
uint64_t count = 0;
|
||||
};
|
||||
|
||||
std::map<std::string, KernelStat> 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<uint8_t*>(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<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 (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<std::mutex> lock(g_mutex);
|
||||
g_stats.clear();
|
||||
}
|
||||
CUptiResult r = cuptiActivityRegisterCallbacks(bufferRequested, bufferCompleted);
|
||||
if (r != CUPTI_SUCCESS) return static_cast<int>(r);
|
||||
r = cuptiActivityEnable(CUPTI_ACTIVITY_KIND_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);
|
||||
return static_cast<int>(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<std::mutex> 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<double>(total_ns) / static_cast<double>(count) / 1000.0 : 0.0;
|
||||
}
|
||||
|
||||
// Number of recorded kernel instances whose name contains `substr`.
|
||||
long kt_get_count(const char* substr) {
|
||||
std::lock_guard<std::mutex> 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<long>(count);
|
||||
}
|
||||
|
||||
} // extern "C"
|
||||
610
test/python/ep/ep_bench_ll.py
Normal file
610
test/python/ep/ep_bench_ll.py
Normal file
@@ -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
|
||||
377
test/python/ep/mscclpp_ep_bench.cu
Normal file
377
test/python/ep/mscclpp_ep_bench.cu
Normal file
@@ -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 <cuda.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cupti.h>
|
||||
#include <mpi.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <map>
|
||||
#include <mscclpp/core.hpp>
|
||||
#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
|
||||
|
||||
#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<std::string, KernStat> 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<uint8_t*>(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<CUpti_ActivityKernel10*>(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<CUpti_ActivityKind>(g_activity_kind));
|
||||
}
|
||||
void stop() {
|
||||
CUPTI_CHECK(cuptiActivityFlushAll(1));
|
||||
CUPTI_CHECK(cuptiActivityDisable(static_cast<CUpti_ActivityKind>(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<double>(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<double>& 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<mscclpp::TcpBootstrap>(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<int64_t>(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<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;
|
||||
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<cudaEvent_t> 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<double> 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;
|
||||
}
|
||||
483
test/python/ep/run_ep_bench.py
Normal file
483
test/python/ep/run_ep_bench.py
Normal file
@@ -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: <nccl>/build/test/nccl_ep/ep_bench ->
|
||||
# <nccl>/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()
|
||||
Reference in New Issue
Block a user