From 8ee2580a38dbf0f2a068c3f0fafab8ad6a1d2024 Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Thu, 27 Aug 2026 02:18:10 +0000 Subject: [PATCH 01/14] perf(ep): trim per-call host overhead on the dispatch/combine path torch.cuda.current_stream() re-resolves the device and builds a Stream object every call (~4.9us measured vs ~0.16us for the raw binding it wraps). _launch_multi and _resolve_launch_params each ran an import statement per call; both moved to module level. At small token counts the host submission path is what paces the GPU, so this is latency, not bookkeeping. --- python/mori/ops/dispatch_combine.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/python/mori/ops/dispatch_combine.py b/python/mori/ops/dispatch_combine.py index 7b53bcc0f..304d3fb02 100644 --- a/python/mori/ops/dispatch_combine.py +++ b/python/mori/ops/dispatch_combine.py @@ -21,6 +21,12 @@ # SOFTWARE. from mori import cpp as mori_cpp from mori.tensor_utils import from_gpu_ptr, dtype_to_int + +# Imported here rather than inside the per-call helpers: both sit on the +# dispatch/combine hot path, where a repeated `from ... import ...` is pure +# interpreter overhead. +from mori.jit.hip_driver import launch_multi +from mori.ops.tuning_config import TuningConfigManager import logging import os from dataclasses import dataclass @@ -183,7 +189,12 @@ def _normalize_quant_type(quant_type): def _current_stream(): - return torch.cuda.current_stream().cuda_stream + # torch.cuda.current_stream() re-resolves the device index and builds a + # Stream object on every call (~4.9us measured); the raw binding it wraps + # costs ~0.16us and returns (stream_ptr, device_index, device_type). At small + # token counts the host submission path is what paces the GPU, so this is + # real latency rather than bookkeeping. + return torch._C._cuda_getCurrentStream(torch.cuda.current_device())[0] @dataclass @@ -774,8 +785,6 @@ def _resolve_launch_params( is_push_transport=False, ): if tuning_rules and dtype is not None: - from mori.ops.tuning_config import TuningConfigManager - params = TuningConfigManager.lookup( tuning_rules, dtype, @@ -996,8 +1005,6 @@ def _launch(self, func_name, grid, block, shared_mem, stream, args_ptr): func.launch_struct(grid, block, shared_mem, stream, args_ptr) def _launch_multi(self, func_names, grids, blocks, shared_mems, stream, args_ptr): - from mori.jit.hip_driver import launch_multi - funcs = [self._get_func(name)._func for name in func_names] launch_multi(funcs, grids, blocks, shared_mems, stream, args_ptr) From e1476844bd0597af7ba885b389ae5f0abcabc2b6 Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Thu, 27 Aug 2026 02:18:42 +0000 Subject: [PATCH 02/14] perf(ep): use the 16B load-first gather in the v1_ll combine Gathering a token reads from up to numExpertPerToken peer GPUs over xGMI, where latency rather than bandwidth is the cap. WarpAccumLF issues AccumNum*Unroll of those reads before accumulating any, so they overlap; WarpAccum keeps only AccumNum in flight and moves 4B/lane. The intra-node combine path has used the 16B load-first form for a while, the v1 internode path had not. Falls back to the 4B path when the staging stride is not 16B-aligned (CombineVecAligned), and sizes each warp's slice to a whole vector step instead of the old fixed warpsPerToken=4, since a slice shorter than one step is slower than not vectorizing at all. Measured: EpCombineInterNodeV1KernelLowLatency mean 56.9 -> 46.2us (-19%) on EP16 at 4 tokens, hidden 6144. --- src/ops/dispatch_combine/common.hpp | 12 +++- src/ops/dispatch_combine/internode_v1.cpp | 74 +++++++++++++++++++---- 2 files changed, 73 insertions(+), 13 deletions(-) diff --git a/src/ops/dispatch_combine/common.hpp b/src/ops/dispatch_combine/common.hpp index 1bb61827b..a4e19f497 100644 --- a/src/ops/dispatch_combine/common.hpp +++ b/src/ops/dispatch_combine/common.hpp @@ -68,10 +68,20 @@ struct MultiWarpIter { size_t dimPerWarp; size_t dimSize; - inline __device__ MultiWarpIter(int globalWarpNum, int numItems, size_t dimSize_) + // dimGranularity rounds dimPerWarp up to a multiple of itself, so callers doing + // vectorized loads get every warp's slice starting on a vector boundary and + // sized in whole vector steps. + inline __device__ MultiWarpIter(int globalWarpNum, int numItems, size_t dimSize_, + size_t dimGranularity = 1) : dimSize(dimSize_) { warpsPerItem = (globalWarpNum + numItems - 1) / numItems; dimPerWarp = (dimSize + warpsPerItem - 1) / warpsPerItem; + if (dimGranularity > 1) { + dimPerWarp = ((dimPerWarp + dimGranularity - 1) / dimGranularity) * dimGranularity; + // A coarser slice means fewer warps are actually needed; keep warpsPerItem + // consistent with dimPerWarp or the tail warps decode to empty ranges. + warpsPerItem = static_cast((dimSize + dimPerWarp - 1) / dimPerWarp); + } } inline __device__ void Decode(int i, int& itemId, int& inItemPartId, size_t& dimOffset, diff --git a/src/ops/dispatch_combine/internode_v1.cpp b/src/ops/dispatch_combine/internode_v1.cpp index 698037b52..7b83d9d7a 100644 --- a/src/ops/dispatch_combine/internode_v1.cpp +++ b/src/ops/dispatch_combine/internode_v1.cpp @@ -764,6 +764,43 @@ inline __device__ void CombineSync(EpDispatchCombineArgs& args) { namespace combine_impl { +// Gathering a token from its experts reads from up to numExpertPerToken peer +// GPUs over xGMI, and peer-read *latency* -- not bandwidth -- is what caps it. +// WarpAccumLF issues AccumNum*Unroll of those reads before accumulating any of +// them so they overlap; WarpAccum keeps only AccumNum in flight and moves 4B per +// lane. The intra-node combine path (intranode.hpp) has used the 16B load-first +// form for a while; the v1 internode path had not. +// +// Two constraints come with it: +// - Both ends must be 16B-aligned. A combine staging slot interleaves the +// hidden payload with the per-token weights, so its stride is only aligned +// for some topk/dtype combinations; CombineVecAligned() decides per launch +// and the caller falls back to the 4B path when it cannot. +// - The vector loop advances CombineVecStep() elements per iteration and drops +// to a per-lane scalar tail below that. A slice shorter than one step is +// *slower* than not vectorizing at all, so slices must be a whole multiple +// of it. +constexpr size_t kCombineVecBytes = 16; + +inline __device__ bool CombineVecAligned(size_t tokHiddenBytes, size_t tokCombXferBytes) { + return ((tokHiddenBytes % kCombineVecBytes) == 0) && ((tokCombXferBytes % kCombineVecBytes) == 0); +} + +template +inline __device__ size_t CombineVecStep(int warpSizeRt) { + return static_cast(WARP_ACCUM_UNROLL) * warpSizeRt * (kCombineVecBytes / sizeof(TokT)); +} + +template +inline __device__ void CombineGather(TokT* dest, TokT** srcPtrs, int accumNum, size_t nelems, + bool vecAligned) { + if (vecAligned) { + core::WarpAccumLF(dest, srcPtrs, nullptr, accumNum, nelems); + } else { + core::WarpAccum(dest, srcPtrs, nullptr, accumNum, nelems); + } +} + template __forceinline__ __device__ void CombineIntraNodeTyped(EpDispatchCombineArgs& args, size_t tokHiddenBytes, @@ -827,7 +864,11 @@ __forceinline__ __device__ void CombineIntraNodeLLTyped(EpDispatchCombineArgs uint8_t* stagingPtr = args.interNodeV1TokBufs.staging->template GetAs() + SendBufSlotOffset(config, nNodes + myNode, 0) * tokCombXferBytes; - MultiWarpIter mwIter(xgmiWarpNum, args.curRankNumToken, hiddenDim); + // Slices are snapped to a whole vector step so the gather below stays on + // WarpAccumLF's vector path instead of its scalar tail. + MultiWarpIter mwIter(xgmiWarpNum, args.curRankNumToken, hiddenDim, + CombineVecStep(warpSize)); + const bool vecAligned = CombineVecAligned(tokHiddenBytes, tokCombXferBytes); for (int i = globalWarpId - blockOffset * warpNum; i < (args.curRankNumToken * mwIter.warpsPerItem); i += xgmiWarpNum) { @@ -849,9 +890,9 @@ __forceinline__ __device__ void CombineIntraNodeLLTyped(EpDispatchCombineArgs destLocalTokId * config.numExpertPerToken; } } - core::WarpAccum( + CombineGather( reinterpret_cast(stagingPtr + tokenId * tokCombXferBytes) + hiddenDimOffset, srcPtrs, - nullptr, config.numExpertPerToken, hiddenDimSize); + config.numExpertPerToken, hiddenDimSize, vecAligned); if (args.weightsBuf && (inTokenPartId == mwIter.warpsPerItem - 1)) { core::WarpAccum( reinterpret_cast(stagingPtr + tokenId * tokCombXferBytes + tokHiddenBytes), @@ -1060,11 +1101,21 @@ __forceinline__ __device__ void CombineInterNodeLLTyped(EpDispatchCombineArgs if (nodeCount > 0) nodeCount -= 1; if (nodeCount == 0) continue; - // int warpsPerToken = (rdmaWarpNum + nodeCount - 1) / nodeCount; - // NOTE: Using a fixed value of 4 for warpsPerToken instead of the dynamic formula above is - // an intentional tuning choice. - int warpsPerToken = 4; - size_t hiddenDimPerWarp = (hiddenDim + warpsPerToken - 1) / warpsPerToken; + // One whole vector step per warp. warpsPerToken was a fixed 4, which for + // hidden 6144 bf16 gives a 1536-element slice -- one full 1024-element vector + // step plus a 512-element scalar tail, and that tail costs more than the + // vector part saves. Sizing the split by the step instead keeps every warp on + // the vector path. + // + // This has to be a static function of the config: chunkFlag is cleared by + // whichever warp completes a chunk, so anything derived from the live counts + // can differ between two warps, and they must agree on the completion target. + const size_t vecStep = CombineVecStep(warpSize); + int warpsPerToken = static_cast(hiddenDim / vecStep); + if (warpsPerToken < 1) warpsPerToken = 1; + size_t hiddenDimPerWarp = core::CeilDiv(hiddenDim, static_cast(warpsPerToken)); + hiddenDimPerWarp = core::CeilDiv(hiddenDimPerWarp, vecStep) * vecStep; + const bool vecAligned = CombineVecAligned(tokHiddenBytes, tokCombXferBytes); for (int i = globalWarpId; i < (nodeCount * warpsPerToken); i += rdmaWarpNum) { int tokenId = i / warpsPerToken; @@ -1095,10 +1146,9 @@ __forceinline__ __device__ void CombineInterNodeLLTyped(EpDispatchCombineArgs destLocalTokId * config.numExpertPerToken; } } - core::WarpAccum( - reinterpret_cast(stagingPtr + globalTokenId * tokCombXferBytes) + - hiddenDimOffset, - srcPtrs, nullptr, config.numExpertPerToken, hiddenDimSize); + CombineGather(reinterpret_cast(stagingPtr + globalTokenId * tokCombXferBytes) + + hiddenDimOffset, + srcPtrs, config.numExpertPerToken, hiddenDimSize, vecAligned); if (args.weightsBuf && (inTokenPartId == 0)) { core::WarpAccum( reinterpret_cast(stagingPtr + globalTokenId * tokCombXferBytes + From 464dfdeaeee8dfdbb46d1ce8ef36c9c15ea1bebd Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Mon, 31 Aug 2026 11:40:12 +0800 Subject: [PATCH 03/14] fix(profiler): use the intranode_ll slot enum in intranode_ll.hpp The generator groups slots by source file, so this file's slots live in IntranodeLlSlot; initializing the intranode context made Slot:: resolve to IntranodeSlot and every ENABLE_PROFILER=ON build failed to compile. --- src/ops/dispatch_combine/intranode_ll.hpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/ops/dispatch_combine/intranode_ll.hpp b/src/ops/dispatch_combine/intranode_ll.hpp index d169f7876..76d0ce037 100644 --- a/src/ops/dispatch_combine/intranode_ll.hpp +++ b/src/ops/dispatch_combine/intranode_ll.hpp @@ -260,9 +260,12 @@ __device__ void EpDispatchIntraNodeLLKernel_body(EpDispatchCombineArgs args) const int numTokens = args.curRankNumToken; const bool hasScales = args.scalesBuf && (config.scaleDim > 0) && (config.scaleTypeSize > 0); + // The generator groups slots by source file, so this file's slots live in + // IntranodeLlSlot, not IntranodeSlot -- initializing the intranode context + // here makes Slot:: resolve to the wrong enum and fails to compile. IF_ENABLE_PROFILER( int globalWarpId = blockIdx.x * warpNum + warpId; - INTRANODE_PROFILER_INIT_CONTEXT(profiler, args.profilerConfig, globalWarpId, laneId)); + INTRANODE_LL_PROFILER_INIT_CONTEXT(profiler, args.profilerConfig, globalWarpId, laneId)); MORI_TRACE_SEQ(seq, profiler); MORI_TRACE_NEXT(seq, Slot::DispatchSendTokens); From 37ea75a011e16453ac4e79565da5bc934c65d43b Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Mon, 31 Aug 2026 11:40:24 +0800 Subject: [PATCH 04/14] perf(profiler): decode trace buffers with numpy _parse_trace_events walked the full worst-case buffer (16384 events x 4096 warps = 134M int64) one .item() at a time, 141s per rank. Same output, 0.07s. --- python/mori/kernel_profiler/__init__.py | 50 ++++++++++++------------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/python/mori/kernel_profiler/__init__.py b/python/mori/kernel_profiler/__init__.py index cd319738a..6194fbddb 100644 --- a/python/mori/kernel_profiler/__init__.py +++ b/python/mori/kernel_profiler/__init__.py @@ -23,6 +23,9 @@ import json import warnings from collections import defaultdict + +import numpy as np + from mori import cpp as mori_cpp @@ -30,39 +33,36 @@ def _parse_trace_events(trace_buffer): """Parse trace event stream: [ts0, meta0, ts1, meta1, ...] Meta encoding: [warpId:16][slot:14][type:2] Returns list of (ts, warp_id, slot, event_type) - """ - events = [] + The buffer is sized for the worst case (MAX_TRACE_EVENTS_PER_WARP * + PROFILER_WARPS_PER_RANK, i.e. ~134M int64 = 1 GiB) and is almost entirely + zeros in practice, so this decodes with numpy rather than element-wise: + at one .item() per element a real buffer takes hours to parse. + """ if trace_buffer.is_cuda: trace_buffer = trace_buffer.cpu() - num_elements = trace_buffer.numel() - warp_stride = 32768 # C++ uses 16384 events * 2 int64 = 32768 - - for base in range(0, num_elements, warp_stride): - warp_buffer = trace_buffer[base : base + warp_stride] + # (num_events, 2) view over the flat [ts, meta] pairs. Warp boundaries do + # not matter here: the caller only wants a single globally ordered stream. + pairs = trace_buffer.contiguous().numpy().reshape(-1, 2) + used = pairs[:, 0] != 0 + ts = pairs[:, 0][used] + meta = pairs[:, 1][used] - warp_events = [] - for i in range(0, warp_stride, 2): - ts = warp_buffer[i].item() - meta = warp_buffer[i + 1].item() - - if ts == 0: - continue - - warp_events.append((ts, meta)) - - warp_events.sort(key=lambda x: x[0]) + if ts.size == 0: + return [] - for ts, meta in warp_events: - event_type = meta & 0x3 - slot = (meta >> 2) & 0x3FFF - warp_id = (meta >> 16) & 0xFFFF + # Stable sort so that events sharing a timestamp keep buffer order, which + # is what the previous per-warp-then-global sort produced. + order = np.argsort(ts, kind="stable") + ts = ts[order] + meta = meta[order] - events.append((ts, warp_id, slot, event_type)) + event_type = meta & 0x3 + slot = (meta >> 2) & 0x3FFF + warp_id = (meta >> 16) & 0xFFFF - events.sort(key=lambda x: x[0]) - return events + return list(zip(ts.tolist(), warp_id.tolist(), slot.tolist(), event_type.tolist())) def _sanitize_events(raw_events, drop_orphan_ends=True, drop_orphan_begins=True): From ce03c1977198dd1cc532424b7b2b8c9d40b565a9 Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Mon, 31 Aug 2026 11:40:24 +0800 Subject: [PATCH 05/14] tools(profiler): aggregate EP kernel traces across ranks analyze_ep_kernel_trace.py shows one rank's timeline; this rolls all ranks up into per-phase duration and share of the dispatch/combine window. --- tools/profiler/aggregate_ep_kernel_trace.py | 203 ++++++++++++++++++++ 1 file changed, 203 insertions(+) create mode 100644 tools/profiler/aggregate_ep_kernel_trace.py diff --git a/tools/profiler/aggregate_ep_kernel_trace.py b/tools/profiler/aggregate_ep_kernel_trace.py new file mode 100644 index 000000000..33b5355a8 --- /dev/null +++ b/tools/profiler/aggregate_ep_kernel_trace.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Aggregate MORI-VIZ kernel traces across all EP ranks. + +analyze_ep_kernel_trace.py answers "what did rank N's timeline look like". +This answers "where does the time go, across the whole job" -- for each +instrumented phase, the wall-clock duration seen on every rank/iteration, plus +what fraction of the dispatch and combine windows it accounts for. + +Usage: aggregate_ep_kernel_trace.py 'traces/trace_rank_*.json' [--per-rank] +""" + +import argparse +import glob +import re +import statistics +import sys +from collections import defaultdict +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from analyze_ep_kernel_trace import _merge_spans, parse_events # noqa: E402 + +# The v1_ll launch order, from src/ops/dispatch_combine/launch.cpp:485-489 +# (dispatch: 2 kernels) and :614-621 (combine: 4 kernels). +DISPATCH_PHASES = [ + "ep_dispatch_copy_to_staging", + "dispatch_inter_node_ll_send", + "dispatch_inter_node_ll_recv", + "dispatch_intra", + "dispatch_sync", +] +COMBINE_PHASES = [ + "combine_sync", + "ep_combine_sync_barrier", + "combine_inter_node_ll", + "combine_intra_node_ll", + "ep_combine_all", +] +PHASE_ORDER = DISPATCH_PHASES + COMBINE_PHASES + +# One iteration is anchored on the first and last kernel of each half, so the +# windows line up with what the torch.cuda.Event timers in run_bench_once +# measure: dispatch = copy_to_staging .. end of the fused LL dispatch kernel, +# combine = EpCombineSync .. end of EpCombineAll. +DISPATCH_ANCHOR = "ep_dispatch_copy_to_staging" +COMBINE_ANCHOR = "combine_sync" +COMBINE_END_ANCHOR = "ep_combine_all" + + +def collect(path): + """Yield one dict per iteration in this rank's trace.""" + intervals = parse_events(path) + by_phase = defaultdict(list) + for name, ts, te, _ in intervals: + by_phase[name].append((ts, te)) + + disp_starts = [s for s, _ in _merge_spans(by_phase[DISPATCH_ANCHOR])] + comb_starts = [s for s, _ in _merge_spans(by_phase[COMBINE_ANCHOR])] + comb_ends = [e for _, e in _merge_spans(by_phase[COMBINE_END_ANCHOR])] + n = min(len(disp_starts), len(comb_starts), len(comb_ends)) + + # Assign every span to the iteration it *starts* in. Overlap-based binning + # double-counts: a warp still spinning in dispatch_sync when the next + # iteration's copy_to_staging begins would otherwise land in both. + bins = [defaultdict(list) for _ in range(n)] + bounds = disp_starts[:n] + [float("inf")] + for name, ts, te, _ in intervals: + for i in range(n): + if bounds[i] <= ts < bounds[i + 1]: + bins[i][name].append((ts, te)) + break + + for i in range(n): + merged = {k: _merge_spans(v) for k, v in bins[i].items()} + # Union of the per-warp spans: the wall-clock time during which any warp + # was inside this phase. Summing the disjoint pieces rather than taking + # last_end - first_start avoids charging a phase for gaps in which every + # warp had already moved on. + durs = {k: sum(e - s for s, e in spans) for k, spans in merged.items()} + disp_end = max( + (e for p in DISPATCH_PHASES if p in merged for _, e in merged[p]), + default=None, + ) + yield { + "durs": durs, + "disp_us": disp_end - disp_starts[i] if disp_end else None, + "comb_us": comb_ends[i] - comb_starts[i], + "gap_us": comb_starts[i] - disp_end if disp_end else None, + } + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("pattern", help="glob matching trace_rank_*.json") + ap.add_argument( + "--skip", + type=int, + default=1, + help="drop the first N iterations per rank (still warming up)", + ) + ap.add_argument( + "--per-rank", + action="store_true", + help="also print each rank's mean dispatch/combine window", + ) + args = ap.parse_args() + + files = sorted(glob.glob(args.pattern)) + if not files: + sys.exit(f"no traces matched {args.pattern!r}") + + per_phase = defaultdict(list) + disp_totals, comb_totals, gaps = [], [], [] + by_rank = [] + n_iters = 0 + for path in files: + its = list(collect(path))[args.skip :] + m = re.search(r"rank_(\d+)_", Path(path).name) + by_rank.append( + ( + int(m.group(1)) if m else -1, + statistics.mean(i["disp_us"] for i in its), + statistics.mean(i["comb_us"] for i in its), + ) + ) + for it in its: + n_iters += 1 + for k, v in it["durs"].items(): + per_phase[k].append(v) + if it["disp_us"] is not None: + disp_totals.append(it["disp_us"]) + gaps.append(it["gap_us"]) + comb_totals.append(it["comb_us"]) + + print( + f"{len(files)} rank traces, {n_iters} rank-iterations " + f"(first {args.skip} per rank dropped)\n" + ) + disp_mean = statistics.mean(disp_totals) + comb_mean = statistics.mean(comb_totals) + print( + f" dispatch window (kernel-side): {disp_mean:7.2f} us " + f"[{min(disp_totals):.1f} .. {max(disp_totals):.1f}]" + ) + print( + f" combine window (kernel-side): {comb_mean:7.2f} us " + f"[{min(comb_totals):.1f} .. {max(comb_totals):.1f}]" + ) + print( + f" gap between the two (host-side convert + launch): " + f"{statistics.mean(gaps):.2f} us" + ) + print() + + hdr = f"{'phase':<32}{'mean':>9}{'min':>9}{'max':>9}{'p90':>9}{'share':>8} n" + print(hdr) + print("-" * len(hdr)) + known = [p for p in PHASE_ORDER if p in per_phase] + rest = sorted(k for k in per_phase if k not in PHASE_ORDER) + for phase in known + rest: + d = sorted(per_phase[phase]) + base = disp_mean if phase in DISPATCH_PHASES else comb_mean + p90 = d[min(len(d) - 1, int(0.9 * len(d)))] + if phase == COMBINE_ANCHOR: + print("-" * len(hdr)) + print( + f"{phase:<32}{statistics.mean(d):9.2f}{d[0]:9.2f}{d[-1]:9.2f}" + f"{p90:9.2f}{statistics.mean(d) / base * 100:7.1f}% {len(d)}" + ) + print("-" * len(hdr)) + print("share = phase mean / (dispatch or combine) window mean; phases that") + print("run concurrently on different blocks make the shares sum past 100%.") + + if args.per_rank: + print(f"\n{'rank':>4}{'dispatch':>10}{'combine':>10}{'sum':>10}") + print("-" * 34) + for r, d, c in sorted(by_rank): + print(f"{r:>4}{d:10.2f}{c:10.2f}{d + c:10.2f}") + + +if __name__ == "__main__": + main() From 4e0577e34ce65bfa86319dad5b7a25cb17bfee33 Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Mon, 31 Aug 2026 11:40:24 +0800 Subject: [PATCH 06/14] tune(ep): add mi308x InterNodeV1LL ep16 rules for hidden 6144 Covers 4/8/16/32 tokens. Only mi300x carried this shape, so AUTO found no rule on mi308x and fell back to the hard-coded (256, 128, 8). --- ...942_mi308x_InterNodeV1LL_ep16_combine.json | 64 +++++++++++++++++++ ...42_mi308x_InterNodeV1LL_ep16_dispatch.json | 56 ++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_combine.json b/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_combine.json index 83ab394ae..8501ecf62 100644 --- a/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_combine.json +++ b/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_combine.json @@ -6,6 +6,70 @@ "ep_size": 16, "phase": "combine", "rules": [ + { + "dtype": "bf16", + "num_tokens": 4, + "hidden_dim": 6144, + "zero_copy": false, + "quant_type": "none", + "block_num": 64, + "rdma_block_num": 32, + "warp_per_block": 4, + "bandwidth_gbps": 4.65, + "avg_rdma_bandwidth_gbps": 1.22, + "avg_xgmi_bandwidth_gbps": 4.07, + "avg_ll_bandwidth_gbps": 4.65, + "avg_latency_us": 80.87, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "bf16", + "num_tokens": 8, + "hidden_dim": 6144, + "zero_copy": false, + "quant_type": "none", + "block_num": 64, + "rdma_block_num": 32, + "warp_per_block": 4, + "bandwidth_gbps": 9.87, + "avg_rdma_bandwidth_gbps": 2.42, + "avg_xgmi_bandwidth_gbps": 7.86, + "avg_ll_bandwidth_gbps": 9.87, + "avg_latency_us": 81.46, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "bf16", + "num_tokens": 16, + "hidden_dim": 6144, + "zero_copy": false, + "quant_type": "none", + "block_num": 64, + "rdma_block_num": 32, + "warp_per_block": 4, + "bandwidth_gbps": 18.96, + "avg_rdma_bandwidth_gbps": 4.58, + "avg_xgmi_bandwidth_gbps": 14.96, + "avg_ll_bandwidth_gbps": 18.96, + "avg_latency_us": 86.07, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "bf16", + "num_tokens": 32, + "hidden_dim": 6144, + "zero_copy": false, + "quant_type": "none", + "block_num": 80, + "rdma_block_num": 53, + "warp_per_block": 4, + "bandwidth_gbps": 32.77, + "avg_rdma_bandwidth_gbps": 7.87, + "avg_xgmi_bandwidth_gbps": 25.86, + "avg_ll_bandwidth_gbps": 32.77, + "avg_latency_us": 100.05, + "bandwidth_metric": "grand_mean" + }, { "dtype": "bf16", "num_tokens": 64, diff --git a/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_dispatch.json b/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_dispatch.json index 643d77255..b8b7a80a4 100644 --- a/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_dispatch.json +++ b/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_dispatch.json @@ -97,6 +97,62 @@ "avg_ll_bandwidth_gbps": 84.69, "avg_latency_us": 1396.41 }, + { + "dtype": "fp8_e4m3_fnuz", + "num_tokens": 4, + "hidden_dim": 6144, + "block_num": 32, + "rdma_block_num": 16, + "warp_per_block": 6, + "bandwidth_gbps": 3.5, + "avg_rdma_bandwidth_gbps": 0.92, + "avg_xgmi_bandwidth_gbps": 3.06, + "avg_ll_bandwidth_gbps": 3.5, + "avg_latency_us": 53.52, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "fp8_e4m3_fnuz", + "num_tokens": 8, + "hidden_dim": 6144, + "block_num": 64, + "rdma_block_num": 16, + "warp_per_block": 4, + "bandwidth_gbps": 7.43, + "avg_rdma_bandwidth_gbps": 1.83, + "avg_xgmi_bandwidth_gbps": 5.92, + "avg_ll_bandwidth_gbps": 7.43, + "avg_latency_us": 53.98, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "fp8_e4m3_fnuz", + "num_tokens": 16, + "hidden_dim": 6144, + "block_num": 64, + "rdma_block_num": 32, + "warp_per_block": 4, + "bandwidth_gbps": 14.49, + "avg_rdma_bandwidth_gbps": 3.51, + "avg_xgmi_bandwidth_gbps": 11.43, + "avg_ll_bandwidth_gbps": 14.49, + "avg_latency_us": 56.27, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "fp8_e4m3_fnuz", + "num_tokens": 32, + "hidden_dim": 6144, + "block_num": 80, + "rdma_block_num": 40, + "warp_per_block": 8, + "bandwidth_gbps": 25.46, + "avg_rdma_bandwidth_gbps": 6.11, + "avg_xgmi_bandwidth_gbps": 20.09, + "avg_ll_bandwidth_gbps": 25.46, + "avg_latency_us": 64.46, + "bandwidth_metric": "grand_mean" + }, { "dtype": "fp8_e4m3_fnuz", "num_tokens": 64, From 8a1be905f48bf677df032f4b2e50be808c4afb4e Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Mon, 31 Aug 2026 08:29:17 +0000 Subject: [PATCH 07/14] fix(ep): get the real stream pointer for kernel launch, not its packed id _current_stream() called torch._C._cuda_getCurrentStream(device)[0] to avoid building a torch.cuda.Stream object (~4.8us) on every dispatch/combine call. That call does not return a stream pointer despite what its result was documented as here: it returns CUDAStream's packed stream_id (pool index + per-pool index + priority folded into an int), not the cudaStream_t/hipStream_t address that _launch's hipModuleLaunchKernel needs. This went unnoticed because the default stream's packed id happens to be 0, which coincides with the null-stream sentinel HIP already treats as "the current stream" -- so ordinary (non-graph) dispatch/combine calls kept working by accident. Entering torch.cuda.graph() switches capture onto a real non-default stream, whose packed id is a small nonzero int (e.g. 3). Passed to hipModuleLaunchKernel as if it were a stream pointer, that value addresses unmapped memory: the launch goes nowhere, capture_end() reports "The CUDA Graph is empty", and later using the captured (empty) graph corrupts the context (HIP error 709, hipErrorContextIsDestroyed). Reproduced with: PYTHONPATH=$(pwd)/python:$(pwd) python3 \ tests/python/ops/bench_dispatch_combine.py --world-size 8 --cmd bench whose default path captures dispatch/combine into CUDA graphs (_capture_split_graphs). Verified torch._C._cuda_getCurrentRawStream(device) returns the same address as the previously-correct torch.cuda.current_stream().cuda_stream in both the default-stream and active-graph-capture cases, and is still ~8x cheaper (~0.6us) than reconstructing the Stream wrapper -- so this keeps the intended host-latency win while launching kernels on the stream that was actually asked for. The same bench command now runs to completion (10 rounds, dispatch/combine/e2e all report). --- python/mori/ops/dispatch_combine.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/python/mori/ops/dispatch_combine.py b/python/mori/ops/dispatch_combine.py index 304d3fb02..067d41926 100644 --- a/python/mori/ops/dispatch_combine.py +++ b/python/mori/ops/dispatch_combine.py @@ -190,11 +190,26 @@ def _normalize_quant_type(quant_type): def _current_stream(): # torch.cuda.current_stream() re-resolves the device index and builds a - # Stream object on every call (~4.9us measured); the raw binding it wraps - # costs ~0.16us and returns (stream_ptr, device_index, device_type). At small - # token counts the host submission path is what paces the GPU, so this is - # real latency rather than bookkeeping. - return torch._C._cuda_getCurrentStream(torch.cuda.current_device())[0] + # Stream object on every call (~4.8us measured). _cuda_getCurrentRawStream + # skips that and returns the same raw cudaStream_t/hipStream_t pointer + # (~0.6us), which is what _launch's hipModuleLaunchKernel call needs. + # + # This used to call _cuda_getCurrentStream(...)[0] instead, which is + # *not* the raw pointer: it is CUDAStream's packed stream_id (pool index + + # per-pool stream index + priority, not an address). That happened to work + # outside CUDA graph capture because the default stream's packed id is 0, + # which coincides with the null-stream sentinel HIP already treats as "the + # current stream". Inside torch.cuda.graph(), the capture stream is a real + # non-default stream with a non-zero packed id (e.g. 3), and passing that + # to hipModuleLaunchKernel as a stream pointer launches on garbage address + # 0x3 instead of the capture stream -- the capture then sees no kernels + # ("UserWarning: The CUDA Graph is empty"), and replaying/using that + # invalid handle afterwards corrupts the context (HIP error 709, + # hipErrorContextIsDestroyed). Reproduced with + # PYTHONPATH=$(pwd) python3 tests/python/ops/bench_dispatch_combine.py + # --world-size 8 --cmd bench, whose default path captures dispatch/combine + # into CUDA graphs. + return torch._C._cuda_getCurrentRawStream(torch.cuda.current_device()) @dataclass From a9116d0e6d72012d7d14e3f828de87a01e4c656e Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Mon, 31 Aug 2026 08:29:34 +0000 Subject: [PATCH 08/14] tune(ep): re-tune MI300X EP16 v1_ll at hidden 6144, tokens 4/8/16/32 Full sweep on the current kernel (16B load-first combine gather + the hierarchical dispatch grid barrier), same hardware/topology/config as the rules being replaced: EP16, gfx942 MI300X, fp8_e4m3_fnuz->bf16, num_qp=1, MORI_RDMA_TC=41. dispatch: 48.42 -> 40.70us (4 tok, -15.9%), 48.52 -> 42.17us (8 tok, -13.1%), 49.84 -> 43.20us (16 tok, -13.3%), 53.54 -> 46.58us (32 tok, -13.0%) combine: 70.32 -> 57.54us (4 tok, -18.2%), 69.57 -> 58.04us (8 tok, -16.6%), 69.93 -> 60.95us (16 tok, -12.8%), 75.33 -> 68.50us (32 tok, -9.1%) Unaffected by the stream-pointer fix in the preceding commit: internode --cmd tuning runs through run_bench_once, which does not use torch.cuda.graph() (the only CUDA Graph capture in this file is under --cmd stress), so these numbers were never exposed to that bug. --- ...942_mi300x_InterNodeV1LL_ep16_combine.json | 58 +++++++++---------- ...42_mi300x_InterNodeV1LL_ep16_dispatch.json | 50 ++++++++-------- 2 files changed, 54 insertions(+), 54 deletions(-) diff --git a/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_combine.json b/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_combine.json index 790f0b4b5..5b58e3433 100644 --- a/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_combine.json +++ b/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_combine.json @@ -12,14 +12,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 128, - "rdma_block_num": 32, + "block_num": 256, + "rdma_block_num": 64, "warp_per_block": 4, - "bandwidth_gbps": 5.33, - "avg_rdma_bandwidth_gbps": 1.4, - "avg_xgmi_bandwidth_gbps": 4.66, - "avg_ll_bandwidth_gbps": 5.33, - "avg_latency_us": 70.32, + "bandwidth_gbps": 6.53, + "avg_rdma_bandwidth_gbps": 1.72, + "avg_xgmi_bandwidth_gbps": 5.71, + "avg_ll_bandwidth_gbps": 6.53, + "avg_latency_us": 57.54, "bandwidth_metric": "grand_mean" }, { @@ -28,14 +28,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 64, + "block_num": 128, "rdma_block_num": 32, - "warp_per_block": 4, - "bandwidth_gbps": 11.53, - "avg_rdma_bandwidth_gbps": 2.83, - "avg_xgmi_bandwidth_gbps": 9.18, - "avg_ll_bandwidth_gbps": 11.53, - "avg_latency_us": 69.57, + "warp_per_block": 6, + "bandwidth_gbps": 13.85, + "avg_rdma_bandwidth_gbps": 3.4, + "avg_xgmi_bandwidth_gbps": 11.04, + "avg_ll_bandwidth_gbps": 13.85, + "avg_latency_us": 58.04, "bandwidth_metric": "grand_mean" }, { @@ -44,14 +44,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 32, - "rdma_block_num": 21, + "block_num": 128, + "rdma_block_num": 64, "warp_per_block": 4, - "bandwidth_gbps": 23.31, - "avg_rdma_bandwidth_gbps": 5.64, - "avg_xgmi_bandwidth_gbps": 18.4, - "avg_ll_bandwidth_gbps": 23.31, - "avg_latency_us": 69.93, + "bandwidth_gbps": 26.79, + "avg_rdma_bandwidth_gbps": 6.47, + "avg_xgmi_bandwidth_gbps": 21.14, + "avg_ll_bandwidth_gbps": 26.79, + "avg_latency_us": 60.95, "bandwidth_metric": "grand_mean" }, { @@ -60,14 +60,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 64, - "rdma_block_num": 32, - "warp_per_block": 4, - "bandwidth_gbps": 43.54, - "avg_rdma_bandwidth_gbps": 10.46, - "avg_xgmi_bandwidth_gbps": 34.36, - "avg_ll_bandwidth_gbps": 43.54, - "avg_latency_us": 75.33, + "block_num": 128, + "rdma_block_num": 64, + "warp_per_block": 6, + "bandwidth_gbps": 47.9, + "avg_rdma_bandwidth_gbps": 11.5, + "avg_xgmi_bandwidth_gbps": 37.79, + "avg_ll_bandwidth_gbps": 47.9, + "avg_latency_us": 68.5, "bandwidth_metric": "grand_mean" } ] diff --git a/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_dispatch.json b/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_dispatch.json index 3f3a12432..5e57727cb 100644 --- a/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_dispatch.json +++ b/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_dispatch.json @@ -10,28 +10,28 @@ "dtype": "fp8_e4m3_fnuz", "num_tokens": 4, "hidden_dim": 6144, - "block_num": 64, - "rdma_block_num": 16, + "block_num": 16, + "rdma_block_num": 10, "warp_per_block": 4, - "bandwidth_gbps": 3.87, - "avg_rdma_bandwidth_gbps": 1.02, - "avg_xgmi_bandwidth_gbps": 3.38, - "avg_ll_bandwidth_gbps": 3.87, - "avg_latency_us": 48.42, + "bandwidth_gbps": 4.6, + "avg_rdma_bandwidth_gbps": 1.21, + "avg_xgmi_bandwidth_gbps": 4.03, + "avg_ll_bandwidth_gbps": 4.6, + "avg_latency_us": 40.7, "bandwidth_metric": "grand_mean" }, { "dtype": "fp8_e4m3_fnuz", "num_tokens": 8, "hidden_dim": 6144, - "block_num": 64, - "rdma_block_num": 32, + "block_num": 32, + "rdma_block_num": 16, "warp_per_block": 4, - "bandwidth_gbps": 8.27, - "avg_rdma_bandwidth_gbps": 2.03, - "avg_xgmi_bandwidth_gbps": 6.59, - "avg_ll_bandwidth_gbps": 8.27, - "avg_latency_us": 48.52, + "bandwidth_gbps": 9.51, + "avg_rdma_bandwidth_gbps": 2.34, + "avg_xgmi_bandwidth_gbps": 7.58, + "avg_ll_bandwidth_gbps": 9.51, + "avg_latency_us": 42.17, "bandwidth_metric": "grand_mean" }, { @@ -39,13 +39,13 @@ "num_tokens": 16, "hidden_dim": 6144, "block_num": 64, - "rdma_block_num": 32, + "rdma_block_num": 42, "warp_per_block": 4, - "bandwidth_gbps": 16.37, - "avg_rdma_bandwidth_gbps": 3.96, - "avg_xgmi_bandwidth_gbps": 12.92, - "avg_ll_bandwidth_gbps": 16.37, - "avg_latency_us": 49.84, + "bandwidth_gbps": 18.88, + "avg_rdma_bandwidth_gbps": 4.57, + "avg_xgmi_bandwidth_gbps": 14.9, + "avg_ll_bandwidth_gbps": 18.88, + "avg_latency_us": 43.2, "bandwidth_metric": "grand_mean" }, { @@ -55,11 +55,11 @@ "block_num": 128, "rdma_block_num": 64, "warp_per_block": 4, - "bandwidth_gbps": 30.68, - "avg_rdma_bandwidth_gbps": 7.37, - "avg_xgmi_bandwidth_gbps": 24.21, - "avg_ll_bandwidth_gbps": 30.68, - "avg_latency_us": 53.54, + "bandwidth_gbps": 35.31, + "avg_rdma_bandwidth_gbps": 8.48, + "avg_xgmi_bandwidth_gbps": 27.86, + "avg_ll_bandwidth_gbps": 35.31, + "avg_latency_us": 46.58, "bandwidth_metric": "grand_mean" } ] From 2fb0f875e0b548556f334bfc96aa353affdf94e0 Mon Sep 17 00:00:00 2001 From: jhchouuu Date: Mon, 31 Aug 2026 17:25:20 +0800 Subject: [PATCH 09/14] feat(bench): socket bootstrap for shmem p2p benchmarks (no MPI) PerfInit uses a socket bootstrap (SocketBootstrapNetwork + ShmemInit) when MASTER_ADDR is set, mirroring the EP tests' RANK/WORLD_SIZE/ MASTER_ADDR launch. This lets p2p_{put,get}_{latency,bw} run across two nodes with one process per node (no mpirun / MPI orchestration), which is what internode IBGDA latency measurement needs. Falls back to MPI_Init when MASTER_ADDR is unset; PerfFinalize guards MPI_Comm_free for the socket path (local_comm == MPI_COMM_NULL). Co-Authored-By: Claude Opus 4.8 (1M context) --- benchmark/shmem/util.cpp | 43 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/benchmark/shmem/util.cpp b/benchmark/shmem/util.cpp index ba30eef99..fbf2e9a3c 100644 --- a/benchmark/shmem/util.cpp +++ b/benchmark/shmem/util.cpp @@ -28,6 +28,7 @@ #include "hip/hip_runtime.h" #include "mori/application/utils/check.hpp" +#include "mori/application/bootstrap/socket_bootstrap.hpp" #include "mori/shmem/shmem_api.hpp" namespace mori::shmem::benchmark { @@ -189,6 +190,44 @@ int PerfInit(int argc, char** argv, struct PerfContext* ctx) { ctx->args = PerfArgs{}; PerfArgs& args = ctx->args; + // Socket bootstrap (no MPI) when MASTER_ADDR is set: launch one process per node + // with RANK/WORLD_SIZE/LOCAL_RANK/MASTER_ADDR/MASTER_PORT env, like the EP tests. + const char* master_addr = std::getenv("MASTER_ADDR"); + if (master_addr != nullptr) { + ctx->world_rank = std::atoi(std::getenv("RANK")); + const int ws = std::atoi(std::getenv("WORLD_SIZE")); + ctx->local_rank = std::getenv("LOCAL_RANK") ? std::atoi(std::getenv("LOCAL_RANK")) : 0; + const int port = std::getenv("MASTER_PORT") ? std::atoi(std::getenv("MASTER_PORT")) : 29500; + ctx->local_comm = MPI_COMM_NULL; + + rc = ParseArgs(argc, argv, &args); + if (rc) { + if (ctx->world_rank == 0) PrintUsage(argv[0]); + return rc; + } + if (args.min_size % sizeof(double) != 0) { + args.min_size = (args.min_size + sizeof(double) - 1) / sizeof(double) * sizeof(double); + } + HIP_RUNTIME_CHECK(hipGetDeviceCount(&ctx->device_count)); + assert(ctx->device_count); + const int device_id = ctx->local_rank % ctx->device_count; + HIP_RUNTIME_CHECK(hipSetDevice(device_id)); + HIP_RUNTIME_CHECK( + hipDeviceGetAttribute(&ctx->device_warp_size, hipDeviceAttributeWarpSize, device_id)); + + auto* bootNet = new application::SocketBootstrapNetwork( + application::SocketBootstrapNetwork::GenerateUniqueId(master_addr, port), ctx->world_rank, + ws); + rc = ShmemInit(bootNet); // takes ownership + initializes internally + if (rc) { + std::fprintf(stderr, "ShmemInit(socket) failed: %d\n", rc); + return 1; + } + ctx->my_pe = ShmemMyPe(); + ctx->npes = ShmemNPes(); + return 0; + } + MPI_Init(&argc, &argv); MPI_Comm_rank(MPI_COMM_WORLD, &ctx->world_rank); @@ -245,7 +284,9 @@ int PerfInit(int argc, char** argv, struct PerfContext* ctx) { } void PerfFinalize(struct PerfContext* ctx) { - MPI_Comm_free(&ctx->local_comm); + if (ctx->local_comm != MPI_COMM_NULL) { + MPI_Comm_free(&ctx->local_comm); + } ShmemFinalize(); } From e626a1b0d1bdfef95ec9da1118c4e6ffb9065931 Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Tue, 1 Sep 2026 17:11:52 +0800 Subject: [PATCH 10/14] perf(ep): optimize the v1_ll combine reduce (#627) --- .../core/transport/p2p/device_primitives.hpp | 13 +++-- ...942_mi300x_InterNodeV1LL_ep16_combine.json | 52 +++++++++---------- ...942_mi308x_InterNodeV1LL_ep16_combine.json | 50 +++++++++--------- src/ops/dispatch_combine/internode_v1.cpp | 17 +++++- 4 files changed, 74 insertions(+), 58 deletions(-) diff --git a/include/mori/core/transport/p2p/device_primitives.hpp b/include/mori/core/transport/p2p/device_primitives.hpp index 63ec64b89..c313ca566 100644 --- a/include/mori/core/transport/p2p/device_primitives.hpp +++ b/include/mori/core/transport/p2p/device_primitives.hpp @@ -858,16 +858,19 @@ __forceinline__ __device__ void WarpAccumLFImpl(T* __restrict__ dest, T* const* } } -template +// Unroll is a template parameter so a caller can override it without changing +// the default for every other user of this primitive. +template __forceinline__ __device__ void WarpAccumLF(T* __restrict__ dest, T* const* __restrict__ srcs, const float* __restrict__ srcScales, size_t accumNum, size_t nelems) { static_assert((VecBytes <= 16) && (VecBytes >= 4) && IsPowerOf2(VecBytes)); + static_assert(Unroll >= 1); size_t offset = 0; -#define WARP_ACCUM_LF_CASE(AccumNum) \ - case AccumNum: \ - WarpAccumLFImpl(dest, srcs, srcScales, offset, \ - nelems); \ +#define WARP_ACCUM_LF_CASE(AccumNum) \ + case AccumNum: \ + WarpAccumLFImpl(dest, srcs, srcScales, offset, \ + nelems); \ break; switch (accumNum) { WARP_ACCUM_LF_CASE(1) diff --git a/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_combine.json b/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_combine.json index 5b58e3433..7740e6aaf 100644 --- a/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_combine.json +++ b/python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_combine.json @@ -12,14 +12,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 256, - "rdma_block_num": 64, + "block_num": 304, + "rdma_block_num": 76, "warp_per_block": 4, - "bandwidth_gbps": 6.53, - "avg_rdma_bandwidth_gbps": 1.72, - "avg_xgmi_bandwidth_gbps": 5.71, - "avg_ll_bandwidth_gbps": 6.53, - "avg_latency_us": 57.54, + "bandwidth_gbps": 7.24, + "avg_rdma_bandwidth_gbps": 1.9, + "avg_xgmi_bandwidth_gbps": 6.34, + "avg_ll_bandwidth_gbps": 7.24, + "avg_latency_us": 51.89, "bandwidth_metric": "grand_mean" }, { @@ -28,14 +28,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 128, - "rdma_block_num": 32, - "warp_per_block": 6, - "bandwidth_gbps": 13.85, - "avg_rdma_bandwidth_gbps": 3.4, - "avg_xgmi_bandwidth_gbps": 11.04, - "avg_ll_bandwidth_gbps": 13.85, - "avg_latency_us": 58.04, + "block_num": 256, + "rdma_block_num": 64, + "warp_per_block": 4, + "bandwidth_gbps": 14.92, + "avg_rdma_bandwidth_gbps": 3.66, + "avg_xgmi_bandwidth_gbps": 11.89, + "avg_ll_bandwidth_gbps": 14.92, + "avg_latency_us": 53.89, "bandwidth_metric": "grand_mean" }, { @@ -44,14 +44,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 128, + "block_num": 256, "rdma_block_num": 64, "warp_per_block": 4, - "bandwidth_gbps": 26.79, - "avg_rdma_bandwidth_gbps": 6.47, - "avg_xgmi_bandwidth_gbps": 21.14, - "avg_ll_bandwidth_gbps": 26.79, - "avg_latency_us": 60.95, + "bandwidth_gbps": 30.53, + "avg_rdma_bandwidth_gbps": 7.37, + "avg_xgmi_bandwidth_gbps": 24.09, + "avg_ll_bandwidth_gbps": 30.53, + "avg_latency_us": 53.56, "bandwidth_metric": "grand_mean" }, { @@ -63,11 +63,11 @@ "block_num": 128, "rdma_block_num": 64, "warp_per_block": 6, - "bandwidth_gbps": 47.9, - "avg_rdma_bandwidth_gbps": 11.5, - "avg_xgmi_bandwidth_gbps": 37.79, - "avg_ll_bandwidth_gbps": 47.9, - "avg_latency_us": 68.5, + "bandwidth_gbps": 53.96, + "avg_rdma_bandwidth_gbps": 12.95, + "avg_xgmi_bandwidth_gbps": 42.58, + "avg_ll_bandwidth_gbps": 53.96, + "avg_latency_us": 60.91, "bandwidth_metric": "grand_mean" } ] diff --git a/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_combine.json b/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_combine.json index 8501ecf62..d706a7485 100644 --- a/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_combine.json +++ b/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_combine.json @@ -12,14 +12,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 64, - "rdma_block_num": 32, - "warp_per_block": 4, - "bandwidth_gbps": 4.65, - "avg_rdma_bandwidth_gbps": 1.22, - "avg_xgmi_bandwidth_gbps": 4.07, - "avg_ll_bandwidth_gbps": 4.65, - "avg_latency_us": 80.87, + "block_num": 32, + "rdma_block_num": 21, + "warp_per_block": 6, + "bandwidth_gbps": 6.59, + "avg_rdma_bandwidth_gbps": 1.73, + "avg_xgmi_bandwidth_gbps": 5.77, + "avg_ll_bandwidth_gbps": 6.59, + "avg_latency_us": 57.03, "bandwidth_metric": "grand_mean" }, { @@ -31,11 +31,11 @@ "block_num": 64, "rdma_block_num": 32, "warp_per_block": 4, - "bandwidth_gbps": 9.87, - "avg_rdma_bandwidth_gbps": 2.42, - "avg_xgmi_bandwidth_gbps": 7.86, - "avg_ll_bandwidth_gbps": 9.87, - "avg_latency_us": 81.46, + "bandwidth_gbps": 13.89, + "avg_rdma_bandwidth_gbps": 3.4, + "avg_xgmi_bandwidth_gbps": 11.07, + "avg_ll_bandwidth_gbps": 13.89, + "avg_latency_us": 57.93, "bandwidth_metric": "grand_mean" }, { @@ -44,14 +44,14 @@ "hidden_dim": 6144, "zero_copy": false, "quant_type": "none", - "block_num": 64, - "rdma_block_num": 32, + "block_num": 80, + "rdma_block_num": 40, "warp_per_block": 4, - "bandwidth_gbps": 18.96, - "avg_rdma_bandwidth_gbps": 4.58, - "avg_xgmi_bandwidth_gbps": 14.96, - "avg_ll_bandwidth_gbps": 18.96, - "avg_latency_us": 86.07, + "bandwidth_gbps": 26.28, + "avg_rdma_bandwidth_gbps": 6.35, + "avg_xgmi_bandwidth_gbps": 20.74, + "avg_ll_bandwidth_gbps": 26.28, + "avg_latency_us": 62.14, "bandwidth_metric": "grand_mean" }, { @@ -63,11 +63,11 @@ "block_num": 80, "rdma_block_num": 53, "warp_per_block": 4, - "bandwidth_gbps": 32.77, - "avg_rdma_bandwidth_gbps": 7.87, - "avg_xgmi_bandwidth_gbps": 25.86, - "avg_ll_bandwidth_gbps": 32.77, - "avg_latency_us": 100.05, + "bandwidth_gbps": 41.98, + "avg_rdma_bandwidth_gbps": 10.08, + "avg_xgmi_bandwidth_gbps": 33.12, + "avg_ll_bandwidth_gbps": 41.98, + "avg_latency_us": 78.11, "bandwidth_metric": "grand_mean" }, { diff --git a/src/ops/dispatch_combine/internode_v1.cpp b/src/ops/dispatch_combine/internode_v1.cpp index 7b83d9d7a..2e510791b 100644 --- a/src/ops/dispatch_combine/internode_v1.cpp +++ b/src/ops/dispatch_combine/internode_v1.cpp @@ -782,20 +782,33 @@ namespace combine_impl { // of it. constexpr size_t kCombineVecBytes = 16; +// How many vector steps of a token's hidden dimension go into one warp's slice. +// +// This only sets the split -- CombineVecStep() feeds warpsPerToken below, and +// the slice is rounded up to a whole number of steps so no warp gets less than +// one. It does not reach inside the gather: WarpAccum advances exactly one step +// (warpSize * kCombineVecBytes) per inner iteration regardless of what is set +// here, so a slice of 2 steps simply means each warp runs two iterations. +// +// 2 measures faster than 1 at every token count tried, which is a statement +// about how wide to spread a token, not about the gather's inner loop. +constexpr int kCombineStepsPerWarpSlice = 2; + inline __device__ bool CombineVecAligned(size_t tokHiddenBytes, size_t tokCombXferBytes) { return ((tokHiddenBytes % kCombineVecBytes) == 0) && ((tokCombXferBytes % kCombineVecBytes) == 0); } template inline __device__ size_t CombineVecStep(int warpSizeRt) { - return static_cast(WARP_ACCUM_UNROLL) * warpSizeRt * (kCombineVecBytes / sizeof(TokT)); + return static_cast(kCombineStepsPerWarpSlice) * warpSizeRt * + (kCombineVecBytes / sizeof(TokT)); } template inline __device__ void CombineGather(TokT* dest, TokT** srcPtrs, int accumNum, size_t nelems, bool vecAligned) { if (vecAligned) { - core::WarpAccumLF(dest, srcPtrs, nullptr, accumNum, nelems); + core::WarpAccum(dest, srcPtrs, nullptr, accumNum, nelems); } else { core::WarpAccum(dest, srcPtrs, nullptr, accumNum, nelems); } From 2e8dca31c7d83cd356d38c4b043219d2bc6103b5 Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Wed, 2 Sep 2026 09:48:06 +0800 Subject: [PATCH 11/14] style: clang-format util.cpp and device_primitives.hpp Both fell out of formatting during the EP16 work: an #include went in out of alphabetical order, and a macro's continuation backslashes drifted out of column alignment when a line above it was edited. --- benchmark/shmem/util.cpp | 2 +- include/mori/core/transport/p2p/device_primitives.hpp | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/benchmark/shmem/util.cpp b/benchmark/shmem/util.cpp index fbf2e9a3c..6c5ca8372 100644 --- a/benchmark/shmem/util.cpp +++ b/benchmark/shmem/util.cpp @@ -27,8 +27,8 @@ #include #include "hip/hip_runtime.h" -#include "mori/application/utils/check.hpp" #include "mori/application/bootstrap/socket_bootstrap.hpp" +#include "mori/application/utils/check.hpp" #include "mori/shmem/shmem_api.hpp" namespace mori::shmem::benchmark { diff --git a/include/mori/core/transport/p2p/device_primitives.hpp b/include/mori/core/transport/p2p/device_primitives.hpp index c313ca566..6fbfed1c6 100644 --- a/include/mori/core/transport/p2p/device_primitives.hpp +++ b/include/mori/core/transport/p2p/device_primitives.hpp @@ -867,10 +867,9 @@ __forceinline__ __device__ void WarpAccumLF(T* __restrict__ dest, T* const* __re static_assert((VecBytes <= 16) && (VecBytes >= 4) && IsPowerOf2(VecBytes)); static_assert(Unroll >= 1); size_t offset = 0; -#define WARP_ACCUM_LF_CASE(AccumNum) \ - case AccumNum: \ - WarpAccumLFImpl(dest, srcs, srcScales, offset, \ - nelems); \ +#define WARP_ACCUM_LF_CASE(AccumNum) \ + case AccumNum: \ + WarpAccumLFImpl(dest, srcs, srcScales, offset, nelems); \ break; switch (accumNum) { WARP_ACCUM_LF_CASE(1) From cacda17f7f565a1c3f46537fc46c1ad4373f8bef Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Wed, 2 Sep 2026 11:15:26 +0800 Subject: [PATCH 12/14] tune(ep): add mi308x bf16 dispatch rules at hidden 6144, tokens 4/8/16/32 test_low_latency.py dispatches in bf16; AUTO had no rule to match and fell back to the hardcoded (256, 128, 8). --- ...42_mi308x_InterNodeV1LL_ep16_dispatch.json | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_dispatch.json b/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_dispatch.json index b8b7a80a4..56b1c7422 100644 --- a/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_dispatch.json +++ b/python/mori/ops/tuning_configs/gfx942_mi308x_InterNodeV1LL_ep16_dispatch.json @@ -6,6 +6,62 @@ "ep_size": 16, "phase": "dispatch", "rules": [ + { + "dtype": "bf16", + "num_tokens": 4, + "hidden_dim": 6144, + "block_num": 32, + "rdma_block_num": 21, + "warp_per_block": 4, + "bandwidth_gbps": 8.03, + "avg_rdma_bandwidth_gbps": 2.12, + "avg_xgmi_bandwidth_gbps": 7.03, + "avg_ll_bandwidth_gbps": 8.03, + "avg_latency_us": 46.63, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "bf16", + "num_tokens": 8, + "hidden_dim": 6144, + "block_num": 64, + "rdma_block_num": 42, + "warp_per_block": 4, + "bandwidth_gbps": 16.46, + "avg_rdma_bandwidth_gbps": 4.05, + "avg_xgmi_bandwidth_gbps": 13.11, + "avg_ll_bandwidth_gbps": 16.46, + "avg_latency_us": 48.78, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "bf16", + "num_tokens": 16, + "hidden_dim": 6144, + "block_num": 80, + "rdma_block_num": 40, + "warp_per_block": 4, + "bandwidth_gbps": 30.16, + "avg_rdma_bandwidth_gbps": 7.29, + "avg_xgmi_bandwidth_gbps": 23.8, + "avg_ll_bandwidth_gbps": 30.16, + "avg_latency_us": 54.06, + "bandwidth_metric": "grand_mean" + }, + { + "dtype": "bf16", + "num_tokens": 32, + "hidden_dim": 6144, + "block_num": 64, + "rdma_block_num": 32, + "warp_per_block": 8, + "bandwidth_gbps": 48.28, + "avg_rdma_bandwidth_gbps": 11.61, + "avg_xgmi_bandwidth_gbps": 38.09, + "avg_ll_bandwidth_gbps": 48.28, + "avg_latency_us": 69.11, + "bandwidth_metric": "grand_mean" + }, { "dtype": "fp4", "num_tokens": 64, From 23ba0a8901a33e8a4104f01186fb8c6dcc564ed1 Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Wed, 2 Sep 2026 11:17:30 +0800 Subject: [PATCH 13/14] feat(ep): support MORI_EP_LAUNCH_CONFIG_MODE=AUTO and --num-qp in test_low_latency AUTO previously had no effect here: the per-call block/warp args were still the old hardcoded numbers, which silently override half of what AUTO would have picked (a tuning rule wins over an explicit arg, but the built-in fallback doesn't). -1 hands both back to the op. QP count isn't part of the tuning lookup key, so it needs its own flag. --- .../ops/dispatch_combine/test_low_latency.py | 39 +++++++++++++++++-- 1 file changed, 36 insertions(+), 3 deletions(-) diff --git a/examples/ops/dispatch_combine/test_low_latency.py b/examples/ops/dispatch_combine/test_low_latency.py index 65bccbc5e..8437a2279 100644 --- a/examples/ops/dispatch_combine/test_low_latency.py +++ b/examples/ops/dispatch_combine/test_low_latency.py @@ -313,6 +313,7 @@ def test_main( seed: int = 0, enable_dedup: bool = True, fused_moe_adaption: bool = True, + num_qp_per_pe: int = -1, ): torch.manual_seed(seed + rank) random.seed(seed + rank) @@ -370,6 +371,22 @@ def test_main( combine_block_num, combine_warp_per_block = block_num, warp_num_per_block rdma_block_num = 0 + # Under AUTO the op resolves block/warp/rdma itself (tuning rule, else its + # own fallback); -1 means "you decide". Passing the hardcoded numbers above + # would only half-apply that, since a tuning rule wins over an explicit arg + # but the fallback doesn't. + auto_launch = os.getenv("MORI_EP_LAUNCH_CONFIG_MODE", "MANUAL").upper() == "AUTO" + if auto_launch: + dispatch_block_num = dispatch_warp_per_block = -1 + combine_block_num = combine_warp_per_block = -1 + + # QP count isn't in the tuning lookup key, so AUTO can't pick it. -1 keeps + # the long-standing default; at 4-32 tokens/rank, hidden 6144, 1 measures + # ~20% faster than the default 4 on 2x8 MI308X (small batches don't have + # enough chunks in flight to earn back the extra QPs' per-QP overhead). + if num_qp_per_pe <= 0: + num_qp_per_pe = 4 if multi_node else 1 + mori.shmem.shmem_torch_process_group_init("default") config = mori.ops.EpDispatchCombineConfig( @@ -389,7 +406,7 @@ def test_main( kernel_type=kernel_type, gpu_per_node=num_ranks // num_nodes, rdma_block_num=rdma_block_num, - num_qp_per_pe=4 if multi_node else 1, + num_qp_per_pe=num_qp_per_pe, ) op = mori.ops.EpDispatchCombineOp(config) @@ -513,7 +530,9 @@ def test_main( topk_idx, block_num=combine_block_num, warp_per_block=( - 4 if zero_copy and not multi_node else combine_warp_per_block + 4 + if zero_copy and not multi_node and not auto_launch + else combine_warp_per_block ), use_external_inp_buf=not zero_copy, ) @@ -554,7 +573,9 @@ def test_func(zero_copy: bool, use_fp8: bool): topk_idx, block_num=combine_block_num, warp_per_block=( - 4 if zero_copy and not multi_node else combine_warp_per_block + 4 + if zero_copy and not multi_node and not auto_launch + else combine_warp_per_block ), use_external_inp_buf=not zero_copy, ) @@ -664,6 +685,7 @@ def test_loop( num_topk: int = 8, num_experts: int = 288, do_pressure_test: bool = False, + num_qp_per_pe: int = -1, ): rank, num_ranks, group, num_nodes = init_dist(local_rank, num_local_ranks) @@ -677,6 +699,7 @@ def test_loop( num_nodes, group, seed=1, + num_qp_per_pe=num_qp_per_pe, ) for seed in range(int(1e9) if do_pressure_test else 0): @@ -692,6 +715,7 @@ def test_loop( num_nodes, group, seed=seed, + num_qp_per_pe=num_qp_per_pe, ) for i in range(20): assert ( @@ -705,6 +729,7 @@ def test_loop( num_nodes, group, seed=seed, + num_qp_per_pe=num_qp_per_pe, ) == ref_hash ), f"Error: seed={seed}" @@ -726,6 +751,13 @@ def parse_args(): p.add_argument( "--num-processes", type=int, default=8, help="ranks per node (GPUs to use)" ) + p.add_argument( + "--num-qp", + type=int, + default=-1, + help="QPs per peer. -1 keeps the default (4 multi-node, 1 single-node); " + "AUTO doesn't cover this, only block/warp/rdma", + ) p.add_argument( "--pressure-test", action="store_true", @@ -746,6 +778,7 @@ def parse_args(): args.num_topk, args.num_experts, args.pressure_test, + args.num_qp, ), nprocs=args.num_processes, ) From 74503348cd2c6bfc4f94c4d57d98cf0f7fccbd34 Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Wed, 2 Sep 2026 12:16:27 +0800 Subject: [PATCH 14/14] docs(ep): correct why -1 matters under AUTO for InterNodeV1LL Previous comment (23ba0a89) said passing the hardcoded launch numbers would 'half-apply' AUTO, i.e. that a matched tuning rule wins but the built-in fallback would still use the caller's value. Verified that's wrong: the fallback tuple for InterNodeV1LL is (256, 128, 8) in dispatch_combine.py's __init__, all non-zero, so it always wins over the caller's argument in _resolve_launch_params -- matched rule or not. -1 vs the old hardcoded values measures identically (104.6us vs 105.1-105.6us) once the mi308x bf16 rule exists; the earlier improvement was entirely the tuning config, not this code path. --- examples/ops/dispatch_combine/test_low_latency.py | 11 +++++++---- python/mori/ops/dispatch_combine.py | 7 +++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/examples/ops/dispatch_combine/test_low_latency.py b/examples/ops/dispatch_combine/test_low_latency.py index 8437a2279..0f049aaee 100644 --- a/examples/ops/dispatch_combine/test_low_latency.py +++ b/examples/ops/dispatch_combine/test_low_latency.py @@ -371,10 +371,13 @@ def test_main( combine_block_num, combine_warp_per_block = block_num, warp_num_per_block rdma_block_num = 0 - # Under AUTO the op resolves block/warp/rdma itself (tuning rule, else its - # own fallback); -1 means "you decide". Passing the hardcoded numbers above - # would only half-apply that, since a tuning rule wins over an explicit arg - # but the fallback doesn't. + # Under AUTO the op resolves block/warp/rdma itself; -1 means "you decide". + # For InterNodeV1LL (the kernel type below) this happens to be a no-op + # today -- see _resolve_launch_params in dispatch_combine.py, whose AUTO + # fallback for this kernel type is fully non-zero and so always wins over + # whatever's passed here, matched tuning rule or not. Kept anyway: it's + # what "let AUTO decide" is supposed to mean, and stops these numbers from + # silently doing something once that fallback logic changes. auto_launch = os.getenv("MORI_EP_LAUNCH_CONFIG_MODE", "MANUAL").upper() == "AUTO" if auto_launch: dispatch_block_num = dispatch_warp_per_block = -1 diff --git a/python/mori/ops/dispatch_combine.py b/python/mori/ops/dispatch_combine.py index 067d41926..2a60b4568 100644 --- a/python/mori/ops/dispatch_combine.py +++ b/python/mori/ops/dispatch_combine.py @@ -811,6 +811,13 @@ def _resolve_launch_params( ) if params is not None: return params.block_num, params.rdma_block_num, params.warp_per_block + # No matching rule: fall back to the per-kernel-type AUTO default set in + # __init__. For kernel types whose default is fully non-zero (currently + # InterNodeV1 and InterNodeV1LL), that default always wins here -- the + # caller's block_num/rdma_block_num/warp_per_block is never reached, + # matched rule or not. Passing an explicit value under AUTO for those + # kernel types is a no-op; only IntraNode's zero rdma default actually + # falls through to the caller's argument. bn = self.auto_block_num if self.auto_block_num else block_num rbn = self.auto_rdma_block_num if self.auto_rdma_block_num else rdma_block_num wpb = self.auto_warp_per_block if self.auto_warp_per_block else warp_per_block