diff --git a/benchmark/shmem/util.cpp b/benchmark/shmem/util.cpp index ba30eef99..6c5ca8372 100644 --- a/benchmark/shmem/util.cpp +++ b/benchmark/shmem/util.cpp @@ -27,6 +27,7 @@ #include #include "hip/hip_runtime.h" +#include "mori/application/bootstrap/socket_bootstrap.hpp" #include "mori/application/utils/check.hpp" #include "mori/shmem/shmem_api.hpp" @@ -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(); } diff --git a/examples/ops/dispatch_combine/test_low_latency.py b/examples/ops/dispatch_combine/test_low_latency.py index 65bccbc5e..0f049aaee 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,25 @@ 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; -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 + 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 +409,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 +533,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 +576,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 +688,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 +702,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 +718,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 +732,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 +754,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 +781,7 @@ def parse_args(): args.num_topk, args.num_experts, args.pressure_test, + args.num_qp, ), nprocs=args.num_processes, ) diff --git a/include/mori/core/transport/p2p/device_primitives.hpp b/include/mori/core/transport/p2p/device_primitives.hpp index 63ec64b89..6fbfed1c6 100644 --- a/include/mori/core/transport/p2p/device_primitives.hpp +++ b/include/mori/core/transport/p2p/device_primitives.hpp @@ -858,16 +858,18 @@ __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/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): diff --git a/python/mori/ops/dispatch_combine.py b/python/mori/ops/dispatch_combine.py index 7b53bcc0f..2a60b4568 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,27 @@ 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.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 @@ -774,8 +800,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, @@ -787,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 @@ -996,8 +1027,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) 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..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": 128, - "rdma_block_num": 32, + "block_num": 304, + "rdma_block_num": 76, "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": 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": 64, - "rdma_block_num": 32, + "block_num": 256, + "rdma_block_num": 64, "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, + "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": 32, - "rdma_block_num": 21, + "block_num": 256, + "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": 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" }, { @@ -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": 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_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" } ] 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..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 @@ -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": 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" + }, + { + "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": 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" + }, + { + "dtype": "bf16", + "num_tokens": 16, + "hidden_dim": 6144, + "zero_copy": false, + "quant_type": "none", + "block_num": 80, + "rdma_block_num": 40, + "warp_per_block": 4, + "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" + }, + { + "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": 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" + }, { "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..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, @@ -97,6 +153,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, 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..2e510791b 100644 --- a/src/ops/dispatch_combine/internode_v1.cpp +++ b/src/ops/dispatch_combine/internode_v1.cpp @@ -764,6 +764,56 @@ 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; + +// 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(kCombineStepsPerWarpSlice) * warpSizeRt * + (kCombineVecBytes / sizeof(TokT)); +} + +template +inline __device__ void CombineGather(TokT* dest, TokT** srcPtrs, int accumNum, size_t nelems, + bool vecAligned) { + if (vecAligned) { + core::WarpAccum(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 +877,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 +903,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 +1114,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 +1159,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 + 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); 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()