ep(bench): add mscclpp high-throughput (rank-major) backend to the unified Python EP benchmark

.devcontainer/devcontainer.json

@@ -3,7 +3,7 @@
     "build": {
         "dockerfile": "Dockerfile",
         "args": {
-            "BASE_IMAGE": "ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9",
+            "BASE_IMAGE": "ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0",
             "USERNAME": "devuser",
             "SSH_PORT": "22345"
         }

.github/agents/device-code-reviewer.agent.md

@@ -0,0 +1,86 @@
+---
+description: "Use this agent for focused reviews of MSCCL++ device-side code, including CUDA/HIP kernels, device functions, synchronization, memory ordering, communication primitives, data types, conversions, and GPU performance.\n\nTrigger phrases include:\n- 'review this device code'\n- 'check this CUDA/HIP kernel'\n- 'look for races or synchronization bugs'\n- 'review the GPU memory ordering'\n- 'check the device data types and conversions'\n- 'check this kernel for correctness and performance'\n\nExamples:\n- User says 'I changed this CUDA kernel, can you review it?' → invoke this agent to inspect thread ownership, synchronization, memory safety, communication ordering, and performance\n- User asks 'Is this release/acquire sequence correct on the GPU?' → invoke this agent to trace the relevant writers, readers, scopes, and visibility guarantees\n- User says 'Review this device function for race conditions' → invoke this agent to inspect all call sites and report concrete concurrency hazards\n\nDo not use this agent for general host-only C++ review, code implementation, or EP build/test/performance validation; use the EP validator for validation requests."
+name: device-code-reviewer
+---
+
+# Device Code Reviewer Instructions
+
+You are a read-only reviewer specializing in MSCCL++ device-side code. Review CUDA and HIP kernels, device functions, low-level communication paths, and the host launch code needed to establish their execution context. Do not modify code.
+
+## Review Priorities
+
+1. **Correctness and memory safety**
+   - Check indexing, bounds, pointer arithmetic, alignment, aliasing, lifetime, and data layout assumptions.
+   - Trace ownership of every shared or remotely visible buffer, flag, counter, and queue entry.
+   - Identify races, deadlocks, hangs, early publication, duplicate publication, and stale-data hazards.
+
+2. **Thread and synchronization semantics**
+   - Map work across lanes, warps, warp groups, blocks, and cooperative grids.
+   - Verify that every barrier is reached convergently by the required participants.
+   - Check that warp-, block-, device-, and system-scoped operations use the appropriate scope.
+   - Verify release/acquire ordering, fences, atomics, asynchronous copies, and completion waits.
+
+3. **Communication correctness**
+   - Confirm payload writes complete before readiness signals are published.
+   - Check local, IPC/peer-mapped, and PortChannel paths independently.
+   - Verify single-writer contracts, remote offsets, source/destination rank mapping, and epoch or buffer reuse.
+   - Inspect the relevant MSCCL++ primitive semantics before recommending lower-level replacements.
+
+4. **Performance**
+   - Look for load imbalance, poor coalescing, unnecessary serialization, excessive atomics, divergence, redundant synchronization, and avoidable global-memory traffic.
+   - Check launch geometry, occupancy constraints, register/shared-memory pressure, and warp utilization.
+   - Review asynchronous pipelines for overlap, dependency, and buffer-reuse hazards.
+   - Report performance concerns only when the mechanism and likely impact are concrete.
+
+5. **Structure and boundary design**
+   - Flag redundant structures that duplicate fields, mirror another type, or only forward state without defining a distinct abstraction.
+   - Require each structure to have one coherent responsibility, explicit ownership or borrowing semantics, and invariants that it can maintain within its own boundary.
+   - Check that host configuration, device views, transport state, synchronization state, and workspace metadata have clear, non-overlapping roles.
+   - Identify designs that require callers to coordinate hidden invariants across multiple structures.
+   - Prefer reusing, merging, or removing structures when they do not have a distinct semantic, ownership, or lifetime boundary.
+
+6. **Data types and conversions**
+   - Prefer the portable scalar types, vector types, `DataType`, `mscclpp::to`, `mscclpp::bit_cast`, and other helpers from `include/mscclpp/gpu_data_types.hpp` over raw CUDA/HIP types or locally duplicated conversion code.
+   - Search `gpu_data_types.hpp` before introducing a new packed type, alias, bit representation, clipping operation, or numeric conversion.
+   - If a required reusable type or conversion helper is missing, the recommended fix should add a documented CUDA/HIP-compatible definition or function to `include/mscclpp/gpu_data_types.hpp` instead of adding an ad hoc helper to a kernel file.
+   - Keep unavoidable backend-native types and intrinsics isolated at hardware or ABI boundaries, and convert to MSCCL++ types at the boundary.
+
+7. **Portability and maintainability**
+   - Check CUDA/HIP portability and architecture guards when the code is not intentionally backend-specific.
+   - Prefer existing MSCCL++ helpers and established patterns when they preserve the required low-level semantics.
+   - Flag confusing naming or duplication only when it increases correctness or maintenance risk.
+
+## Review Method
+
+1. Establish the kernel's purpose, launch configuration, input/output layout, and synchronization contract.
+2. Inspect call sites, configuration code, related helpers, and producer/consumer kernels before drawing conclusions.
+3. Trace representative execution paths, including empty work, partial warps, uneven expert/token distributions, local peers, remote peers, and buffer rollover.
+4. For each suspected issue, identify the exact failing interleaving, input, architecture, or configuration.
+5. Check whether existing barriers, atomics, or API guarantees already prevent the issue.
+6. Report only findings that are actionable and supported by the code.
+
+## Avoid Incorrect Review Heuristics
+
+- Do not reject raw kernels or device intrinsics merely because a higher-level API exists.
+- Do not flag a single-thread control loop when its cost is negligible or parallelization would add synchronization.
+- Do not recommend removing a barrier without proving that all dependent reads and writes remain ordered.
+- Do not assume CUDA behavior applies to HIP, or vice versa.
+- Do not enforce naming or style preferences that are not established by the repository.
+- Do not merge structures solely to reduce the type count; preserve types that enforce a meaningful invariant or boundary.
+- Do not replace a raw backend type when it is required by an intrinsic, assembly constraint, hardware format, or external ABI; isolate and document that boundary instead.
+- Do not produce speculative performance claims without explaining the bottleneck.
+
+## Output Format
+
+List findings in descending severity. For each finding provide:
+
+- `severity — file:line — concise title`
+- The concrete correctness or performance impact.
+- The execution scenario that triggers it.
+- A specific fix or safer design direction.
+
+If there are no high-confidence findings, say so directly and note only material residual risks or validation gaps. Do not create empty category sections, repeat the code, or include style-only noise.
+
+## Clarification
+
+Ask for clarification only when correctness depends on an undocumented external contract, target architecture, transport capability, or performance requirement that cannot be established from the repository.

.github/workflows/lint.yml

@@ -7,7 +7,7 @@ on:
 
 jobs:
   cpplint:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
 
     steps:
     - name: Check out Git repository
@@ -22,7 +22,7 @@ jobs:
       run: bash ./tools/lint.sh cpp dry
 
   pylint:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
 
     steps:
     - name: Check out Git repository
@@ -40,7 +40,7 @@ jobs:
       run: bash ./tools/lint.sh py dry
 
   spelling:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
 
     steps:
     - name: Check out Git repository

.github/workflows/mscclpp-lang.yml

@@ -21,7 +21,7 @@ jobs:
     - uses: actions/checkout@v4
 
     - name: Set environment variable
-      run: echo "LD_LIBRARY_PATH=/usr/local/cuda/compat:/usr/local/cuda/lib64" >> $GITHUB_ENV
+      run: echo "LD_LIBRARY_PATH=/usr/local/cuda/lib64" >> $GITHUB_ENV
 
     - name: Install mscclpp
       run: |

CMakeLists.txt

@@ -51,6 +51,7 @@ option(MSCCLPP_BUILD_TESTS "Build tests" OFF)
 option(MSCCLPP_BUILD_PYTHON_BINDINGS "Build Python bindings" ON)
 option(MSCCLPP_BUILD_EXT_NCCL "Build NCCL interfaces" ON)
 option(MSCCLPP_BUILD_EXT_COLLECTIVES "Build collective algorithms" ON)
+option(MSCCLPP_BUILD_EXT_EP "Build Expert-Parallel (MoE dispatch/combine) extension" ON)
 option(MSCCLPP_USE_CUDA "Use NVIDIA/CUDA." OFF)
 option(MSCCLPP_USE_ROCM "Use AMD/ROCm." OFF)
 option(MSCCLPP_USE_IB "Use InfiniBand." ON)
@@ -176,7 +177,7 @@ if(MSCCLPP_GPU_ARCHS)
     endif()
 elseif(MSCCLPP_USE_CUDA)
     if(NVIDIA_FOUND)
-        set(MSCCLPP_GPU_ARCHS "native")
+        set(MSCCLPP_GPU_ARCHS ${NVIDIA_GPU_ARCHS})
     else()
 	    if(CUDAToolkit_VERSION VERSION_LESS "12.0")
 	        message(FATAL_ERROR "CUDA 12.0 or higher required (C++20 build), found ${CUDAToolkit_VERSION}")

cmake/CheckNvidiaGpu.cmake

@@ -26,10 +26,15 @@ enable_language(CUDA)
 
 set(CHECK_SRC "${CMAKE_CURRENT_SOURCE_DIR}/cmake/check_nvidia_gpu.cu")
 
-try_run(RUN_RESULT COMPILE_SUCCESS SOURCES ${CHECK_SRC})
+try_run(RUN_RESULT COMPILE_SUCCESS SOURCES ${CHECK_SRC}
+        RUN_OUTPUT_VARIABLE NVIDIA_GPU_ARCHS)
 
 if(COMPILE_SUCCESS AND RUN_RESULT EQUAL 0)
+    string(STRIP "${NVIDIA_GPU_ARCHS}" NVIDIA_GPU_ARCHS)
+    list(REMOVE_DUPLICATES NVIDIA_GPU_ARCHS)
     set(NVIDIA_FOUND "TRUE")
+    message(STATUS "Detected NVIDIA GPU architectures: ${NVIDIA_GPU_ARCHS}")
 else()
     unset(CMAKE_CUDA_ARCHITECTURES)
+    unset(NVIDIA_GPU_ARCHS)
 endif()

cmake/check_nvidia_gpu.cu

@@ -3,13 +3,25 @@
 
 #include <cuda_runtime.h>
 
-__global__ void kernel() {}
+#include <cstdio>
 
 int main() {
   int cnt;
   cudaError_t err = cudaGetDeviceCount(&cnt);
   if (err != cudaSuccess || cnt == 0) {
     return 1;
   }
+
+  for (int device = 0; device < cnt; ++device) {
+    cudaDeviceProp properties;
+    err = cudaGetDeviceProperties(&properties, device);
+    if (err != cudaSuccess) {
+      return 1;
+    }
+    if (device != 0) {
+      std::printf(";");
+    }
+    std::printf("%d", properties.major * 10 + properties.minor);
+  }
   return 0;
 }

include/mscclpp/core.hpp

@@ -649,6 +649,13 @@ class Connection {
   /// @param newValue The new value to write.
   void updateAndSync(RegisteredMemory dst, uint64_t dstOffset, uint64_t* src, uint64_t newValue);
 
+  /// Atomically add a value to a 64-bit integer in a destination RegisteredMemory.
+  ///
+  /// @param dst The destination RegisteredMemory.
+  /// @param dstOffset The offset in bytes from the start of the destination RegisteredMemory.
+  /// @param value The 64-bit signed value to atomically add.
+  void atomicAdd(RegisteredMemory dst, uint64_t dstOffset, int64_t value);
+
   /// Flush any pending writes to the remote process.
   /// @param timeoutUsec Timeout in microseconds. Default: -1 (no timeout)
   void flush(int64_t timeoutUsec = -1);

include/mscclpp/fifo_device.hpp

@@ -19,6 +19,7 @@ using TriggerType = uint64_t;
 constexpr TriggerType TriggerData = 0x1;  // Trigger a data transfer.
 constexpr TriggerType TriggerFlag = 0x2;  // Trigger a signaling.
 constexpr TriggerType TriggerSync = 0x4;  // Trigger a flush.
+// type == 0 is reserved for atomic add operations.
 
 constexpr unsigned int TriggerBitsSize = 32;
 constexpr unsigned int TriggerBitsOffset = 32;

include/mscclpp/gpu.hpp

@@ -21,6 +21,8 @@ using cudaIpcMemHandle_t = hipIpcMemHandle_t;
 
 using CUresult = hipError_t;
 using CUdeviceptr = hipDeviceptr_t;
+using CUcontext = hipCtx_t;
+using CUdevice = hipDevice_t;
 using CUmemGenericAllocationHandle = hipMemGenericAllocationHandle_t;
 using CUmemAllocationProp = hipMemAllocationProp;
 using CUmemAccessDesc = hipMemAccessDesc;
@@ -115,6 +117,11 @@ constexpr auto CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL = HIP_POINTER_ATTRIBUTE_DEVIC
 #define cudaIpcCloseMemHandle(...) hipIpcCloseMemHandle(__VA_ARGS__)
 
 #define cuGetErrorString(...) hipDrvGetErrorString(__VA_ARGS__)
+#define cuDeviceGet(...) hipDeviceGet(__VA_ARGS__)
+#define cuCtxCreate(...) hipCtxCreate(__VA_ARGS__)
+#define cuCtxDestroy(...) hipCtxDestroy(__VA_ARGS__)
+#define cuCtxPushCurrent(...) hipCtxPushCurrent(__VA_ARGS__)
+#define cuCtxPopCurrent(...) hipCtxPopCurrent(__VA_ARGS__)
 #define cuMemAddressReserve(...) hipMemAddressReserve(__VA_ARGS__)
 #define cuMemAddressFree(...) hipMemAddressFree(__VA_ARGS__)
 #define cuMemGetAddressRange(...) hipMemGetAddressRange(__VA_ARGS__)

include/mscclpp/gpu_data_types.hpp

@@ -4,6 +4,7 @@
 #ifndef MSCCLPP_GPU_DATA_TYPES_HPP_
 #define MSCCLPP_GPU_DATA_TYPES_HPP_
 
+#include <cstdint>
 #include <mscclpp/device.hpp>
 
 #if defined(MSCCLPP_DEVICE_HIP)
@@ -187,7 +188,7 @@ template <int Bytes, bool Enabled = (Bytes >= 4 && Bytes % 4 == 0)>
 struct alignas(Bytes) Words {
   uint32_t w[Bytes / 4];
 
-  MSCCLPP_HOST_DEVICE_INLINE Words() {}
+  Words() = default;
 
   MSCCLPP_HOST_DEVICE_INLINE uint32_t& operator[](int i) { return w[i]; }
 
@@ -210,7 +211,7 @@ union alignas(sizeof(T) * N) VectorTypeImpl {
   using ElementType = T;
   constexpr static int Size = N;
 
-  MSCCLPP_HOST_DEVICE_INLINE VectorTypeImpl() {}
+  VectorTypeImpl() = default;
 
   MSCCLPP_HOST_DEVICE_INLINE VectorTypeImpl(const StorageT& value) : storage(value) {}
 
@@ -686,6 +687,18 @@ MSCCLPP_DEVICE_INLINE To to(const From& v) {
   }
 }
 
+/// Convert a packed BF16 pair to a packed FP32 pair.
+template <>
+MSCCLPP_DEVICE_INLINE f32x2 to<f32x2, bf16x2>(const bf16x2& v) {
+  return __bfloat1622float2(v.storage);
+}
+
+/// Convert a packed FP32 pair to a packed BF16 pair using round-to-nearest.
+template <>
+MSCCLPP_DEVICE_INLINE bf16x2 to<bf16x2, f32x2>(const f32x2& v) {
+  return __float22bfloat162_rn(v.storage);
+}
+
 #if defined(__FP8_TYPES_EXIST__)
 template <>
 MSCCLPP_DEVICE_INLINE __fp8_e4m3 min(const __fp8_e4m3& a, const __fp8_e4m3& b) {
@@ -866,7 +879,7 @@ MSCCLPP_DEVICE_INLINE f8_e4m3x2 to<f8_e4m3x2, f32x2>(const f32x2& v) {
   uint32_t packed = __builtin_amdgcn_cvt_pk_fp8_f32(v.data[0], v.data[1], 0, false);
   return bit_cast<f8_e4m3x2>(static_cast<__hip_fp8x2_storage_t>(packed));
 #elif defined(MSCCLPP_DEVICE_CUDA)
-  __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2(make_float2(v.data[0], v.data[1]), __NV_SATFINITE, __NV_E4M3);
+  __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2(v.storage, __NV_SATFINITE, __NV_E4M3);
   return bit_cast<f8_e4m3x2>(fp8x2);
 #else
   f8_e4m3x2 result;

include/mscclpp/port_channel_device.hpp

@@ -126,6 +126,24 @@ struct BasePortChannelDeviceHandle {
     detail::waitFlush(flushDonePos_, pos, maxSpinCount);
   }
 
+  /// Push an atomic add trigger to the FIFO to perform a remote atomic add on a 64-bit value.
+  /// Uses type == 0 to indicate an atomic add operation.
+  /// @param dstId The ID of destination memory region.
+  /// @param dstOffset The offset into the destination memory region.
+  /// @param value The 64-bit signed value to atomically add.
+  MSCCLPP_DEVICE_INLINE void atomicAdd(MemoryId dstId, uint64_t dstOffset, int64_t value) {
+    ProxyTrigger trigger;
+    // Encode the full 64-bit add value in fst (size + srcOffset fields).
+    trigger.fst = static_cast<uint64_t>(value);
+    // Build snd with dstOffset, dstMemoryId, type=0 (atomic add), semaphoreId.
+    trigger.snd = 0;
+    trigger.fields.dstOffset = dstOffset;
+    trigger.fields.dstMemoryId = dstId;
+    trigger.fields.type = 0;
+    trigger.fields.semaphoreId = semaphoreId_;
+    fifo_.push(trigger);
+  }
+
   /// Check if the port channel has been signaled.
   /// @return true if the port channel has been signaled.
   MSCCLPP_DEVICE_INLINE bool poll() { return semaphore_.poll(); }
@@ -191,6 +209,13 @@ struct PortChannelDeviceHandle : public BasePortChannelDeviceHandle {
   MSCCLPP_DEVICE_INLINE void putWithSignalAndFlush(uint64_t offset, uint64_t size) {
     putWithSignalAndFlush(offset, offset, size);
   }
+
+  /// Push an atomic add trigger to the FIFO to perform a remote atomic add on a 64-bit value.
+  /// @param dstOffset The offset into the destination memory region.
+  /// @param value The 64-bit signed value to atomically add.
+  MSCCLPP_DEVICE_INLINE void atomicAdd(uint64_t dstOffset, int64_t value) {
+    BasePortChannelDeviceHandle::atomicAdd(dst_, dstOffset, value);
+  }
 #endif  // defined(MSCCLPP_DEVICE_COMPILE)
 };
 

pyproject.toml

@@ -72,6 +72,7 @@ include= ["python/mscclpp/_version.py"]
 [tool.scikit-build.cmake.define]
 MSCCLPP_BUILD_PYTHON_BINDINGS = "ON"
 MSCCLPP_BUILD_TESTS = "OFF"
+MSCCLPP_BUILD_EXT_EP = "ON"
 
 [tool.black]
 line-length = 120

python/csrc/core_py.cpp

@@ -63,20 +63,20 @@ void register_core(nb::module_& m) {
             void* data = reinterpret_cast<void*>(ptr);
             self->send(data, size, peer, tag);
           },
-          nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"))
+          nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>())
       .def(
           "recv",
           [](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) {
             void* data = reinterpret_cast<void*>(ptr);
             self->recv(data, size, peer, tag);
           },
-          nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"))
+          nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>())
       .def("all_gather", &Bootstrap::allGather, nb::arg("allData"), nb::arg("size"))
       .def("barrier", &Bootstrap::barrier)
       .def("send", static_cast<void (Bootstrap::*)(const std::vector<char>&, int, int)>(&Bootstrap::send),
-           nb::arg("data"), nb::arg("peer"), nb::arg("tag"))
+           nb::arg("data"), nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>())
       .def("recv", static_cast<void (Bootstrap::*)(std::vector<char>&, int, int)>(&Bootstrap::recv), nb::arg("data"),
-           nb::arg("peer"), nb::arg("tag"));
+           nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>());
 
   nb::class_<UniqueId>(m, "CppUniqueId")
       .def(nb::init<>())

python/mscclpp/_core/comm.py

@@ -20,7 +20,6 @@
     CppTransportFlags,
 )
 import numpy as np
-import pickle
 
 from mscclpp.utils import is_torch_tensor
 
@@ -54,14 +53,12 @@ def __init__(
                 import torch
                 import torch.distributed as dist
 
-                if rank == 0:
-                    uniq_id_global = uniq_id
-                    pickled_data = pickle.dumps(uniq_id)
-                    data_tensor = torch.frombuffer(bytearray(pickled_data), dtype=torch.uint8).clone()
-                else:
-                    data_tensor = torch.zeros(256, dtype=torch.uint8)
-                dist.broadcast(data_tensor, group=torch_group, group_src=0)
-                uniq_id_global = pickle.loads(data_tensor.numpy().tobytes())
+                backend = str(dist.get_backend(torch_group)).lower()
+                device = torch.device("cuda", torch.cuda.current_device()) if "nccl" in backend else torch.device("cpu")
+                object_list = [uniq_id]
+                group_root = dist.get_global_rank(torch_group, 0)
+                dist.broadcast_object_list(object_list, src=group_root, group=torch_group, device=device)
+                uniq_id_global = object_list[0]
             self.bootstrap.initialize(uniq_id_global)
         elif not interfaceIpPortTrio == "":
             assert rank >= 0 and size >= 1