perf(ep): optimize InterNodeV1LL small-token latency on EP16

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();
 }
 

python/mori/kernel_profiler/__init__.py

@@ -23,46 +23,46 @@
 import json
 import warnings
 from collections import defaultdict
+
+import numpy as np
+
 from mori import cpp as mori_cpp
 
 
 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):

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,
@@ -996,8 +1020,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)
 

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"
     }
   ]

python/mori/ops/tuning_configs/gfx942_mi300x_InterNodeV1LL_ep16_dispatch.json

@@ -10,42 +10,42 @@
       "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"
     },
     {
       "dtype": "fp8_e4m3_fnuz",
       "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"
     }
   ]