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

benchmark/shmem/util.cpp

@@ -27,6 +27,7 @@
 #include <string>
 
 #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();
 }
 

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,
     )

include/mori/core/transport/p2p/device_primitives.hpp

@@ -858,16 +858,18 @@ __forceinline__ __device__ void WarpAccumLFImpl(T* __restrict__ dest, T* const*
   }
 }
 
-template <typename T, int VecBytes>
+// Unroll is a template parameter so a caller can override it without changing
+// the default for every other user of this primitive.
+template <typename T, int VecBytes, int Unroll = WARP_ACCUM_UNROLL>
 __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<T, VecBytes, AccumNum, WARP_ACCUM_UNROLL>(dest, srcs, srcScales, offset, \
-                                                              nelems);                       \
+#define WARP_ACCUM_LF_CASE(AccumNum)                                                       \
+  case AccumNum:                                                                           \
+    WarpAccumLFImpl<T, VecBytes, AccumNum, Unroll>(dest, srcs, srcScales, offset, nelems); \
     break;
   switch (accumNum) {
     WARP_ACCUM_LF_CASE(1)

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,
@@ -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)