Repository: microsoft/mscclpp
Pull request #505: Fix performance issue introduced in PR: 499
Author: login: Binyang2014
Created: 2025-04-18T02:43:41Z | Merged: 2025-04-22T21:03:37Z

1. use `fence+relaxed` to replace `release` for fifo. `fence+relax` is more efficient on A100
2. Update the deviceSyncer. Previous one cannot handle threadBlock number change correctly. Use three counters to solve this issue. Reset previous counter before sync on current counter.
3. Introduce relaxedWait which can be used with relaxedSignal for case doesn't need guarantee the memory visibility


Diff:
diff --git a/include/mscclpp/concurrency_device.hpp b/include/mscclpp/concurrency_device.hpp
index e0ca9d94b..e1c055a35 100644
--- a/include/mscclpp/concurrency_device.hpp
+++ b/include/mscclpp/concurrency_device.hpp
@@ -7,6 +7,8 @@
 #include "atomic_device.hpp"
 #include "poll_device.hpp"
 
+#define NUM_DEVICE_SYNCER_COUNTER 3
+
 namespace mscclpp {
 
 /// A device-wide barrier.
@@ -24,15 +26,17 @@ struct DeviceSyncer {
   /// @param blockNum The number of blocks that will synchronize.
   /// @param maxSpinCount The maximum number of spin counts before asserting. Never assert if negative.
   MSCCLPP_DEVICE_INLINE void sync(int blockNum, int64_t maxSpinCount = 100000000) {
-    int targetCnt = blockNum;
+    unsigned int targetCnt = blockNum;
     __syncthreads();
     if (blockNum == 1) return;
     if (threadIdx.x == 0) {
-      unsigned int tmp = preFlag_ ^ 1;
-      int val = (tmp << 1) - 1;
-      targetCnt = val == 1 ? targetCnt : 0;
-      atomicFetchAdd<int, scopeDevice>(&count_, val, memoryOrderRelease);
-      POLL_MAYBE_JAILBREAK((atomicLoad<int, scopeDevice>(&count_, memoryOrderAcquire) != targetCnt), maxSpinCount);
+      unsigned int tmp = (preFlag_ + 1) % NUM_DEVICE_SYNCER_COUNTER;
+      unsigned int next = (tmp + 1) % NUM_DEVICE_SYNCER_COUNTER;
+      unsigned int* count = &count_[tmp];
+      count_[next] = 0;
+      atomicFetchAdd<unsigned int, scopeDevice>(count, 1U, memoryOrderRelease);
+      POLL_MAYBE_JAILBREAK((atomicLoad<unsigned int, scopeDevice>(count, memoryOrderAcquire) != targetCnt),
+                           maxSpinCount);
       preFlag_ = tmp;
     }
     // We need sync here because only a single thread is checking whether
@@ -43,7 +47,7 @@ struct DeviceSyncer {
 
  private:
   /// The counter of synchronized blocks.
-  int count_;
+  unsigned int count_[NUM_DEVICE_SYNCER_COUNTER];
   /// The flag to indicate whether to increase or decrease @ref flag_.
   unsigned int preFlag_;
 };
diff --git a/include/mscclpp/fifo_device.hpp b/include/mscclpp/fifo_device.hpp
index 03ca63395..f431b1d8e 100644
--- a/include/mscclpp/fifo_device.hpp
+++ b/include/mscclpp/fifo_device.hpp
@@ -67,7 +67,13 @@ struct FifoDeviceHandle {
 
     // Make sure the data is visible to the host before we update the tail.
 #if defined(MSCCLPP_DEVICE_CUDA)
+#if __CUDA_ARCH__ == 800
+    // For A100, threadfence_system is more efficient than release
+    __threadfence_system();
+    asm volatile("st.global.relaxed.sys.v2.u64 [%0], {%1,%2};" ::"l"(triggerPtr), "l"(trigger.fst), "l"(trigger.snd));
+#else
     asm volatile("st.global.release.sys.v2.u64 [%0], {%1,%2};" ::"l"(triggerPtr), "l"(trigger.fst), "l"(trigger.snd));
+#endif
 #else   // !defined(MSCCLPP_DEVICE_CUDA)
     // store snd no later than fst.
     atomicStore(&(triggerPtr->snd), trigger.snd, memoryOrderRelaxed);
diff --git a/include/mscclpp/memory_channel_device.hpp b/include/mscclpp/memory_channel_device.hpp
index d49eb4def..1cf6a853e 100644
--- a/include/mscclpp/memory_channel_device.hpp
+++ b/include/mscclpp/memory_channel_device.hpp
@@ -273,6 +273,14 @@ struct MemoryChannelDeviceHandle {
   /// Wait for the remote semaphore to send a signal.
   /// @param maxSpinCount The maximum number of spins before asserting. Never assert if negative.
   MSCCLPP_DEVICE_INLINE void wait(int64_t maxSpinCount = 10000000) { semaphore_.wait(maxSpinCount); }
+
+  /// Wait for the remote semaphore to send a signal.
+  ///
+  /// This function is a relaxed version of signal() and provides no guarantee on the completion of memory operations.
+  /// User requires to call proper fencing before using this function.
+  ///
+  /// @param maxSpinCount The maximum number of spins before asserting. Never assert if negative.
+  MSCCLPP_DEVICE_INLINE void relaxedWait() { semaphore_.relaxedWait(); }
 #endif  // defined(MSCCLPP_DEVICE_COMPILE)
 };
 
diff --git a/include/mscclpp/semaphore_device.hpp b/include/mscclpp/semaphore_device.hpp
index 31dff14d1..88ca98823 100644
--- a/include/mscclpp/semaphore_device.hpp
+++ b/include/mscclpp/semaphore_device.hpp
@@ -54,6 +54,17 @@ struct MemoryDevice2DeviceSemaphoreDeviceHandle {
     POLL_MAYBE_JAILBREAK((atomicLoad(inboundSemaphoreId, memoryOrderAcquire) < flag), maxSpinCount);
   }
 
+  /// Wait for the remote device to signal.
+  ///
+  /// This function is a relaxed version of Wait() and provides no guarantee on the completion of memory operations.
+  /// User requires to call proper fencing before using this function.
+  ///
+  MSCCLPP_DEVICE_INLINE void relaxedWait(int64_t maxSpinCount = 100000000) {
+    (*expectedInboundSemaphoreId) += 1;
+    uint64_t flag = (*expectedInboundSemaphoreId);
+    POLL_MAYBE_JAILBREAK((atomicLoad(inboundSemaphoreId, memoryOrderRelaxed) < flag), maxSpinCount);
+  }
+
   /// Signal the remote device.
   ///
   /// This function guarantees that all the memory operation before this function is completed before the remote
@@ -65,7 +76,11 @@ struct MemoryDevice2DeviceSemaphoreDeviceHandle {
     semaphoreIncrement();
     // use memoryOrderSeqCst instead of memoryOrderRelease since memoryOrderSeqCst
     // is more efficient on A100.
+#if __CUDA_ARCH__ == 800
     atomicStore(remoteInboundSemaphoreId, semaphoreGetLocal(), memoryOrderSeqCst);
+#else
+    atomicStore(remoteInboundSemaphoreId, semaphoreGetLocal(), memoryOrderRelease);
+#endif
   }
 
   /// Signal the remote device.
diff --git a/python/mscclpp_benchmark/allreduce.cu b/python/mscclpp_benchmark/allreduce.cu
index dbe376a3a..ce124da5d 100644
--- a/python/mscclpp_benchmark/allreduce.cu
+++ b/python/mscclpp_benchmark/allreduce.cu
@@ -134,16 +134,11 @@ __device__ void allreduce1_helper(mscclpp::MemoryChannelDeviceHandle* memChans,
   int4* buff4 = (int4*)buff;
   const int tid = threadIdx.x + blockIdx.x * blockDim.x;
 
-  // synchronize everyone
-  if (tid == 0) {
-    __threadfence_system();
-  }
-  __syncthreads();
   if (tid < nPeer) {
     memChans[tid].relaxedSignal();
   }
   if (tid >= nPeer && tid < nPeer * 2) {
-    memChans[tid - nPeer].wait();
+    memChans[tid - nPeer].relaxedWait();
   }
   deviceSyncer.sync(gridDim.x);
 
@@ -193,15 +188,11 @@ __device__ void allreduce1_helper(mscclpp::MemoryChannelDeviceHandle* memChans,
 
   // synchronize everyone again
   deviceSyncer.sync(gridDim.x);
-  if (tid == 0) {
-    __threadfence_system();
-  }
-  __syncthreads();
   if (tid < nPeer) {
     memChans[tid].relaxedSignal();
   }
   if (tid >= nPeer && tid < nPeer * 2) {
-    memChans[tid - nPeer].wait();
+    memChans[tid - nPeer].relaxedWait();
   }
 
   if (READ_ONLY) {
@@ -437,7 +428,7 @@ __device__ void localReduceScatterMem(mscclpp::MemoryChannelDeviceHandle* memCha
     memChans[peerIdx].relaxedSignal();
   }
   for (int peerIdx = threadIdx.x + blockIdx.x * blockDim.x; peerIdx < nPeer; peerIdx += blockDim.x * nBlocks) {
-    memChans[peerIdx].wait();
+    memChans[peerIdx].relaxedWait();
   }
   reduceScatterDeviceSyncer.sync(nBlocks);
 
@@ -497,7 +488,7 @@ __device__ void localAllGatherMem(mscclpp::MemoryChannelDeviceHandle* memChans,
   }
   if (threadIdx.x == 0 && peerLocalBlockIdx == 0) {
     memChans[peerIdx].relaxedSignal();
-    memChans[peerIdx].wait();
+    memChans[peerIdx].relaxedWait();
   }
   allGatherDeviceSyncer.sync(nBlocks);
   size_t offset = rankChunkSize * (startRankChunkIndex + remoteRankLocalIndex) + offsetInRankChunk;
diff --git a/test/deploy/perf_ndmv4.jsonl b/test/deploy/perf_ndmv4.jsonl
index 757998851..ed8bc17a4 100644
--- a/test/deploy/perf_ndmv4.jsonl
+++ b/test/deploy/perf_ndmv4.jsonl
@@ -3,7 +3,7 @@
 {"name":"allgather", "kernel":3, "ranks":8, "ranksPerNode":8, "algBw":0.1112, "busBw":0.0973, "size":8192,       "time":73.63,   "target":"latency"}
 {"name":"allreduce", "kernel":1, "ranks":8, "ranksPerNode":8, "algBw":139.41, "busBw":243.96, "size":1073741824, "time":7701.98, "target":"throughput"}
 {"name":"allreduce", "kernel":2, "ranks":8, "ranksPerNode":8, "algBw":1.25,   "busBw":2.19,   "size":8192,       "time":6.51,    "target":"latency"}
-{"name":"allreduce", "kernel":2, "ranks":16,"ranksPerNode":8, "algBw":0.51,   "busBw":0.96,   "size":8192,       "time":15.96,   "target":"latency"}
+{"name":"allreduce", "kernel":2, "ranks":16,"ranksPerNode":8, "algBw":0.44,   "busBw":0.83,   "size":8192,       "time":18.42,   "target":"latency"}
 {"name":"allreduce", "kernel":3, "ranks":8, "ranksPerNode":8, "algBw":139.08, "busBw":243.40, "size":1073741824, "time":7719.85, "target":"throughput"}
 {"name":"allreduce", "kernel":4, "ranks":8, "ranksPerNode":8, "algBw":106.98, "busBw":187.22, "size":16777216,   "time":156.81,  "target":"throughput"}
 {"name":"allreduce", "kernel":4, "ranks":8, "ranksPerNode":8, "algBw":116.24, "busBw":203.42, "size":33554432,   "time":288.65,  "target":"throughput"}


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR addresses a performance issue by replacing the release fence with a more efficient fence+relaxed approach for A100 and improves the device synchronization mechanism to properly handle thread block number changes.  
- Added a relaxed wait function in semaphore_device.hpp.  
- Introduced a relaxed wait in memory_channel_device.hpp that delegates to the semaphore’s relaxed wait.  
- Updated fifo_device.hpp to use threadfence_system with a relaxed store for better performance on A100.  
- Revised DeviceSyncer in concurrency_device.hpp to use three counters for handling thread block number changes.

### Reviewed Changes

Copilot reviewed 4 out of 5 changed files in this pull request and generated no comments.

| File                                   | Description                                                                              |
| -------------------------------------- | ---------------------------------------------------------------------------------------- |
| include/mscclpp/semaphore_device.hpp   | Added a relaxedWait function with relaxed memory ordering.                              |
| include/mscclpp/memory_channel_device.hpp | Added a new relaxedWait function that forwards the call to the semaphore implementation. |
| include/mscclpp/fifo_device.hpp         | Replaced the release store with __threadfence_system() call and a relaxed store for A100.   |
| include/mscclpp/concurrency_device.hpp  | Modified DeviceSyncer to use an array of unsigned counters and adjusted synchronization logic to handle thread block changes. |

<details>
<summary>Files not reviewed (1)</summary>

* **python/mscclpp_benchmark/allreduce.cu**: Language not supported
</details>



state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-18T02:47:28Z
user: login: liangyuRain
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-18T06:52:30Z
user: login: Binyang2014
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-18T16:25:57Z
user: login: chhwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-20T23:18:14Z
user: login: chhwang
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-20T23:19:10Z
user: login: chhwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-21T20:23:01Z
user: login: chhwang
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T21:03:26Z

Review comments:
path: include/mscclpp/fifo_device.hpp
user: login: liangyuRain
body: Surprised to learn this, given threadfence is offering the stronger sequential consistency than release. Maybe also consider using `asm volatile("fence.acq_rel.gpu;":: : "memory");`? PTX documentation [link](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar-fence) says
> `fence.acq_rel` is a light-weight fence that is sufficient for memory synchronization in most programs.
> ...
> `fence.sc` is a slower fence that can restore sequential consistency when used in sufficient places, at the cost of performance.
created_at: 2025-04-18T06:52:30Z
author_association: NONE
path: include/mscclpp/fifo_device.hpp
user: login: Binyang2014
body: I think it's a bug on A100. H100 doesn't have such issue, they provide similar performance.
created_at: 2025-04-18T16:25:57Z
author_association: CONTRIBUTOR
path: include/mscclpp/fifo_device.hpp
user: login: chhwang
body: I would prefer making the current version specific to older CUDA architectures and keep the `release` version as the default.
created_at: 2025-04-20T23:18:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -28,11 +30,13 @@ struct DeviceSyncer {
     __syncthreads();
     if (blockNum == 1) return;
     if (threadIdx.x == 0) {
-      unsigned int tmp = preFlag_ ^ 1;
-      int val = (tmp << 1) - 1;
-      targetCnt = val == 1 ? targetCnt : 0;
-      atomicFetchAdd<int, scopeDevice>(&count_, val, memoryOrderRelease);
-      POLL_MAYBE_JAILBREAK((atomicLoad<int, scopeDevice>(&count_, memoryOrderAcquire) != targetCnt), maxSpinCount);
+      unsigned int tmp = (preFlag_ + 1) % NUM_DEVICE_SYNCER_COUNTER;
+      unsigned int next = (tmp + 1) % NUM_DEVICE_SYNCER_COUNTER;
+      unsigned int* count = &count_[tmp];
+      count_[next] = 0;
+      atomicFetchAdd<unsigned int, scopeDevice>(count, 1U, memoryOrderRelease);
+      POLL_MAYBE_JAILBREAK((atomicLoad<unsigned int, scopeDevice>(count, memoryOrderAcquire) != targetCnt),
path: include/mscclpp/concurrency_device.hpp
user: login: chhwang
body: Please change `targetCnt` data type to `unsigned int`
created_at: 2025-04-21T20:23:01Z
author_association: CONTRIBUTOR

PR comments:
user: login: Binyang2014
created_at: 2025-04-18T02:49:04Z
body: @liangyuRain pls help to review

author_association: CONTRIBUTOR
user: login: liangyuRain
created_at: 2025-04-18T06:54:47Z
body: Left a comment. Otherwise looks good to me.
author_association: NONE
user: login: Binyang2014
created_at: 2025-04-22T04:29:19Z
body: /azp run
author_association: CONTRIBUTOR
user: login: azure-pipelines[bot]
created_at: 2025-04-22T04:29:36Z
body: <samp>
Azure Pipelines successfully started running 3 pipeline(s).<br>

</samp>
author_association: NONE