mscclpp/test/python/ep/ep_bench_ll.py

#!/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