#!/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")), choices=(4096, 7168, 8192, 9216), help="hidden dimension", ) p.add_argument( "-k", "--num-topk", type=int, default=int(os.environ.get("MSCCLPP_EP_BENCH_TOPK", "8")), choices=range(1, 10), help="top-k experts per token", ) p.add_argument( "-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( "--dispatch-dtype", choices=("bf16", "fp8_e4m3"), default="bf16", help="low-latency dispatch payload format", ) p.add_argument( "--combine-mode", choices=("rank_local_reduce", "direct_send"), default="rank_local_reduce", help="low-latency combine algorithm", ) p.add_argument("--num-blocks", type=int, default=130, help="total low-latency dispatch blocks") p.add_argument( "--no-kernel-timing", dest="kernel_timing", 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. Matches mangled dispatch/combine kernel names and " "replaces the torch.profiler pass.", ) p.add_argument("--seed", type=int, default=0xB3C4, help="per-rank RNG seed base") args = p.parse_args() if args.hidden not in (4096, 7168, 8192, 9216): p.error("--hidden must be one of 4096, 7168, 8192, 9216") if not 1 <= args.num_topk <= 9: p.error("--num-topk must be in [1, 9]") if args.num_tokens <= 0 or args.num_experts <= 0: p.error("--num-tokens and --num-experts must be positive") if args.num_topk > args.num_experts: p.error("--num-topk must not exceed --num-experts") if args.num_warmup < 0 or args.num_iters <= 0: p.error("--num-warmup must be non-negative and --num-iters must be positive") return args def init_dist(): 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.environ.get( "MSCCLPP_EP_CUPTI_TIMER_LIB", _os.path.join(_os.path.dirname(_os.path.abspath(__file__)), "build", "libcupti_kernel_timer.so"), ) self.lib = ctypes.CDLL(so) self.lib.kt_start.restype = ctypes.c_int self.lib.kt_stop.restype = ctypes.c_int 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 dispatch_data_type = { "bf16": ep.DispatchDataType.BF16, "fp8_e4m3": ep.DispatchDataType.FP8_E4M3, }[args.dispatch_dtype] combine_mode = { "rank_local_reduce": ep.CombineMode.RANK_LOCAL_REDUCE, "direct_send": ep.CombineMode.DIRECT_SEND, }[args.combine_mode] dispatch_quant = ( None if dispatch_data_type == ep.DispatchDataType.BF16 else ep.QuantConfig(format=dispatch_data_type) ) dispatch_dtype = torch.bfloat16 if dispatch_quant is None else torch.float8_e4m3fn dispatch_label = "BF16" if dispatch_quant is None else "FP8_E4M3" # bf16 precision anchor (same convention as test_low_latency_multirank.py). rank_offset = 128 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()) dispatch_bytes_per_token = hidden * 2 if dispatch_quant is None else hidden + hidden // 128 * 4 disp_bytes = num_valid_selections * dispatch_bytes_per_token comb_bytes = num_valid_selections * hidden * 2 # BF16 (symmetric, per ep_bench) num_rdma_bytes = get_low_latency_rdma_size_hint(num_tokens, hidden, num_ranks, num_experts, num_topk) if rank == 0: print( f"[cfg] algorithm=LOW_LATENCY num_ranks={num_ranks} tokens/rank={num_tokens} hidden={hidden} " f"num_experts={num_experts} top_k={num_topk} warmup={warmup} iters={iters} " f"dispatch_dtype={args.dispatch_dtype} combine_mode={args.combine_mode} " f"num_rdma_bytes={num_rdma_bytes}", flush=True, ) # 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), low_latency_num_blocks=args.num_blocks, low_latency_combine_mode=combine_mode, quant=dispatch_quant, ) 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=dispatch_dtype, device="cuda" ) expert_output = ( None if dispatch_quant is None else torch.zeros((num_local_experts, num_ranks * num_tokens, hidden), dtype=torch.bfloat16, device="cuda") ) out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") 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 if expert_output is None else expert_output, 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 inproc_requested = bool(getattr(args, "cupti_inproc", False)) local_inproc_ready = False if inproc_requested: try: _inproc = _InProcCupti() local_inproc_ready = True except Exception as exc: if rank == 0: print(f"[warn] in-proc CUPTI unavailable ({exc}); host-observed only", flush=True) _inproc = None ready = torch.tensor(int(local_inproc_ready), dtype=torch.int32, device="cuda") dist.all_reduce(ready, op=dist.ReduceOp.MIN, group=group) if ready.item() == 0: if rank == 0: print("[warn] in-proc CUPTI unavailable on at least one rank; disabling globally", flush=True) _inproc = None else: torch.cuda.synchronize() dist.barrier(group=group) try: _rc = _inproc.start() except Exception: _rc = -1 started = torch.tensor(int(_rc == 0), dtype=torch.int32, device="cuda") dist.all_reduce(started, op=dist.ReduceOp.MIN, group=group) if started.item() == 0: if _rc == 0: _inproc.stop() if rank == 0: print("[warn] in-proc CUPTI failed to start on at least one rank; disabling globally", flush=True) _inproc = None dist.barrier(group=group) d_start = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] d_end = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] 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 ({dispatch_label}): avg={g_d_avg:.2f} us, min={g_d_min:.2f} us, max={g_d_max:.2f} us") print( f" throughput: avg={avg_d_tp:.2f} GB/s, " f"min={d_lo:.2f} GB/s (rank {d_lo_r}), max={d_hi:.2f} GB/s (rank {d_hi_r})" ) 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: min={g_ck_min:.2f} us (representative) " f"[avg={g_ck_avg:.2f}, max={g_ck_max:.2f} us -- inflated by profiler rank skew]" ) print(f" throughput @min: {(comb_bytes / 1e9) / (g_ck_min * 1e-6):.2f} GB/s") print( " NOTE: for an authoritative low-perturbation kernel-only number, run under " "nsys (as ep_bench's CUPTI path does); torch.profiler perturbs the LL recv-spin." ) 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: min={g_ik_c_min:.2f} us (representative) " f"[avg={g_ik_c_avg:.2f}, max={g_ik_c_max:.2f} us -- rank skew on lagging ranks]" ) print(f" throughput @min: {(comb_bytes / 1e9) / (g_ik_c_min * 1e-6):.2f} GB/s") else: print(" NOTE: in-process CUPTI collector unavailable (see [warn] above).") print( f"\nByte counts: dispatch={disp_bytes / 1e6:.2f} MB ({dispatch_label}), " 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