From c1071318c84f968bc292e2ef9b8296ba837d06af Mon Sep 17 00:00:00 2001
From: Caio Rocha <164253795+caiomcbr@users.noreply.github.com>
Date: Tue, 19 May 2026 13:06:53 -0700
Subject: [PATCH 01/16] Include a static synchronization check in the DSL.
 (#806)

---
 python/mscclpp/language/channel.py |  6 ++++++
 python/mscclpp/language/program.py | 32 ++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/python/mscclpp/language/channel.py b/python/mscclpp/language/channel.py
index 23d76eda..de0f65c5 100644
--- a/python/mscclpp/language/channel.py
+++ b/python/mscclpp/language/channel.py
@@ -78,6 +78,7 @@ class MemoryChannel:
         tb_channel_ids = get_program().setup_channel(tb, self)
         op = SignalOperation(tb_channel_ids, self.channel_type, data_sync, relaxed)
         get_program().add_operation(self.src_rank, tb, op)
+        get_program().register_signal(self.src_rank, self.dst_rank, self.channel_type)
 
     def wait(self, tb: int, data_sync: SyncType = SyncType.both, relaxed: bool = False):
         """Wait for a signal through the memory channel.
@@ -99,6 +100,7 @@ class MemoryChannel:
         tb_channel_ids = get_program().setup_channel(tb, self)
         op = WaitOperation(tb_channel_ids, self.channel_type, data_sync, relaxed)
         get_program().add_operation(self.src_rank, tb, op)
+        get_program().register_wait(self.src_rank, self.dst_rank, self.channel_type)
 
     def get(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int = None, tb_group: ThreadBlockGroup = None):
         """Retrieve data from remote memory to local memory.
@@ -508,6 +510,7 @@ class PortChannel:
         tb_channel_ids = get_program().setup_channel(tb, self)
         op = SignalOperation(tb_channel_ids, self.channel_type, data_sync)
         get_program().add_operation(self.src_rank, tb, op)
+        get_program().register_signal(self.src_rank, self.dst_rank, self.channel_type)
 
     def wait(self, tb: int, data_sync: SyncType = SyncType.both):
         """Wait for a signal through the port channel.
@@ -527,6 +530,7 @@ class PortChannel:
         tb_channel_ids = get_program().setup_channel(tb, self)
         op = WaitOperation(tb_channel_ids, self.channel_type, data_sync)
         get_program().add_operation(self.src_rank, tb, op)
+        get_program().register_wait(self.src_rank, self.dst_rank, self.channel_type)
 
     def flush(self, tb: int, data_sync: SyncType = SyncType.both):
         """Flush pending operations through the port channel.
@@ -636,6 +640,7 @@ class PortChannel:
         )
 
         get_program().add_operation(self.src_rank, tb, op)
+        get_program().register_signal(self.src_rank, self.dst_rank, self.channel_type)
 
     def put_with_signal_and_flush(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int):
         """Send data from local memory to remote memory with signal and flush.
@@ -681,6 +686,7 @@ class PortChannel:
         )
 
         get_program().add_operation(self.src_rank, tb, op)
+        get_program().register_signal(self.src_rank, self.dst_rank, self.channel_type)
 
     def put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int):
         """Transfer data from local buffer to remote scratch buffer in packet format.
diff --git a/python/mscclpp/language/program.py b/python/mscclpp/language/program.py
index c29e9ab7..825a9d40 100644
--- a/python/mscclpp/language/program.py
+++ b/python/mscclpp/language/program.py
@@ -10,6 +10,7 @@ from mscclpp.language.rank import Semaphore
 from mscclpp.language.collectives import *
 from mscclpp.language.utils import AlgoSpec, ReplicationPolicy
 from typing import List
+from collections import defaultdict
 import json
 
 
@@ -112,6 +113,9 @@ class CollectiveProgram:
 
         self.loop_context = None
 
+        self._signal_counts = defaultdict(int)
+        self._wait_counts = defaultdict(int)
+
     @classmethod
     def from_spec(cls, spec: AlgoSpec):
         """Initialize a new CollectiveProgram from an algorithm specification.
@@ -206,7 +210,35 @@ class CollectiveProgram:
         else:
             self.gpus[rank].add_operation(tb, operation)
 
+    def register_signal(self, src_rank, dst_rank, channel_type):
+        """Record that `src_rank` issued a signal targeting `dst_rank` over `channel_type`."""
+        self._signal_counts[(src_rank, dst_rank, channel_type)] += 1
+
+    def register_wait(self, src_rank, dst_rank, channel_type):
+        """Record that `src_rank` performed a wait for `dst_rank` over `channel_type`."""
+        self._wait_counts[(src_rank, dst_rank, channel_type)] += 1
+
+    def validate_signal_wait_pairing(self):
+        """Validate that every signal issued by a rank is matched by a wait on the peer rank.
+
+        For each (src_rank, dst_rank, channel_type) triple, the number of signals sent by
+        `src_rank` to `dst_rank` must equal the number of waits performed by `dst_rank`
+        for `src_rank` on a channel of the same type. Raises RuntimeError on mismatch.
+        """
+        keys = set(self._signal_counts.keys()) | {(dst, src, t) for (src, dst, t) in self._wait_counts.keys()}
+        for src_rank, dst_rank, channel_type in keys:
+            signals = self._signal_counts.get((src_rank, dst_rank, channel_type), 0)
+            waits = self._wait_counts.get((dst_rank, src_rank, channel_type), 0)
+            if signals != waits:
+                raise RuntimeError(
+                    f"Signal/Wait mismatch on {channel_type}: rank {src_rank} issues {signals} "
+                    f"signal(s) to rank {dst_rank}, but rank {dst_rank} performs {waits} wait(s) "
+                    f"for rank {src_rank}. Every signal must be matched by a corresponding wait "
+                    f"on the peer rank over a channel of the same type."
+                )
+
     def post_process_operations(self):
+        self.validate_signal_wait_pairing()
         for gpu in self.gpus:
             if self.instr_fusion:
                 gpu.optimize_operations()

From 72621e72216c15ac9d636e00704dadc212e110aa Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Wed, 20 May 2026 09:29:55 -0700
Subject: [PATCH 02/16] add nBlocks check for allreduce_allpair_packet algo
 (#807)

- Fix the correctness issue for allreduce_allpair_packet algo. Make sure
no overwrite for input buffer. Use same tb for send/reduce/write-back.
- Check if nBlocks/nthreads validate for packet algorithm.
- Add more logs
- Modify flag update logic, make it work for the case: nthreadPerNBlock
< nflags

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../allreduce/allreduce_allpair_packet.cu     | 60 +++++++++++--------
 .../collectives/allreduce/allreduce_packet.cu |  5 ++
 .../allreduce/allreduce_allpair_packet.hpp    |  2 +-
 3 files changed, 42 insertions(+), 25 deletions(-)

diff --git a/src/ext/collectives/allreduce/allreduce_allpair_packet.cu b/src/ext/collectives/allreduce/allreduce_allpair_packet.cu
index 17bcfc33..faef5459 100644
--- a/src/ext/collectives/allreduce/allreduce_allpair_packet.cu
+++ b/src/ext/collectives/allreduce/allreduce_allpair_packet.cu
@@ -7,7 +7,7 @@
 #include "allreduce/allreduce_allpair_packet.hpp"
 #include "allreduce/common.hpp"
 #include "collective_utils.hpp"
-#include "debug.h"
+#include "logger.hpp"
 
 namespace mscclpp {
 namespace collective {
@@ -27,22 +27,30 @@ __global__ void allreduceAllPairs(T* buff, T* scratch, T* resultBuff, DeviceHand
   size_t scratchBaseOffset = (flag % numScratchBuff) ? (scratchBufferSize / numScratchBuff) : 0;
   size_t channelScratchOffset = scratchBaseOffset;
 
-  const int nBlocksPerPeer = gridDim.x / nPeers;
-  const int localBlockIdx = blockIdx.x % nBlocksPerPeer;
-  const int tid = threadIdx.x + localBlockIdx * blockDim.x;
-  const int peerIdx = blockIdx.x / nBlocksPerPeer;
-  size_t srcOffset = channelDataOffset;
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
   size_t scratchOffset = channelScratchOffset + rank * nelems * sizeof(LL8Packet);
   void* scratchBuff = (void*)((char*)scratch + channelScratchOffset);
   uint32_t* src = (uint32_t*)((char*)buff);
   uint32_t* dst = (uint32_t*)((char*)resultBuff);
 
-  // step 1: write data to each peer's scratch buffer
-  memoryChannels[peerIdx].putPackets<LL8Packet>(scratchOffset, srcOffset, nelems * sizeof(uint32_t), tid,
-                                                blockDim.x * nBlocksPerPeer, flag);
+  const int warpId = threadIdx.x / WARP_SIZE;
+  const int lane = threadIdx.x % WARP_SIZE;
+  const int nWarpsPerBlock = blockDim.x / WARP_SIZE;
+  // Assign one warp in every block to each peer. Each peer warp sends the
+  // same block-owned stripe, so nBlocks only partitions data and no longer
+  // needs to be grouped by nPeers.
+  if (warpId < nPeers) {
+    memoryChannels[warpId].putPackets<LL8Packet>(scratchOffset, channelDataOffset, nelems * sizeof(uint32_t),
+                                                 lane + blockIdx.x * WARP_SIZE, gridDim.x * WARP_SIZE, flag);
+  }
+  // Safe for in-place allreduce: all peer warps must finish reading src for
+  // this block's stripe before any warp writes reduced data back to dst/src.
+  __syncthreads();
 
-  // step 2: Reduce Data
-  for (size_t idx = threadIdx.x + blockIdx.x * blockDim.x; idx < nelems; idx += blockDim.x * gridDim.x) {
+  // Split the same sent stream across all warps for reduction. warpId selects
+  // which strided subset to reduce while lane preserves coalesced packet reads.
+  for (size_t idx = lane + blockIdx.x * WARP_SIZE + warpId * WARP_SIZE * gridDim.x; idx < nelems;
+       idx += nWarpsPerBlock * WARP_SIZE * gridDim.x) {
     uint32_t data = src[idx];
     using AccRaw = std::conditional_t<std::is_same_v<T, AccumT>, uint32_t,
                                       mscclpp::VectorType<AccumT, sizeof(uint32_t) / sizeof(T)>>;
@@ -59,14 +67,14 @@ __global__ void allreduceAllPairs(T* buff, T* scratch, T* resultBuff, DeviceHand
   if (threadIdx.x == 0) {
     ((uint32_t*)flags)[blockIdx.x] = flag + 1;
   }
-  if (blockIdx.x == 0 && threadIdx.x >= gridDim.x && threadIdx.x < flagSize / sizeof(uint32_t)) {
-    ((uint32_t*)flags)[threadIdx.x] = flag + 1;
+  if (tid >= gridDim.x && tid < flagSize / sizeof(uint32_t)) {
+    ((uint32_t*)flags)[tid] = flag + 1;
   }
 }
 
 inline std::pair<int, int> getDefaultBlockNumAndThreadNum(size_t inputSize, int worldSize) {
   if (inputSize < worldSize * sizeof(int)) {
-    return {worldSize - 1, 32};
+    return {worldSize - 1, (worldSize - 1) * WARP_SIZE};
   }
   return {(worldSize - 1) * 4, 512};
 }
@@ -80,11 +88,6 @@ struct AllpairAdapter {
                           int nThreadsPerBlock = 0) {
     using ChannelType = DeviceHandle<MemoryChannel>;
     const size_t nelems = inputSize / sizeof(T);
-    // Round nBlocks to multiple of nPeers so every block maps to a valid peer.
-    const int nPeers = worldSize - 1;
-    if (nPeers > 0) {
-      nBlocks = (nBlocks / nPeers) * nPeers;
-    }
     allreduceAllPairs<OpType, T, AccumT><<<nBlocks, nThreadsPerBlock, 0, stream>>>(
         (T*)buff, (T*)scratch, (T*)resultBuff, (ChannelType*)memoryChannels, channelInOffset, scratchBufferSize, rank,
         nRanksPerNode, worldSize, nelems, numScratchBuff, flags, flagSize);
@@ -110,9 +113,17 @@ CommResult AllreduceAllpairPacket::allreduceKernelFunc(const std::shared_ptr<voi
   if (blockAndThreadNum.first == 0 || blockAndThreadNum.second == 0) {
     blockAndThreadNum = getDefaultBlockNumAndThreadNum(inputSize, algoCtx->workSize);
   }
-  // nBlocks must be at least nPeers for allpair — each block maps to one peer.
+  if (blockAndThreadNum.first > maxBlockNum_) {
+    WARN(ALGO, "Requested block number ", blockAndThreadNum.first, " exceeds the maximum supported block number ",
+         maxBlockNum_, ".");
+    return CommResult::CommInvalidArgument;
+  }
   const int nPeers = algoCtx->nRanksPerNode - 1;
-  if (nPeers > 0 && blockAndThreadNum.first < nPeers) {
+  // The kernel maps peer sends by warpId, so every peer needs a full warp.
+  if (blockAndThreadNum.second % WARP_SIZE != 0 || blockAndThreadNum.second / WARP_SIZE < nPeers) {
+    WARN(ALGO,
+         "Allpair packet requires at least one full warp per peer, but got nThreadsPerBlock=", blockAndThreadNum.second,
+         " and nPeers=", nPeers, ".");
     return CommResult::CommInvalidArgument;
   }
   size_t sendBytes;
@@ -122,7 +133,8 @@ CommResult AllreduceAllpairPacket::allreduceKernelFunc(const std::shared_ptr<voi
 
   AllreduceFunc allreduce = dispatch<AllpairAdapter>(op, dtype, accumDtype);
   if (!allreduce) {
-    WARN("Unsupported operation or data type for allreduce: op=%d, dtype=%d", op, static_cast<int>(dtype));
+    WARN(ALGO, "Unsupported operation or data type for allreduce: op=", static_cast<int>(op),
+         ", dtype=", static_cast<int>(dtype));
     return CommResult::CommInvalidArgument;
   }
   cudaError_t error =
@@ -131,7 +143,7 @@ CommResult AllreduceAllpairPacket::allreduceKernelFunc(const std::shared_ptr<voi
                 algoCtx->workSize, inputSize, stream, (void*)flagBuffer_, (uint32_t)flagBufferSize_,
                 this->nSegmentsForScratchBuffer_, blockAndThreadNum.first, blockAndThreadNum.second);
   if (error != cudaSuccess) {
-    WARN("AllreducePacket failed with error: %s", cudaGetErrorString(error));
+    WARN(ALGO, "AllreducePacket failed with error: ", cudaGetErrorString(error));
     return CommResult::CommUnhandledCudaError;
   }
   return CommResult::CommSuccess;
@@ -189,4 +201,4 @@ std::shared_ptr<Algorithm> AllreduceAllpairPacket::build() {
       });
 }
 }  // namespace collective
-}  // namespace mscclpp
\ No newline at end of file
+}  // namespace mscclpp
diff --git a/src/ext/collectives/allreduce/allreduce_packet.cu b/src/ext/collectives/allreduce/allreduce_packet.cu
index 6199f192..3c75a746 100644
--- a/src/ext/collectives/allreduce/allreduce_packet.cu
+++ b/src/ext/collectives/allreduce/allreduce_packet.cu
@@ -235,6 +235,11 @@ CommResult AllreducePacket::allreduceKernelFunc(const std::shared_ptr<void> ctx_
   if (blockAndThreadNum.first == 0 || blockAndThreadNum.second == 0) {
     blockAndThreadNum = getDefaultBlockNumAndThreadNum(inputSize, ctx->workSize, ctx->nRanksPerNode, dtype);
   }
+  if (blockAndThreadNum.first > maxBlockNum_) {
+    WARN(ALGO, "Requested block number ", blockAndThreadNum.first, " exceeds the maximum supported block number ",
+         maxBlockNum_, ".");
+    return CommResult::CommInvalidArgument;
+  }
 
   size_t sendBytes;
   CUdeviceptr sendBasePtr;
diff --git a/src/ext/collectives/include/allreduce/allreduce_allpair_packet.hpp b/src/ext/collectives/include/allreduce/allreduce_allpair_packet.hpp
index 362308b2..64f5ec54 100644
--- a/src/ext/collectives/include/allreduce/allreduce_allpair_packet.hpp
+++ b/src/ext/collectives/include/allreduce/allreduce_allpair_packet.hpp
@@ -29,7 +29,7 @@ class AllreduceAllpairPacket : public AlgorithmBuilder {
   void* scratchBuffer_;
   size_t scratchBufferSize_;
   const int nSegmentsForScratchBuffer_ = 2;
-  const int maxBlockNum_ = 28;
+  const int maxBlockNum_ = 64;
   std::vector<Connection> conns_;
   std::vector<std::shared_ptr<MemoryDevice2DeviceSemaphore>> memorySemaphores_;
   std::vector<RegisteredMemory> registeredMemories_;

From 9e177b388c1ace99be176c9553512849e2df6ae7 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Wed, 20 May 2026 16:49:49 -0700
Subject: [PATCH 03/16] remove useless sync (#809)

---
 src/ext/collectives/allreduce/allreduce_allpair_packet.cu | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/ext/collectives/allreduce/allreduce_allpair_packet.cu b/src/ext/collectives/allreduce/allreduce_allpair_packet.cu
index faef5459..49058f59 100644
--- a/src/ext/collectives/allreduce/allreduce_allpair_packet.cu
+++ b/src/ext/collectives/allreduce/allreduce_allpair_packet.cu
@@ -63,7 +63,6 @@ __global__ void allreduceAllPairs(T* buff, T* scratch, T* resultBuff, DeviceHand
     }
     dst[idx] = mscclpp::downcastVector<T, AccumT, uint32_t>(acc);
   }
-  __syncthreads();
   if (threadIdx.x == 0) {
     ((uint32_t*)flags)[blockIdx.x] = flag + 1;
   }

From 08ee18be64248ebdbd94bc758e2cbb59e3ac9bf5 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 22 May 2026 09:18:41 -0700
Subject: [PATCH 04/16] Add check to filter invalid nblock/nthread candidates
 (#811)

Add check for invalid nblock/nthread candidate
---
 .../customized_comm_with_tuning.py            |  4 +--
 .../allgather/allgather_fullmesh.cu           | 30 ++++++++++++-------
 .../allgather/allgather_fullmesh_2.cu         | 26 ++++++++++++++--
 3 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/examples/torch-integration/customized_comm_with_tuning.py b/examples/torch-integration/customized_comm_with_tuning.py
index 060a0097..b96087c2 100644
--- a/examples/torch-integration/customized_comm_with_tuning.py
+++ b/examples/torch-integration/customized_comm_with_tuning.py
@@ -70,12 +70,12 @@ class CustomizedComm:
     _TUNE_N_WARMUP = 5
     _TUNE_N_GRAPH_LAUNCHES = 10
     _TUNE_N_OPS_PER_GRAPH = 100
-    _CANDIDATE_NBLOCKS = [4, 8, 16, 24, 32, 48, 64, 128]
+    _CANDIDATE_NBLOCKS = [4, 8, 16, 24, 32, 48, 56, 64, 128]
     _CANDIDATE_NTHREADS = [512, 768, 1024]
     _NBLOCKS_LIMIT = {
         "default_allreduce_nvls_packet": 16,
         "default_allreduce_packet": 56,
-        "default_allreduce_allpair_packet": 56,
+        "default_allreduce_allpair_packet": 64,
         "default_allreduce_fullmesh": 64,
         "default_allgather_fullmesh2": 32,
     }
diff --git a/src/ext/collectives/allgather/allgather_fullmesh.cu b/src/ext/collectives/allgather/allgather_fullmesh.cu
index fb51a342..d1b4e731 100644
--- a/src/ext/collectives/allgather/allgather_fullmesh.cu
+++ b/src/ext/collectives/allgather/allgather_fullmesh.cu
@@ -8,6 +8,11 @@
 namespace mscclpp {
 namespace collective {
 
+namespace {
+constexpr int kMaxBlocks = 56;
+constexpr int kMaxThreadsPerBlock = 1024;
+}  // namespace
+
 template <bool IsOutOfPlace>
 __global__ void __launch_bounds__(1024, 1)
     allgatherFullmesh(void* buff, void* scratch, void* resultBuff, DeviceHandle<MemoryChannel>* memoryChannels,
@@ -116,12 +121,19 @@ CommResult AllgatherFullmesh::allgatherKernelFunc(const std::shared_ptr<void> ct
   int rank = ctx->rank;
   const size_t nElem = inputSize / sizeof(int);
   std::pair<int, int> numBlocksAndThreads = {nBlocks, nThreadsPerBlock};
-  if (numBlocksAndThreads.first > 56) {
-    WARN("AllgatherFullmesh: number of blocks exceeds maximum supported blocks, which is 56");
-    return mscclpp::CommResult::CommInvalidArgument;
-  }
   if (numBlocksAndThreads.first == 0 || numBlocksAndThreads.second == 0) {
-    numBlocksAndThreads = {56, 1024};
+    numBlocksAndThreads = {kMaxBlocks, kMaxThreadsPerBlock};
+  }
+  if (numBlocksAndThreads.first > kMaxBlocks || numBlocksAndThreads.second > kMaxThreadsPerBlock) {
+    WARN(
+        "AllgatherFullmesh: number of blocks must be no more than %d and threads per block must be no more than %d; "
+        "got nBlocks=%d, nThreadsPerBlock=%d",
+        kMaxBlocks, kMaxThreadsPerBlock, numBlocksAndThreads.first, numBlocksAndThreads.second);
+    return CommResult::CommInvalidArgument;
+  }
+  if (numBlocksAndThreads.second % WARP_SIZE != 0) {
+    WARN("AllgatherFullmesh: threads per block must be a multiple of warp size %d", WARP_SIZE);
+    return CommResult::CommInvalidArgument;
   }
   if ((char*)input == (char*)output + rank * inputSize) {
     allgatherFullmesh<false><<<numBlocksAndThreads.first, numBlocksAndThreads.second, 0, stream>>>(
@@ -142,15 +154,13 @@ CommResult AllgatherFullmesh::allgatherKernelFunc(const std::shared_ptr<void> ct
 
 std::shared_ptr<void> AllgatherFullmesh::initAllgatherContext(std::shared_ptr<Communicator> comm, const void* input,
                                                               void*, size_t inputSize, DataType) {
-  constexpr int nChannelsPerConnection = 56;
-
   auto ctx = std::make_shared<AlgorithmCtx>();
   ctx->rank = comm->bootstrap()->getRank();
   ctx->workSize = comm->bootstrap()->getNranks();
   ctx->nRanksPerNode = comm->bootstrap()->getNranksPerNode();
 
   // setup semaphores
-  ctx->memorySemaphores = setupMemorySemaphores(comm, this->conns_, nChannelsPerConnection);
+  ctx->memorySemaphores = setupMemorySemaphores(comm, this->conns_, kMaxBlocks);
 
   // register the memory for the broadcast operation
   RegisteredMemory localMemory = comm->registerMemory((void*)input, inputSize, Transport::CudaIpc);
@@ -159,7 +169,7 @@ std::shared_ptr<void> AllgatherFullmesh::initAllgatherContext(std::shared_ptr<Co
 
   // setup channels
   ctx->memoryChannels =
-      setupMemoryChannels(this->conns_, ctx->memorySemaphores, remoteMemories, localMemory, nChannelsPerConnection);
+      setupMemoryChannels(this->conns_, ctx->memorySemaphores, remoteMemories, localMemory, kMaxBlocks);
   ctx->memoryChannelDeviceHandles = setupMemoryChannelDeviceHandles(ctx->memoryChannels);
 
   // keep registered memories reference
@@ -196,4 +206,4 @@ std::shared_ptr<Algorithm> AllgatherFullmesh::build() {
       });
 }
 }  // namespace collective
-}  // namespace mscclpp
\ No newline at end of file
+}  // namespace mscclpp
diff --git a/src/ext/collectives/allgather/allgather_fullmesh_2.cu b/src/ext/collectives/allgather/allgather_fullmesh_2.cu
index 9d169d68..89581822 100644
--- a/src/ext/collectives/allgather/allgather_fullmesh_2.cu
+++ b/src/ext/collectives/allgather/allgather_fullmesh_2.cu
@@ -18,7 +18,11 @@ __global__ void __launch_bounds__(1024, 1)
   const size_t lid = tid % WARP_SIZE;
   const size_t wid = tid / WARP_SIZE;
 
-  const size_t nThread = blockDim.x * gridDim.x;
+  // Round down to multiple of warp size
+  const size_t nThread = (blockDim.x * gridDim.x) / WARP_SIZE * WARP_SIZE;
+  if (tid >= nThread) {
+    return;
+  }
   const size_t nWarp = nThread / WARP_SIZE;
   const size_t nPeer = nRanksPerNode - 1;
   const size_t chanOffset = nPeer * blockIdx.x;
@@ -135,6 +139,24 @@ CommResult AllgatherFullmesh2::allgatherKernelFunc(const std::shared_ptr<void> c
       numBlocksAndThreads.first = 35;
     }
   }
+  const int nPeer = ctx->nRanksPerNode - 1;
+  const int nWarp = numBlocksAndThreads.first * numBlocksAndThreads.second / WARP_SIZE;
+  if (numBlocksAndThreads.first > nChannelsPerConnection_ || numBlocksAndThreads.first <= 0 ||
+      numBlocksAndThreads.second <= 0) {
+    WARN(
+        "AllgatherFullmesh2: number of blocks must be a positive multiple of peer count and no more than %d, threads "
+        "per block must be positive; got nBlocks=%d, nThreadsPerBlock=%d, nPeers=%d",
+        nChannelsPerConnection_, numBlocksAndThreads.first, numBlocksAndThreads.second, nPeer);
+    return CommResult::CommInvalidArgument;
+  }
+  if (nWarp < nPeer) {
+    WARN(
+        "AllgatherFullmesh2: total number of warps must be no less than peer count; got nBlocks=%d, "
+        "nThreadsPerBlock=%d, "
+        "nPeers=%d",
+        numBlocksAndThreads.first, numBlocksAndThreads.second, nPeer);
+    return CommResult::CommInvalidArgument;
+  }
 
   size_t channelOutOffset = *static_cast<size_t*>(ctx->extras["channel_out_offset"].get());
   if ((char*)input == (char*)output + rank * inputSize) {
@@ -226,4 +248,4 @@ std::shared_ptr<Algorithm> AllgatherFullmesh2::build() {
 }
 
 }  // namespace collective
-}  // namespace mscclpp
\ No newline at end of file
+}  // namespace mscclpp

From 379d0e51e4fbcaf28d604067fa5ed045fefc9eb5 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Tue, 26 May 2026 12:34:13 -0700
Subject: [PATCH 05/16] Fix for allgather_fullmesh algo (#813)

---
 src/ext/collectives/allgather/allgather_fullmesh_2.cu | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/ext/collectives/allgather/allgather_fullmesh_2.cu b/src/ext/collectives/allgather/allgather_fullmesh_2.cu
index 89581822..72a2be9d 100644
--- a/src/ext/collectives/allgather/allgather_fullmesh_2.cu
+++ b/src/ext/collectives/allgather/allgather_fullmesh_2.cu
@@ -17,14 +17,14 @@ __global__ void __launch_bounds__(1024, 1)
   const size_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   const size_t lid = tid % WARP_SIZE;
   const size_t wid = tid / WARP_SIZE;
+  const size_t nPeer = nRanksPerNode - 1;
 
-  // Round down to multiple of warp size
-  const size_t nThread = (blockDim.x * gridDim.x) / WARP_SIZE * WARP_SIZE;
+  // Round down to multiple of peer count.
+  const size_t nThread = (blockDim.x * gridDim.x) / WARP_SIZE / nPeer * nPeer * WARP_SIZE;
   if (tid >= nThread) {
     return;
   }
   const size_t nWarp = nThread / WARP_SIZE;
-  const size_t nPeer = nRanksPerNode - 1;
   const size_t chanOffset = nPeer * blockIdx.x;
   auto memChans = memoryChannels + chanOffset;
 

From 29d5beb3489ad030b228ad19875c3af2e16b87d7 Mon Sep 17 00:00:00 2001
From: RJ Souza <rjsouza1995@gmail.com>
Date: Wed, 27 May 2026 14:57:21 -0700
Subject: [PATCH 06/16] Adding Support for SGLang CI Tests (#800)

Adds Azure DevOps pipelines, templates, and supporting scripts to run
SGLang end-to-end and benchmark tests against MSCCL++ on H100 GPU nodes,
plus the Docker image and small infrastructure tweaks needed to make
those pipelines runnable.

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 .azure-pipelines/integration-test.yml         |  22 +--
 .azure-pipelines/multi-nodes-test.yml         |  11 +-
 .azure-pipelines/sglang-multi-node-test.yml   | 141 ++++++++++++++++++
 .azure-pipelines/sglang-test.yml              |  63 ++++++++
 .azure-pipelines/templates/deploy.yml         |   8 +-
 .../templates/integration-test.yml            |   2 +-
 .azure-pipelines/templates/nccl-test.yml      |   2 +-
 .../templates/sglang-multi-test.yml           |  95 ++++++++++++
 .azure-pipelines/templates/sglang-test.yml    |  87 +++++++++++
 .azure-pipelines/templates/ut-no-ib-env.yml   |   2 +-
 .azure-pipelines/templates/ut-npkit.yml       |   2 +-
 docker/build.sh                               |   3 +-
 test/deploy/deploy.sh                         |  68 ++++++---
 test/deploy/run-remote.sh                     |  11 +-
 14 files changed, 473 insertions(+), 44 deletions(-)
 create mode 100644 .azure-pipelines/sglang-multi-node-test.yml
 create mode 100644 .azure-pipelines/sglang-test.yml
 create mode 100644 .azure-pipelines/templates/sglang-multi-test.yml
 create mode 100644 .azure-pipelines/templates/sglang-test.yml

diff --git a/.azure-pipelines/integration-test.yml b/.azure-pipelines/integration-test.yml
index d5d5f9bd..45bb1e96 100644
--- a/.azure-pipelines/integration-test.yml
+++ b/.azure-pipelines/integration-test.yml
@@ -19,11 +19,11 @@ pr:
   drafts: false
   paths:
     exclude:
-      - .devcontainer/**
-      - .github/**
-      - docker/**
-      - docs/**
-      - '**/*.md'
+    - .devcontainer/**
+    - .github/**
+    - docker/**
+    - docs/**
+    - '**/*.md'
 
 jobs:
 - job: IntegrationTestA100
@@ -43,9 +43,9 @@ jobs:
   steps:
   - template: templates/integration-test.yml
     parameters:
-      subscription:     mscclpp-ci
-      vmssName:         mscclpp-ci
-      gpuArch:          '80'
+      subscription: mscclpp-ci
+      vmssName: mscclpp-ci
+      gpuArch: '80'
 
 - job: IntegrationTestH100
   displayName: Integration test H100
@@ -62,7 +62,7 @@ jobs:
   steps:
   - template: templates/integration-test.yml
     parameters:
-      subscription:     mscclpp-ci-h100
-      vmssName:         mscclpp-h100-ci
+      subscription: mscclpp-ci-h100
+      vmssName: mscclpp-h100-ci
       perfBaselineFile: test/deploy/perf_ndmv5.jsonl
-      gpuArch:          '90'
+      gpuArch: '90'
diff --git a/.azure-pipelines/multi-nodes-test.yml b/.azure-pipelines/multi-nodes-test.yml
index 3b3ebe1f..ee2766fd 100644
--- a/.azure-pipelines/multi-nodes-test.yml
+++ b/.azure-pipelines/multi-nodes-test.yml
@@ -14,7 +14,6 @@ trigger:
 # Do not run multi-nodes-test for PR, we can trigger it manually
 pr: none
 
-
 parameters:
 - name: vmssName
   type: string
@@ -79,10 +78,10 @@ jobs:
 
   - template: templates/deploy.yml
     parameters:
-      subscription:  mscclpp-ci-h100
-      vmssName:      ${{ parameters.vmssName }}
+      subscription: mscclpp-ci-h100
+      vmssName: ${{ parameters.vmssName }}
       resourceGroup: mscclpp
-      gpuArch:       '90'
+      gpuArch: '90'
 
   - template: templates/run-remote-task.yml
     parameters:
@@ -119,6 +118,6 @@ jobs:
 
   - template: templates/stop.yml
     parameters:
-      subscription:  mscclpp-ci-h100
-      vmssName:      ${{ parameters.vmssName }}
+      subscription: mscclpp-ci-h100
+      vmssName: ${{ parameters.vmssName }}
       resourceGroup: mscclpp
diff --git a/.azure-pipelines/sglang-multi-node-test.yml b/.azure-pipelines/sglang-multi-node-test.yml
new file mode 100644
index 00000000..bf640db2
--- /dev/null
+++ b/.azure-pipelines/sglang-multi-node-test.yml
@@ -0,0 +1,141 @@
+# =============================================================================
+# Multi-node SGLang integration test pipeline.
+#
+# This pipeline runs MSCCL++ SGLang tests across two H100 VMSS GPU nodes.
+# High-level flow:
+#   1. The pipeline agent runs inside a container on the `mscclpp-multi-node`
+#      pool. The agent itself has no GPUs.
+#   2. SSH/host configuration is generated so the agent can reach the two
+#      pre-provisioned VMSS GPU nodes.
+#   3. `templates/deploy.yml` builds and ships MSCCL++ to the GPU nodes.
+#   4. `templates/sglang-multi-test.yml` runs the SGLang multi-node tests.
+#   5. `templates/stop.yml` tears down / stops the VMSS nodes.
+#
+# Docs / non-code changes are excluded from triggering this pipeline.
+# =============================================================================
+
+trigger:
+  branches:
+    include:
+    - main
+    - release/*
+  paths:
+    exclude:
+    - .devcontainer/**
+    - .github/**
+    - docker/**
+    - docs/**
+    - '**/*.md'
+
+pr:
+  branches:
+    include:
+    - main
+    - release/*
+  drafts: false
+  paths:
+    exclude:
+    - .devcontainer/**
+    - .github/**
+    - docker/**
+    - docs/**
+    - '**/*.md'
+
+parameters:
+# Name of the pre-provisioned Azure VMSS that hosts the GPU test nodes.
+# Node hostnames are derived as "${vmssName}000000" and "${vmssName}000001".
+- name: vmssName
+  type: string
+  default: mscclpp-h100-multinode-ci
+# Static /etc/hosts entries mapping VMSS node hostnames to their private IPs.
+# These IPs are tied to the specific VMSS above; update both together if the
+# VMSS is reprovisioned or renamed.
+- name: hostEntries
+  type: string
+  default: |
+    10.0.0.5 mscclpp-h100-multinode-ci000000
+    10.0.0.4 mscclpp-h100-multinode-ci000001
+# Docker image used for the SGLang test container on the GPU nodes.
+- name: sglangImage
+  type: string
+  default: lmsysorg/sglang:latest-cu129
+
+jobs:
+- job: SGLangTestMultiNode
+  displayName: SGLang Test Multi Node
+  strategy:
+    matrix:
+      cuda12:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+  pool:
+    name: mscclpp-multi-node
+  container:
+    image: $(containerImage)
+
+  steps:
+  # Ensure the VMSS node hostnames resolve from the pipeline agent container.
+  # Idempotent: only appends lines that are not already present in /etc/hosts.
+  - task: Bash@3
+    displayName: Add HostEntry
+    inputs:
+      targetType: 'inline'
+      script: |
+        while IFS= read -r line; do
+          [ -z "$line" ] && continue
+          if ! grep -qxF "$line" /etc/hosts; then
+            echo "Adding to /etc/hosts: $line"
+            echo "$line" | sudo tee -a /etc/hosts
+          else
+            echo "Entry already exists: $line"
+          fi
+        done <<< "${{ parameters.hostEntries }}"
+
+  # Generate the SSH config and hostfile consumed by the deploy / test
+  # templates below:
+  #   - config   : SSH client config (custom port + key) for each node
+  #   - hostfile : user@host list used by deploy / test scripts (parallel-ssh)
+  - task: Bash@3
+    displayName: Generate deploy files
+    inputs:
+      targetType: 'inline'
+      script: |
+        set -e
+        VMSS="${{ parameters.vmssName }}"
+        DEPLOY_DIR="$(System.DefaultWorkingDirectory)/test/deploy"
+        NODE0="${VMSS}000000"
+        NODE1="${VMSS}000001"
+
+        echo "Host ${NODE0}
+          Port 22345
+          IdentityFile /root/mscclpp/sshkey
+          StrictHostKeyChecking no
+        Host ${NODE1}
+          Port 22345
+          IdentityFile /root/mscclpp/sshkey
+          StrictHostKeyChecking no" > "${DEPLOY_DIR}/config"
+
+        printf '%s\n%s\n' "azureuser@${NODE0}" "azureuser@${NODE1}" > "${DEPLOY_DIR}/hostfile"
+
+  # Build MSCCL++ and deploy it onto the VMSS GPU nodes.
+  - template: templates/deploy.yml
+    parameters:
+      subscription: mscclpp-ci-h100
+      vmssName: ${{ parameters.vmssName }}
+      resourceGroup: mscclpp
+      gpuArch: '90'
+      deployArgs: 'multi-node-test true cuda'
+      containerName: 'sglang-mscclpp-test'
+      sglangImage: ${{ parameters.sglangImage }}
+
+  # Run the SGLang multi-node tests across the two GPU nodes.
+  - template: templates/sglang-multi-test.yml
+    parameters:
+      subscription: mscclpp-ci-h100
+      vmssName: ${{ parameters.vmssName }}
+
+  # Stop/deallocate the VMSS GPU nodes to release resources.
+  - template: templates/stop.yml
+    parameters:
+      subscription: mscclpp-ci-h100
+      vmssName: ${{ parameters.vmssName }}
+      resourceGroup: mscclpp
diff --git a/.azure-pipelines/sglang-test.yml b/.azure-pipelines/sglang-test.yml
new file mode 100644
index 00000000..70e30d35
--- /dev/null
+++ b/.azure-pipelines/sglang-test.yml
@@ -0,0 +1,63 @@
+# =============================================================================
+# Single-node SGLang integration test pipeline.
+#
+# Runs MSCCL++ SGLang tests on a single H100 GPU node from the `msccl-ci-h100`
+# pool. All deploy / run / teardown logic is delegated to
+# `templates/sglang-test.yml`.
+#
+# Docs / non-code changes are excluded from triggering this pipeline.
+# =============================================================================
+
+trigger:
+  branches:
+    include:
+    - main
+    - release/*
+  paths:
+    exclude:
+    - .devcontainer/**
+    - .github/**
+    - docker/**
+    - docs/**
+    - '**/*.md'
+
+pr:
+  branches:
+    include:
+    - main
+    - release/*
+  drafts: false
+  paths:
+    exclude:
+    - .devcontainer/**
+    - .github/**
+    - docker/**
+    - docs/**
+    - '**/*.md'
+
+parameters:
+# Docker image used for the SGLang test container on the GPU node.
+- name: sglangImage
+  type: string
+  default: lmsysorg/sglang:latest-cu129
+
+jobs:
+- job: SGLangTest
+  displayName: SGLang Test
+  strategy:
+    matrix:
+      cuda12:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+  pool:
+    name: msccl-ci-h100
+  container:
+    image: $(containerImage)
+
+  steps:
+  # Deploy MSCCL++ to the GPU node and run the SGLang single-node tests.
+  - template: templates/sglang-test.yml
+    parameters:
+      subscription: mscclpp-ci-h100
+      vmssName: mscclpp-h100-ci
+      gpuArch: '90'
+      sglangImage: ${{ parameters.sglangImage }}
diff --git a/.azure-pipelines/templates/deploy.yml b/.azure-pipelines/templates/deploy.yml
index 2f642f1d..9eb46d8b 100644
--- a/.azure-pipelines/templates/deploy.yml
+++ b/.azure-pipelines/templates/deploy.yml
@@ -32,6 +32,12 @@ parameters:
 - name: deployArgs
   type: string
   default: ''
+- name: containerName
+  type: string
+  default: 'mscclpp-test'
+- name: sglangImage
+  type: string
+  default: ''
 
 steps:
 # 0. Ensure Azure CLI exists before running AzureCLI@2 tasks.
@@ -147,5 +153,5 @@ steps:
   inputs:
     targetType: filePath
     filePath: test/deploy/deploy.sh
-    arguments: ${{ parameters.deployArgs }}
+    arguments: ${{ parameters.deployArgs }} ${{ parameters.containerName }} ${{ parameters.sglangImage }}
     workingDirectory: '$(System.DefaultWorkingDirectory)'
diff --git a/.azure-pipelines/templates/integration-test.yml b/.azure-pipelines/templates/integration-test.yml
index b686e4f2..ad95cbc2 100644
--- a/.azure-pipelines/templates/integration-test.yml
+++ b/.azure-pipelines/templates/integration-test.yml
@@ -15,7 +15,7 @@ steps:
     subscription:     ${{ parameters.subscription }}
     vmssName:         ${{ parameters.vmssName }}
     gpuArch:          ${{ parameters.gpuArch }}
-    deployArgs:       'single-node-test'
+    deployArgs:       'single-node-test true cuda'
 
 - template: run-remote-task.yml
   parameters:
diff --git a/.azure-pipelines/templates/nccl-test.yml b/.azure-pipelines/templates/nccl-test.yml
index fa3900f1..585f5b48 100644
--- a/.azure-pipelines/templates/nccl-test.yml
+++ b/.azure-pipelines/templates/nccl-test.yml
@@ -23,7 +23,7 @@ steps:
     subscription:     ${{ parameters.subscription }}
     vmssName:         ${{ parameters.vmssName }}
     gpuArch:          ${{ parameters.gpuArch }}
-    deployArgs:       'nccltest-single-node'
+    deployArgs:       'nccltest-single-node true cuda'
 
 - template: run-remote-task.yml
   parameters:
diff --git a/.azure-pipelines/templates/sglang-multi-test.yml b/.azure-pipelines/templates/sglang-multi-test.yml
new file mode 100644
index 00000000..80e72926
--- /dev/null
+++ b/.azure-pipelines/templates/sglang-multi-test.yml
@@ -0,0 +1,95 @@
+# =============================================================================
+# SGLang multi-node test template.
+#
+# Runs on the pipeline agent and dispatches remote steps to the two VMSS GPU
+# nodes (via run-remote-task.yml + the SSH config / hostfile produced by the
+# caller pipeline). Steps:
+#   1. Build and install MSCCL++ on each node.
+#   2. Install a (currently forked) SGLang on each node, replacing any
+#      pre-baked copy from the base image.
+#   3. Run a 2-node sglang.bench_one_batch smoke test with MSCCL++ enabled.
+#   4. Run the MSCCL++ all-reduce micro-benchmark via torchrun across both
+#      nodes.
+# =============================================================================
+
+parameters:
+- name: subscription
+  type: string
+- name: vmssName
+  type: string
+- name: containerName
+  type: string
+  default: 'sglang-mscclpp-test'
+
+steps:
+# TODO: Switch to the official upstream sglang repo once Caio's PR is merged.
+# Tracking: the fork below (`caiomcbr/sglang` @ caiorocha/mscclpp) is a personal
+# branch and should not remain a long-term CI dependency.
+- template: run-remote-task.yml
+  parameters:
+    name: InstallSGLang
+    displayName: Install SGLang
+    runRemoteArgs: '--container ${{ parameters.containerName }} --hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --user azureuser'
+    remoteScript: |
+      git clone -b main https://github.com/caiomcbr/sglang.git
+      cd sglang/python
+      pip install -e .
+
+# Smoke test: 2-node tensor-parallel benchmark of Qwen3-8B with MSCCL++.
+# Port 20003 is the SGLang distributed-init rendezvous port (arbitrary, must
+# match across ranks and be free on node 0).
+- template: run-remote-task.yml
+  parameters:
+    name: RunSGLangMultiBenchOneBatch
+    displayName: Run SGLang Multi-Node Bench One Batch
+    runRemoteArgs: '--container ${{ parameters.containerName }} --hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --user azureuser'
+    remoteScript: |
+      export FLASHINFER_DISABLE_VERSION_CHECK=1
+      VMSS="${{ parameters.vmssName }}"
+      HOSTNAME=$(hostname)
+      # Explicit 2-node mapping: hostname suffix -> SGLang node rank.
+      if [ "$HOSTNAME" = "${VMSS}000000" ]; then
+        NODE_RANK=0
+      elif [ "$HOSTNAME" = "${VMSS}000001" ]; then
+        NODE_RANK=1
+      else
+        echo "Unknown hostname: $HOSTNAME"
+        exit 1
+      fi
+      python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 1 2 4 8 16 32 64 128 256 512 --input-len 256 --output-len 256 --tp-size 16 --dist-init-addr ${VMSS}000000:20003 --nnodes 2 --node-rank $NODE_RANK --enable-mscclpp
+
+# Depends on the `sglang/` source tree cloned by the InstallSGLang step above
+# (steps on the same remote share a working directory).
+- template: run-remote-task.yml
+  parameters:
+    name: RunSGLangMultiTestAllReduce
+    displayName: Run SGLang Multi-Node Test All Reduce
+    runRemoteArgs: '--container ${{ parameters.containerName }} --hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --user azureuser'
+    remoteScript: |
+      export FLASHINFER_DISABLE_VERSION_CHECK=1
+      VMSS="${{ parameters.vmssName }}"
+      HOSTNAME=$(hostname)
+      # Explicit 2-node mapping: hostname suffix -> torchrun node rank.
+      if [ "$HOSTNAME" = "${VMSS}000000" ]; then
+        NODE_RANK=0
+      elif [ "$HOSTNAME" = "${VMSS}000001" ]; then
+        NODE_RANK=1
+      else
+        echo "Unknown hostname: $HOSTNAME"
+        exit 1
+      fi
+
+      export NODE_SIZE=2
+      export WORLD_SIZE=8
+
+      cd sglang
+
+      # Port 20004 is the torchrun rendezvous port (arbitrary, must match
+      # across ranks and be free on node 0). Distinct from 20003 used by
+      # sglang.bench_one_batch above.
+      torchrun --nproc_per_node $WORLD_SIZE \
+        --nnodes $NODE_SIZE \
+        --node_rank $NODE_RANK \
+        --master_addr ${VMSS}000000 \
+        --master_port 20004 \
+        benchmark/kernels/all_reduce/benchmark_mscclpp.py
diff --git a/.azure-pipelines/templates/sglang-test.yml b/.azure-pipelines/templates/sglang-test.yml
new file mode 100644
index 00000000..0d663b71
--- /dev/null
+++ b/.azure-pipelines/templates/sglang-test.yml
@@ -0,0 +1,87 @@
+# =============================================================================
+# SGLang single-node test template.
+#
+# Runs on the pipeline agent and dispatches remote steps to a single VMSS GPU
+# node (via run-remote-task.yml). Steps:
+#   1. Deploy: build the test container and bring the VMSS node online.
+#   2. Build and install MSCCL++ on the node.
+#   3. Install a (currently forked) SGLang.
+#   4. Run sglang.bench_one_batch at several batch sizes.
+#   5. Run a longer end-to-end validation: bring up an sglang server and
+#      drive it with sglang.bench_serving.
+#   6. Run the MSCCL++ all-reduce micro-benchmark via torchrun.
+#   7. Stop / deallocate the VMSS node.
+# =============================================================================
+
+parameters:
+- name: subscription
+  type: string
+- name: vmssName
+  type: string
+- name: gpuArch
+  type: string
+- name: containerName
+  type: string
+  default: 'sglang-mscclpp-test'
+- name: sglangImage
+  type: string
+  default: 'lmsysorg/sglang:latest'
+
+steps:
+# deployArgs positional fields: <test-mode> <use-gpu> <cuda|rocm>
+- template: deploy.yml
+  parameters:
+    subscription: ${{ parameters.subscription }}
+    vmssName: ${{ parameters.vmssName }}
+    gpuArch: ${{ parameters.gpuArch }}
+    deployArgs: 'single-node-test true cuda'
+    containerName: ${{ parameters.containerName }}
+    sglangImage: ${{ parameters.sglangImage }}
+
+# TODO: Switch to the official upstream sglang repo once Caio's PR is merged.
+# Tracking: the fork below (`caiomcbr/sglang` @ caiorocha/mscclpp) is a personal branch and
+# should not remain a long-term CI dependency. Also consider pinning to a
+# release branch or commit SHA for reproducibility.
+- template: run-remote-task.yml
+  parameters:
+    name: InstallSGLang
+    displayName: Install SGLang
+    runRemoteArgs: '--container ${{ parameters.containerName }}'
+    remoteScript: |
+      git clone -b main https://github.com/caiomcbr/sglang.git
+      cd sglang/python
+      pip install -e .
+
+- template: run-remote-task.yml
+  parameters:
+    name: RunSGLangBenchOneBatch
+    displayName: Run SGLang Bench One Batch
+    runRemoteArgs: '--container ${{ parameters.containerName }}'
+    remoteScript: |
+      export FLASHINFER_DISABLE_VERSION_CHECK=1
+      python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 1 2 4 8 16 32 64 128 256 512 --input-len 256 --output-len 256 --tp-size 8 --enable-mscclpp
+
+# Depends on the `sglang/` source tree cloned by the InstallSGLang step above
+# (steps on the same remote share a working directory).
+- template: run-remote-task.yml
+  parameters:
+    name: RunSGLangTestAllReduce
+    displayName: Run SGLang Test All Reduce
+    runRemoteArgs: '--container ${{ parameters.containerName }}'
+    remoteScript: |
+      export FLASHINFER_DISABLE_VERSION_CHECK=1
+      export NODE_SIZE=1
+      export WORLD_SIZE=8
+      export RANK=0
+
+      cd sglang
+
+      torchrun --nproc_per_node $WORLD_SIZE \
+      --nnodes $NODE_SIZE \
+      --node_rank $RANK \
+      benchmark/kernels/all_reduce/benchmark_mscclpp.py
+
+- template: stop.yml
+  parameters:
+    subscription: ${{ parameters.subscription }}
+    vmssName: ${{ parameters.vmssName }}
diff --git a/.azure-pipelines/templates/ut-no-ib-env.yml b/.azure-pipelines/templates/ut-no-ib-env.yml
index a62f1a77..cc7d2018 100644
--- a/.azure-pipelines/templates/ut-no-ib-env.yml
+++ b/.azure-pipelines/templates/ut-no-ib-env.yml
@@ -13,7 +13,7 @@ steps:
     vmssName:         ${{ parameters.vmssName }}
     gpuArch:          ${{ parameters.gpuArch }}
     cmakeArgs:        '-DMSCCLPP_USE_IB=OFF'
-    deployArgs:       'single-node-test false'
+    deployArgs:       'single-node-test false cuda'
 
 - template: run-remote-task.yml
   parameters:
diff --git a/.azure-pipelines/templates/ut-npkit.yml b/.azure-pipelines/templates/ut-npkit.yml
index 1bd89caf..18934e6b 100644
--- a/.azure-pipelines/templates/ut-npkit.yml
+++ b/.azure-pipelines/templates/ut-npkit.yml
@@ -14,7 +14,7 @@ steps:
     vmssName:         ${{ parameters.vmssName }}
     gpuArch:          ${{ parameters.gpuArch }}
     cmakeArgs:        '-DMSCCLPP_NPKIT_FLAGS="-DENABLE_NPKIT -DENABLE_NPKIT_EVENT_TIME_SYNC_CPU -DENABLE_NPKIT_EVENT_TIME_SYNC_GPU -DENABLE_NPKIT_EVENT_EXECUTOR_INIT_ENTRY -DENABLE_NPKIT_EVENT_EXECUTOR_INIT_EXIT -DENABLE_NPKIT_EVENT_EXECUTOR_OP_BASE_ENTRY -DENABLE_NPKIT_EVENT_EXECUTOR_OP_BASE_EXIT"'
-    deployArgs:       'single-node-test'
+    deployArgs:       'single-node-test true cuda'
 
 - template: run-remote-task.yml
   parameters:
diff --git a/docker/build.sh b/docker/build.sh
index 651a6122..b84eac9a 100755
--- a/docker/build.sh
+++ b/docker/build.sh
@@ -75,6 +75,7 @@ docker build -t ${TAG_BASE_DEV} \
     --build-arg BASE_IMAGE=${TAG_BASE} \
     --build-arg TARGET=${TARGET} .
 
+
 GHCR="ghcr.io/microsoft/mscclpp/mscclpp"
 GHCR_TAG_BASE_DEV=${GHCR}:base-dev-${TARGET}
 GHCR_TAG_BASE_DEV_ARCH=${GHCR}:base-dev-${TARGET}-${OS_ARCH}
@@ -107,4 +108,4 @@ echo ""
 echo "    docker buildx imagetools create \\"
 echo "        --tag ${GHCR_TAG_BASE_DEV} \\"
 echo "        ${GHCR_TAG_BASE_DEV_ARCH}"
-echo ""
+echo ""
\ No newline at end of file
diff --git a/test/deploy/deploy.sh b/test/deploy/deploy.sh
index 6358787b..02fe4fd2 100644
--- a/test/deploy/deploy.sh
+++ b/test/deploy/deploy.sh
@@ -1,17 +1,34 @@
+#!/bin/bash
+# deploy.sh — Provisions remote hosts, copies sources, and launches Docker containers
+# for mscclpp CI/CD test environments.
+#
+# Usage: deploy.sh <test_name> [ib_environment] [platform] [container_name] [sglang_image]
+#   test_name       : Test suite to deploy (e.g. single-node-test, nccltest-single-node)
+#   ib_environment  : Enable InfiniBand networking (default: true)
+#   platform        : Target GPU platform — "cuda" or "rocm" (default: cuda)
+#   container_name  : Docker container name (default: mscclpp-test)
+#   sglang_image    : Docker image used for the SGLang test container
+#                     (default: lmsysorg/sglang:latest). Only used when
+#                     container_name is "sglang-mscclpp-test".
+
 set -ex
 
 TEST_NAME=$1
 IB_ENVIRONMENT="${2:-true}"
 PLATFORM="${3:-cuda}"
+CONTAINER_NAME="${4:-mscclpp-test}"
+SGLANG_IMAGE="${5:-lmsysorg/sglang:latest}"
 
 KeyFilePath=${SSHKEYFILE_SECUREFILEPATH}
 ROOT_DIR="${SYSTEM_DEFAULTWORKINGDIRECTORY}/"
 DST_DIR="/tmp/mscclpp"
+
 if [ "${TEST_NAME}" == "nccltest-single-node" ] || [ "${TEST_NAME}" == "single-node-test" ]; then
   HOSTFILE="${SYSTEM_DEFAULTWORKINGDIRECTORY}/test/deploy/hostfile_ci"
 else
   HOSTFILE="${SYSTEM_DEFAULTWORKINGDIRECTORY}/test/deploy/hostfile"
 fi
+
 SSH_OPTION="StrictHostKeyChecking=no"
 
 chmod 400 ${KeyFilePath}
@@ -26,8 +43,8 @@ while true; do
   echo "Waiting for sshd to start..."
   sleep 5
 done
-
 set -e
+
 parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION "sudo rm -rf ${DST_DIR}"
 tar czf /tmp/mscclpp.tar.gz -C ${ROOT_DIR} .
 parallel-scp -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION /tmp/mscclpp.tar.gz /tmp/mscclpp.tar.gz
@@ -57,25 +74,38 @@ if [ "${PLATFORM}" == "cuda" ]; then
     fi"
 fi
 
-# force to pull the latest image
-parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \
-  "sudo docker pull ${CONTAINERIMAGE}"
-
-LAUNCH_OPTION="--gpus=all"
-if [ "${PLATFORM}" == "rocm" ]; then
-  LAUNCH_OPTION="--device=/dev/kfd --device=/dev/dri --group-add=video"
-fi
-if [ "${IB_ENVIRONMENT}" == "true" ]; then
+if [ "${CONTAINER_NAME}" == "sglang-mscclpp-test" ]; then
+  # force to pull the latest image
   parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \
-    "sudo docker run --rm -itd --privileged --net=host --ipc=host ${LAUNCH_OPTION} \
-    -w /root -v ${DST_DIR}:/root/mscclpp -v /opt/microsoft:/opt/microsoft --ulimit memlock=-1:-1 --name=mscclpp-test \
-    --entrypoint /bin/bash ${CONTAINERIMAGE}"
+    "sudo docker pull ${SGLANG_IMAGE}"
+
+  parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \
+    "sudo docker run --rm -itd --name=${CONTAINER_NAME} --privileged --net=host --ipc=host --gpus=all -w /root -v ${DST_DIR}:/root/mscclpp --entrypoint /bin/bash ${SGLANG_IMAGE}"
 else
+  # force to pull the latest image
   parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \
-    "sudo docker run --rm -itd --net=host --ipc=host ${LAUNCH_OPTION} --cap-add=SYS_ADMIN --security-opt seccomp=unconfined \
-    -w /root -v ${DST_DIR}:/root/mscclpp -v /opt/microsoft:/opt/microsoft --ulimit memlock=-1:-1 --name=mscclpp-test \
-    --entrypoint /bin/bash ${CONTAINERIMAGE}"
-fi
-parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \
-  "sudo docker exec -t --user root mscclpp-test bash '/root/mscclpp/test/deploy/setup.sh' ${PLATFORM}"
+    "sudo docker pull ${CONTAINERIMAGE}"
 
+  # Set GPU passthrough flags based on platform
+  LAUNCH_OPTION="--gpus=all"
+  if [ "${PLATFORM}" == "rocm" ]; then
+    LAUNCH_OPTION="--device=/dev/kfd --device=/dev/dri --group-add=video"
+  fi
+
+  if [ "${IB_ENVIRONMENT}" == "true" ]; then
+    # InfiniBand: use --privileged for RDMA device access
+    parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \
+      "sudo docker run --rm -itd --privileged --net=host --ipc=host ${LAUNCH_OPTION} \
+      -w /root -v ${DST_DIR}:/root/mscclpp -v /opt/microsoft:/opt/microsoft --ulimit memlock=-1:-1 --name=${CONTAINER_NAME} \
+      --entrypoint /bin/bash ${CONTAINERIMAGE}"
+  else
+    # Non-IB: grant SYS_ADMIN and disable seccomp instead of full --privileged
+    parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \
+      "sudo docker run --rm -itd --net=host --ipc=host ${LAUNCH_OPTION} --cap-add=SYS_ADMIN --security-opt seccomp=unconfined \
+      -w /root -v ${DST_DIR}:/root/mscclpp -v /opt/microsoft:/opt/microsoft --ulimit memlock=-1:-1 --name=${CONTAINER_NAME} \
+      --entrypoint /bin/bash ${CONTAINERIMAGE}"
+  fi
+fi
+
+parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \
+  "sudo docker exec -t --user root ${CONTAINER_NAME} bash '/root/mscclpp/test/deploy/setup.sh' ${PLATFORM}"
diff --git a/test/deploy/run-remote.sh b/test/deploy/run-remote.sh
index 2468243e..9607664f 100755
--- a/test/deploy/run-remote.sh
+++ b/test/deploy/run-remote.sh
@@ -11,6 +11,7 @@
 #   --hostfile    Override hostfile path (default: test/deploy/hostfile_ci)
 #   --host        Run command on a single host (uses parallel-ssh -H)
 #   --user        SSH user when using --host or custom hostfile
+#   --container   Docker container name to exec into (default: mscclpp-test)
 
 set -e
 
@@ -23,9 +24,10 @@ USE_DOCKER=true
 USE_LOG=true
 TARGET_HOST=""
 REMOTE_USER=""
+CONTAINER_NAME="mscclpp-test"
 
 usage() {
-    echo "Usage: $0 [--no-docker] [--no-log] [--hostfile <path>] [--host <name>] [--user <name>] < <command_script>" >&2
+    echo "Usage: $0 [--no-docker] [--no-log] [--hostfile <path>] [--host <name>] [--user <name>] [--container <name>] < <command_script>" >&2
 }
 
 require_value() {
@@ -56,6 +58,11 @@ while [[ "$1" == --* ]]; do
             REMOTE_USER="$2"
             shift 2
             ;;
+        --container)
+            require_value "--container" "${2-}"
+            CONTAINER_NAME="$2"
+            shift 2
+            ;;
         *) echo "Unknown option: $1" >&2; exit 1 ;;
     esac
 done
@@ -103,7 +110,7 @@ if $USE_DOCKER; then
     INNER+=" rm -f \\\"\\\$TMP\\\""
 
     parallel-ssh -i "${PSSH_COMMON[@]}" \
-        "sudo docker exec mscclpp-test bash -c \"${INNER}\""
+        "sudo docker exec ${CONTAINER_NAME} bash -c \"${INNER}\""
 else
     parallel-ssh -i "${PSSH_COMMON[@]}" \
         "set -euxo pipefail; CMD_B64='${CMD_B64}'; TMP=\$(mktemp); printf '%s' \"\$CMD_B64\" | base64 -d > \"\$TMP\"; bash -euxo pipefail \"\$TMP\"; rm -f \"\$TMP\""

From c9f8be64bbc01eaa5b6d1f4051cf5136b4dafbf2 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Thu, 4 Jun 2026 09:22:10 -0700
Subject: [PATCH 07/16] Add collective benchmark and correctness check (#814)

- Add unit-test for float8_e4m3b15 data type.
- And tuner and benchmark for allreduce/allgather algo, make sure the
correctness and performance.
---
 .azure-pipelines/templates/nccl-test.yml      |   9 +
 .azure-pipelines/templates/rccl-test.yml      |   9 +
 docs/quickstart.md                            |  37 +-
 include/mscclpp/gpu_data_types.hpp            |  69 +-
 pyproject.toml                                |  20 +-
 python/mscclpp_benchmark/__init__.py          |  16 +-
 python/mscclpp_benchmark/bench_collective.py  | 645 ++++++++++++++++++
 python/mscclpp_benchmark/comm.py              | 409 +++++++++++
 python/mscclpp_benchmark/correctness.py       | 402 +++++++++++
 python/mscclpp_benchmark/gpu.py               | 187 +++++
 python/mscclpp_benchmark/tuner.py             |  84 +++
 python/mscclpp_benchmark/tuning_config.py     | 242 +++++++
 python/requirements_cuda11.txt                |   1 +
 python/requirements_cuda12.txt                |   1 +
 python/requirements_cuda13.txt                |   1 +
 python/requirements_rocm6.txt                 |   3 +-
 python/test/test_fp8_accum.py                 |  10 +-
 .../allgather/allgather_fullmesh_2.cu         | 125 ++--
 .../collectives/allreduce/allreduce_packet.cu |   7 +
 test/unit/CMakeLists.txt                      |   1 +
 test/unit/gpu_data_types_tests.cu             | 175 +++++
 21 files changed, 2327 insertions(+), 126 deletions(-)
 create mode 100644 python/mscclpp_benchmark/bench_collective.py
 create mode 100644 python/mscclpp_benchmark/comm.py
 create mode 100644 python/mscclpp_benchmark/correctness.py
 create mode 100644 python/mscclpp_benchmark/gpu.py
 create mode 100644 python/mscclpp_benchmark/tuner.py
 create mode 100644 python/mscclpp_benchmark/tuning_config.py
 create mode 100644 test/unit/gpu_data_types_tests.cu

diff --git a/.azure-pipelines/templates/nccl-test.yml b/.azure-pipelines/templates/nccl-test.yml
index 585f5b48..550f5690 100644
--- a/.azure-pipelines/templates/nccl-test.yml
+++ b/.azure-pipelines/templates/nccl-test.yml
@@ -74,6 +74,15 @@ steps:
       mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN -x MSCCLPP_ENABLE_NCCL_FALLBACK=TRUE -x MSCCLPP_NCCL_LIB_PATH=/root/nccl/build/lib/libnccl.so -x MSCCLPP_FORCE_NCCL_FALLBACK_OPERATION="broadcast" /root/nccl-tests/build/broadcast_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
       mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN -x MSCCLPP_ENABLE_NCCL_FALLBACK=TRUE -x MSCCLPP_NCCL_LIB_PATH=/root/nccl/build/lib/libnccl.so -x MSCCLPP_FORCE_NCCL_FALLBACK_OPERATION="allreduce" /root/nccl-tests/build/broadcast_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
 
+- template: run-remote-task.yml
+  parameters:
+    name: PyBench
+    displayName: Run Collective Benchmarks
+    remoteScript: |
+      mpirun --allow-run-as-root -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3b15 --accum-type float32 --autotune --symmetric-memory
+      mpirun --allow-run-as-root -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3fn --accum-type float16 --autotune --symmetric-memory
+      mpirun --allow-run-as-root -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float16 --symmetric-memory --autotune
+
 - template: stop.yml
   parameters:
     subscription: ${{ parameters.subscription }}
diff --git a/.azure-pipelines/templates/rccl-test.yml b/.azure-pipelines/templates/rccl-test.yml
index 8e247161..63788ac2 100644
--- a/.azure-pipelines/templates/rccl-test.yml
+++ b/.azure-pipelines/templates/rccl-test.yml
@@ -57,6 +57,15 @@ steps:
       mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
       mpirun -np 8 --bind-to numa --allow-run-as-root /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
 
+- template: run-remote-task.yml
+  parameters:
+    name: PyBench
+    displayName: Run Collective Benchmarks
+    remoteScript: |
+      mpirun --allow-run-as-root -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3b15 --accum-type float32 --autotune
+      mpirun --allow-run-as-root -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3fnuz --accum-type float32 --autotune
+      mpirun --allow-run-as-root -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allgather --dtype float8_e4m3b15 --autotune --buffer-mode out-of-place
+
 - template: stop.yml
   parameters:
     subscription: ${{ parameters.subscription }}
diff --git a/docs/quickstart.md b/docs/quickstart.md
index 716fcf61..320a2db7 100644
--- a/docs/quickstart.md
+++ b/docs/quickstart.md
@@ -110,12 +110,12 @@ $ CXX=/opt/rocm/bin/hipcc python -m pip install ".[rocm6]"
 ```
 
 > **Note:** A platform extra (`cuda11`, `cuda12`, `cuda13`, or `rocm6`) is required to install CuPy.
-> The CUDA extras install pre-built CuPy wheels. The `rocm6` extra installs CuPy from source,
-> which requires ROCm and may take longer. Running `pip install .` without an extra will not install CuPy.
+> The CUDA extras install pre-built CuPy wheels and CUDA Python bindings. The `rocm6` extra installs CuPy from source
+> and HIP Python 6.x, which require ROCm and may take longer. Running `pip install .` without an extra will not install CuPy.
 
 Optional extras can be installed by specifying them in brackets. Available extras:
-- **`cuda11`**, **`cuda12`**, **`cuda13`**: Install a pre-built CuPy package for your CUDA version.
-- **`rocm6`**: Install CuPy from source for AMD ROCm platforms.
+- **`cuda11`**, **`cuda12`**, **`cuda13`**: Install a pre-built CuPy package and CUDA Python bindings for your CUDA version.
+- **`rocm6`**: Install CuPy from source and HIP Python 6.x for AMD ROCm platforms.
 - **`benchmark`**: Install benchmark dependencies (mpi4py, prettytable, netifaces, matplotlib).
 - **`test`**: Install test dependencies (pytest, mpi4py, netifaces).
 
@@ -209,15 +209,37 @@ $ mpirun -np 16 -npernode 8 -hostfile hostfile ./bin/mp_unit_tests -ip_port 10.0
 
 ## Performance Benchmark
 
-### Python Benchmark
+### Python Benchmark and Tuning
 
-[Install the MSCCL++ Python package](#install-from-source-python-module) and run our Python AllReduce benchmark as follows. It requires MPI on the system.
+[Install the MSCCL++ Python package](#install-from-source-python-module) and run the Python collective benchmark as follows. It requires MPI on the system.
 
 ```bash
 # Install with benchmark dependencies and the appropriate CUDA/ROCm extras.
 # Replace `cuda12` with your platform: cuda11, cuda12, cuda13, or rocm6.
 $ python3 -m pip install ".[cuda12,benchmark,test]"
-$ mpirun -tag-output -np 8 python3 ./python/mscclpp_benchmark/allreduce_bench.py
+
+```
+
+To autotune launch parameters and save a tuned config:
+
+```bash
+$ PYTHONPATH=$PWD/python mpirun -np 8 --allow-run-as-root \
+    python3 -m mscclpp_benchmark.bench_collective \
+    --collective allreduce \
+    --dtype float16 \
+    --batch-sizes 1,2,4,8 \
+    --autotune \
+    --write-config /tmp/mscclpp_tuned_configs.json
+```
+
+Use the tuned config in a benchmark:
+
+```bash
+$ PYTHONPATH=$PWD/python mpirun -np 8 --allow-run-as-root \
+    python3 -m mscclpp_benchmark.bench_collective \
+    --collective allreduce \
+    --dtype float16 \
+    --config-path /tmp/mscclpp_tuned_configs.json
 ```
 
 (nccl-benchmark)=
@@ -291,4 +313,3 @@ Version: 0.8.0.post1.dev0+gc632fee37.d20251007
 mscclpp.version
 {'version': '0.8.0.post1.dev0+gc632fee37.d20251007', 'git_commit': 'g50382c567'}
 ```
-
diff --git a/include/mscclpp/gpu_data_types.hpp b/include/mscclpp/gpu_data_types.hpp
index 672434f9..4a16628e 100644
--- a/include/mscclpp/gpu_data_types.hpp
+++ b/include/mscclpp/gpu_data_types.hpp
@@ -71,7 +71,7 @@ using __bfloat162 = __nv_bfloat162;
 
 /// Software float8 with 4 exponent bits, 3 mantissa bits, exponent bias = 15.
 /// Format (MSB first): [sign:1][exponent:4][mantissa:3]
-/// No infinities, no NaN. Encode saturates to ±1.75 (0x7e/0xfe).
+/// No infinities, no NaN. Encode saturates to ±1.875 (0x7f/0xff).
 /// Adapted from the Triton compiler's fp8e4b15 format.
 struct alignas(1) __fp8_e4m3b15 {
   uint8_t __x;
@@ -103,7 +103,7 @@ struct alignas(1) __fp8_e4m3b15 {
   /// then convert fp16 → float32.
   static MSCCLPP_HOST_DEVICE_INLINE float toFloat(uint8_t bits) {
     // Branch-free decode: fp8 → fp16 → fp32, no special-case handling.
-    // Encode saturates to ±1.75, so 0x7f/0xff are never produced.
+    // Every byte maps to a finite value; encode saturates at ±1.875, so 0x7f/0xff decode to ±1.875.
     // Refer:
     // https://github.com/triton-lang/triton/blob/cf34004b8a67d290a962da166f5aa2fc66751326/python/triton/language/extra/cuda/utils.py#L34
     uint16_t h = (uint16_t)bits << 8;             // place fp8 in upper byte of fp16
@@ -132,10 +132,9 @@ struct alignas(1) __fp8_e4m3b15 {
     } cvt = {h_val};
     uint16_t fp16_bits = cvt.u;
 
-    // Clamp abs to max encodable value: 1.75 → fp16 = 0x3F00.
-    // Matches Triton: encode saturates, 0x7f/0xff are never produced.
+    // Clamp abs to max encodable value: 1.875 → fp16 = 0x3F80 (largest byte 0x7f/0xff).
     uint16_t abs_fp16 = fp16_bits & 0x7FFFu;
-    if (abs_fp16 > 0x3F00u) abs_fp16 = 0x3F00u;
+    if (abs_fp16 > 0x3F80u) abs_fp16 = 0x3F80u;
 
     // Reconstruct with sign.
     uint16_t sign16 = fp16_bits & 0x8000u;
@@ -852,27 +851,17 @@ MSCCLPP_DEVICE_INLINE f32x4 to<f32x4, f8_e5m2x4>(const f8_e5m2x4& v) {
 
 /// f32x2 -> f8_e4m3x2.
 /// HIP gfx942: float -> fp8 (via __builtin_amdgcn_cvt_pk_fp8_f32).
-/// NVIDIA SM90+: float -> half -> fp8 (via __nv_cvt_halfraw2_to_fp8x2).
-/// NVIDIA pre-SM90: float -> half -> fp8 (via __nv_cvt_halfraw_to_fp8, element-wise).
+/// NVIDIA: float -> fp8 directly (via __nv_cvt_float2_to_fp8x2). On SM89+ this maps to a
+/// single hardware round-to-nearest-even instruction; on older arch it falls back to a
+/// software direct conversion.
 template <>
 MSCCLPP_DEVICE_INLINE f8_e4m3x2 to<f8_e4m3x2, f32x2>(const f32x2& v) {
 #if defined(MSCCLPP_DEVICE_HIP) && defined(__gfx942__)
   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) && __CUDA_ARCH__ >= 900
-  __half2_raw h2;
-  h2.x = bit_cast<unsigned short>(__float2half_rn(v.data[0]));
-  h2.y = bit_cast<unsigned short>(__float2half_rn(v.data[1]));
-  __nv_fp8x2_storage_t fp8x2 = __nv_cvt_halfraw2_to_fp8x2(h2, __NV_SATFINITE, __NV_E4M3);
-  return bit_cast<f8_e4m3x2>(fp8x2);
 #elif defined(MSCCLPP_DEVICE_CUDA)
-  __half_raw h0, h1;
-  h0.x = bit_cast<unsigned short>(__float2half_rn(v.data[0]));
-  h1.x = bit_cast<unsigned short>(__float2half_rn(v.data[1]));
-  f8_e4m3x2 result;
-  result.data[0] = bit_cast<__fp8_e4m3>(__nv_cvt_halfraw_to_fp8(h0, __NV_SATFINITE, __NV_E4M3));
-  result.data[1] = bit_cast<__fp8_e4m3>(__nv_cvt_halfraw_to_fp8(h1, __NV_SATFINITE, __NV_E4M3));
-  return result;
+  __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2(make_float2(v.data[0], v.data[1]), __NV_SATFINITE, __NV_E4M3);
+  return bit_cast<f8_e4m3x2>(fp8x2);
 #else
   f8_e4m3x2 result;
   result.data[0] = static_cast<__fp8_e4m3>(v.data[0]);
@@ -909,27 +898,17 @@ MSCCLPP_DEVICE_INLINE f8_e4m3x4 to<f8_e4m3x4, f32x4>(const f32x4& v) {
 
 /// f32x2 -> f8_e5m2x2.
 /// HIP gfx942: float -> bf8 (via __builtin_amdgcn_cvt_pk_bf8_f32).
-/// NVIDIA SM90+: float -> half -> fp8 (via __nv_cvt_halfraw2_to_fp8x2 with __NV_E5M2).
-/// NVIDIA pre-SM90: float -> half -> fp8 (via __nv_cvt_halfraw_to_fp8, element-wise).
+/// NVIDIA: float -> fp8 directly (via __nv_cvt_float2_to_fp8x2 with __NV_E5M2). On SM89+ this
+/// maps to a single hardware round-to-nearest-even instruction; on older arch it falls back to a
+/// software direct conversion.
 template <>
 MSCCLPP_DEVICE_INLINE f8_e5m2x2 to<f8_e5m2x2, f32x2>(const f32x2& v) {
 #if defined(MSCCLPP_DEVICE_HIP) && defined(__gfx942__)
   uint32_t packed = __builtin_amdgcn_cvt_pk_bf8_f32(v.data[0], v.data[1], 0, false);
   return bit_cast<f8_e5m2x2>(static_cast<__hip_fp8x2_storage_t>(packed));
-#elif defined(MSCCLPP_DEVICE_CUDA) && __CUDA_ARCH__ >= 900
-  __half2_raw h2;
-  h2.x = bit_cast<unsigned short>(__float2half_rn(v.data[0]));
-  h2.y = bit_cast<unsigned short>(__float2half_rn(v.data[1]));
-  __nv_fp8x2_storage_t fp8x2 = __nv_cvt_halfraw2_to_fp8x2(h2, __NV_SATFINITE, __NV_E5M2);
-  return bit_cast<f8_e5m2x2>(fp8x2);
 #elif defined(MSCCLPP_DEVICE_CUDA)
-  __half_raw h0, h1;
-  h0.x = bit_cast<unsigned short>(__float2half_rn(v.data[0]));
-  h1.x = bit_cast<unsigned short>(__float2half_rn(v.data[1]));
-  f8_e5m2x2 result;
-  result.data[0] = bit_cast<__fp8_e5m2>(__nv_cvt_halfraw_to_fp8(h0, __NV_SATFINITE, __NV_E5M2));
-  result.data[1] = bit_cast<__fp8_e5m2>(__nv_cvt_halfraw_to_fp8(h1, __NV_SATFINITE, __NV_E5M2));
-  return result;
+  __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2(make_float2(v.data[0], v.data[1]), __NV_SATFINITE, __NV_E5M2);
+  return bit_cast<f8_e5m2x2>(fp8x2);
 #else
   f8_e5m2x2 result;
   result.data[0] = static_cast<__fp8_e5m2>(v.data[0]);
@@ -1103,11 +1082,11 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x2 to<f8_e4m3b15x2, f16x2>(const f16x2& v) {
 #if defined(MSCCLPP_DEVICE_CUDA)
   uint32_t in0;
   asm("mov.b32 %0, %1;" : "=r"(in0) : "r"(*reinterpret_cast<const uint32_t*>(&v)));
-  // Clamp abs to max encodable e4m3b15 (0x3F00 = 1.75 in fp16).
+  // Clamp abs to max encodable e4m3b15 (0x3F80 = 1.875 in fp16).
   uint32_t lo = in0 & 0xFFFFu, hi = in0 >> 16;
   uint32_t alo = lo & 0x7FFFu, ahi = hi & 0x7FFFu;
-  alo = alo < 0x3F00u ? alo : 0x3F00u;
-  ahi = ahi < 0x3F00u ? ahi : 0x3F00u;
+  alo = alo < 0x3F80u ? alo : 0x3F80u;
+  ahi = ahi < 0x3F80u ? ahi : 0x3F80u;
   uint32_t a0 = alo | (ahi << 16);
   a0 = a0 * 2u + 0x00800080u;
   uint32_t b0 = a0 | (in0 & 0x80008000u);
@@ -1118,7 +1097,7 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x2 to<f8_e4m3b15x2, f16x2>(const f16x2& v) {
   uint32_t in0 = v.words[0];
   uint32_t abs0 = in0 & 0x7fff7fffu;
   uint32_t a0;
-  asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a0) : "v"(abs0), "v"(0x3F003F00u));
+  asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a0) : "v"(abs0), "v"(0x3F803F80u));
   a0 = a0 * 2u + 0x00800080u;
   uint32_t b0 = a0 | (in0 & 0x80008000u);
   uint16_t packed = (uint16_t)(((b0 >> 8) & 0xFFu) | ((b0 >> 16) & 0xFF00u));
@@ -1141,8 +1120,8 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x4 to<f8_e4m3b15x4, f16x4>(const f16x4& v) {
   asm("mov.b32 %0, %1;" : "=r"(in1) : "r"(v.words[1]));
   uint32_t abs0 = in0 & 0x7fff7fffu;
   uint32_t abs1 = in1 & 0x7fff7fffu;
-  uint32_t a0 = __vminu2(abs0, 0x3F003F00u);
-  uint32_t a1 = __vminu2(abs1, 0x3F003F00u);
+  uint32_t a0 = __vminu2(abs0, 0x3F803F80u);
+  uint32_t a1 = __vminu2(abs1, 0x3F803F80u);
   a0 = a0 * 2u + 0x00800080u;
   a1 = a1 * 2u + 0x00800080u;
   uint32_t b0, b1;
@@ -1155,8 +1134,8 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x4 to<f8_e4m3b15x4, f16x4>(const f16x4& v) {
   uint32_t in0 = v.words[0], in1 = v.words[1];
   uint32_t abs0 = in0 & 0x7fff7fffu, abs1 = in1 & 0x7fff7fffu;
   uint32_t a0, a1;
-  asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a0) : "v"(abs0), "v"(0x3F003F00u));
-  asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a1) : "v"(abs1), "v"(0x3F003F00u));
+  asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a0) : "v"(abs0), "v"(0x3F803F80u));
+  asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a1) : "v"(abs1), "v"(0x3F803F80u));
   a0 = a0 * 2u + 0x00800080u;
   a1 = a1 * 2u + 0x00800080u;
   uint32_t b0 = a0 | (in0 & 0x80008000u);
@@ -1268,8 +1247,8 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x4 to<f8_e4m3b15x4, f32x4>(const f32x4& v) {
   return to<f8_e4m3b15x4, f16x4>(h);
 #elif defined(MSCCLPP_DEVICE_HIP) && defined(__gfx942__)
   f16x4 h;
-  h.words[0] = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(v.data[0], v.data[1]));
-  h.words[1] = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(v.data[2], v.data[3]));
+  h.words[0] = __builtin_bit_cast(uint32_t, __floats2half2_rn(v.data[0], v.data[1]));
+  h.words[1] = __builtin_bit_cast(uint32_t, __floats2half2_rn(v.data[2], v.data[3]));
   return to<f8_e4m3b15x4, f16x4>(h);
 #else
   f8_e4m3b15x4 result;
diff --git a/pyproject.toml b/pyproject.toml
index 0ea569cb..b35b1b3a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,10 +21,22 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-cuda11 = ["cupy-cuda11x"]
-cuda12 = ["cupy-cuda12x"]
-cuda13 = ["cupy-cuda13x"]
-rocm6 = ["cupy"]
+cuda11 = [
+    "cupy-cuda11x",
+    "cuda-bindings>=11.8,<12",
+]
+cuda12 = [
+    "cupy-cuda12x",
+    "cuda-bindings>=12,<13",
+]
+cuda13 = [
+    "cupy-cuda13x",
+    "cuda-bindings>=13,<14",
+]
+rocm6 = [
+    "cupy",
+    "hip-python>=6,<7",
+]
 benchmark = [
     "mpi4py",
     "prettytable",
diff --git a/python/mscclpp_benchmark/__init__.py b/python/mscclpp_benchmark/__init__.py
index 1ee3f3bf..11e08c9b 100644
--- a/python/mscclpp_benchmark/__init__.py
+++ b/python/mscclpp_benchmark/__init__.py
@@ -1,4 +1,18 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from .mscclpp_op import MscclppAllReduce1, MscclppAllReduce2, MscclppAllReduce3, MscclppAllReduce4, MscclppAllReduce5
+__all__ = [
+    "MscclppAllReduce1",
+    "MscclppAllReduce2",
+    "MscclppAllReduce3",
+    "MscclppAllReduce4",
+    "MscclppAllReduce5",
+]
+
+
+def __getattr__(name):
+    if name in __all__:
+        from . import mscclpp_op
+
+        return getattr(mscclpp_op, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/python/mscclpp_benchmark/bench_collective.py b/python/mscclpp_benchmark/bench_collective.py
new file mode 100644
index 00000000..c526438d
--- /dev/null
+++ b/python/mscclpp_benchmark/bench_collective.py
@@ -0,0 +1,645 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from typing import Any
+
+import cupy as cp
+from mpi4py import MPI
+
+_mscclpp_module = None
+
+from mscclpp_benchmark.comm import Comm
+from mscclpp_benchmark.correctness import (
+    CorrectnessStats,
+    check_correctness as _check_correctness,
+    fill_case_for_benchmark as _fill_case_for_benchmark,
+)
+from mscclpp_benchmark.gpu import capture_graph, init_runtime
+from mscclpp_benchmark.tuner import OfflineTuner
+from mscclpp_benchmark.tuning_config import HardwareProfile, TunedConfig, TunedConfigStore, normalize_sku
+
+_ALLREDUCE = "allreduce"
+_ALLGATHER = "allgather"
+_DEFAULT_BATCH_SIZES = (
+    1,
+    2,
+    3,
+    4,
+    8,
+    16,
+    24,
+    32,
+    48,
+    64,
+    96,
+    128,
+    256,
+    512,
+    1024,
+    1280,
+    1536,
+    1792,
+    2048,
+    2560,
+    3072,
+    3584,
+    4096,
+)
+_DEFAULT_CANDIDATE_NBLOCKS = (1, 4, 8, 16, 24, 32, 48, 56, 64)
+_DEFAULT_CANDIDATE_NTHREADS = (256, 512, 768, 1024)
+
+
+def _mscclpp():
+    global _mscclpp_module
+    if _mscclpp_module is None:
+        import mscclpp
+        import mscclpp.ext
+
+        _mscclpp_module = mscclpp
+    return _mscclpp_module
+
+
+@dataclass(frozen=True)
+class DTypeSpec:
+    name: str
+    cupy_dtype: Any
+    mscclpp_dtype: Any
+    accum_dtype: Any | None = None
+    fp8_format: str | None = None
+
+
+@dataclass(frozen=True)
+class CandidateSpec:
+    algorithm: str
+    min_message_size: int | None = None
+    max_message_size: int | None = None
+    max_nblocks: int | None = None
+    supported_skus: tuple[str, ...] | None = None
+    requires_nvls: bool = False
+    requires_symmetric_memory: bool = False
+
+
+@dataclass
+class BenchmarkCase:
+    collective: str
+    message_size: int
+    total_size: int
+    input: cp.ndarray
+    output: cp.ndarray
+    dtype_spec: DTypeSpec
+    symmetric_memory: bool = False
+
+
+def _device_name() -> str:
+    props = cp.cuda.runtime.getDeviceProperties(cp.cuda.Device().id)
+    name = props.get("name", "UNKNOWN")
+    if isinstance(name, bytes):
+        return name.decode("utf-8")
+    return str(name)
+
+
+def _detect_hardware_profile(scale: int) -> HardwareProfile:
+    return HardwareProfile(sku=normalize_sku(_device_name()), scale=scale)
+
+
+def _parse_dtype(dtype_name: str) -> DTypeSpec:
+    mscclpp = _mscclpp()
+    normalized = dtype_name.strip().lower().replace("-", "_")
+    if normalized in {"float16", "fp16", "half"}:
+        return DTypeSpec("float16", cp.float16, mscclpp.DataType.float16)
+    if normalized in {"float32", "fp32", "float"}:
+        return DTypeSpec("float32", cp.float32, mscclpp.DataType.float32)
+    if normalized in {"int32", "i32"}:
+        return DTypeSpec("int32", cp.int32, mscclpp.DataType.int32)
+    if normalized in {"uint8", "u8"}:
+        return DTypeSpec("uint8", cp.uint8, mscclpp.DataType.uint8)
+    if normalized in {"float8_e4m3fn", "fp8_e4m3fn"}:
+        return DTypeSpec(
+            "float8_e4m3fn",
+            cp.uint8,
+            mscclpp.DataType.float8_e4m3fn,
+            accum_dtype=mscclpp.DataType.float16,
+            fp8_format="e4m3fn",
+        )
+    if normalized in {"float8_e4m3fnuz", "fp8_e4m3fnuz"}:
+        return DTypeSpec(
+            "float8_e4m3fnuz",
+            cp.uint8,
+            mscclpp.DataType.float8_e4m3fnuz,
+            accum_dtype=mscclpp.DataType.float16,
+            fp8_format="e4m3fnuz",
+        )
+    if normalized in {"float8_e4m3b15", "fp8_e4m3b15"}:
+        return DTypeSpec(
+            "float8_e4m3b15",
+            cp.uint8,
+            mscclpp.DataType.float8_e4m3b15,
+            accum_dtype=mscclpp.DataType.float32,
+            fp8_format="e4m3b15",
+        )
+    raise ValueError(
+        f"Unsupported dtype {dtype_name!r}; use float16, float32, int32, uint8, "
+        "float8_e4m3fn, float8_e4m3fnuz, or float8_e4m3b15"
+    )
+
+
+def _with_accum_type(dtype_spec: DTypeSpec, accum_type: str | None) -> DTypeSpec:
+    if accum_type is None:
+        return dtype_spec
+
+    mscclpp = _mscclpp()
+    normalized = accum_type.strip().lower().replace("-", "_")
+    if normalized in {"native", "same", "auto"}:
+        accum_dtype = dtype_spec.mscclpp_dtype
+    elif normalized in {"float16", "fp16", "half"}:
+        accum_dtype = mscclpp.DataType.float16
+    elif normalized in {"float32", "fp32", "float"}:
+        accum_dtype = mscclpp.DataType.float32
+    else:
+        raise ValueError(f"Unsupported accum type {accum_type!r}; use native, float16, or float32")
+
+    return DTypeSpec(
+        name=dtype_spec.name,
+        cupy_dtype=dtype_spec.cupy_dtype,
+        mscclpp_dtype=dtype_spec.mscclpp_dtype,
+        accum_dtype=accum_dtype,
+        fp8_format=dtype_spec.fp8_format,
+    )
+
+
+def _human_size(size: int) -> str:
+    value = float(size)
+    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
+        if value < 1024.0 or unit == "TiB":
+            return f"{value:.1f} {unit}"
+        value /= 1024.0
+    raise AssertionError("unreachable")
+
+
+def _parse_int_list(raw: str | None, default: tuple[int, ...]) -> tuple[int, ...]:
+    if raw is None:
+        return default
+    values = tuple(sorted({int(item.strip()) for item in raw.split(",") if item.strip()}))
+    if not values or values[0] <= 0:
+        raise ValueError(f"Expected a comma-separated list of positive integers, got {raw!r}")
+    return values
+
+
+def _candidate_specs(collective: str, *, symmetric_memory: bool = False) -> tuple[CandidateSpec, ...]:
+    if collective == _ALLGATHER:
+        return (CandidateSpec("default_allgather_fullmesh2", max_nblocks=64, supported_skus=("MI300X",)),)
+    if collective != _ALLREDUCE:
+        raise ValueError(f"Unsupported collective: {collective}")
+    candidates = (
+        CandidateSpec(
+            "default_allreduce_nvls_packet",
+            max_message_size=512 * 1024,
+            max_nblocks=16,
+            supported_skus=("H100", "GB300"),
+            requires_nvls=True,
+        ),
+        CandidateSpec(
+            "default_allreduce_packet",
+            max_message_size=4 * 1024 * 1024,
+            max_nblocks=56,
+        ),
+        CandidateSpec(
+            "default_allreduce_allpair_packet",
+            max_message_size=4 * 1024 * 1024,
+            max_nblocks=56,
+        ),
+        CandidateSpec(
+            "default_allreduce_rsag_zero_copy",
+            min_message_size=512 * 1024 + 1,
+        ),
+        CandidateSpec(
+            "default_allreduce_fullmesh",
+            min_message_size=512 * 1024 + 1,
+            max_nblocks=64,
+            supported_skus=("MI300X",),
+        ),
+    )
+    if symmetric_memory:
+        return (
+            CandidateSpec(
+                "default_allreduce_nvls_zero_copy",
+                max_nblocks=32,
+                supported_skus=("H100", "GB300"),
+                requires_nvls=True,
+                requires_symmetric_memory=True,
+            ),
+            *candidates,
+        )
+    return candidates
+
+
+def _candidate_algorithms(comm: Comm, case: BenchmarkCase) -> list[tuple[Any, CandidateSpec]]:
+    available = comm.algorithms.get(case.collective, {})
+    candidates: list[tuple[Any, CandidateSpec]] = []
+    seen: set[str] = set()
+    symmetric_memory = case.symmetric_memory
+    profile = getattr(comm, "hardware_profile", None)
+    filtered_out = False
+    for candidate in _candidate_specs(case.collective, symmetric_memory=symmetric_memory):
+        if not _candidate_supports_profile(candidate, profile):
+            filtered_out = True
+            continue
+        if not _candidate_supports_message_size(candidate, case.message_size):
+            filtered_out = True
+            continue
+        if candidate.requires_nvls and not _mscclpp().is_nvls_supported():
+            filtered_out = True
+            continue
+        if candidate.requires_symmetric_memory and not symmetric_memory:
+            filtered_out = True
+            continue
+        algorithm = available.get(candidate.algorithm)
+        if algorithm is None or algorithm.name in seen:
+            continue
+        seen.add(algorithm.name)
+        candidates.append((algorithm, candidate))
+    if candidates:
+        return candidates
+    if filtered_out:
+        return []
+    return [(algorithm, CandidateSpec(algorithm.name)) for algorithm in available.values()]
+
+
+def _candidate_supports_profile(candidate: CandidateSpec, profile: HardwareProfile | None) -> bool:
+    if candidate.supported_skus is None:
+        return True
+    sku = None if profile is None else profile.sku
+    if not sku or sku == "UNKNOWN":
+        return True
+    return sku in candidate.supported_skus
+
+
+def _candidate_supports_message_size(candidate: CandidateSpec, message_size: int) -> bool:
+    if candidate.min_message_size is not None and message_size < candidate.min_message_size:
+        return False
+    if candidate.max_message_size is not None and message_size > candidate.max_message_size:
+        return False
+    return True
+
+
+def _make_case(
+    *,
+    collective: str,
+    nelems: int,
+    dtype_spec: DTypeSpec,
+    comm_group: Any,
+    buffer_mode: str,
+    symmetric_memory: bool = False,
+) -> BenchmarkCase:
+    if buffer_mode not in ("in-place", "out-of-place"):
+        raise ValueError(f"Unsupported buffer mode: {buffer_mode}")
+
+    if collective == _ALLREDUCE:
+        if buffer_mode == "in-place":
+            memory = _mscclpp().GpuBuffer(nelems, dtype=dtype_spec.cupy_dtype)
+            input_buffer = memory
+            output = memory
+        else:
+            input_buffer = _mscclpp().GpuBuffer(nelems, dtype=dtype_spec.cupy_dtype)
+            output = _mscclpp().GpuBuffer(nelems, dtype=dtype_spec.cupy_dtype)
+        return BenchmarkCase(
+            collective=collective,
+            message_size=input_buffer.nbytes,
+            total_size=output.nbytes,
+            input=input_buffer,
+            output=output,
+            dtype_spec=dtype_spec,
+            symmetric_memory=symmetric_memory,
+        )
+
+    if collective != _ALLGATHER:
+        raise ValueError(f"Unsupported collective: {collective}")
+
+    if buffer_mode == "in-place":
+        output = _mscclpp().GpuBuffer(nelems * comm_group.nranks, dtype=dtype_spec.cupy_dtype)
+        start = comm_group.my_rank * nelems
+        input_buffer = output[start : start + nelems]
+    else:
+        input_buffer = _mscclpp().GpuBuffer(nelems, dtype=dtype_spec.cupy_dtype)
+        output = _mscclpp().GpuBuffer(nelems * comm_group.nranks, dtype=dtype_spec.cupy_dtype)
+
+    return BenchmarkCase(
+        collective=collective,
+        message_size=input_buffer.nbytes,
+        total_size=output.nbytes,
+        input=input_buffer,
+        output=output,
+        dtype_spec=dtype_spec,
+        symmetric_memory=symmetric_memory,
+    )
+
+
+def _try_measure_case(
+    comm: Comm,
+    case: BenchmarkCase,
+    config: TunedConfig,
+    *,
+    n_warmup: int,
+    n_graph_launches: int,
+    n_ops_per_graph: int,
+) -> float | None:
+    try:
+        return _measure_case(
+            comm,
+            case,
+            config,
+            n_warmup=n_warmup,
+            n_graph_launches=n_graph_launches,
+            n_ops_per_graph=n_ops_per_graph,
+        )
+    except Exception as exc:
+        if comm.rank == 0:
+            print(
+                f"[skip] {config.algorithm} nb={config.nblocks} nt={config.nthreads} "
+                f"size={case.message_size}: {type(exc).__name__}: {exc}",
+                flush=True,
+            )
+        return None
+
+
+def _measure_case(
+    comm: Comm,
+    case: BenchmarkCase,
+    config: TunedConfig,
+    *,
+    n_warmup: int,
+    n_graph_launches: int,
+    n_ops_per_graph: int,
+) -> float:
+    _fill_case_for_benchmark(case, comm.rank)
+    comm.comm_group.barrier()
+    if comm.run(case, config) != 0:
+        raise RuntimeError("algorithm returned non-zero status")
+    cp.cuda.runtime.deviceSynchronize()
+    comm.comm_group.barrier()
+
+    stream = cp.cuda.Stream(non_blocking=True)
+    graph = None
+
+    def capture_ops() -> None:
+        for _ in range(n_ops_per_graph):
+            ret = comm.run(case, config, stream)
+            if ret != 0:
+                raise RuntimeError("algorithm returned non-zero status during graph capture")
+
+    try:
+        with stream:
+            graph = capture_graph(stream, capture_ops)
+
+        for _ in range(n_warmup):
+            graph.launch(stream)
+        stream.synchronize()
+        comm.comm_group.barrier()
+
+        start = cp.cuda.Event()
+        end = cp.cuda.Event()
+        start.record(stream)
+        for _ in range(n_graph_launches):
+            graph.launch(stream)
+        end.record(stream)
+        end.synchronize()
+
+        elapsed_us = cp.cuda.get_elapsed_time(start, end) * 1000.0 / (n_graph_launches * n_ops_per_graph)
+        return float(MPI.COMM_WORLD.allreduce(elapsed_us, op=MPI.MAX))
+    finally:
+        if graph is not None:
+            graph.close()
+
+
+def _bandwidth_gbps(num_bytes: int, time_us: float) -> float:
+    return num_bytes / time_us / 1e3
+
+
+def _busbw_factor(collective: str, nranks: int) -> float:
+    if nranks <= 1:
+        return 1.0
+    if collective == _ALLREDUCE:
+        return 2 * (nranks - 1) / nranks
+    if collective == _ALLGATHER:
+        return (nranks - 1) / nranks
+    raise ValueError(f"Unsupported collective: {collective}")
+
+
+def _format_table(headers: list[str], rows: list[list[str]]) -> str:
+    widths = [len(header) for header in headers]
+    for row in rows:
+        widths = [max(width, len(cell)) for width, cell in zip(widths, row)]
+    header_line = " | ".join(header.ljust(width) for header, width in zip(headers, widths))
+    sep_line = "-+-".join("-" * width for width in widths)
+    row_lines = [" | ".join(cell.ljust(width) for cell, width in zip(row, widths)) for row in rows]
+    return "\n".join([header_line, sep_line, *row_lines])
+
+
+def _format_stat(value: float | None) -> str:
+    if value is None:
+        return "-"
+    return f"{value:.6g}"
+
+
+def _format_mismatches(stats: CorrectnessStats | None) -> str:
+    if stats is None or stats.total == 0:
+        return "-"
+    return f"{stats.mismatches}/{stats.total}"
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Benchmark MSCCL++ collectives without PyTorch dependencies")
+    parser.add_argument("--collective", choices=(_ALLREDUCE, _ALLGATHER), default=_ALLREDUCE)
+    parser.add_argument("--d-model", type=int, default=5120)
+    parser.add_argument("--dtype", default="float16")
+    parser.add_argument("--accum-type", help="Accumulation type for reductions: native, float16, or float32")
+    parser.add_argument("--batch-sizes", help="Comma-separated batch sizes; default uses the benchmark sweep")
+    parser.add_argument(
+        "--buffer-mode",
+        choices=("in-place", "out-of-place"),
+        default="in-place",
+        help="Buffer layout for the collective: in-place (input aliases output) or out-of-place (separate buffers)",
+    )
+    parser.add_argument("--config-path", help="Optional MSCCL++ tuned config JSON")
+    parser.add_argument("--write-config", help="Write autotuned configs to this JSON path")
+    parser.add_argument("--autotune", action="store_true", help="Tune each benchmark size before timing it")
+    parser.add_argument("--skip-correctness", action="store_true")
+    parser.add_argument("--correctness-iters", type=int, default=1)
+    parser.add_argument("--scratch-buffer-size", type=int, default=1 << 27)
+    parser.add_argument("--warmup", type=int, default=5, help="Warmup graph replays before benchmark timing")
+    parser.add_argument("--graph-launches", type=int, default=10, help="Timed graph replays")
+    parser.add_argument("--iterations", type=int, default=100, help="Collective operations captured per CUDA graph")
+    parser.add_argument("--tune-warmup", type=int, default=2)
+    parser.add_argument("--tune-graph-launches", type=int, default=3)
+    parser.add_argument("--tune-iterations", type=int, default=20)
+    parser.add_argument("--candidate-nblocks", help="Comma-separated nblocks tuning candidates")
+    parser.add_argument("--candidate-nthreads", help="Comma-separated nthreads tuning candidates")
+    parser.add_argument("--symmetric-memory", action="store_true")
+    return parser
+
+
+def _validate_args(args: argparse.Namespace) -> None:
+    for name in (
+        "d_model",
+        "scratch_buffer_size",
+        "graph_launches",
+        "iterations",
+        "tune_graph_launches",
+        "tune_iterations",
+        "correctness_iters",
+    ):
+        if getattr(args, name) <= 0:
+            raise ValueError(f"--{name.replace('_', '-')} must be positive")
+    if args.warmup < 0 or args.tune_warmup < 0:
+        raise ValueError("warmup counts must be non-negative")
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = _build_parser().parse_args(argv)
+    _validate_args(args)
+    init_runtime()
+
+    local_comm = MPI.COMM_WORLD.Split_type(MPI.COMM_TYPE_SHARED, 0, MPI.INFO_NULL)
+    try:
+        visible_devices = cp.cuda.runtime.getDeviceCount()
+        if visible_devices <= 0:
+            raise RuntimeError("MSCCL++ benchmark requires at least one visible GPU")
+        cp.cuda.Device(local_comm.Get_rank() % visible_devices).use()
+    finally:
+        local_comm.Free()
+
+    dtype_spec = _with_accum_type(_parse_dtype(args.dtype), args.accum_type)
+    batch_sizes = _parse_int_list(args.batch_sizes, _DEFAULT_BATCH_SIZES)
+    candidate_nblocks = _parse_int_list(args.candidate_nblocks, _DEFAULT_CANDIDATE_NBLOCKS)
+    candidate_nthreads = _parse_int_list(args.candidate_nthreads, _DEFAULT_CANDIDATE_NTHREADS)
+
+    comm_group = _mscclpp().CommGroup(MPI.COMM_WORLD)
+    setattr(comm_group, "_mpi_comm", MPI.COMM_WORLD)
+    hardware_profile = _detect_hardware_profile(comm_group.nranks)
+    config_store = TunedConfigStore.load_path(args.config_path) if args.config_path else TunedConfigStore.empty()
+    comm = Comm(
+        comm_group,
+        config_store=config_store,
+        hardware_profile=hardware_profile,
+        scratch_buffer_size=args.scratch_buffer_size,
+    )
+    tuner = OfflineTuner(
+        comm,
+        candidate_nblocks=candidate_nblocks,
+        candidate_nthreads=candidate_nthreads,
+        n_warmup=args.tune_warmup,
+        n_graph_launches=args.tune_graph_launches,
+        n_ops_per_graph=args.tune_iterations,
+        candidate_algorithms=_candidate_algorithms,
+        check_correctness=_check_correctness,
+        measure=_try_measure_case,
+    )
+
+    rows: list[list[str]] = []
+    try:
+        if comm.rank == 0:
+            print(
+                f"MSCCL++ {args.collective} benchmark: profile={hardware_profile} dtype={dtype_spec.name} "
+                f"graph_launches={args.graph_launches} iterations={args.iterations}",
+                flush=True,
+            )
+
+        for batch_size in batch_sizes:
+            nelems = batch_size * args.d_model
+            case = _make_case(
+                collective=args.collective,
+                nelems=nelems,
+                dtype_spec=dtype_spec,
+                comm_group=comm_group,
+                buffer_mode=args.buffer_mode,
+                symmetric_memory=args.symmetric_memory,
+            )
+            config = tuner.tune(case) if args.autotune else comm.resolve_config(case)
+            if config is None:
+                continue
+            if args.autotune:
+                config_store.upsert(hardware_profile, args.collective, case.message_size, config)
+
+            correctness = "SKIP"
+            correctness_stats: CorrectnessStats | None = None
+            if not args.skip_correctness:
+                correctness_stats = _check_correctness(comm, case, config, niter=args.correctness_iters)
+                correctness = "PASS" if correctness_stats else "FAIL"
+                comm.reset(config)
+                if correctness != "PASS":
+                    raise RuntimeError(
+                        f"Correctness failed for batch_size={batch_size}, message_size={case.message_size}, "
+                        f"config={config}"
+                    )
+
+            time_us = _measure_case(
+                comm,
+                case,
+                config,
+                n_warmup=args.warmup,
+                n_graph_launches=args.graph_launches,
+                n_ops_per_graph=args.iterations,
+            )
+            comm.reset(config)
+
+            algbw = _bandwidth_gbps(case.total_size, time_us)
+            busbw = algbw * _busbw_factor(args.collective, comm_group.nranks)
+            rows.append(
+                [
+                    str(batch_size),
+                    _human_size(case.message_size),
+                    _human_size(case.total_size),
+                    config.algorithm,
+                    str(config.nblocks or "auto"),
+                    str(config.nthreads or "auto"),
+                    f"{time_us:.2f}",
+                    f"{algbw:.2f}",
+                    f"{busbw:.2f}",
+                    correctness,
+                    _format_stat(None if correctness_stats is None else correctness_stats.max_abs_diff),
+                    _format_stat(None if correctness_stats is None else correctness_stats.mean_abs_diff),
+                    _format_mismatches(correctness_stats),
+                ]
+            )
+            if comm.rank == 0:
+                print(".", end="", flush=True)
+
+        if args.write_config and comm.rank == 0:
+            config_store.write_path(args.write_config)
+            print(f"\nWrote tuned config to {args.write_config}", flush=True)
+
+        if comm.rank == 0:
+            print(
+                "\n"
+                + _format_table(
+                    [
+                        "batch",
+                        "msg",
+                        "total",
+                        "algorithm",
+                        "nblocks",
+                        "nthreads",
+                        "time_us",
+                        "algBW_GB/s",
+                        "busBW_GB/s",
+                        "check",
+                        "max_diff",
+                        "mean_diff",
+                        "mismatch",
+                    ],
+                    rows,
+                ),
+                flush=True,
+            )
+    finally:
+        comm_group.barrier()
+        cp.cuda.runtime.deviceSynchronize()
+        comm.close()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/python/mscclpp_benchmark/comm.py b/python/mscclpp_benchmark/comm.py
new file mode 100644
index 00000000..23770ac2
--- /dev/null
+++ b/python/mscclpp_benchmark/comm.py
@@ -0,0 +1,409 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+_ALLREDUCE_COLLECTIVE = "allreduce"
+_ALLGATHER_COLLECTIVE = "allgather"
+_mscclpp_module = None
+
+from mscclpp_benchmark.gpu import current_device, device_name, set_device
+from mscclpp_benchmark.tuning_config import HardwareProfile, TunedConfig, TunedConfigStore, normalize_sku
+
+
+def _mscclpp():
+    global _mscclpp_module
+    if _mscclpp_module is None:
+        import mscclpp
+        import mscclpp.ext
+
+        _mscclpp_module = mscclpp
+    return _mscclpp_module
+
+
+class Buffer:
+    def __init__(
+        self,
+        nbytes: int | None = None,
+        *,
+        dtype: str | Any = "float16",
+        shape: tuple[int, ...] | None = None,
+        buffer: Any | None = None,
+    ) -> None:
+        self.dtype = dtype
+        self.element_size = _dtype_size(dtype)
+        if buffer is None:
+            if nbytes is None:
+                if shape is None:
+                    raise ValueError("Either nbytes or shape is required")
+                nbytes = _numel(shape) * self.element_size
+            _ensure_device()
+            buffer = _mscclpp().RawGpuBuffer(int(nbytes))
+        self.buffer = buffer
+        self.nbytes = int(buffer.bytes())
+        self.shape = shape if shape is not None else (self.nbytes // self.element_size,)
+
+    @property
+    def ndim(self) -> int:
+        return len(self.shape)
+
+    @property
+    def size(self) -> int:
+        return _numel(self.shape)
+
+    def data_ptr(self) -> int:
+        return int(self.buffer.data())
+
+
+class _AllReduceOp:
+    def __init__(self, comm: "Comm", x: Any, *, symmetric_memory: bool = False) -> None:
+        self._comm = comm
+        self._x = x
+        self._symmetric_memory = symmetric_memory
+
+    def __call__(self, **_: Any) -> Any:
+        self._comm.run(self._x, symmetric_memory=self._symmetric_memory)
+        return self._x
+
+
+class _AllGatherOp:
+    def __init__(self, comm: "Comm", x: Any, *, dim: int, y: Any | None = None, symmetric_memory: bool = False) -> None:
+        shape = _shape(x)
+        if len(shape) == 0:
+            raise ValueError("MSCCL++ allgather requires a non-scalar buffer")
+        if dim % len(shape) != 0:
+            raise NotImplementedError("Raw-buffer allgather currently supports only dim=0")
+        if y is None:
+            y_shape = (comm._scale() * shape[0], *shape[1:])
+            y = Buffer(dtype=_dtype(x), shape=y_shape)
+        self._comm = comm
+        self._x = x
+        self.y = y
+        self._symmetric_memory = symmetric_memory
+
+    def __call__(self, **_: Any) -> Any:
+        self._comm.run(
+            self._x,
+            collective=_ALLGATHER_COLLECTIVE,
+            output_tensor=self.y,
+            symmetric_memory=self._symmetric_memory,
+        )
+        return self.y
+
+
+class Comm:
+    """Runtime MSCCL++ wrapper that owns algorithm handles and execution without Torch/CuPy tensors."""
+
+    def __init__(
+        self,
+        comm_group: Any,
+        scratch_buffer_size: int = 1 << 27,
+        *,
+        config_store: "TunedConfigStore | None" = None,
+        hardware_profile: HardwareProfile | None = None,
+    ) -> None:
+        self._comm_group = comm_group
+        self._mpi_comm = getattr(comm_group, "_mpi_comm", None)
+        self._rank = comm_group.my_rank
+        self._closed = False
+        _ensure_device()
+        self._mscclpp = _mscclpp()
+        self._scratch_buffer = self._mscclpp.RawGpuBuffer(scratch_buffer_size)
+        self._config_store = TunedConfigStore.empty() if config_store is None else config_store
+        self._hardware_profile = (
+            _detect_hardware_profile(scale=self._scale()) if hardware_profile is None else hardware_profile
+        )
+        self._default_config_warning_keys: set[tuple[str, str, str, int]] = set()
+
+        algorithms = self._mscclpp.ext.AlgorithmCollectionBuilder().build_default_algorithms(
+            scratch_buffer=self._scratch_buffer.data(),
+            scratch_buffer_size=self._scratch_buffer.bytes(),
+            rank=self._rank,
+        )
+        self._algorithms_by_collective: dict[str, dict[str, Any]] = {}
+        for algorithm in algorithms:
+            self._algorithms_by_collective.setdefault(algorithm.collective, {})[algorithm.name] = algorithm
+
+    @property
+    def comm_group(self) -> Any:
+        return self._comm_group
+
+    @property
+    def rank(self) -> int:
+        return self._rank
+
+    @property
+    def nranks(self) -> int:
+        return self._comm_group.nranks
+
+    @property
+    def algorithms(self) -> dict[str, dict[str, Any]]:
+        return self._algorithms_by_collective
+
+    @property
+    def hardware_profile(self) -> HardwareProfile:
+        return self._hardware_profile
+
+    def make_allreduce(self, x: Any, *, symmetric_memory: bool = False) -> _AllReduceOp:
+        return _AllReduceOp(self, x, symmetric_memory=symmetric_memory)
+
+    def make_allgather(self, x: Any, dim: int, y: Any | None = None, *, symmetric_memory: bool = False) -> _AllGatherOp:
+        return _AllGatherOp(self, x, dim=dim, y=y, symmetric_memory=symmetric_memory)
+
+    def _scale(self) -> int:
+        if self._mpi_comm is not None:
+            return int(self._mpi_comm.Get_size())
+        return 1
+
+    def resolve_config(self, case: Any, *, symmetric_memory: bool = False) -> TunedConfig:
+        dtype_override = getattr(getattr(case, "dtype_spec", None), "mscclpp_dtype", None)
+        accum_dtype = getattr(getattr(case, "dtype_spec", None), "accum_dtype", None) or dtype_override
+        symmetric_memory = symmetric_memory or bool(getattr(case, "symmetric_memory", False))
+        return self._resolve_config(
+            case.collective,
+            case.input,
+            dtype_override=dtype_override,
+            accum_dtype=accum_dtype,
+            symmetric_memory=symmetric_memory,
+        )
+
+    def _resolve_config(
+        self,
+        collective: str,
+        buffer: Any,
+        *,
+        dtype_override: Any | None = None,
+        accum_dtype: Any | None = None,
+        symmetric_memory: bool = False,
+    ) -> TunedConfig:
+        tuned_config = self._config_store.select(self._hardware_profile, collective, _nbytes(buffer))
+        if tuned_config is not None and tuned_config.algorithm in self._algorithms_by_collective.get(collective, {}):
+            return tuned_config
+
+        if self._rank == 0:
+            dim = int(_shape(buffer)[1]) if len(_shape(buffer)) > 1 else 1
+            warning_key = (
+                collective,
+                str(dtype_override if dtype_override is not None else _dtype(buffer)),
+                str(
+                    accum_dtype
+                    if accum_dtype is not None
+                    else dtype_override if dtype_override is not None else _dtype(buffer)
+                ),
+                dim,
+            )
+            if warning_key not in self._default_config_warning_keys:
+                self._default_config_warning_keys.add(warning_key)
+                logger.warning(
+                    "MSCCL++ default config: no tuning for collective=%s profile=%s dtype=%s accum=%s dim=%s; perf may be poor",
+                    collective,
+                    self._hardware_profile,
+                    warning_key[1],
+                    warning_key[2],
+                    dim,
+                )
+        return _default_tuned_config(
+            collective,
+            _nbytes(buffer),
+            self._algorithms_by_collective,
+            symmetric_memory=symmetric_memory,
+        )
+
+    def run(
+        self,
+        buffer: Any,
+        config: TunedConfig | None = None,
+        stream: Any | None = None,
+        *,
+        collective: str = _ALLREDUCE_COLLECTIVE,
+        output_tensor: Any | None = None,
+        dtype_override: Any | None = None,
+        accum_dtype: Any | None = None,
+        symmetric_memory: bool = False,
+    ) -> int:
+        if self._closed:
+            raise RuntimeError("Cannot use a closed MSCCL++ comm")
+
+        raise_on_error = True
+        if hasattr(buffer, "input") and hasattr(buffer, "output") and hasattr(buffer, "dtype_spec"):
+            case = buffer
+            buffer = case.input
+            output_tensor = case.output
+            collective = case.collective
+            dtype_override = case.dtype_spec.mscclpp_dtype
+            accum_dtype = case.dtype_spec.accum_dtype or dtype_override
+            symmetric_memory = symmetric_memory or bool(getattr(case, "symmetric_memory", False))
+            raise_on_error = False
+
+        if collective not in self._algorithms_by_collective:
+            raise RuntimeError(f"No supported MSCCL++ {collective} algorithm is available")
+
+        if config is None:
+            config = self._resolve_config(
+                collective,
+                buffer,
+                dtype_override=dtype_override,
+                accum_dtype=accum_dtype,
+                symmetric_memory=symmetric_memory,
+            )
+        symmetric_memory = symmetric_memory or config.symmetric_memory
+        algorithm = self._algorithms_by_collective[collective][config.algorithm]
+        output = buffer if output_tensor is None else output_tensor
+        dtype = dtype_override if dtype_override is not None else _dtype_to_mscclpp(_dtype(buffer))
+        accum = accum_dtype if accum_dtype is not None else dtype
+        ret = algorithm.execute(
+            comm=self._comm_group.communicator,
+            input_buffer=_data_ptr(buffer),
+            output_buffer=_data_ptr(output),
+            input_size=_nbytes(buffer),
+            output_size=_nbytes(output),
+            dtype=dtype,
+            op=self._mscclpp.ReduceOp.SUM if collective == _ALLREDUCE_COLLECTIVE else self._mscclpp.ReduceOp.NOP,
+            stream=_stream_ptr(stream),
+            nblocks=config.nblocks or 0,
+            nthreads_per_block=config.nthreads or 0,
+            symmetric_memory=symmetric_memory,
+            accum_dtype=accum,
+        )
+        if ret != 0 and raise_on_error:
+            raise RuntimeError(f"MSCCL++ {collective} failed on rank {self._rank} with error code {ret}")
+        return ret
+
+    def reset(self, config: TunedConfig | None = None) -> None:
+        if config is not None:
+            for algorithms_by_name in self._algorithms_by_collective.values():
+                algorithm = algorithms_by_name.get(config.algorithm)
+                if algorithm is not None:
+                    algorithm.reset()
+                    return
+        for algorithms_by_name in self._algorithms_by_collective.values():
+            for algorithm in algorithms_by_name.values():
+                algorithm.reset()
+
+    def close(self) -> None:
+        self.reset()
+        self._algorithms_by_collective = {}
+        self._scratch_buffer = None
+        self._closed = True
+        self._mscclpp.ext.AlgorithmCollectionBuilder.reset()
+
+
+def _numel(shape: tuple[int, ...]) -> int:
+    out = 1
+    for dim in shape:
+        out *= int(dim)
+    return out
+
+
+def _dtype_size(dtype: Any) -> int:
+    dtype_name = _dtype_name(dtype)
+    if dtype_name in {"float16", "bfloat16"}:
+        return 2
+    if dtype_name in {"float32", "int32", "uint32"}:
+        return 4
+    if dtype_name in {"uint8", "float8_e4m3b15", "float8_e4m3fn", "float8_e4m3fnuz"}:
+        return 1
+    raise ValueError(f"Unknown data type size for {dtype}")
+
+
+def _dtype_name(dtype: Any) -> str:
+    if isinstance(dtype, str):
+        return dtype.strip().lower().replace("-", "_")
+    name = str(dtype).rsplit(".", 1)[-1]
+    return name.strip().lower().replace("-", "_")
+
+
+def _dtype_to_mscclpp(dtype: Any) -> Any:
+    dtype_name = _dtype_name(dtype)
+    mapping = {
+        "float16": _mscclpp().DataType.float16,
+        "float32": _mscclpp().DataType.float32,
+        "int32": _mscclpp().DataType.int32,
+        "uint8": _mscclpp().DataType.uint8,
+        "float8_e4m3b15": _mscclpp().DataType.float8_e4m3b15,
+        "float8_e4m3fn": _mscclpp().DataType.float8_e4m3fn,
+        "float8_e4m3fnuz": _mscclpp().DataType.float8_e4m3fnuz,
+    }
+    try:
+        return mapping[dtype_name]
+    except KeyError as exc:
+        raise ValueError(f"Unknown data type: {dtype}") from exc
+
+
+def _data_ptr(buffer: Any) -> int:
+    if hasattr(buffer, "data_ptr"):
+        data_ptr = buffer.data_ptr
+        return int(data_ptr() if callable(data_ptr) else data_ptr)
+    if hasattr(buffer, "data"):
+        data = buffer.data
+        if callable(data):
+            return int(data())
+        if hasattr(data, "ptr"):
+            return int(data.ptr)
+    raise TypeError(f"Cannot get device pointer from {type(buffer)!r}")
+
+
+def _stream_ptr(stream: Any | None) -> int:
+    if stream is None:
+        return 0
+    return int(getattr(stream, "ptr", stream))
+
+
+def _nbytes(buffer: Any) -> int:
+    if hasattr(buffer, "nbytes"):
+        return int(buffer.nbytes)
+    if hasattr(buffer, "bytes"):
+        value = buffer.bytes
+        return int(value() if callable(value) else value)
+    raise TypeError(f"Cannot get byte size from {type(buffer)!r}")
+
+
+def _shape(buffer: Any) -> tuple[int, ...]:
+    shape = getattr(buffer, "shape", None)
+    if shape is None:
+        return (_nbytes(buffer) // _dtype_size(_dtype(buffer)),)
+    return tuple(int(dim) for dim in shape)
+
+
+def _dtype(buffer: Any) -> Any:
+    dtype = getattr(buffer, "dtype", None)
+    if dtype is None:
+        return "uint8"
+    return dtype
+
+
+def _detect_hardware_profile(*, scale: int) -> HardwareProfile:
+    try:
+        sku = device_name()
+    except Exception:
+        sku = "UNKNOWN"
+    return HardwareProfile(sku=normalize_sku(sku), scale=scale)
+
+
+def _ensure_device() -> None:
+    set_device(current_device())
+
+
+def _default_tuned_config(
+    collective: str,
+    message_size: int,
+    algorithms_by_collective: dict[str, dict[str, Any]],
+    *,
+    symmetric_memory: bool = False,
+) -> TunedConfig:
+    if collective == _ALLGATHER_COLLECTIVE:
+        return TunedConfig("default_allgather_fullmesh2", symmetric_memory=symmetric_memory)
+    available = algorithms_by_collective.get(collective, {})
+    if symmetric_memory and _mscclpp().is_nvls_supported() and "default_allreduce_nvls_zero_copy" in available:
+        return TunedConfig("default_allreduce_nvls_zero_copy", symmetric_memory=True)
+    if message_size <= 512 * 1024 and "default_allreduce_packet" in available:
+        return TunedConfig("default_allreduce_packet", symmetric_memory=symmetric_memory)
+    if "default_allreduce_rsag_zero_copy" in available:
+        return TunedConfig("default_allreduce_rsag_zero_copy", symmetric_memory=symmetric_memory)
+    if available:
+        return TunedConfig(next(iter(available)), symmetric_memory=symmetric_memory)
+    raise RuntimeError(f"No MSCCL++ algorithm is available for {collective}")
diff --git a/python/mscclpp_benchmark/correctness.py b/python/mscclpp_benchmark/correctness.py
new file mode 100644
index 00000000..0d9ab5c1
--- /dev/null
+++ b/python/mscclpp_benchmark/correctness.py
@@ -0,0 +1,402 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import cupy as cp
+from mpi4py import MPI
+
+_mscclpp_module = None
+
+
+def _mscclpp():
+    global _mscclpp_module
+    if _mscclpp_module is None:
+        import mscclpp
+
+        _mscclpp_module = mscclpp
+    return _mscclpp_module
+
+
+@dataclass(frozen=True)
+class CorrectnessStats:
+    ok: bool
+    max_abs_diff: float = 0.0
+    mean_abs_diff: float = 0.0
+    mismatches: int = 0
+    total: int = 0
+
+    def __bool__(self) -> bool:
+        return self.ok
+
+
+def config_accum_dtype(case: Any) -> Any:
+    return case.dtype_spec.accum_dtype or case.dtype_spec.mscclpp_dtype
+
+
+def fill_case_for_benchmark(case: Any, rank: int) -> None:
+    values = _benchmark_input_values(case, rank)
+    encoded = _encode_correctness_input(case, values)
+    if case.collective == "allreduce":
+        case.input[...] = encoded
+        return
+    case.output.fill(0)
+    case.input[...] = encoded
+
+
+def check_correctness(
+    comm: Any,
+    case: Any,
+    config: Any,
+    *,
+    niter: int = 1,
+) -> CorrectnessStats:
+    all_ok = True
+    local_max_abs_diff = 0.0
+    local_sum_abs_diff = 0.0
+    local_mismatches = 0
+    local_total = 0
+    for iteration in range(niter):
+        _fill_case_for_correctness(case, comm.rank, iteration)
+        comm.comm_group.barrier()
+        ret = comm.run(case, config)
+        cp.cuda.runtime.deviceSynchronize()
+        comm.comm_group.barrier()
+        if ret != 0:
+            all_ok = False
+            continue
+
+        expected, stats_expected = _expected_outputs(case, comm.nranks, iteration)
+        iter_stats = _local_diff_stats(case, case.output, expected, comm.nranks, stats_expected=stats_expected)
+        local_ok = _compare_output(case, case.output, expected, comm.nranks)
+        all_ok = all_ok and local_ok
+        local_max_abs_diff = max(local_max_abs_diff, iter_stats.max_abs_diff)
+        local_sum_abs_diff += iter_stats.mean_abs_diff * iter_stats.total
+        local_mismatches += iter_stats.mismatches
+        local_total += iter_stats.total
+
+        if not local_ok:
+            mismatch = _mismatch_mask(case, case.output, expected, comm.nranks)
+            print(
+                "not close: "
+                f"iter={iteration}, rank={comm.rank}, output={case.output[mismatch][0]}, "
+                f"expected={expected[mismatch][0]}, max_abs_diff={iter_stats.max_abs_diff:.6g}, "
+                f"mean_abs_diff={iter_stats.mean_abs_diff:.6g}, mismatches={iter_stats.mismatches}/{iter_stats.total}",
+                flush=True,
+            )
+
+    global_ok = bool(MPI.COMM_WORLD.allreduce(all_ok, op=MPI.LAND))
+    global_max_abs_diff = float(MPI.COMM_WORLD.allreduce(local_max_abs_diff, op=MPI.MAX))
+    global_sum_abs_diff = float(MPI.COMM_WORLD.allreduce(local_sum_abs_diff, op=MPI.SUM))
+    global_mismatches = int(MPI.COMM_WORLD.allreduce(local_mismatches, op=MPI.SUM))
+    global_total = int(MPI.COMM_WORLD.allreduce(local_total, op=MPI.SUM))
+    global_mean_abs_diff = global_sum_abs_diff / global_total if global_total else 0.0
+    return CorrectnessStats(
+        ok=global_ok,
+        max_abs_diff=global_max_abs_diff,
+        mean_abs_diff=global_mean_abs_diff,
+        mismatches=global_mismatches,
+        total=global_total,
+    )
+
+
+def _fill_case_for_correctness(case: Any, rank: int, iteration: int) -> None:
+    values = _correctness_input_values(case, rank, iteration)
+    encoded = _encode_correctness_input(case, values)
+    if case.collective == "allreduce":
+        case.input[...] = encoded
+        return
+    case.output.fill(0)
+    case.input[...] = encoded
+
+
+def _correctness_input_values(case: Any, rank: int, iteration: int):
+    shape = case.input.shape
+    rng = cp.random.RandomState(_correctness_seed(rank, iteration))
+    return _random_input_values(case, rng, shape)
+
+
+def _benchmark_input_values(case: Any, rank: int):
+    rng = cp.random.RandomState(17_000_003 + rank)
+    return _random_input_values(case, rng, case.input.shape)
+
+
+def _random_input_values(case: Any, rng, shape):
+    if case.dtype_spec.fp8_format is not None:
+        value_range = _fp8_correctness_input_range(case)
+        return rng.uniform(-value_range, value_range, size=shape).astype(cp.float32)
+    if case.dtype_spec.cupy_dtype == cp.int32:
+        return rng.randint(-1, 2, size=shape).astype(cp.int32)
+    if case.dtype_spec.cupy_dtype == cp.uint8:
+        return rng.randint(0, 2, size=shape).astype(cp.uint8)
+    return rng.uniform(-1.0, 1.0, size=shape).astype(cp.float32)
+
+
+def _correctness_seed(rank: int, iteration: int) -> int:
+    return (iteration + 1) * 1_000_003 + rank
+
+
+def _fp8_correctness_input_range(case: Any) -> float:
+    if case.collective != "allreduce":
+        return 1.0
+    fp8_format = case.dtype_spec.fp8_format
+    if fp8_format is None:
+        return 1.0
+    return min(1.0, _fp8_max_abs_value(fp8_format) / max(1, MPI.COMM_WORLD.size))
+
+
+def _encode_correctness_input(case: Any, values):
+    if case.dtype_spec.fp8_format is not None:
+        # FP8 buffers are stored as uint8 raw bytes, so a normal astype(uint8) cast would not produce FP8 bits.
+        return _encode_fp8_values(case.dtype_spec.fp8_format, values)
+    return values.astype(case.dtype_spec.cupy_dtype)
+
+
+def _local_diff_stats(case: Any, output, expected, nranks: int, *, stats_expected=None) -> CorrectnessStats:
+    mismatch = _mismatch_mask(case, output, expected, nranks)
+    mismatches = int(cp.count_nonzero(mismatch).item())
+    total = int(output.size)
+    if total == 0:
+        return CorrectnessStats(ok=mismatches == 0)
+
+    output_values = _stats_values(case, output)
+    expected_values = _stats_values(case, expected) if stats_expected is None else stats_expected.astype(cp.float64)
+    abs_diff = cp.abs(output_values - expected_values)
+    return CorrectnessStats(
+        ok=mismatches == 0,
+        max_abs_diff=float(cp.max(abs_diff).item()),
+        mean_abs_diff=float(cp.mean(abs_diff).item()),
+        mismatches=mismatches,
+        total=total,
+    )
+
+
+def _stats_values(case: Any, values):
+    # Convert storage buffers into numeric values before computing max/mean diff.
+    if case.dtype_spec.fp8_format is not None:
+        return _decode_fp8_array(case.dtype_spec.fp8_format, values)
+    if cp.issubdtype(values.dtype, cp.floating):
+        return values.astype(cp.float64)
+    return values.astype(cp.int64)
+
+
+def _expected_outputs(case: Any, nranks: int, iteration: int):
+    if case.collective == "allreduce":
+        encoded_inputs = _encoded_rank_inputs(case, nranks, iteration)
+        if case.dtype_spec.fp8_format is not None:
+            stats_expected = _expected_fp8_accum_values(case, encoded_inputs)
+            return _encode_reduced_output(case, stats_expected), stats_expected
+        return _encode_reduced_output(case, sum(values.astype(cp.float32) for values in encoded_inputs)), None
+
+    expected = cp.empty_like(case.output)
+    chunk = case.input.size
+    for rank, values in enumerate(_encoded_rank_inputs(case, nranks, iteration)):
+        expected[rank * chunk : (rank + 1) * chunk] = values.reshape(-1)
+    return expected, None
+
+
+def _encoded_rank_inputs(case: Any, nranks: int, iteration: int) -> list[Any]:
+    return [_encode_correctness_input(case, _correctness_input_values(case, rank, iteration)) for rank in range(nranks)]
+
+
+def _expected_fp8_accum_values(case: Any, encoded_inputs: list[Any]):
+    fp8_format = case.dtype_spec.fp8_format
+    if fp8_format is None:
+        raise ValueError("FP8 format is required")
+
+    accum_dtype = config_accum_dtype(case)
+    if accum_dtype == _mscclpp().DataType.float16:
+        acc = cp.zeros_like(_decode_fp8_array(fp8_format, encoded_inputs[0]), dtype=cp.float16)
+        for values in encoded_inputs:
+            acc = (acc + _decode_fp8_array(fp8_format, values).astype(cp.float16)).astype(cp.float16)
+        return acc.astype(cp.float32)
+
+    if accum_dtype == _mscclpp().DataType.float32:
+        acc = cp.zeros_like(_decode_fp8_array(fp8_format, encoded_inputs[0]), dtype=cp.float32)
+        for values in encoded_inputs:
+            acc += _decode_fp8_array(fp8_format, values).astype(cp.float32)
+        return acc
+
+    acc = encoded_inputs[0]
+    for values in encoded_inputs[1:]:
+        acc = _encode_fp8_values(fp8_format, _decode_fp8_array(fp8_format, acc) + _decode_fp8_array(fp8_format, values))
+    return _decode_fp8_array(fp8_format, acc).astype(cp.float32)
+
+
+def _encode_reduced_output(case: Any, values):
+    if case.dtype_spec.fp8_format is not None:
+        return _encode_fp8_values(case.dtype_spec.fp8_format, values)
+    return values.astype(case.output.dtype)
+
+
+def _compare_output(case: Any, output, expected, nranks: int) -> bool:
+    return bool(cp.all(~_mismatch_mask(case, output, expected, nranks)).item())
+
+
+def _mismatch_mask(case: Any, output, expected, nranks: int):
+    tolerance = _comparison_tolerance(case, nranks)
+    if tolerance is None:
+        return output != expected
+    rtol, atol = tolerance
+    return ~cp.isclose(_stats_values(case, output), _stats_values(case, expected), rtol=rtol, atol=atol)
+
+
+def _comparison_tolerance(case: Any, nranks: int) -> tuple[float, float] | None:
+    scale = max(1, nranks) if case.collective == "allreduce" else 1
+    if case.dtype_spec.fp8_format is not None:
+        accum_dtype = config_accum_dtype(case)
+        if accum_dtype == _mscclpp().DataType.float32:
+            return None
+        atol = _max_fp8_spacing(case.dtype_spec.fp8_format, float(scale))
+        if accum_dtype == _mscclpp().DataType.float16:
+            return (0.0, atol)
+        return (0.0, atol * 2)
+    if case.dtype_spec.cupy_dtype == cp.float16:
+        return (1.0e-2, 5.0e-4 * scale)
+    if case.dtype_spec.cupy_dtype == cp.float32:
+        return (1.0e-5 * scale, 1.0e-6 * scale)
+    return None
+
+
+_FP8_TABLES: dict[str, list[tuple[int, float]]] = {}
+_FP8_LOOKUP_CACHE: dict[str, tuple[Any, Any]] = {}
+_FP8_SPACING_CACHE: dict[tuple[str, float], float] = {}
+
+
+def _encode_fp8_values(fp8_format: str, values):
+    values = values.astype(cp.float32)
+    if fp8_format == "e4m3b15":
+        return _encode_e4m3b15_values(values)
+
+    # Round each value to the nearest representable FP8 value (ties to even).
+    table_values, table_bytes = _fp8_lookup_arrays(fp8_format)
+    flat_values = values.ravel()
+
+    # For each value find its two surrounding table entries: lower <= value <= upper.
+    upper = cp.clip(cp.searchsorted(table_values, flat_values), 1, table_values.size - 1)
+    lower = upper - 1
+
+    # Pick the closer neighbor; on an exact tie pick the one with an even byte.
+    dist_to_upper = table_values[upper] - flat_values
+    dist_to_lower = flat_values - table_values[lower]
+    upper_is_even = (table_bytes[upper] & cp.uint8(1)) == 0
+    pick_upper = (dist_to_upper < dist_to_lower) | ((dist_to_upper == dist_to_lower) & upper_is_even)
+
+    return cp.where(pick_upper, table_bytes[upper], table_bytes[lower]).reshape(values.shape)
+
+
+def _fp8_lookup_arrays(fp8_format: str):
+    # Cache a sorted (value -> byte) table per format for fast nearest-value lookup.
+    if fp8_format in _FP8_LOOKUP_CACHE:
+        return _FP8_LOOKUP_CACHE[fp8_format]
+
+    # Different bytes can decode to the same value (e.g. +0 and -0); keep one byte per value.
+    byte_for_value: dict[float, int] = {}
+    for byte, value in _FP8_TABLES.setdefault(fp8_format, _build_fp8_table(fp8_format)):
+        if value not in byte_for_value or byte < byte_for_value[value]:
+            byte_for_value[value] = byte
+
+    table = sorted(byte_for_value.items())
+    table_values = cp.asarray([value for value, _ in table], dtype=cp.float32)
+    table_bytes = cp.asarray([byte for _, byte in table], dtype=cp.uint8)
+    _FP8_LOOKUP_CACHE[fp8_format] = (table_values, table_bytes)
+    return _FP8_LOOKUP_CACHE[fp8_format]
+
+
+def _max_fp8_spacing(fp8_format: str, max_abs_value: float) -> float:
+    cache_key = (fp8_format, max_abs_value)
+    if cache_key in _FP8_SPACING_CACHE:
+        return _FP8_SPACING_CACHE[cache_key]
+
+    values = sorted(
+        {
+            value
+            for _, value in _FP8_TABLES.setdefault(fp8_format, _build_fp8_table(fp8_format))
+            if abs(value) <= max_abs_value
+        }
+    )
+    if len(values) < 2:
+        spacing = 0.0
+    else:
+        spacing = max(right - left for left, right in zip(values, values[1:]))
+    _FP8_SPACING_CACHE[cache_key] = spacing
+    return spacing
+
+
+def _fp8_max_abs_value(fp8_format: str) -> float:
+    return max(abs(value) for _, value in _FP8_TABLES.setdefault(fp8_format, _build_fp8_table(fp8_format)))
+
+
+def _encode_e4m3b15_values(values):
+    # Mirrors the device e4m3b15 encode (gpu_data_types.hpp): clamp the fp16 intermediate
+    # to 0x3F80 (+/-1.875) so the max encodable byte is 0x7F/0xFF.
+    fp16_bits = values.astype(cp.float16).view(cp.uint16)
+    abs_fp16 = fp16_bits & cp.uint16(0x7FFF)
+    abs_fp16 = cp.minimum(abs_fp16, cp.uint16(0x3F80)).astype(cp.uint32)
+    sign16 = (fp16_bits & cp.uint16(0x8000)).astype(cp.uint32)
+    adjusted = abs_fp16 * cp.uint32(2) + cp.uint32(0x0080)
+    return (((sign16 | adjusted) >> cp.uint32(8)) & cp.uint32(0xFF)).astype(cp.uint8)
+
+
+def _build_fp8_table(fp8_format: str) -> list[tuple[int, float]]:
+    table = []
+    for byte in range(256):
+        value = _decode_fp8_scalar(fp8_format, byte)
+        if not math.isnan(value):
+            table.append((byte, value))
+    return table
+
+
+def _decode_fp8_scalar(fp8_format: str, byte: int) -> float:
+    if fp8_format == "e4m3fnuz" and byte == 0x80:
+        return float("nan")
+    sign = -1.0 if byte & 0x80 else 1.0
+    return sign * _decode_fp8_positive(fp8_format, byte & 0x7F)
+
+
+def _decode_fp8_positive(fp8_format: str, byte: int) -> float:
+    exp = (byte >> 3) & 0xF
+    mant = byte & 0x7
+    if fp8_format == "e4m3fn" and exp == 0xF and mant == 0x7:
+        return float("nan")
+    if exp == 0 and mant == 0:
+        return 0.0
+    if fp8_format == "e4m3fn":
+        return math.ldexp(mant / 8.0, -6) if exp == 0 else math.ldexp(1.0 + mant / 8.0, exp - 7)
+    if fp8_format == "e4m3fnuz":
+        return math.ldexp(mant / 8.0, -7) if exp == 0 else math.ldexp(1.0 + mant / 8.0, exp - 8)
+    if fp8_format == "e4m3b15":
+        return math.ldexp(mant / 8.0, -14) if exp == 0 else math.ldexp(1.0 + mant / 8.0, exp - 15)
+    raise ValueError(f"Unknown FP8 format: {fp8_format}")
+
+
+def _decode_fp8_array(fp8_format: str, values):
+    bits = values.astype(cp.int32)
+    sign = (bits >> 7) & 1
+    exp = (bits >> 3) & 0xF
+    mant = bits & 0x7
+
+    if fp8_format == "e4m3fn":
+        subnormal = cp.ldexp(mant.astype(cp.float32) / cp.float32(8.0), cp.int32(-6))
+        normal = cp.ldexp(cp.float32(1.0) + mant.astype(cp.float32) / cp.float32(8.0), exp.astype(cp.int32) - 7)
+        decoded = cp.where(exp == 0, subnormal, normal)
+        decoded = cp.where((exp == 0xF) & (mant == 0x7), cp.nan, decoded)
+    elif fp8_format == "e4m3fnuz":
+        subnormal = cp.ldexp(mant.astype(cp.float32) / cp.float32(8.0), cp.int32(-7))
+        normal = cp.ldexp(cp.float32(1.0) + mant.astype(cp.float32) / cp.float32(8.0), exp.astype(cp.int32) - 8)
+        decoded = cp.where(exp == 0, subnormal, normal)
+    elif fp8_format == "e4m3b15":
+        subnormal = cp.ldexp(mant.astype(cp.float32) / cp.float32(8.0), cp.int32(-14))
+        normal = cp.ldexp(cp.float32(1.0) + mant.astype(cp.float32) / cp.float32(8.0), exp.astype(cp.int32) - 15)
+        decoded = cp.where(exp == 0, subnormal, normal)
+    else:
+        raise ValueError(f"Unknown FP8 format: {fp8_format}")
+
+    result = cp.where(sign == 1, -decoded, decoded)
+    if fp8_format == "e4m3fnuz":
+        result = cp.where(bits == 0x80, cp.float32(float("nan")), result)
+    return result
diff --git a/python/mscclpp_benchmark/gpu.py b/python/mscclpp_benchmark/gpu.py
new file mode 100644
index 00000000..1309a504
--- /dev/null
+++ b/python/mscclpp_benchmark/gpu.py
@@ -0,0 +1,187 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+_API_NAMES = {
+    "get_device_count": ("hipGetDeviceCount", "cudaGetDeviceCount"),
+    "get_device": ("hipGetDevice", "cudaGetDevice"),
+    "get_device_properties": ("hipGetDeviceProperties", "cudaGetDeviceProperties"),
+    "set_device": ("hipSetDevice", "cudaSetDevice"),
+    "stream_begin_capture": ("hipStreamBeginCapture", "cudaStreamBeginCapture"),
+    "stream_end_capture": ("hipStreamEndCapture", "cudaStreamEndCapture"),
+    "graph_instantiate": ("hipGraphInstantiate", "cudaGraphInstantiate"),
+    "graph_launch": ("hipGraphLaunch", "cudaGraphLaunch"),
+    "graph_destroy": ("hipGraphDestroy", "cudaGraphDestroy"),
+    "graph_exec_destroy": ("hipGraphExecDestroy", "cudaGraphExecDestroy"),
+    "get_error_string": ("hipGetErrorString", "cudaGetErrorString"),
+}
+
+
+@dataclass(frozen=True)
+class _Runtime:
+    name: str
+    success: Any
+    capture_mode_relaxed: Any
+    funcs: dict[str, Callable[..., Any] | None]
+
+    @classmethod
+    def create(cls, name: str, module: Any, success: Any, capture_mode_relaxed: Any) -> "_Runtime":
+        index = 0 if name == "hip" else 1
+        funcs = {
+            attr: (None if names[index] is None else getattr(module, names[index]))
+            for attr, names in _API_NAMES.items()
+        }
+        return cls(name=name, success=success, capture_mode_relaxed=capture_mode_relaxed, funcs=funcs)
+
+    def call(self, name: str, *args: Any) -> tuple[Any, ...]:
+        fn = self.funcs[name]
+        if fn is None:
+            raise RuntimeError(f"{name} is not available for {self.name}")
+        result = fn(*args)
+        if not isinstance(result, tuple):
+            result = (result,)
+        self.check(result[0], name)
+        return result[1:]
+
+    def check(self, error: Any, api: str) -> None:
+        if error == self.success:
+            return
+        result = self.funcs["get_error_string"](error)
+        if not isinstance(result, tuple):
+            result = (result,)
+        err, message = result
+        if err != self.success:
+            raise RuntimeError(f"{api} failed with error {int(error)}")
+        decoded = message.decode("utf-8") if isinstance(message, bytes) else str(message)
+        raise RuntimeError(f"{api} failed: {decoded} ({int(error)})")
+
+
+def _load_runtime() -> _Runtime:
+    errors: list[str] = []
+
+    try:
+        from hip import hip
+
+        runtime = _Runtime.create(
+            name="hip",
+            module=hip,
+            success=hip.hipError_t.hipSuccess,
+            capture_mode_relaxed=hip.hipStreamCaptureMode.hipStreamCaptureModeRelaxed,
+        )
+        count = runtime.call("get_device_count")[0]
+        if count and count > 0:
+            return runtime
+        errors.append(f"hipGetDeviceCount returned count={count}")
+    except ImportError as exc:
+        errors.append(f"hip-python unavailable: {exc}")
+
+    try:
+        from cuda.bindings import runtime as cuda_runtime
+
+        runtime = _Runtime.create(
+            name="cuda",
+            module=cuda_runtime,
+            success=cuda_runtime.cudaError_t.cudaSuccess,
+            capture_mode_relaxed=cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeRelaxed,
+        )
+        count = runtime.call("get_device_count")[0]
+        if count and count > 0:
+            return runtime
+        errors.append(f"cudaGetDeviceCount returned count={count}")
+    except ImportError as exc:
+        errors.append(f"cuda-bindings unavailable: {exc}")
+
+    raise RuntimeError("No usable CUDA/HIP Python runtime found: " + "; ".join(errors))
+
+
+_RUNTIME = _load_runtime()
+
+
+class Graph:
+    def __init__(self, graph_exec: Any) -> None:
+        self._graph_exec = graph_exec
+
+    def launch(self, stream: Any) -> None:
+        _api("graph_launch")(self._graph_exec, _stream_ptr(stream))
+
+    def close(self) -> None:
+        if self._graph_exec is not None:
+            _api("graph_exec_destroy")(self._graph_exec)
+            self._graph_exec = None
+
+
+def init_runtime() -> None:
+    return None
+
+
+def capture_graph(stream: Any, capture_fn: Callable[[], None]) -> Graph:
+    _api("set_device")(current_device())
+    stream_ptr = _stream_ptr(stream)
+    _api("stream_begin_capture")(stream_ptr, _RUNTIME.capture_mode_relaxed)
+
+    graph = None
+    try:
+        capture_fn()
+        graph = _api("stream_end_capture")(stream_ptr)[0]
+    except Exception:
+        try:
+            _api("stream_end_capture")(stream_ptr)
+        except Exception:
+            pass
+        raise
+
+    try:
+        graph_exec = _instantiate_graph(graph)
+        return Graph(graph_exec)
+    finally:
+        if graph is not None:
+            _api("graph_destroy")(graph)
+
+
+def current_device() -> int:
+    return int(_api("get_device")()[0])
+
+
+def device_name(device_id: int | None = None) -> str:
+    if device_id is None:
+        device_id = current_device()
+    prop = _api("get_device_properties")(int(device_id))[0]
+    name = getattr(prop, "name", "UNKNOWN")
+    return name.decode("utf-8") if isinstance(name, bytes) else str(name)
+
+
+def _stream_ptr(stream: Any) -> int:
+    return int(getattr(stream, "ptr", stream))
+
+
+def _instantiate_graph(graph: Any) -> Any:
+    if _RUNTIME.name == "hip":
+        return _api("graph_instantiate")(graph, None, 0)[0]
+    return _api("graph_instantiate")(graph, 0)[0]
+
+
+def _api(name: str) -> Callable[..., tuple[Any, ...]]:
+    api = globals().get(name)
+    if api is None:
+        api = __getattr__(name)
+    return api
+
+
+def _make_api(name: str) -> Callable[..., tuple[Any, ...]]:
+    def api(*args: Any) -> tuple[Any, ...]:
+        return _RUNTIME.call(name, *args)
+
+    api.__name__ = name
+    return api
+
+
+def __getattr__(name: str) -> Callable[..., tuple[Any, ...]]:
+    if name in _API_NAMES:
+        api = _make_api(name)
+        globals()[name] = api
+        return api
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/python/mscclpp_benchmark/tuner.py b/python/mscclpp_benchmark/tuner.py
new file mode 100644
index 00000000..8df3259b
--- /dev/null
+++ b/python/mscclpp_benchmark/tuner.py
@@ -0,0 +1,84 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from typing import Any, Callable, Iterable
+
+from mscclpp_benchmark.tuning_config import TunedConfig
+
+
+class OfflineTuner:
+    def __init__(
+        self,
+        comm: Any,
+        *,
+        candidate_nblocks: Iterable[int],
+        candidate_nthreads: Iterable[int],
+        n_warmup: int,
+        n_graph_launches: int,
+        n_ops_per_graph: int,
+        candidate_algorithms: Callable[[Any, Any], list[tuple[Any, Any]]],
+        check_correctness: Callable[..., bool],
+        measure: Callable[..., float | None],
+    ) -> None:
+        self.comm = comm
+        self.candidate_nblocks = tuple(candidate_nblocks)
+        self.candidate_nthreads = tuple(candidate_nthreads)
+        self.n_warmup = n_warmup
+        self.n_graph_launches = n_graph_launches
+        self.n_ops_per_graph = n_ops_per_graph
+        self._candidate_algorithms = candidate_algorithms
+        self._check_correctness = check_correctness
+        self._measure = measure
+
+    def tune(self, case: Any) -> TunedConfig | None:
+        best_config: TunedConfig | None = None
+        best_time_us = float("inf")
+        symmetric_memory = bool(getattr(case, "symmetric_memory", False))
+        candidates = self._candidate_algorithms(self.comm, case)
+        if not candidates:
+            if self.comm.rank == 0:
+                print(
+                    f"[skip] no supported tuning candidates for collective={case.collective} "
+                    f"size={case.message_size}",
+                    flush=True,
+                )
+            return None
+        for algorithm, candidate_spec in candidates:
+            for nblocks in self.candidate_nblocks:
+                if candidate_spec.max_nblocks is not None and nblocks > candidate_spec.max_nblocks:
+                    continue
+                for nthreads in self.candidate_nthreads:
+                    config = TunedConfig(
+                        algorithm=algorithm.name,
+                        nblocks=nblocks,
+                        nthreads=nthreads,
+                        symmetric_memory=symmetric_memory,
+                    )
+                    if not self._check_correctness(self.comm, case, config):
+                        self.comm.reset(config)
+                        continue
+                    self.comm.reset(config)
+                    time_us = self._measure(
+                        self.comm,
+                        case,
+                        config,
+                        n_warmup=self.n_warmup,
+                        n_graph_launches=self.n_graph_launches,
+                        n_ops_per_graph=self.n_ops_per_graph,
+                    )
+                    self.comm.reset(config)
+                    if time_us is None or time_us >= best_time_us:
+                        continue
+                    best_time_us = time_us
+                    best_config = TunedConfig(
+                        algorithm=algorithm.name,
+                        nblocks=nblocks,
+                        nthreads=nthreads,
+                        symmetric_memory=symmetric_memory,
+                        time_us=time_us,
+                    )
+        if best_config is None:
+            return self.comm.resolve_config(case)
+        return best_config
diff --git a/python/mscclpp_benchmark/tuning_config.py b/python/mscclpp_benchmark/tuning_config.py
new file mode 100644
index 00000000..2a914ec9
--- /dev/null
+++ b/python/mscclpp_benchmark/tuning_config.py
@@ -0,0 +1,242 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import json
+import re
+from bisect import bisect_left
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+_KNOWN_GPU_SKUS = ("GB300", "MI300X", "H100", "A100")
+
+
+@dataclass(frozen=True)
+class HardwareProfile:
+    sku: str | None = None
+    scale: int | None = None
+
+
+@dataclass(frozen=True)
+class TunedConfig:
+    algorithm: str
+    nblocks: int | None = None
+    nthreads: int | None = None
+    symmetric_memory: bool = False
+    time_us: float | None = None
+
+
+@dataclass(order=True, frozen=True)
+class TunedConfigBySize:
+    message_size: int
+    config: TunedConfig
+
+
+class TunedConfigStore:
+    def __init__(self, profiles: dict[HardwareProfile, dict[str, list[TunedConfigBySize]]]) -> None:
+        self._profiles = profiles
+
+    @classmethod
+    def empty(cls) -> "TunedConfigStore":
+        return cls({})
+
+    @classmethod
+    def load_path(cls, path: str | Path) -> "TunedConfigStore":
+        with Path(path).open("r", encoding="utf-8") as handle:
+            return cls.from_payload(json.load(handle))
+
+    @classmethod
+    def from_payload(cls, payload: Any) -> "TunedConfigStore":
+        if not isinstance(payload, dict):
+            raise ValueError("MSCCL++ tuned config must be a JSON object")
+        raw_profiles = payload.get("profiles")
+        if not isinstance(raw_profiles, list):
+            raise ValueError("MSCCL++ tuned config must contain a 'profiles' list")
+        profiles: dict[HardwareProfile, dict[str, list[TunedConfigBySize]]] = {}
+        for raw_profile in raw_profiles:
+            profile = _profile_from_payload(raw_profile)
+            profiles[profile] = _configs_by_collective_from_payload(raw_profile.get("collectives", {}))
+        return cls(profiles)
+
+    def select(self, profile: HardwareProfile, collective: str, message_size: int) -> TunedConfig | None:
+        for _, configs_by_collective in _matching_profiles(self._profiles, profile):
+            config = _select_config(configs_by_collective, collective, message_size)
+            if config is not None:
+                return config
+        return None
+
+    def upsert(self, profile: HardwareProfile, collective: str, message_size: int, config: TunedConfig) -> None:
+        configs = self._profiles.setdefault(profile, {}).setdefault(collective, [])
+        for index, existing in enumerate(configs):
+            if existing.message_size == message_size:
+                configs[index] = TunedConfigBySize(message_size, config)
+                break
+        else:
+            configs.append(TunedConfigBySize(message_size, config))
+        configs.sort(key=lambda item: item.message_size)
+
+    def write_path(self, path: str | Path) -> None:
+        profiles_payload: list[dict[str, Any]] = []
+        for profile, configs_by_collective in sorted(
+            self._profiles.items(),
+            key=lambda item: (item[0].sku is None, item[0].sku or "", item[0].scale is None, item[0].scale or 0),
+        ):
+            collectives: dict[str, list[dict[str, Any]]] = {}
+            for collective, configs in sorted(configs_by_collective.items()):
+                collectives[collective] = [_config_entry_payload(item) for item in sorted(configs)]
+            profile_payload: dict[str, Any] = {}
+            if profile.sku is not None:
+                profile_payload["sku"] = profile.sku
+            if profile.scale is not None:
+                profile_payload["scale"] = profile.scale
+            profile_payload["collectives"] = collectives
+            profiles_payload.append(profile_payload)
+
+        with Path(path).open("w", encoding="utf-8") as handle:
+            handle.write(_format_tuned_config_json({"version": 1, "profiles": profiles_payload}))
+
+
+def normalize_sku(raw_sku: str) -> str:
+    upper_sku = raw_sku.upper()
+    for known_sku in _KNOWN_GPU_SKUS:
+        if known_sku in upper_sku:
+            return known_sku
+    normalized = re.sub(r"[^A-Z0-9]+", "_", upper_sku).strip("_")
+    return normalized or "UNKNOWN"
+
+
+def _profile_from_payload(raw_profile: Any) -> HardwareProfile:
+    if not isinstance(raw_profile, dict):
+        raise ValueError(f"Invalid tuned config profile: {raw_profile!r}")
+    raw_sku = raw_profile.get("sku")
+    return HardwareProfile(
+        sku=None if raw_sku is None else normalize_sku(str(raw_sku)),
+        scale=_optional_positive_int(raw_profile.get("scale"), "scale"),
+    )
+
+
+def _matching_profiles(
+    profiles: dict[HardwareProfile, dict[str, list[TunedConfigBySize]]],
+    runtime_profile: HardwareProfile,
+) -> list[tuple[int, dict[str, list[TunedConfigBySize]]]]:
+    matches: list[tuple[int, dict[str, list[TunedConfigBySize]]]] = []
+    for profile, configs_by_collective in profiles.items():
+        specificity = _profile_match_specificity(profile, runtime_profile)
+        if specificity is not None:
+            matches.append((specificity, configs_by_collective))
+    return sorted(matches, key=lambda item: item[0], reverse=True)
+
+
+def _profile_match_specificity(profile: HardwareProfile, runtime_profile: HardwareProfile) -> int | None:
+    specificity = 0
+    if profile.sku is not None:
+        if profile.sku != runtime_profile.sku:
+            return None
+        specificity += 1
+    if profile.scale is not None:
+        if profile.scale != runtime_profile.scale:
+            return None
+        specificity += 1
+    return specificity
+
+
+def _select_config(
+    configs_by_collective: dict[str, list[TunedConfigBySize]], collective: str, message_size: int
+) -> TunedConfig | None:
+    configs = configs_by_collective.get(collective, [])
+    if not configs:
+        return None
+    sizes = [item.message_size for item in configs]
+    index = bisect_left(sizes, message_size)
+    if index == len(sizes):
+        return configs[-1].config
+    if sizes[index] == message_size or index == 0:
+        return configs[index].config
+    return configs[index - 1].config
+
+
+def _configs_by_collective_from_payload(payload: Any) -> dict[str, list[TunedConfigBySize]]:
+    if not isinstance(payload, dict):
+        raise ValueError("MSCCL++ tuned config collectives must be an object")
+
+    result: dict[str, list[TunedConfigBySize]] = {}
+    for collective, raw_entries in payload.items():
+        if isinstance(raw_entries, dict):
+            raw_entries = raw_entries.get("configs", [])
+        if not isinstance(raw_entries, list):
+            continue
+        configs = []
+        for raw_entry in raw_entries:
+            if not isinstance(raw_entry, dict):
+                raise ValueError(f"Invalid tuned config entry for {collective}: {raw_entry!r}")
+            configs.append(
+                TunedConfigBySize(
+                    message_size=_parse_positive_int(raw_entry.get("message_size"), "message_size"),
+                    config=TunedConfig(
+                        algorithm=str(raw_entry["algorithm"]),
+                        nblocks=_optional_int(raw_entry.get("nblocks")),
+                        nthreads=_optional_int(raw_entry.get("nthreads")),
+                        symmetric_memory=_optional_bool(raw_entry.get("symmetric_memory", False)),
+                        time_us=_optional_float(raw_entry.get("time_us")),
+                    ),
+                )
+            )
+        result[str(collective)] = sorted(configs)
+    return result
+
+
+def _config_entry_payload(item: TunedConfigBySize) -> dict[str, Any]:
+    payload: dict[str, Any] = {"message_size": item.message_size, "algorithm": item.config.algorithm}
+    if item.config.nblocks is not None:
+        payload["nblocks"] = item.config.nblocks
+    if item.config.nthreads is not None:
+        payload["nthreads"] = item.config.nthreads
+    if item.config.symmetric_memory:
+        payload["symmetric_memory"] = item.config.symmetric_memory
+    if item.config.time_us is not None:
+        payload["time_us"] = item.config.time_us
+    return payload
+
+
+def _format_tuned_config_json(payload: dict[str, Any]) -> str:
+    text = json.dumps(payload, indent=2)
+    pattern = re.compile(
+        r"(?m)^(?P<indent> +)\{\n"
+        r'(?P<body>(?P=indent)  "message_size": [^\n]+,?\n(?:(?P=indent)  "[^"]+": [^\n]+,?\n)*)'
+        r"(?P=indent)\}(?P<comma>,?)$"
+    )
+
+    def compact(match: re.Match[str]) -> str:
+        body = " ".join(line.strip() for line in match.group("body").splitlines())
+        return f"{match.group('indent')}{{{body}}}{match.group('comma')}"
+
+    return pattern.sub(compact, text) + "\n"
+
+
+def _optional_int(value: Any | None) -> int | None:
+    return None if value is None else int(value)
+
+
+def _optional_float(value: Any | None) -> float | None:
+    return None if value is None else float(value)
+
+
+def _optional_positive_int(value: Any | None, name: str) -> int | None:
+    return None if value is None else _parse_positive_int(value, name)
+
+
+def _optional_bool(value: Any | None) -> bool | None:
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return value
+    raise ValueError(f"Expected boolean value, got {value!r}")
+
+
+def _parse_positive_int(value: Any, name: str) -> int:
+    parsed = int(value)
+    if parsed <= 0:
+        raise ValueError(f"{name} must be positive, got {parsed}")
+    return parsed
diff --git a/python/requirements_cuda11.txt b/python/requirements_cuda11.txt
index a9786071..1f575f67 100644
--- a/python/requirements_cuda11.txt
+++ b/python/requirements_cuda11.txt
@@ -1,5 +1,6 @@
 mpi4py
 cupy-cuda11x
+cuda-bindings>=11.8,<12
 prettytable
 netifaces
 pytest
diff --git a/python/requirements_cuda12.txt b/python/requirements_cuda12.txt
index 71572714..fcc59660 100644
--- a/python/requirements_cuda12.txt
+++ b/python/requirements_cuda12.txt
@@ -1,5 +1,6 @@
 mpi4py
 cupy-cuda12x
+cuda-bindings>=12,<13
 prettytable
 netifaces
 pytest
diff --git a/python/requirements_cuda13.txt b/python/requirements_cuda13.txt
index 95e99533..19ad93d7 100644
--- a/python/requirements_cuda13.txt
+++ b/python/requirements_cuda13.txt
@@ -1,5 +1,6 @@
 mpi4py
 cupy-cuda13x
+cuda-bindings>=13,<14
 prettytable
 netifaces
 pytest
diff --git a/python/requirements_rocm6.txt b/python/requirements_rocm6.txt
index 757d4e26..bcc22dfb 100644
--- a/python/requirements_rocm6.txt
+++ b/python/requirements_rocm6.txt
@@ -7,4 +7,5 @@ numpy
 matplotlib
 sortedcontainers
 blake3
-pybind11
\ No newline at end of file
+pybind11
+hip-python>=6,<7
\ No newline at end of file
diff --git a/python/test/test_fp8_accum.py b/python/test/test_fp8_accum.py
index ba33c085..554e131a 100644
--- a/python/test/test_fp8_accum.py
+++ b/python/test/test_fp8_accum.py
@@ -167,7 +167,7 @@ else:
 
 
 # ---------------------------------------------------------------------------
-# FP8 E4M3B15 helpers (bias=15, encode saturates to ±1.75, no NaN)
+# FP8 E4M3B15 helpers (bias=15, float source saturates to ±1.875, no NaN)
 # Matches Triton's fp8e4b15: all 256 bit patterns are finite.
 # ---------------------------------------------------------------------------
 
@@ -193,7 +193,7 @@ def float_to_e4m3b15(f32_array, chunk_size=65536):
     """Encode a cupy float32 array to uint8 E4M3B15 bit patterns.
 
     Same lookup-table approach as float_to_e4m3fn.
-    Saturates to ±1.75 (0x7e/0xfe), matching Triton's fp8e4b15.
+    Saturates to ±1.875 (0x7f/0xff), matching the device float32 → e4m3b15 path.
     """
     # Build lookup table of all 128 positive E4M3B15 values (0x00..0x7F)
     all_bytes = cp.arange(128, dtype=cp.uint8)
@@ -203,7 +203,7 @@ def float_to_e4m3b15(f32_array, chunk_size=65536):
     values = f32_array.astype(cp.float32)
     signs = cp.signbit(values).astype(cp.uint8)
     absval = cp.abs(values)
-    absval = cp.clip(absval, cp.float32(0.0), cp.float32(1.75))
+    absval = cp.clip(absval, cp.float32(0.0), cp.float32(1.875))
 
     result = cp.zeros(absval.shape, dtype=cp.uint8)
     n = absval.size
@@ -442,8 +442,8 @@ def test_fp8_e4m3b15_accum(mpi_group: MpiGroup, algo_name: str, size: int):
             bits_r = cp.asarray(rng_r.randint(0, 256, (size,)).astype(np.uint8))
             ref_f32 += e4m3b15_to_float(bits_r)
 
-        # Clamp reference to e4m3b15 representable range
-        ref_f32 = cp.clip(ref_f32, -1.75, 1.75)
+        # Clamp reference to e4m3b15 representable range (float source saturates at ±1.875)
+        ref_f32 = cp.clip(ref_f32, -1.875, 1.875)
 
         # Compute errors
         abs_err = cp.abs(result_f32 - ref_f32)
diff --git a/src/ext/collectives/allgather/allgather_fullmesh_2.cu b/src/ext/collectives/allgather/allgather_fullmesh_2.cu
index 72a2be9d..3500c0c4 100644
--- a/src/ext/collectives/allgather/allgather_fullmesh_2.cu
+++ b/src/ext/collectives/allgather/allgather_fullmesh_2.cu
@@ -8,7 +8,6 @@
 namespace mscclpp {
 namespace collective {
 
-__device__ DeviceSyncer deviceSyncer;
 template <bool IsOutOfPlace>
 __global__ void __launch_bounds__(1024, 1)
     allgatherFullmesh2(void* sendbuff, mscclpp::DeviceHandle<mscclpp::MemoryChannel>* memoryChannels,
@@ -21,9 +20,7 @@ __global__ void __launch_bounds__(1024, 1)
 
   // Round down to multiple of peer count.
   const size_t nThread = (blockDim.x * gridDim.x) / WARP_SIZE / nPeer * nPeer * WARP_SIZE;
-  if (tid >= nThread) {
-    return;
-  }
+  bool isWorker = tid < nThread;
   const size_t nWarp = nThread / WARP_SIZE;
   const size_t chanOffset = nPeer * blockIdx.x;
   auto memChans = memoryChannels + chanOffset;
@@ -34,76 +31,80 @@ __global__ void __launch_bounds__(1024, 1)
   }
   __syncthreads();
 
-  const size_t bytesPerGPU = nelemsPerGPU * sizeof(int);
-  const size_t bytes = bytesPerGPU * nPeer;
-  size_t unitBytesPerThread;
-  if (bytes >= nThread * 64) {
-    unitBytesPerThread = 64;
-  } else {
-    unitBytesPerThread = 16;
-  }
-  const size_t unitBytesPerWarp = unitBytesPerThread * WARP_SIZE;
-  const size_t unitBytes = unitBytesPerWarp * nWarp;
-  const size_t nLoop = bytes / unitBytes;
-
-  if (nLoop > 0) {
-    // First loop unrolling
-    const size_t peerIdx = wid % nPeer;
-    const size_t offset = bytesPerGPU * rank + (wid / nPeer) * unitBytesPerWarp;
-    if constexpr (IsOutOfPlace) {
-      char* dst = reinterpret_cast<char*>(memChans[peerIdx].dst_);
-      char* src = reinterpret_cast<char*>(memChans[peerIdx].src_);
-      char* buff = reinterpret_cast<char*>(sendbuff);
-      const size_t offsetWithinRank = (wid / nPeer) * unitBytesPerWarp;
-      mscclpp::copy<16, false>(src + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid,
-                               WARP_SIZE);
-      mscclpp::copy<16, false>(dst + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid,
-                               WARP_SIZE);
+  if (isWorker) {
+    const size_t bytesPerGPU = nelemsPerGPU * sizeof(int);
+    const size_t bytes = bytesPerGPU * nPeer;
+    size_t unitBytesPerThread;
+    if (bytes >= nThread * 64) {
+      unitBytesPerThread = 64;
     } else {
-      memChans[peerIdx].put<16, false>(offset + channelOutOffset, unitBytesPerWarp, lid, WARP_SIZE);
+      unitBytesPerThread = 16;
     }
-  }
+    const size_t unitBytesPerWarp = unitBytesPerThread * WARP_SIZE;
+    const size_t unitBytes = unitBytesPerWarp * nWarp;
+    const size_t nLoop = bytes / unitBytes;
 
-  for (size_t i = 1; i < nLoop; ++i) {
-    const size_t gWid = wid + i * nWarp;
-    const size_t peerIdx = gWid % nPeer;
-    const size_t offset = bytesPerGPU * rank + (gWid / nPeer) * unitBytesPerWarp;
-    if constexpr (IsOutOfPlace) {
-      char* dst = reinterpret_cast<char*>(memChans[peerIdx].dst_);
-      char* src = reinterpret_cast<char*>(memChans[peerIdx].src_);
-      char* buff = reinterpret_cast<char*>(sendbuff);
-      const size_t offsetWithinRank = (gWid / nPeer) * unitBytesPerWarp;
-      mscclpp::copy<16, false>(src + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid,
-                               WARP_SIZE);
-      mscclpp::copy<16, false>(dst + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid,
-                               WARP_SIZE);
-    } else {
-      memChans[peerIdx].put<16, false>(offset + channelOutOffset, unitBytesPerWarp, lid, WARP_SIZE);
-    }
-  }
-
-  if (bytes % unitBytes > 0) {
-    const size_t gWid = wid + nLoop * nWarp;
-    const size_t peerIdx = gWid % nPeer;
-    const size_t offsetWithinRank = (gWid / nPeer) * unitBytesPerWarp;
-    const size_t offset = bytesPerGPU * rank + offsetWithinRank;
-    const size_t remainBytes = (offsetWithinRank + unitBytesPerWarp > bytesPerGPU)
-                                   ? ((bytesPerGPU > offsetWithinRank) ? (bytesPerGPU - offsetWithinRank) : 0)
-                                   : unitBytesPerWarp;
-    if (remainBytes > 0) {
+    if (nLoop > 0) {
+      // First loop unrolling
+      const size_t peerIdx = wid % nPeer;
+      const size_t offset = bytesPerGPU * rank + (wid / nPeer) * unitBytesPerWarp;
       if constexpr (IsOutOfPlace) {
         char* dst = reinterpret_cast<char*>(memChans[peerIdx].dst_);
         char* src = reinterpret_cast<char*>(memChans[peerIdx].src_);
         char* buff = reinterpret_cast<char*>(sendbuff);
-        mscclpp::copy<16, true>(src + offset + channelOutOffset, buff + offsetWithinRank, remainBytes, lid, WARP_SIZE);
-        mscclpp::copy<16, true>(dst + offset + channelOutOffset, buff + offsetWithinRank, remainBytes, lid, WARP_SIZE);
+        const size_t offsetWithinRank = (wid / nPeer) * unitBytesPerWarp;
+        mscclpp::copy<16, false>(src + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid,
+                                 WARP_SIZE);
+        mscclpp::copy<16, false>(dst + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid,
+                                 WARP_SIZE);
       } else {
-        memChans[peerIdx].put<16, true>(offset + channelOutOffset, remainBytes, lid, WARP_SIZE);
+        memChans[peerIdx].put<16, false>(offset + channelOutOffset, unitBytesPerWarp, lid, WARP_SIZE);
+      }
+    }
+
+    for (size_t i = 1; i < nLoop; ++i) {
+      const size_t gWid = wid + i * nWarp;
+      const size_t peerIdx = gWid % nPeer;
+      const size_t offset = bytesPerGPU * rank + (gWid / nPeer) * unitBytesPerWarp;
+      if constexpr (IsOutOfPlace) {
+        char* dst = reinterpret_cast<char*>(memChans[peerIdx].dst_);
+        char* src = reinterpret_cast<char*>(memChans[peerIdx].src_);
+        char* buff = reinterpret_cast<char*>(sendbuff);
+        const size_t offsetWithinRank = (gWid / nPeer) * unitBytesPerWarp;
+        mscclpp::copy<16, false>(src + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid,
+                                 WARP_SIZE);
+        mscclpp::copy<16, false>(dst + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid,
+                                 WARP_SIZE);
+      } else {
+        memChans[peerIdx].put<16, false>(offset + channelOutOffset, unitBytesPerWarp, lid, WARP_SIZE);
+      }
+    }
+
+    if (bytes % unitBytes > 0) {
+      const size_t gWid = wid + nLoop * nWarp;
+      const size_t peerIdx = gWid % nPeer;
+      const size_t offsetWithinRank = (gWid / nPeer) * unitBytesPerWarp;
+      const size_t offset = bytesPerGPU * rank + offsetWithinRank;
+      const size_t remainBytes = (offsetWithinRank + unitBytesPerWarp > bytesPerGPU)
+                                     ? ((bytesPerGPU > offsetWithinRank) ? (bytesPerGPU - offsetWithinRank) : 0)
+                                     : unitBytesPerWarp;
+      if (remainBytes > 0) {
+        if constexpr (IsOutOfPlace) {
+          char* dst = reinterpret_cast<char*>(memChans[peerIdx].dst_);
+          char* src = reinterpret_cast<char*>(memChans[peerIdx].src_);
+          char* buff = reinterpret_cast<char*>(sendbuff);
+          mscclpp::copy<16, true>(src + offset + channelOutOffset, buff + offsetWithinRank, remainBytes, lid,
+                                  WARP_SIZE);
+          mscclpp::copy<16, true>(dst + offset + channelOutOffset, buff + offsetWithinRank, remainBytes, lid,
+                                  WARP_SIZE);
+        } else {
+          memChans[peerIdx].put<16, true>(offset + channelOutOffset, remainBytes, lid, WARP_SIZE);
+        }
       }
     }
   }
 
-  deviceSyncer.sync(gridDim.x);
+  __syncthreads();
 
   if (threadIdx.x < nPeer) {
     memChans[threadIdx.x].signal();
diff --git a/src/ext/collectives/allreduce/allreduce_packet.cu b/src/ext/collectives/allreduce/allreduce_packet.cu
index 3c75a746..414c2b1f 100644
--- a/src/ext/collectives/allreduce/allreduce_packet.cu
+++ b/src/ext/collectives/allreduce/allreduce_packet.cu
@@ -240,6 +240,13 @@ CommResult AllreducePacket::allreduceKernelFunc(const std::shared_ptr<void> ctx_
          maxBlockNum_, ".");
     return CommResult::CommInvalidArgument;
   }
+  const int nPeers = ctx->nRanksPerNode - 1;
+  if (blockAndThreadNum.first < nPeers) {
+    WARN(ALGO,
+         "AllreducePacket requires block number to be at least peer count, but got nBlocks=", blockAndThreadNum.first,
+         " and nPeers=", nPeers, ".");
+    return CommResult::CommInvalidArgument;
+  }
 
   size_t sendBytes;
   CUdeviceptr sendBasePtr;
diff --git a/test/unit/CMakeLists.txt b/test/unit/CMakeLists.txt
index a345effc..0d075ac8 100644
--- a/test/unit/CMakeLists.txt
+++ b/test/unit/CMakeLists.txt
@@ -14,5 +14,6 @@ target_sources(unit_tests PRIVATE
     utils_tests.cc
     utils_internal_tests.cc
     compile_tests.cu
+    gpu_data_types_tests.cu
     local_channel_tests.cu
 )
diff --git a/test/unit/gpu_data_types_tests.cu b/test/unit/gpu_data_types_tests.cu
new file mode 100644
index 00000000..5f91c684
--- /dev/null
+++ b/test/unit/gpu_data_types_tests.cu
@@ -0,0 +1,175 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include <array>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <mscclpp/gpu_data_types.hpp>
+#include <mscclpp/gpu_utils.hpp>
+#include <vector>
+
+#include "../framework.hpp"
+
+namespace {
+
+constexpr int kConversionPaths = 3;
+
+template <class T, class... Args>
+std::array<T, sizeof...(Args)> makeArray(Args... args) {
+  return {static_cast<T>(args)...};
+}
+
+__device__ uint32_t floatToBitsDevice(float value) {
+  union {
+    float f;
+    uint32_t u;
+  } cvt = {value};
+  return cvt.u;
+}
+
+uint32_t floatToBitsHost(float value) {
+  uint32_t bits;
+  std::memcpy(&bits, &value, sizeof(bits));
+  return bits;
+}
+
+__global__ void kernelE4m3b15TypeConvert(const float* input, int encodeCases, const uint8_t* raw, int decodeCases,
+                                         uint8_t* encoded, uint32_t* decodedBits) {
+  if (threadIdx.x != 0 || blockIdx.x != 0) return;
+
+  for (int offset = 0; offset < encodeCases; offset += 4) {
+    mscclpp::f32x4 inputX4;
+    for (int i = 0; i < 4; ++i) {
+      inputX4.data[i] = input[offset + i];
+    }
+
+    mscclpp::f8_e4m3b15x4 encodedX4 = mscclpp::to<mscclpp::f8_e4m3b15x4>(inputX4);
+    for (int i = 0; i < 4; ++i) {
+      encoded[offset + i] = encodedX4.data[i].__x;
+    }
+
+    for (int pair = 0; pair < 2; ++pair) {
+      mscclpp::f32x2 inputX2;
+      inputX2.data[0] = input[offset + pair * 2];
+      inputX2.data[1] = input[offset + pair * 2 + 1];
+      mscclpp::f8_e4m3b15x2 encodedX2 = mscclpp::to<mscclpp::f8_e4m3b15x2>(inputX2);
+      encoded[encodeCases + offset + pair * 2] = encodedX2.data[0].__x;
+      encoded[encodeCases + offset + pair * 2 + 1] = encodedX2.data[1].__x;
+    }
+  }
+
+  for (int i = 0; i < encodeCases; ++i) {
+    encoded[2 * encodeCases + i] = __fp8_e4m3b15(input[i]).__x;
+  }
+
+  for (int offset = 0; offset < decodeCases; offset += 4) {
+    mscclpp::f8_e4m3b15x4 rawX4;
+    for (int i = 0; i < 4; ++i) {
+      rawX4.data[i] = __fp8_e4m3b15::fromRaw(raw[offset + i]);
+    }
+
+    mscclpp::f32x4 decodedX4 = mscclpp::to<mscclpp::f32x4>(rawX4);
+    for (int i = 0; i < 4; ++i) {
+      decodedBits[offset + i] = floatToBitsDevice(decodedX4.data[i]);
+    }
+
+    for (int pair = 0; pair < 2; ++pair) {
+      mscclpp::f8_e4m3b15x2 rawX2;
+      rawX2.data[0] = __fp8_e4m3b15::fromRaw(raw[offset + pair * 2]);
+      rawX2.data[1] = __fp8_e4m3b15::fromRaw(raw[offset + pair * 2 + 1]);
+      mscclpp::f32x2 decodedX2 = mscclpp::to<mscclpp::f32x2>(rawX2);
+      decodedBits[decodeCases + offset + pair * 2] = floatToBitsDevice(decodedX2.data[0]);
+      decodedBits[decodeCases + offset + pair * 2 + 1] = floatToBitsDevice(decodedX2.data[1]);
+    }
+  }
+
+  for (int i = 0; i < decodeCases; ++i) {
+    decodedBits[2 * decodeCases + i] = floatToBitsDevice(float(__fp8_e4m3b15::fromRaw(raw[i])));
+  }
+}
+
+}  // namespace
+
+TEST(GpuDataTypesTest, E4m3b15TypeConvert) {
+  const float inf = std::numeric_limits<float>::infinity();
+  const float nan = std::numeric_limits<float>::quiet_NaN();
+  const float maxFloat = std::numeric_limits<float>::max();
+
+  // Each input value maps to the byte at the same index in expectedEncoded. The fp8_e4m3b15 format has no
+  // NaN/Inf encoding, so NaN, Inf, and overflow inputs saturate to +/-1.875 (max byte 0x7f/0xff).
+  const auto input = makeArray<float>(0.0f, -0.0f,                // +/-0
+                                      0x1.0p-19f, -0x1.0p-19f,    // +/-2^-19: underflows to signed 0
+                                      0x1.0p-18f, -0x1.0p-18f,    // +/-2^-18: rounds to min subnormal
+                                      0x1.0p-17f, -0x1.0p-17f,    // +/-2^-17: min subnormal
+                                      0x1.0p-14f, -0x1.0p-14f,    // +/-2^-14: min normal
+                                      0x1.0fcp-2f, -0x1.0fcp-2f,  // Boundary rounds down in magnitude
+                                      0x1.0fep-2f, -0x1.0fep-2f,  // Boundary rounds up in magnitude
+                                      0x1.cfep-2f, -0x1.cfep-2f,  // Boundary rounds to +/-0.46875
+                                      0x1.cp0f, -0x1.cp0f,        // +/-1.75: max finite
+                                      2.0f, -2.0f,                // Overflow saturation
+                                      inf, -inf,                  // +/-Inf saturation
+                                      nan, -maxFloat);            // NaN / large negative saturation
+
+  const auto expectedEncoded = makeArray<uint8_t>(0x00, 0x80,   // +/-0
+                                                  0x00, 0x80,   // Underflow to signed zero
+                                                  0x01, 0x81,   // Round to min signed subnormal
+                                                  0x01, 0x81,   // Min signed subnormal
+                                                  0x08, 0x88,   // Min signed normal
+                                                  0x68, 0xe8,   // Boundary rounds to +/-0.25
+                                                  0x69, 0xe9,   // Boundary rounds to +/-0.28125
+                                                  0x6f, 0xef,   // Boundary rounds to +/-0.46875
+                                                  0x7e, 0xfe,   // Max finite at fp16 grid (1.75)
+                                                  0x7f, 0xff,   // Overflow saturation (1.875)
+                                                  0x7f, 0xff,   // Inf saturation (1.875)
+                                                  0x7f, 0xff);  // NaN / large negative saturation (1.875)
+
+  // Raw bytes to decode, with expectedDecoded giving the exact float value at the same index.
+  const auto raw = makeArray<uint8_t>(0x00, 0x80,                         // +/-0
+                                      0x01, 0x81,                         // +/-2^-17: min subnormal
+                                      0x08, 0x88,                         // +/-2^-14: min normal
+                                      0x68, 0xe8,                         // +/-0.25
+                                      0x69, 0xe9,                         // +/-0.28125
+                                      0x7e, 0xfe);                        // +/-1.75: max finite
+  const auto expectedDecoded = makeArray<float>(0.0f, -0.0f,              // +/-0
+                                                0x1.0p-17f, -0x1.0p-17f,  // +/-2^-17: min subnormal
+                                                0x1.0p-14f, -0x1.0p-14f,  // +/-2^-14: min normal
+                                                0x1.0p-2f, -0x1.0p-2f,    // +/-0.25
+                                                0x1.2p-2f, -0x1.2p-2f,    // +/-0.28125
+                                                0x1.cp0f, -0x1.cp0f);     // +/-1.75: max finite
+
+  ASSERT_EQ(input.size(), expectedEncoded.size());
+  ASSERT_EQ(raw.size(), expectedDecoded.size());
+  ASSERT_EQ(input.size() % 4, size_t(0));
+  ASSERT_EQ(raw.size() % 4, size_t(0));
+
+  auto inputDev = mscclpp::detail::gpuCallocShared<float>(input.size());
+  auto rawDev = mscclpp::detail::gpuCallocShared<uint8_t>(raw.size());
+  auto encodedDev = mscclpp::detail::gpuCallocShared<uint8_t>(input.size() * kConversionPaths);
+  auto decodedBitsDev = mscclpp::detail::gpuCallocShared<uint32_t>(raw.size() * kConversionPaths);
+
+  mscclpp::gpuMemcpy(inputDev.get(), input.data(), input.size(), cudaMemcpyHostToDevice);
+  mscclpp::gpuMemcpy(rawDev.get(), raw.data(), raw.size(), cudaMemcpyHostToDevice);
+
+  kernelE4m3b15TypeConvert<<<1, 1>>>(inputDev.get(), static_cast<int>(input.size()), rawDev.get(),
+                                     static_cast<int>(raw.size()), encodedDev.get(), decodedBitsDev.get());
+  MSCCLPP_CUDATHROW(cudaGetLastError());
+  MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
+
+  std::vector<uint8_t> encoded(input.size() * kConversionPaths);
+  std::vector<uint32_t> decodedBits(raw.size() * kConversionPaths);
+  mscclpp::gpuMemcpy(encoded.data(), encodedDev.get(), encoded.size(), cudaMemcpyDeviceToHost);
+  mscclpp::gpuMemcpy(decodedBits.data(), decodedBitsDev.get(), decodedBits.size(), cudaMemcpyDeviceToHost);
+
+  for (int path = 0; path < kConversionPaths; ++path) {
+    for (size_t i = 0; i < input.size(); ++i) {
+      EXPECT_EQ(static_cast<int>(encoded[path * input.size() + i]), static_cast<int>(expectedEncoded[i]));
+    }
+  }
+
+  for (int path = 0; path < kConversionPaths; ++path) {
+    for (size_t i = 0; i < raw.size(); ++i) {
+      EXPECT_EQ(decodedBits[path * raw.size() + i], floatToBitsHost(expectedDecoded[i]));
+    }
+  }
+}

From 7c390fffd607bcd11d822428afcd2f992c017ca1 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Thu, 4 Jun 2026 13:16:18 -0700
Subject: [PATCH 08/16] Expose NVLS multicast granularity option for GpuBuffer
 (#815)

Add a public Granularity enum (MultiCastMinimum, MultiCastRecommended)
and let GpuBuffer choose the NVLS multicast allocation granularity via a
constructor argument, defaulting to MultiCastMinimum to minimize memory
usage. Expose the same option through the C++ and Python (nanobind)
APIs.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 docs/cpp_api.rst               |  2 ++
 include/mscclpp/gpu_utils.hpp  | 21 +++++++++++++++++++--
 python/csrc/gpu_utils_py.cpp   |  7 ++++++-
 python/mscclpp/__init__.py     |  1 +
 python/mscclpp/_core/buffer.py | 15 +++++++++++----
 5 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/docs/cpp_api.rst b/docs/cpp_api.rst
index 6b5f24c0..a7ebaaf9 100644
--- a/docs/cpp_api.rst
+++ b/docs/cpp_api.rst
@@ -128,6 +128,8 @@ Utilities
 .. doxygenclass:: mscclpp::GpuBuffer
    :members:
 
+.. doxygenenum:: mscclpp::GpuBufferGranularity
+
 .. doxygenclass:: mscclpp::GpuStream
    :members:
 
diff --git a/include/mscclpp/gpu_utils.hpp b/include/mscclpp/gpu_utils.hpp
index b079e0fd..82fa3ec0 100644
--- a/include/mscclpp/gpu_utils.hpp
+++ b/include/mscclpp/gpu_utils.hpp
@@ -317,6 +317,16 @@ bool isNvlsSupported();
 /// @return True if the pointer is allocated by cuMemMap, false otherwise.
 bool isCuMemMapAllocated(void* ptr);
 
+/// Granularity used to size a `GpuBuffer` allocation so that it is compatible with the multicast (NVLS) API.
+enum class GpuBufferGranularity {
+  /// Minimum multicast granularity. Rounds the allocation up to the minimum granularity required for multicast
+  /// compatibility, minimizing memory footprint. This is the default.
+  MultiCastMinimum,
+  /// Recommended multicast granularity. Rounds the allocation up to the granularity recommended by the driver,
+  /// which may be larger than the minimum but can yield better performance.
+  MultiCastRecommended,
+};
+
 /// Allocates a GPU memory space specialized for communication. The memory is zeroed out. Get the device pointer by
 /// `GpuBuffer::data()`.
 ///
@@ -334,7 +344,11 @@ class GpuBuffer {
  public:
   /// Constructs a GpuBuffer with the specified number of elements.
   /// @param nelems Number of elements to allocate. If it is zero, `data()` will return a null pointer.
-  GpuBuffer(size_t nelems) : nelems_(nelems) {
+  /// @param granularity Granularity used to size the allocation for multicast (NVLS) compatibility. Defaults to
+  /// `GpuBufferGranularity::MultiCastMinimum`, which minimizes memory usage. This is ignored when the buffer is not
+  /// allocated through the multicast-compatible path.
+  GpuBuffer(size_t nelems, [[maybe_unused]] GpuBufferGranularity granularity = GpuBufferGranularity::MultiCastMinimum)
+      : nelems_(nelems) {
     if (nelems == 0) {
       bytes_ = 0;
       return;
@@ -342,7 +356,10 @@ class GpuBuffer {
     MSCCLPP_CUDATHROW(cudaGetDevice(&deviceId_));
 #if (CUDA_NVLS_API_AVAILABLE)
     if (isNvlsSupported()) {
-      size_t gran = detail::getMulticastGranularity(nelems * sizeof(T), CU_MULTICAST_GRANULARITY_RECOMMENDED);
+      CUmulticastGranularity_flags granFlag = (granularity == GpuBufferGranularity::MultiCastRecommended)
+                                                  ? CU_MULTICAST_GRANULARITY_RECOMMENDED
+                                                  : CU_MULTICAST_GRANULARITY_MINIMUM;
+      size_t gran = detail::getMulticastGranularity(nelems * sizeof(T), granFlag);
       bytes_ = (nelems * sizeof(T) + gran - 1) / gran * gran / sizeof(T) * sizeof(T);
       memory_ = detail::gpuCallocPhysicalShared<T>(nelems, gran);
       return;
diff --git a/python/csrc/gpu_utils_py.cpp b/python/csrc/gpu_utils_py.cpp
index 60880456..d6527502 100644
--- a/python/csrc/gpu_utils_py.cpp
+++ b/python/csrc/gpu_utils_py.cpp
@@ -114,8 +114,13 @@ static nb::capsule toDlpack(GpuBuffer<char> buffer, std::string dataType, std::v
 void register_gpu_utils(nb::module_& m) {
   m.def("is_nvls_supported", &isNvlsSupported);
 
+  nb::enum_<GpuBufferGranularity>(m, "CppGpuBufferGranularity")
+      .value("MultiCastMinimum", GpuBufferGranularity::MultiCastMinimum)
+      .value("MultiCastRecommended", GpuBufferGranularity::MultiCastRecommended);
+
   nb::class_<GpuBuffer<char>>(m, "CppRawGpuBuffer")
-      .def(nb::init<size_t>(), nb::arg("nelems"))
+      .def(nb::init<size_t, GpuBufferGranularity>(), nb::arg("nelems"),
+           nb::arg("granularity") = GpuBufferGranularity::MultiCastMinimum)
       .def("nelems", &GpuBuffer<char>::nelems)
       .def("bytes", &GpuBuffer<char>::bytes)
       .def("data", [](GpuBuffer<char>& self) { return reinterpret_cast<uintptr_t>(self.data()); })
diff --git a/python/mscclpp/__init__.py b/python/mscclpp/__init__.py
index 5f3a2302..09408171 100644
--- a/python/mscclpp/__init__.py
+++ b/python/mscclpp/__init__.py
@@ -100,6 +100,7 @@ __all__ = [
     "AlgorithmCollection",
     "CommGroup",
     "GpuBuffer",
+    "GpuBufferGranularity",
 ]
 
 
diff --git a/python/mscclpp/_core/buffer.py b/python/mscclpp/_core/buffer.py
index 0575ca68..e07424f5 100644
--- a/python/mscclpp/_core/buffer.py
+++ b/python/mscclpp/_core/buffer.py
@@ -6,14 +6,21 @@ from typing import Union, Tuple
 
 import cupy as cp
 import numpy as np
-from mscclpp._mscclpp import CppRawGpuBuffer
+from mscclpp._mscclpp import CppRawGpuBuffer, CppGpuBufferGranularity
 
-__all__ = ["GpuBuffer"]
+__all__ = ["GpuBuffer", "GpuBufferGranularity"]
+
+GpuBufferGranularity = CppGpuBufferGranularity
 
 
 class GpuBuffer(cp.ndarray):
     def __new__(
-        cls, shape: Union[int, Tuple[int]], dtype: cp.dtype = float, strides: Tuple[int] = None, order: str = "C"
+        cls,
+        shape: Union[int, Tuple[int]],
+        dtype: cp.dtype = float,
+        strides: Tuple[int] = None,
+        order: str = "C",
+        granularity: CppGpuBufferGranularity = CppGpuBufferGranularity.MultiCastMinimum,
     ):
         # Check if `shape` is valid
         if isinstance(shape, int):
@@ -25,6 +32,6 @@ class GpuBuffer(cp.ndarray):
         if any(s <= 0 for s in shape):
             raise ValueError("Shape must be positive.")
         # Create the buffer
-        buffer = CppRawGpuBuffer(np.prod(shape) * np.dtype(dtype).itemsize)
+        buffer = CppRawGpuBuffer(np.prod(shape) * np.dtype(dtype).itemsize, granularity)
         memptr = cp.cuda.MemoryPointer(cp.cuda.UnownedMemory(buffer.data(), buffer.bytes(), buffer), 0)
         return cp.ndarray(shape, dtype=dtype, strides=strides, order=order, memptr=memptr)

From dc0b8d75f325e9c21e556f8fd7a3035722e61c07 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 19 Jun 2026 13:19:01 -0700
Subject: [PATCH 09/16] GB200 support: SendRecv DSL collective and per-channel
 executor connections (#810)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

GB200 support work: introduces point-to-point send/receive in the
MSCCL++ DSL
and extends the executor for split-NVL-domain topologies where some
ranks are
NVL-connected within a node and other ranks must communicate across the
network.

 ### DSL
 - New `SendRecv` collective with separate input/output buffers
   (`python/mscclpp/language/collectives.py`).
 - New multi-node sendrecv DSL example
(`python/mscclpp/language/tests/multi_node/send_recv.py`) with
`--split_mask`
(group size − 1) and `--instances` CLI options. Documents the
channel-ordering
   trick that keeps signal tags cross-matched between paired peers when
   `prev == next`.
 - `BaseBuffer.__getitem__` now accepts slices with `None` start/stop
   (e.g., `buf[:]`).

 ### Executor
 - One connection (unique QP) per channel entry instead of one per peer.
Required for HostNoAtomic IB mode where each QP can forward signals to a
single semaphore. Uses per-peer tag counters so paired ranks agree on
tag
ordering regardless of the order peers appear in each rank's
`connected_to`
   list.
- MEMORY channels now unconditionally use `Transport::CudaIpc`; only
PORT
   channels can use IB. Matches the invariant already enforced by
   `getTransportFlags`.
- `ExecutionContext::connections` is now a `vector<Connection>` indexed
by
channel order (was `unordered_map<int, Connection>` keyed by peer).
Removes
   redundant semaphore fields from `ExecutionContext`.
 - TODO: explicit NVL-domain check in `useIB`

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Changho Hwang <changhohwang@microsoft.com>
---
 python/mscclpp/language/collectives.py        |  43 +++++
 python/mscclpp/language/rank.py               |  21 ++-
 .../language/tests/multi_node/send_recv.py    |  95 +++++++++++
 python/test/executor_test.py                  | 105 ++++++++----
 python/test/executor_test_verifier.cu         | 155 +++++++++++-------
 src/core/executor/executor.cc                 |  54 ++++--
 src/core/include/execution_kernel.hpp         |   8 +-
 7 files changed, 366 insertions(+), 115 deletions(-)
 create mode 100644 python/mscclpp/language/tests/multi_node/send_recv.py

diff --git a/python/mscclpp/language/collectives.py b/python/mscclpp/language/collectives.py
index 55c0e6b6..15d41ad1 100644
--- a/python/mscclpp/language/collectives.py
+++ b/python/mscclpp/language/collectives.py
@@ -236,3 +236,46 @@ class AllToAll(Collective):
             }
             rank_buffers.append(buffers)
         return rank_buffers
+
+
+class SendRecv(Collective):
+    """A SendRecv collective communication pattern.
+
+    SendRecv performs a point-to-point send/receive operation.
+    Each rank sends its input buffer to the next rank and receives data from the
+    previous rank into its output buffer.
+
+    This operation creates input and output buffers both sized by chunk_factor,
+    as each rank sends and receives the same amount of data.
+    """
+
+    def __init__(self, num_ranks, chunk_factor, inplace):
+        """Initialize a new SendRecv collective.
+
+        Args:
+            num_ranks (int): The number of ranks participating in the SendRecv.
+            chunk_factor (int): The size factor for data chunks.
+            inplace (bool): Whether the operation should be performed in-place.
+
+        Example:
+            >>> sendrecv = SendRecv(num_ranks=4, chunk_factor=1, inplace=False)
+        """
+        Collective.__init__(self, num_ranks, chunk_factor, inplace)
+        self.name = "sendrecv"
+
+    def init_buffers(self):
+        """Initialize buffers for the SendRecv operation.
+
+        Creates input and output buffers both sized by chunk_factor.
+
+        Returns:
+            list: A list of buffer dictionaries, one for each rank.
+        """
+        rank_buffers = []
+        for rank in range(self.num_ranks):
+            buffers = {
+                BufferType.input: BaseBuffer(rank, BufferType.input, 0, self.chunk_factor),
+                BufferType.output: BaseBuffer(rank, BufferType.output, 0, self.chunk_factor),
+            }
+            rank_buffers.append(buffers)
+        return rank_buffers
diff --git a/python/mscclpp/language/rank.py b/python/mscclpp/language/rank.py
index e5b7aab8..3fd93dc7 100644
--- a/python/mscclpp/language/rank.py
+++ b/python/mscclpp/language/rank.py
@@ -304,11 +304,24 @@ class BaseBuffer:
         self.size = offset + size
 
     def __getitem__(self, key):
-        if self.offset + key.stop > self.size:
-            raise RuntimeError(
-                f"Index range from {self.offset + key.start} - {self.offset + key.stop} is out of bounds for buffer {self.buffer_type}. Buffer size: {self.size}"
+        if not isinstance(key, slice):
+            raise TypeError(f"Buffer indices must be slices, not {type(key).__name__}")
+        if key.step is not None and key.step != 1:
+            raise ValueError(f"Buffer slicing does not support step != 1 (got step={key.step})")
+        buffer_size = self.size - self.offset
+        start = key.start if key.start is not None else 0
+        stop = key.stop if key.stop is not None else buffer_size
+        if start < 0 or stop < 0:
+            raise ValueError(
+                f"Buffer slicing does not support negative indices (got start={key.start}, stop={key.stop})"
             )
-        return Chunk(self.rank, self.buffer_type, self.offset + key.start, key.stop - key.start)
+        if start > stop:
+            raise ValueError(f"Buffer slice start ({start}) must be <= stop ({stop})")
+        if self.offset + stop > self.size:
+            raise RuntimeError(
+                f"Index range from {self.offset + start} - {self.offset + stop} is out of bounds for buffer {self.buffer_type}. Buffer size: {self.size}"
+            )
+        return Chunk(self.rank, self.buffer_type, self.offset + start, stop - start)
 
 
 class Buffer(BaseBuffer):
diff --git a/python/mscclpp/language/tests/multi_node/send_recv.py b/python/mscclpp/language/tests/multi_node/send_recv.py
new file mode 100644
index 00000000..0e898f95
--- /dev/null
+++ b/python/mscclpp/language/tests/multi_node/send_recv.py
@@ -0,0 +1,95 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+from mscclpp.language.channel import *
+from mscclpp.language.rank import *
+from mscclpp.language.general import *
+from mscclpp.language.program import *
+from mscclpp.language.collectives import *
+
+
+def send_recv(name, nnodes, gpus_per_node, split_mask, instances):
+    gpu_size = nnodes * gpus_per_node
+    group_size = split_mask + 1
+    if split_mask < 0 or (split_mask & (split_mask + 1)) != 0 or gpu_size % group_size != 0:
+        raise ValueError(
+            f"split_mask must be of the form 2^k - 1 and gpu_size ({gpu_size}) must be divisible by "
+            f"group_size ({group_size}), got split_mask={hex(split_mask)}"
+        )
+    collective = SendRecv(gpu_size, 1, False)
+    with CollectiveProgram(
+        name,
+        collective,
+        gpu_size,
+        protocol="Simple",
+        num_threads_per_block=1024,
+        use_double_scratch_buffer=False,
+        min_message_size=0,
+        max_message_size=2**64 - 1,
+        instances=instances,
+    ):
+        # Creating separate port channels for next and prev directions.
+        # When prev and next are the same peer (e.g., 2-node ring), both channels go to the same peer
+        # and get distinct tags. To ensure cross-rank tag matching (rank A's prev_channel signal
+        # arrives at rank B's next_channel wait), we create channels in opposite order for the
+        # "higher" rank so that tags cross-match:
+        #   Lower rank:  [next(tag0), prev(tag1)]
+        #   Higher rank:  [prev(tag0), next(tag1)]
+        # Then lower.prev(tag1) == higher.next(tag1) and higher.prev(tag0) == lower.next(tag0)
+        # When prev != next (3+ nodes), each channel targets a different peer so each gets tag 0
+        # and this ordering doesn't matter.
+        group_size = group_size
+        num_groups = gpu_size // group_size
+        next_channels = {}  # channel for sending to next rank
+        prev_channels = {}  # channel for receiving from prev rank
+        prev_next_ids = {}
+        for node in range(nnodes):
+            for gpu in range(gpus_per_node):
+                global_rank_id = gpu + gpus_per_node * node
+                position_in_group = global_rank_id & split_mask
+                group_id = global_rank_id // group_size
+                next_group_id = (group_id + 1) % num_groups
+                next_global_rank_id = next_group_id * group_size + position_in_group
+                prev_group_id = (group_id - 1 + num_groups) % num_groups
+                prev_global_rank_id = prev_group_id * group_size + position_in_group
+                if prev_global_rank_id == next_global_rank_id and global_rank_id > prev_global_rank_id:
+                    # Higher rank: create prev first, then next (swapped order)
+                    prev_channels[global_rank_id] = PortChannel(prev_global_rank_id, global_rank_id)
+                    next_channels[global_rank_id] = PortChannel(next_global_rank_id, global_rank_id)
+                else:
+                    # Lower rank or different peers: create next first, then prev
+                    next_channels[global_rank_id] = PortChannel(next_global_rank_id, global_rank_id)
+                    prev_channels[global_rank_id] = PortChannel(prev_global_rank_id, global_rank_id)
+                prev_next_ids[global_rank_id] = (prev_global_rank_id, next_global_rank_id)
+
+        # sync with the next rank and the previous rank in the group
+        for node in range(nnodes):
+            for gpu in range(gpus_per_node):
+                global_rank_id = gpu + gpus_per_node * node
+                prev_global_rank_id, next_global_rank_id = prev_next_ids[global_rank_id]
+                prev_channels[global_rank_id].signal(tb=0, data_sync=SyncType.none)
+                next_channels[global_rank_id].wait(tb=0, data_sync=SyncType.after)
+
+                src_rank = Rank(global_rank_id)
+                src_buffer = src_rank.get_input_buffer()
+                dst_rank = Rank(next_global_rank_id)
+                dst_buffer = dst_rank.get_output_buffer()
+
+                next_channels[global_rank_id].put_with_signal(dst_buffer[:], src_buffer[:], tb=0)
+                prev_channels[global_rank_id].wait(tb=0, data_sync=SyncType.none)
+
+        print(JSON())
+
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--name", type=str, help="name of the program")
+parser.add_argument("--nnodes", type=int, default=1, help="number of nodes")
+parser.add_argument("--gpus_per_node", type=int, help="number of gpus per node")
+parser.add_argument("--split_mask", type=lambda x: int(x, 0), default=0x0, help="split mask (e.g. 0x3)")
+parser.add_argument("--instances", type=int, default=4, help="number of instances")
+
+args = parser.parse_args()
+
+send_recv(args.name, args.nnodes, args.gpus_per_node, args.split_mask, args.instances)
diff --git a/python/test/executor_test.py b/python/test/executor_test.py
index 8a309de5..0159b8fa 100644
--- a/python/test/executor_test.py
+++ b/python/test/executor_test.py
@@ -14,6 +14,7 @@ from mscclpp import CommGroup, GpuBuffer
 from mscclpp.utils import KernelBuilder, pack
 import os
 import struct
+from typing import Callable
 
 import cupy as cp
 from mpi4py import MPI
@@ -34,13 +35,13 @@ def parse_dtype(dtype_str):
         raise ValueError(f"Unknown data type: {dtype_str}")
 
 
-def bench_time(n_iters: int, n_graph_iters: int, func):
-    # capture cuda graph for n_iters of the kernel launch
+def bench_time(n_iters: int, n_graph_iters: int, funcs: list[Callable]):
+    """Benchmark execution time. `funcs` is a list of callables; iteration i runs funcs[i % len(funcs)]."""
     stream = cp.cuda.Stream(non_blocking=True)
     with stream:
         stream.begin_capture()
         for i in range(n_iters):
-            func(stream)
+            funcs[i % len(funcs)](stream)
         graph = stream.end_capture()
 
     # now run a warm up round
@@ -61,15 +62,17 @@ def bench_time(n_iters: int, n_graph_iters: int, func):
 
 def bench_correctness(
     collective: str,
-    input_buf: cp.ndarray,
-    result_buf: cp.ndarray,
-    test_buf: cp.ndarray,
+    input_bufs: list[cp.ndarray],
+    result_bufs: list[cp.ndarray],
+    test_bufs: list[cp.ndarray],
     dtype_str: str,
     rank: int,
     num_ranks: int,
     n_iters: int,
-    func,
+    funcs: list[Callable],
+    split_mask: int = 0,
 ):
+    """Validate correctness. Buffers and funcs are parallel lists; iteration i uses index i % len(funcs)."""
     type_size = cp.dtype(parse_dtype(dtype_str)).itemsize
 
     fill_data_kernel_name = "fill_data_%s" % dtype_str
@@ -79,8 +82,12 @@ def bench_correctness(
         coll = "reduce_scatter"
     elif "allreduce" in collective:
         coll = "all_reduce"
-    else:
+    elif "alltoall" in collective:
         coll = "all_to_all"
+    elif "sendrecv" in collective:
+        coll = "send_recv"
+    else:
+        raise ValueError(f"Unknown collective: {collective}")
     test_data_kernel_name = "test_data_%s_%s" % (coll, dtype_str)
 
     file_dir = os.path.dirname(os.path.abspath(__file__))
@@ -97,11 +104,20 @@ def bench_correctness(
     with stream:
         stream.begin_capture()
         for i in range(n_iters):
-            fill_data_params = pack(input_buf) + struct.pack("Q", input_buf.nbytes // type_size) + pack(rank, i)
+            idx = i % len(funcs)
+            cur_input = input_bufs[idx]
+            cur_result = result_bufs[idx]
+            cur_test = test_bufs[idx]
+
+            fill_data_params = (
+                pack(cur_input) + struct.pack("Q", cur_input.nbytes // type_size) + pack(rank, i, split_mask)
+            )
             fill_data_kernel.launch_kernel(fill_data_params, nblocks, nthreads, 0, stream)
-            func(stream)
+            funcs[idx](stream)
             test_data_params = (
-                pack(result_buf, test_buf) + struct.pack("Q", input_buf.nbytes // type_size) + pack(num_ranks, rank, i)
+                pack(cur_result, cur_test)
+                + struct.pack("Q", cur_input.nbytes // type_size)
+                + pack(num_ranks, rank, i, split_mask)
             )
             test_data_kernel.launch_kernel(test_data_params, nblocks, nthreads, 0, stream)
         graph = stream.end_capture()
@@ -143,10 +159,20 @@ def build_bufs(
     rank: int,
     num_ranks: int,
 ):
+    """Allocate input/result/test buffers. Returns parallel lists (length 2 for sendrecv double-buffering,
+    length 1 otherwise) so callers can iterate uniformly."""
     type_size = cp.dtype(dtype).itemsize
     assert (size % type_size) == 0, "size %d not multiple of type size %d" % (size, type_size)
     nelems = size // type_size
 
+    # Sendrecv uses double buffering: build two parallel buffer slots.
+    if "sendrecv" in collective:
+        n_slots = 2
+        input_bufs = [GpuBuffer(nelems, dtype=dtype) for _ in range(n_slots)]
+        result_bufs = [GpuBuffer(nelems, dtype=dtype) for _ in range(n_slots)]
+        test_bufs = [cp.zeros(nelems, dtype=dtype) for _ in range(n_slots)]
+        return input_bufs, result_bufs, test_bufs, nelems
+
     if "allgather" in collective:
         assert (nelems % num_ranks) == 0, "nelems %d not multiple of num_ranks %d" % (nelems, num_ranks)
         nelems_input = nelems if in_place else nelems // num_ranks
@@ -173,7 +199,7 @@ def build_bufs(
 
     test_buf = cp.zeros(nelems, dtype=dtype)
 
-    return input_buf, result_buf, test_buf
+    return [input_buf], [result_buf], [test_buf], nelems
 
 
 def main(
@@ -184,8 +210,14 @@ def main(
     packet_type: PacketType = PacketType.LL16,
     n_iters: int = 10,
     n_graph_iters: int = 10,
+    split_mask: int = 0,
 ):
     mscclpp_group = CommGroup(MPI.COMM_WORLD)
+    if split_mask < 0 or (split_mask & (split_mask + 1)) != 0 or mscclpp_group.nranks % (split_mask + 1) != 0:
+        raise ValueError(
+            f"split_mask must be of the form 2^k - 1 and nranks ({mscclpp_group.nranks}) must be divisible "
+            f"by group_size ({split_mask + 1}), got split_mask={hex(split_mask)}"
+        )
     cp.cuda.Device(mscclpp_group.my_rank % mscclpp_group.nranks_per_node).use()
     executor = Executor(mscclpp_group.communicator)
     npkit_dump_dir = env().npkit_dump_dir
@@ -195,7 +227,7 @@ def main(
     collective = execution_plan.collective
 
     dtype = parse_dtype(dtype_str)
-    input_buf, result_buf, test_buf = build_bufs(
+    input_bufs, result_bufs, test_bufs, nelem = build_bufs(
         collective,
         size,
         in_place,
@@ -204,39 +236,48 @@ def main(
         mscclpp_group.nranks,
     )
 
-    executor_func = lambda stream: executor.execute(
-        mscclpp_group.my_rank,
-        input_buf.data.ptr,
-        result_buf.data.ptr,
-        input_buf.nbytes,
-        result_buf.nbytes,
-        dtype_to_mscclpp_dtype(dtype_str),
-        execution_plan,
-        stream.ptr,
-        packet_type,
-    )
+    executor_funcs = [
+        (
+            lambda stream, inp=inp, res=res: executor.execute(
+                mscclpp_group.my_rank,
+                inp.data.ptr,
+                res.data.ptr,
+                inp.nbytes,
+                res.nbytes,
+                dtype_to_mscclpp_dtype(dtype_str),
+                execution_plan,
+                stream.ptr,
+                packet_type,
+            )
+        )
+        for inp, res in zip(input_bufs, result_bufs)
+    ]
 
     mscclpp_group.barrier()
     bench_correctness(
         collective,
-        input_buf,
-        result_buf,
-        test_buf,
+        input_bufs,
+        result_bufs,
+        test_bufs,
         dtype_str,
         mscclpp_group.my_rank,
         mscclpp_group.nranks,
         n_iters,
-        executor_func,
+        executor_funcs,
+        split_mask=split_mask,
     )
 
     mscclpp_group.barrier()
-    execution_time = bench_time(n_iters, n_graph_iters, executor_func)
+    execution_time = bench_time(n_iters, n_graph_iters, executor_funcs)
     if npkit_dump_dir is not None:
         npkit.dump(npkit_dump_dir)
         npkit.shutdown()
+
+    result_nbytes = result_bufs[0].nbytes
     print(
         f"Rank: {mscclpp_group.my_rank} Execution time: {execution_time} us, "
-        f"data size: {result_buf.nbytes} bytes data type: {dtype_str} "
+        f"data size: {result_nbytes} bytes data type: {dtype_str} "
+        f"bandwidth: {result_nbytes / (execution_time * 1e-6) / (1024**3):.2f} GB/s, "
         f"packet type: {packet_type}"
     )
     executor = None
@@ -252,6 +293,9 @@ if __name__ == "__main__":
     parser.add_argument("--packet_type", type=str, default="LL16", help="Choose from LL8, LL16")
     parser.add_argument("--n_iters", type=int, default=10)
     parser.add_argument("--n_graph_iters", type=int, default=10)
+    parser.add_argument(
+        "--split_mask", type=lambda x: int(x, 0), default=0x0, help="split mask for sendrecv (e.g. 0x3)"
+    )
     args = parser.parse_args()
 
     packet_type = PacketType.LL16
@@ -267,4 +311,5 @@ if __name__ == "__main__":
         packet_type,
         args.n_iters,
         args.n_graph_iters,
+        args.split_mask,
     )
diff --git a/python/test/executor_test_verifier.cu b/python/test/executor_test_verifier.cu
index e7749197..96ab25c4 100644
--- a/python/test/executor_test_verifier.cu
+++ b/python/test/executor_test_verifier.cu
@@ -22,14 +22,19 @@ static __device__ unsigned int ranqd1(unsigned int seed) {
 // fill/test kernel pairs must have the same thread block size to
 // match their random number series.
 
-#define FILL_DATA(FuncNameType, DataType)                                                                \
-  extern "C" __global__ void __launch_bounds__(1024, 1)                                                  \
-      fill_data_##FuncNameType(DataType* input_buf, size_t num_elems, int rank, int seq) {               \
-    unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq);              \
-    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
-      seed = ranqd1(seed);                                                                               \
-      input_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x);                                 \
-    }                                                                                                    \
+// `split_mask` groups ranks together: group_size = split_mask + 1, group_id = rank / group_size.
+// Data is seeded by group_id so that all ranks within a group produce the same fill, and ranks
+// in different groups produce different fills. With split_mask == 0 this reduces to per-rank
+// seeding (group_id == rank).
+#define FILL_DATA(FuncNameType, DataType)                                                                  \
+  extern "C" __global__ void __launch_bounds__(1024, 1)                                                    \
+      fill_data_##FuncNameType(DataType* input_buf, size_t num_elems, int rank, int seq, int split_mask) { \
+    int seed_rank = rank / (split_mask + 1);                                                               \
+    unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq);           \
+    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) {   \
+      seed = ranqd1(seed);                                                                                 \
+      input_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x);                                   \
+    }                                                                                                      \
   }
 
 FILL_DATA(bfloat16, __nv_bfloat16)
@@ -37,18 +42,20 @@ FILL_DATA(float16, __half)
 FILL_DATA(float32, float)
 FILL_DATA(int32, int)
 
-#define TEST_DATA_ALL_GATHER(FuncNameType, DataType)                                                       \
-  extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_gather_##FuncNameType(               \
-      DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) {   \
-    for (int rank = 0; rank < num_ranks; rank++) {                                                         \
-      size_t rank_offset = rank * num_elems;                                                               \
-      unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq);              \
-      for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
-        seed = ranqd1(seed);                                                                               \
-        test_buf[rank_offset + i] = DataType(seed % blockDim.x) / DataType(blockDim.x);                    \
-        assert(result_buf[rank_offset + i] == test_buf[rank_offset + i]);                                  \
-      }                                                                                                    \
-    }                                                                                                      \
+#define TEST_DATA_ALL_GATHER(FuncNameType, DataType)                                                                 \
+  extern "C" __global__ void __launch_bounds__(1024, 1)                                                              \
+      test_data_all_gather_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \
+                                          int my_rank, int seq, int split_mask) {                                    \
+    for (int rank = 0; rank < num_ranks; rank++) {                                                                   \
+      size_t rank_offset = rank * num_elems;                                                                         \
+      int seed_rank = rank / (split_mask + 1);                                                                       \
+      unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq);                   \
+      for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) {           \
+        seed = ranqd1(seed);                                                                                         \
+        test_buf[rank_offset + i] = DataType(seed % blockDim.x) / DataType(blockDim.x);                              \
+        assert(result_buf[rank_offset + i] == test_buf[rank_offset + i]);                                            \
+      }                                                                                                              \
+    }                                                                                                                \
   }
 
 TEST_DATA_ALL_GATHER(bfloat16, __nv_bfloat16)
@@ -56,25 +63,27 @@ TEST_DATA_ALL_GATHER(float16, __half)
 TEST_DATA_ALL_GATHER(float32, float)
 TEST_DATA_ALL_GATHER(int32, int)
 
-#define TEST_DATA_ALL_REDUCE(FuncNameType, DataType, Eps)                                                  \
-  extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_reduce_##FuncNameType(               \
-      DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) {   \
-    for (int rank = 0; rank < num_ranks; rank++) {                                                         \
-      unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq);              \
-      for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
-        if (rank == 0) {                                                                                   \
-          test_buf[i] = 0;                                                                                 \
-        }                                                                                                  \
-        seed = ranqd1(seed);                                                                               \
-        test_buf[i] += DataType(seed % blockDim.x) / DataType(blockDim.x);                                 \
-      }                                                                                                    \
-    }                                                                                                      \
-    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) {   \
-      float expected = float(test_buf[i]);                                                                 \
-      float result = float(result_buf[i]);                                                                 \
-      float tol = Eps * num_ranks * (1.0f + abs(expected));                                                \
-      assert(abs(result - expected) <= tol);                                                               \
-    }                                                                                                      \
+#define TEST_DATA_ALL_REDUCE(FuncNameType, DataType, Eps)                                                            \
+  extern "C" __global__ void __launch_bounds__(1024, 1)                                                              \
+      test_data_all_reduce_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \
+                                          int my_rank, int seq, int split_mask) {                                    \
+    for (int rank = 0; rank < num_ranks; rank++) {                                                                   \
+      int seed_rank = rank / (split_mask + 1);                                                                       \
+      unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq);                   \
+      for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) {           \
+        if (rank == 0) {                                                                                             \
+          test_buf[i] = 0;                                                                                           \
+        }                                                                                                            \
+        seed = ranqd1(seed);                                                                                         \
+        test_buf[i] += DataType(seed % blockDim.x) / DataType(blockDim.x);                                           \
+      }                                                                                                              \
+    }                                                                                                                \
+    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) {             \
+      float expected = float(test_buf[i]);                                                                           \
+      float result = float(result_buf[i]);                                                                           \
+      float tol = Eps * num_ranks * (1.0f + abs(expected));                                                          \
+      assert(abs(result - expected) <= tol);                                                                         \
+    }                                                                                                                \
   }
 
 TEST_DATA_ALL_REDUCE(bfloat16, __nv_bfloat16, 7.8125e-3f)
@@ -83,12 +92,14 @@ TEST_DATA_ALL_REDUCE(float32, float, 1.1920929e-7f)
 TEST_DATA_ALL_REDUCE(int32, int, 0.0f)
 
 #define TEST_DATA_REDUCE_SCATTER(FuncNameType, DataType, Eps)                                              \
-  extern "C" __global__ void __launch_bounds__(1024, 1) test_data_reduce_scatter_##FuncNameType(           \
-      DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) {   \
+  extern "C" __global__ void __launch_bounds__(1024, 1)                                                    \
+      test_data_reduce_scatter_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems,  \
+                                              int num_ranks, int my_rank, int seq, int split_mask) {       \
     int nem_elems_per_rank = num_elems / num_ranks;                                                        \
     int offset = nem_elems_per_rank * my_rank;                                                             \
     for (int rank = 0; rank < num_ranks; rank++) {                                                         \
-      unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq);              \
+      int seed_rank = rank / (split_mask + 1);                                                             \
+      unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq);         \
       for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
         if (rank == 0) {                                                                                   \
           test_buf[i] = 0;                                                                                 \
@@ -112,25 +123,51 @@ TEST_DATA_REDUCE_SCATTER(float16, __half, 9.765625e-4f)
 TEST_DATA_REDUCE_SCATTER(float32, float, 1.1920929e-7f)
 TEST_DATA_REDUCE_SCATTER(int32, int, 0.0f)
 
-#define TEST_DATA_ALL_TO_ALL(FuncNameType, DataType)                                                       \
-  extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_to_all_##FuncNameType(               \
-      DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) {   \
-    int nem_elems_per_rank = num_elems / num_ranks;                                                        \
-    int offset = nem_elems_per_rank * my_rank;                                                             \
-    for (int rank = 0; rank < num_ranks; rank++) {                                                         \
-      size_t rank_offset = rank * nem_elems_per_rank;                                                      \
-      unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq);              \
-      for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
-        seed = ranqd1(seed);                                                                               \
-        if (i >= my_rank * nem_elems_per_rank && i < (my_rank + 1) * nem_elems_per_rank) {                 \
-          test_buf[rank_offset + i - offset] = DataType(seed % blockDim.x) / DataType(blockDim.x);         \
-          assert(result_buf[rank_offset + i - offset] == test_buf[rank_offset + i - offset]);              \
-        }                                                                                                  \
-      }                                                                                                    \
-    }                                                                                                      \
+#define TEST_DATA_ALL_TO_ALL(FuncNameType, DataType)                                                                 \
+  extern "C" __global__ void __launch_bounds__(1024, 1)                                                              \
+      test_data_all_to_all_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \
+                                          int my_rank, int seq, int split_mask) {                                    \
+    int nem_elems_per_rank = num_elems / num_ranks;                                                                  \
+    int offset = nem_elems_per_rank * my_rank;                                                                       \
+    for (int rank = 0; rank < num_ranks; rank++) {                                                                   \
+      size_t rank_offset = rank * nem_elems_per_rank;                                                                \
+      int seed_rank = rank / (split_mask + 1);                                                                       \
+      unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq);                   \
+      for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) {           \
+        seed = ranqd1(seed);                                                                                         \
+        if (i >= my_rank * nem_elems_per_rank && i < (my_rank + 1) * nem_elems_per_rank) {                           \
+          test_buf[rank_offset + i - offset] = DataType(seed % blockDim.x) / DataType(blockDim.x);                   \
+          assert(result_buf[rank_offset + i - offset] == test_buf[rank_offset + i - offset]);                        \
+        }                                                                                                            \
+      }                                                                                                              \
+    }                                                                                                                \
   }
 
 TEST_DATA_ALL_TO_ALL(bfloat16, __nv_bfloat16)
 TEST_DATA_ALL_TO_ALL(float16, __half)
 TEST_DATA_ALL_TO_ALL(float32, float)
-TEST_DATA_ALL_TO_ALL(int32, int)
\ No newline at end of file
+TEST_DATA_ALL_TO_ALL(int32, int)
+
+// Sendrecv verification: receive from the prev group in the ring.
+// fill_data seeds by group_id (rank / (split_mask + 1)); the receiver in group g expects the
+// data produced by group (g - 1 + num_groups) % num_groups, so we recompute that seed here.
+#define TEST_DATA_SEND_RECV(FuncNameType, DataType)                                                                 \
+  extern "C" __global__ void __launch_bounds__(1024, 1)                                                             \
+      test_data_send_recv_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \
+                                         int my_rank, int seq, int split_mask) {                                    \
+    int group_size = split_mask + 1;                                                                                \
+    int num_groups = num_ranks / group_size;                                                                        \
+    int my_group_id = my_rank / group_size;                                                                         \
+    int prev_group_id = (my_group_id - 1 + num_groups) % num_groups;                                                \
+    unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + prev_group_id + seq);                \
+    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) {            \
+      seed = ranqd1(seed);                                                                                          \
+      test_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x);                                             \
+      assert(result_buf[i] == test_buf[i]);                                                                         \
+    }                                                                                                               \
+  }
+
+TEST_DATA_SEND_RECV(bfloat16, __nv_bfloat16)
+TEST_DATA_SEND_RECV(float16, __half)
+TEST_DATA_SEND_RECV(float32, float)
+TEST_DATA_SEND_RECV(int32, int)
diff --git a/src/core/executor/executor.cc b/src/core/executor/executor.cc
index fcecc4dd..9ef59bc1 100644
--- a/src/core/executor/executor.cc
+++ b/src/core/executor/executor.cc
@@ -94,6 +94,7 @@ struct hash<mscclpp::DeviceExecutionPlanKey> {
 namespace {
 auto hasIBDevices = []() { return mscclpp::getIBDeviceCount() > 0; };
 
+// TODO(binyli): Need to add NVL domain check.
 auto useIB = [](int rank1, int rank2, int nranksPerNode) {
   bool inSameNode = rank1 / nranksPerNode == rank2 / nranksPerNode;
   return hasIBDevices() && !inSameNode;
@@ -108,7 +109,7 @@ namespace mscclpp {
 
 struct ExecutionContext {
   std::shared_ptr<ProxyService> proxyService;
-  std::unordered_map<int, Connection> connections;
+  std::vector<Connection> connections;
   std::vector<std::shared_ptr<NvlsConnection>> nvlsConnections;
   MemoryId localMemoryIdBegin = MemoryId(0);
 
@@ -120,8 +121,6 @@ struct ExecutionContext {
   // local registered memories to keep resources alive
   std::vector<mscclpp::RegisteredMemory> localRegisteredMemories;
 
-  std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> memorySemaphores;
-  std::vector<mscclpp::SemaphoreId> proxySemaphores;
   std::vector<mscclpp::BaseMemoryChannel> memoryChannels;
   std::vector<mscclpp::BasePortChannel> portChannels;
   std::vector<mscclpp::SwitchChannel> nvlsChannels;
@@ -265,15 +264,36 @@ struct Executor::Impl {
       }
     };
 
-    std::vector<int> connectedPeers = plan.impl_->getConnectedPeers();
-    std::vector<std::shared_future<mscclpp::Connection>> connectionFutures;
-    for (int peer : connectedPeers) {
-      Transport transport =
-          !useIB(rank, peer, this->nranksPerNode) ? Transport::CudaIpc : IBs[rank % this->nranksPerNode];
-      connectionFutures.push_back(this->comm->connect(transport, peer));
+    // Create one connection (unique QP) per channel entry. Each channel gets its own
+    // QP — no shared connections.
+    // Use per-peer tag counters so that matched connections between pairs of ranks use
+    // the same tag, regardless of the order peers appear in each rank's connected_to list.
+    std::unordered_map<int, int> peerTagCounters;
+    Transport ibTransport = IBs[rank % this->nranksPerNode];
+    std::vector<std::shared_future<Connection>> connFutures;
+    for (ChannelType channelType : {ChannelType::MEMORY, ChannelType::PORT}) {
+      std::vector<ChannelInfo> channelInfos = plan.impl_->getChannelInfos(channelType);
+      for (const auto& info : channelInfos) {
+        for (int peer : info.connectedPeers) {
+          Transport transport = channelType == ChannelType::PORT && useIB(rank, peer, this->nranksPerNode)
+                                    ? ibTransport
+                                    : Transport::CudaIpc;
+          connFutures.push_back(this->comm->connect(transport, peer, peerTagCounters[peer]++));
+        }
+      }
+      channelInfos = plan.impl_->getUnpairedChannelInfos(nranks, channelType);
+      for (const auto& info : channelInfos) {
+        for (int peer : info.connectedPeers) {
+          Transport transport = channelType == ChannelType::PORT && useIB(rank, peer, this->nranksPerNode)
+                                    ? ibTransport
+                                    : Transport::CudaIpc;
+          connFutures.push_back(this->comm->connect(transport, peer, peerTagCounters[peer]++));
+        }
+      }
     }
-    for (size_t i = 0; i < connectionFutures.size(); i++) {
-      context.connections[connectedPeers[i]] = connectionFutures[i].get();
+
+    for (auto& future : connFutures) {
+      context.connections.push_back(future.get());
     }
 
     std::vector<NvlsInfo> nvlsInfos = plan.impl_->nvlsInfos.at(rank);
@@ -327,10 +347,11 @@ struct Executor::Impl {
     std::vector<std::shared_future<Semaphore>> futureProxySemaphores;
     std::vector<std::shared_ptr<MemoryDevice2DeviceSemaphore>> memorySemaphores;
     std::vector<mscclpp::SemaphoreId> proxySemaphores;
+    int connIdx = 0;
     auto processChannelInfos = [&](std::vector<ChannelInfo>& channelInfos) {
       for (ChannelInfo& info : channelInfos) {
-        for (int peer : info.connectedPeers) {
-          auto connection = context.connections.at(peer);
+        for (size_t i = 0; i < info.connectedPeers.size(); i++) {
+          auto& connection = context.connections[connIdx++];
           if (info.channelType == ChannelType::MEMORY) {
             futureMemorySemaphores.push_back(this->comm->buildSemaphore(
                 connection, this->comm->remoteRankOf(connection), this->comm->tagOf(connection)));
@@ -359,18 +380,15 @@ struct Executor::Impl {
       proxySemaphores.push_back(context.proxyService->addSemaphore(sem.get()));
     }
 
-    context.memorySemaphores = std::move(memorySemaphores);
-    context.proxySemaphores = std::move(proxySemaphores);
-
     for (ChannelType channelType : channelTypes) {
       std::vector<ChannelInfo> channelInfos = plan.impl_->getChannelInfos(channelType);
       int index = 0;
       for (ChannelInfo& info : channelInfos) {
         for (size_t i = 0; i < info.connectedPeers.size(); i++) {
           if (channelType == ChannelType::MEMORY) {
-            context.memoryChannels.emplace_back(context.memorySemaphores[index++]);
+            context.memoryChannels.emplace_back(memorySemaphores[index++]);
           } else if (channelType == ChannelType::PORT) {
-            context.portChannels.emplace_back(context.proxyService->basePortChannel(context.proxySemaphores[index++]));
+            context.portChannels.emplace_back(context.proxyService->basePortChannel(proxySemaphores[index++]));
           }
         }
       }
diff --git a/src/core/include/execution_kernel.hpp b/src/core/include/execution_kernel.hpp
index cb808bc8..12223181 100644
--- a/src/core/include/execution_kernel.hpp
+++ b/src/core/include/execution_kernel.hpp
@@ -174,11 +174,11 @@ MSCCLPP_DEVICE_INLINE void handlePut(const Operation& op, void* input, void* out
       uint32_t dstOffset =
           dstOffsets[tid] + getOffset<ReuseScratch>(portChannelBufferTypes_[op.outputBufferRefs[tid].id], offset);
       uint32_t srcOffset = srcOffsets[tid] + getOffset<ReuseScratch>(op.inputBufferRefs[tid].type, offset);
-      if constexpr (PutWithSignal) {
-        portChannels_[channelIndexes[tid]].putWithSignal(dstMemoryId, dstOffset, srcMemoryId, srcOffset, size);
-      } else if constexpr (PutWithSignalAndFlush) {
+      if constexpr (PutWithSignalAndFlush) {
         portChannels_[channelIndexes[tid]].putWithSignalAndFlush(dstMemoryId, (uint64_t)dstOffset, srcMemoryId,
-                                                                 (uint64_t)srcOffsets, size);
+                                                                 (uint64_t)srcOffset, size);
+      } else if constexpr (PutWithSignal) {
+        portChannels_[channelIndexes[tid]].putWithSignal(dstMemoryId, dstOffset, srcMemoryId, srcOffset, size);
       } else {
         portChannels_[channelIndexes[tid]].put(dstMemoryId, dstOffset, srcMemoryId, srcOffset, size);
       }

From 9aab9cacc0e576827eaa9c3bfbf3d509f4d1bf4c Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Wed, 24 Jun 2026 16:09:34 -0700
Subject: [PATCH 10/16] support rocm7.2 (#819)

This pull request introduces support for ROCm 7.2 across the build
system, CI pipelines, Docker images, and documentation, while also
improving ROCm FP8 type selection and CUDA IPC memory handle management.
It updates dependencies and configurations to ensure compatibility with
ROCm 7.2, adds new options for native FP8 variants, and refines some
benchmarking and internal memory handling logic.

Pls notice: there is an issue in rocm7.2 (rocm7.2 user lib + rocm6.2
driver) when execution code in this order: allocating memory -> ipc
communication -> allocate new memory -> free old memory.

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 .azure-pipelines/codecov.yml                  |  2 +
 .azure-pipelines/rccl-api-test.yml            |  2 +
 .azure-pipelines/templates/rccl-test.yml      | 14 +--
 .azure-pipelines/ut.yml                       |  2 +
 .github/workflows/codeql-analysis.yml         |  2 +-
 CMakeLists.txt                                | 12 ++-
 docker/base-dev-x.dockerfile                  |  4 +-
 docker/build.sh                               |  5 +-
 docs/quickstart.md                            | 14 +--
 include/mscclpp/gpu_data_types.hpp            |  5 +-
 pyproject.toml                                |  4 +
 python/mscclpp_benchmark/bench_collective.py  |  7 +-
 python/mscclpp_benchmark/correctness.py       |  1 -
 python/mscclpp_benchmark/gpu.py               | 16 ++++
 python/mscclpp_benchmark/tuner.py             |  1 -
 python/requirements_rocm7.txt                 | 11 +++
 src/core/gpu_ipc_mem.cc                       | 86 ++++++++++---------
 src/core/registered_memory.cc                 |  5 +-
 .../allreduce/allreduce_fullmesh.cu           | 53 ++++++------
 .../include/allreduce/allreduce_fullmesh.hpp  |  2 -
 test/deploy/setup.sh                          |  4 +-
 21 files changed, 158 insertions(+), 94 deletions(-)
 create mode 100644 python/requirements_rocm7.txt

diff --git a/.azure-pipelines/codecov.yml b/.azure-pipelines/codecov.yml
index c4abeaa7..5f3e531e 100644
--- a/.azure-pipelines/codecov.yml
+++ b/.azure-pipelines/codecov.yml
@@ -80,6 +80,8 @@ jobs:
     matrix:
       rocm6_2:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm6.2
+      rocm7_2:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm7.2
 
   container:
     image: $(containerImage)
diff --git a/.azure-pipelines/rccl-api-test.yml b/.azure-pipelines/rccl-api-test.yml
index 43841079..fc793c88 100644
--- a/.azure-pipelines/rccl-api-test.yml
+++ b/.azure-pipelines/rccl-api-test.yml
@@ -35,6 +35,8 @@ jobs:
     matrix:
       rocm6_2:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm6.2
+      rocm7_2:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm7.2
 
   container:
     image: $(containerImage)
diff --git a/.azure-pipelines/templates/rccl-test.yml b/.azure-pipelines/templates/rccl-test.yml
index 63788ac2..65f1c984 100644
--- a/.azure-pipelines/templates/rccl-test.yml
+++ b/.azure-pipelines/templates/rccl-test.yml
@@ -46,25 +46,25 @@ steps:
     name: RunRcclAllGatherTest
     displayName: Run RCCL AllGather Test with or without MSCCLPP Lib
     remoteScript: |
-      mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_gather_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
-      mpirun -np 8 --bind-to numa --allow-run-as-root /root/rocm-systems/projects/rccl-tests/build/all_gather_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
+      mpirun -np 8 --bind-to numa --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_gather_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
+      mpirun -np 8 --bind-to numa --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 /root/rocm-systems/projects/rccl-tests/build/all_gather_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
 
 - template: run-remote-task.yml
   parameters:
     name: RunRcclAllReduceTest
     displayName: Run RCCL AllReduce Test with or without MSCCLPP Lib
     remoteScript: |
-      mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
-      mpirun -np 8 --bind-to numa --allow-run-as-root /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
+      mpirun -np 8 --bind-to numa --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
+      mpirun -np 8 --bind-to numa --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20
 
 - template: run-remote-task.yml
   parameters:
     name: PyBench
     displayName: Run Collective Benchmarks
     remoteScript: |
-      mpirun --allow-run-as-root -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3b15 --accum-type float32 --autotune
-      mpirun --allow-run-as-root -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3fnuz --accum-type float32 --autotune
-      mpirun --allow-run-as-root -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allgather --dtype float8_e4m3b15 --autotune --buffer-mode out-of-place
+      mpirun --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3b15 --accum-type float32 --autotune
+      mpirun --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3fnuz --accum-type float32 --autotune
+      mpirun --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allgather --dtype float8_e4m3b15 --autotune --buffer-mode out-of-place
 
 - template: stop.yml
   parameters:
diff --git a/.azure-pipelines/ut.yml b/.azure-pipelines/ut.yml
index 6b8c9eda..0cabdaeb 100644
--- a/.azure-pipelines/ut.yml
+++ b/.azure-pipelines/ut.yml
@@ -137,6 +137,8 @@ jobs:
     matrix:
       rocm6_2:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm6.2
+      rocm7_2:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm7.2
 
   container:
     image: $(containerImage)
diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml
index fb065141..982d8633 100644
--- a/.github/workflows/codeql-analysis.yml
+++ b/.github/workflows/codeql-analysis.yml
@@ -85,7 +85,7 @@ jobs:
       fail-fast: false
       matrix:
         language: [ 'cpp', 'python' ]
-        version: [ 'rocm6.2' ]
+        version: [ 'rocm6.2', 'rocm7.2' ]
 
     steps:
     - name: Checkout repository
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 49154e0b..92822012 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -59,6 +59,7 @@ option(MSCCLPP_BYPASS_GPU_CHECK "Bypass GPU check." OFF)
 option(MSCCLPP_NPKIT_FLAGS "Set NPKIT flags" OFF)
 option(MSCCLPP_ENABLE_COVERAGE "Enable code coverage" OFF)
 option(MSCCLPP_DISABLE_NB_LEAK_WARNINGS "Disable Nanobind leak warnings" ON)
+option(MSCCLPP_ROCM_USE_FNUZ_FP8 "Use ROCm FNUZ native FP8 types." ON)
 set(MSCCLPP_GPU_ARCHS "" CACHE STRING "Specify GPU architectures with delimiters (comma, space, or semicolon).")
 
 if(MSCCLPP_BYPASS_GPU_CHECK)
@@ -183,11 +184,20 @@ elseif(MSCCLPP_USE_CUDA)
 	    endif()
     endif()
 elseif(MSCCLPP_USE_ROCM)
-    set(MSCCLPP_GPU_ARCHS gfx90a gfx941 gfx942)
+    set(MSCCLPP_GPU_ARCHS gfx90a gfx942)
 endif()
 
 message(STATUS "GPU architectures: ${MSCCLPP_GPU_ARCHS}")
 
+if(MSCCLPP_USE_ROCM)
+    if(MSCCLPP_ROCM_USE_FNUZ_FP8)
+        add_compile_definitions(MSCCLPP_ROCM_FP8_FNUZ)
+        message(STATUS "ROCm native FP8 aliases: FNUZ (MSCCLPP_ROCM_USE_FNUZ_FP8=ON)")
+    else()
+        message(STATUS "ROCm native FP8 aliases: OCP (MSCCLPP_ROCM_USE_FNUZ_FP8=OFF)")
+    endif()
+endif()
+
 # Declare project
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra")
diff --git a/docker/base-dev-x.dockerfile b/docker/base-dev-x.dockerfile
index 47436202..89ea1ec7 100644
--- a/docker/base-dev-x.dockerfile
+++ b/docker/base-dev-x.dockerfile
@@ -70,7 +70,7 @@ RUN if echo "$TARGET" | grep -q "^cuda"; then \
 # Install ROCm-specific packages if building for ROCm
 RUN if echo "$TARGET" | grep -q "^rocm"; then \
         apt-get update -y && \
-        apt-get install -y hipblas hipsparse rocsparse rocrand hiprand rocthrust rocsolver rocfft hipfft hipcub rocprim rccl roctracer-dev && \
+        apt-get install -y hipblas hipsparse rocsparse rocrand hiprand rocthrust-dev rocsolver rocfft hipfft hipcub-dev rocprim-dev rccl roctracer-dev && \
         apt-get autoremove -y && \
         apt-get clean && \
         rm -rf /var/lib/apt/lists/* /tmp/*; \
@@ -84,7 +84,7 @@ ENV PATH="/root/venv/bin:${PATH}"
 # Install Python dependencies
 ADD . /tmp/mscclpp
 WORKDIR /tmp/mscclpp
-RUN target_type=$(echo $TARGET | sed 's/\.[0-9]*$//') && \
+RUN target_type=$(echo "$TARGET" | sed -E 's/^([[:alpha:]]+[0-9]+).*/\1/') && \
     if echo "$TARGET" | grep -q "^rocm"; then \
         export CUPY_INSTALL_USE_HIP=1 && export ROCM_HOME=/opt/rocm; \
     fi && \
diff --git a/docker/build.sh b/docker/build.sh
index b84eac9a..73434cb9 100755
--- a/docker/build.sh
+++ b/docker/build.sh
@@ -10,11 +10,13 @@ baseImageTable=(
     ["cuda12.9"]="nvidia/cuda:12.9.1-devel-ubuntu24.04"
     ["cuda13.0"]="nvidia/cuda:13.0.2-devel-ubuntu24.04"
     ["rocm6.2"]="rocm/dev-ubuntu-22.04:6.2.2"
+    ["rocm7.2"]="rocm/dev-ubuntu-24.04:7.2.4"
 )
 
 declare -A extraLdPathTable
 extraLdPathTable=(
     ["rocm6.2"]="/opt/rocm/lib"
+    ["rocm7.2"]="/opt/rocm/lib"
 )
 
 declare -A ofedVersionTable
@@ -25,13 +27,14 @@ ofedVersionTable=(
     ["cuda12.9"]="24.10-1.1.4.0"
     ["cuda13.0"]="24.10-3.2.5.0"
     ["rocm6.2"]="24.10-1.1.4.0"
+    ["rocm7.2"]="24.10-3.2.5.0"
 )
 
 TARGET=${1}
 OS_ARCH=$(uname -m)
 
 print_usage() {
-    echo "Usage: $0 [cuda11.8|cuda12.4|cuda12.8|cuda12.9|cuda13.0|rocm6.2]"
+    echo "Usage: $0 [cuda11.8|cuda12.4|cuda12.8|cuda12.9|cuda13.0|rocm6.2|rocm7.2]"
 }
 
 if [[ ! -v "baseImageTable[${TARGET}]" ]]; then
diff --git a/docs/quickstart.md b/docs/quickstart.md
index 320a2db7..6bda2877 100644
--- a/docs/quickstart.md
+++ b/docs/quickstart.md
@@ -47,7 +47,7 @@ We provide docker images which package all prerequisites for MSCCL++. You can se
 # For NVIDIA platforms
 $ docker run -it --privileged --net=host --ipc=host --gpus all --name mscclpp-dev ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 bash
 # For AMD platforms
-$ docker run -it --privileged --net=host --ipc=host --security-opt=seccomp=unconfined --group-add=video --name mscclpp-dev ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm6.2 bash
+$ docker run -it --privileged --net=host --ipc=host --security-opt=seccomp=unconfined --group-add=video --name mscclpp-dev ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm7.2 bash
 ```
 
 See all available images [here](https://github.com/microsoft/mscclpp/pkgs/container/mscclpp%2Fmscclpp).
@@ -106,16 +106,16 @@ Python 3.10 or later is required.
 # For NVIDIA platforms (specify your CUDA version)
 $ python -m pip install ".[cuda12]"
 # For AMD platforms
-$ CXX=/opt/rocm/bin/hipcc python -m pip install ".[rocm6]"
+$ CXX=/opt/rocm/bin/hipcc python -m pip install ".[rocm7]"
 ```
 
-> **Note:** A platform extra (`cuda11`, `cuda12`, `cuda13`, or `rocm6`) is required to install CuPy.
-> The CUDA extras install pre-built CuPy wheels and CUDA Python bindings. The `rocm6` extra installs CuPy from source
-> and HIP Python 6.x, which require ROCm and may take longer. Running `pip install .` without an extra will not install CuPy.
+> **Note:** A platform extra (`cuda11`, `cuda12`, `cuda13`, `rocm6`, or `rocm7`) is required to install CuPy.
+> The CUDA extras install pre-built CuPy wheels and CUDA Python bindings. The ROCm extras install CuPy from source
+> and HIP Python for the matching ROCm major version, which require ROCm and may take longer. Running `pip install .` without an extra will not install CuPy.
 
 Optional extras can be installed by specifying them in brackets. Available extras:
 - **`cuda11`**, **`cuda12`**, **`cuda13`**: Install a pre-built CuPy package and CUDA Python bindings for your CUDA version.
-- **`rocm6`**: Install CuPy from source and HIP Python 6.x for AMD ROCm platforms.
+- **`rocm6`**, **`rocm7`**: Install CuPy from source and HIP Python for AMD ROCm platforms.
 - **`benchmark`**: Install benchmark dependencies (mpi4py, prettytable, netifaces, matplotlib).
 - **`test`**: Install test dependencies (pytest, mpi4py, netifaces).
 
@@ -215,7 +215,7 @@ $ mpirun -np 16 -npernode 8 -hostfile hostfile ./bin/mp_unit_tests -ip_port 10.0
 
 ```bash
 # Install with benchmark dependencies and the appropriate CUDA/ROCm extras.
-# Replace `cuda12` with your platform: cuda11, cuda12, cuda13, or rocm6.
+# Replace `cuda12` with your platform: cuda11, cuda12, cuda13, rocm6, or rocm7.
 $ python3 -m pip install ".[cuda12,benchmark,test]"
 
 ```
diff --git a/include/mscclpp/gpu_data_types.hpp b/include/mscclpp/gpu_data_types.hpp
index 4a16628e..0c1e6479 100644
--- a/include/mscclpp/gpu_data_types.hpp
+++ b/include/mscclpp/gpu_data_types.hpp
@@ -16,8 +16,7 @@ using __bfloat16 = __hip_bfloat16;
 using __bfloat162 = __hip_bfloat162;
 #define __CUDA_BF16_TYPES_EXIST__
 
-// AMD FP8 support - Use fnuz types for HIP 6.0 or when HIP_FP8_TYPE_FNUZ is enabled and HIP_FP8_TYPE_OCP is not
-// enabled. Otherwise, use the standard FP8 types.
+// AMD FP8 support. The CMake MSCCLPP_ROCM_USE_FNUZ_FP8 option controls whether native FP8 aliases use FNUZ.
 #if defined(HIP_VERSION_MAJOR) && (HIP_VERSION_MAJOR >= 6)
 #include <hip/hip_fp8.h>
 
@@ -25,7 +24,7 @@ using __bfloat162 = __hip_bfloat162;
 // Define __FP8_E4M3_IS_FNUZ__ / __FP8_E5M2_IS_FNUZ__ when the platform-native FP8 is the
 // "fnuz" variant (no infinities, NaN-only at 0x80, bias differs from OCP). Dispatch layers
 // use these macros to throw on unsupported variants requested via DataType.
-#if (HIP_VERSION_MAJOR == 6) || (HIP_VERSION_MAJOR > 6 && HIP_FP8_TYPE_FNUZ && !HIP_FP8_TYPE_OCP)
+#if defined(MSCCLPP_ROCM_FP8_FNUZ)
 using __fp8_e4m3 = __hip_fp8_e4m3_fnuz;
 using __fp8_e5m2 = __hip_fp8_e5m2_fnuz;
 using __fp8x2_e4m3 = __hip_fp8x2_e4m3_fnuz;
diff --git a/pyproject.toml b/pyproject.toml
index b35b1b3a..89ee4a73 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,10 @@ rocm6 = [
     "cupy",
     "hip-python>=6,<7",
 ]
+rocm7 = [
+    "cupy",
+    "hip-python>=7,<8",
+]
 benchmark = [
     "mpi4py",
     "prettytable",
diff --git a/python/mscclpp_benchmark/bench_collective.py b/python/mscclpp_benchmark/bench_collective.py
index c526438d..0cb11739 100644
--- a/python/mscclpp_benchmark/bench_collective.py
+++ b/python/mscclpp_benchmark/bench_collective.py
@@ -18,7 +18,7 @@ from mscclpp_benchmark.correctness import (
     check_correctness as _check_correctness,
     fill_case_for_benchmark as _fill_case_for_benchmark,
 )
-from mscclpp_benchmark.gpu import capture_graph, init_runtime
+from mscclpp_benchmark.gpu import capture_graph, init_runtime, runtime_name, version
 from mscclpp_benchmark.tuner import OfflineTuner
 from mscclpp_benchmark.tuning_config import HardwareProfile, TunedConfig, TunedConfigStore, normalize_sku
 
@@ -376,7 +376,6 @@ def _measure_case(
     n_ops_per_graph: int,
 ) -> float:
     _fill_case_for_benchmark(case, comm.rank)
-    comm.comm_group.barrier()
     if comm.run(case, config) != 0:
         raise RuntimeError("algorithm returned non-zero status")
     cp.cuda.runtime.deviceSynchronize()
@@ -607,6 +606,10 @@ def main(argv: list[str] | None = None) -> None:
             )
             if comm.rank == 0:
                 print(".", end="", flush=True)
+            if runtime_name() == "hip" and version()[:2] == (7, 2):
+                # TODO: remove this after ROCm 7.2 HIP IPC export issue is fixed.
+                del case
+                comm.comm_group.barrier()
 
         if args.write_config and comm.rank == 0:
             config_store.write_path(args.write_config)
diff --git a/python/mscclpp_benchmark/correctness.py b/python/mscclpp_benchmark/correctness.py
index 0d9ab5c1..15de2f5a 100644
--- a/python/mscclpp_benchmark/correctness.py
+++ b/python/mscclpp_benchmark/correctness.py
@@ -62,7 +62,6 @@ def check_correctness(
     local_total = 0
     for iteration in range(niter):
         _fill_case_for_correctness(case, comm.rank, iteration)
-        comm.comm_group.barrier()
         ret = comm.run(case, config)
         cp.cuda.runtime.deviceSynchronize()
         comm.comm_group.barrier()
diff --git a/python/mscclpp_benchmark/gpu.py b/python/mscclpp_benchmark/gpu.py
index 1309a504..248065ef 100644
--- a/python/mscclpp_benchmark/gpu.py
+++ b/python/mscclpp_benchmark/gpu.py
@@ -18,6 +18,7 @@ _API_NAMES = {
     "graph_destroy": ("hipGraphDestroy", "cudaGraphDestroy"),
     "graph_exec_destroy": ("hipGraphExecDestroy", "cudaGraphExecDestroy"),
     "get_error_string": ("hipGetErrorString", "cudaGetErrorString"),
+    "runtime_get_version": ("hipRuntimeGetVersion", "cudaRuntimeGetVersion"),
 }
 
 
@@ -118,6 +119,21 @@ def init_runtime() -> None:
     return None
 
 
+def runtime_name() -> str:
+    return _RUNTIME.name
+
+
+def _runtime_version_raw() -> int:
+    return int(_api("runtime_get_version")()[0])
+
+
+def version() -> tuple[int, int, int]:
+    version_value = _runtime_version_raw()
+    if _RUNTIME.name == "hip":
+        return version_value // 10_000_000, (version_value // 100_000) % 100, version_value % 100_000
+    return version_value // 1000, (version_value % 1000) // 10, version_value % 10
+
+
 def capture_graph(stream: Any, capture_fn: Callable[[], None]) -> Graph:
     _api("set_device")(current_device())
     stream_ptr = _stream_ptr(stream)
diff --git a/python/mscclpp_benchmark/tuner.py b/python/mscclpp_benchmark/tuner.py
index 8df3259b..42c15ab5 100644
--- a/python/mscclpp_benchmark/tuner.py
+++ b/python/mscclpp_benchmark/tuner.py
@@ -59,7 +59,6 @@ class OfflineTuner:
                     if not self._check_correctness(self.comm, case, config):
                         self.comm.reset(config)
                         continue
-                    self.comm.reset(config)
                     time_us = self._measure(
                         self.comm,
                         case,
diff --git a/python/requirements_rocm7.txt b/python/requirements_rocm7.txt
new file mode 100644
index 00000000..982f9ae6
--- /dev/null
+++ b/python/requirements_rocm7.txt
@@ -0,0 +1,11 @@
+mpi4py
+cupy
+prettytable
+netifaces
+pytest
+numpy
+matplotlib
+sortedcontainers
+blake3
+pybind11
+hip-python>=7,<8
diff --git a/src/core/gpu_ipc_mem.cc b/src/core/gpu_ipc_mem.cc
index 0f58ed20..49dfcf82 100644
--- a/src/core/gpu_ipc_mem.cc
+++ b/src/core/gpu_ipc_mem.cc
@@ -81,8 +81,6 @@ bool isFabricMemHandleAvailable() {
 #endif  // !(CUDA_NVLS_API_AVAILABLE)
 }
 
-#if defined(MSCCLPP_DEVICE_HIP)
-
 // Custom hash and equality for cudaIpcMemHandle_t
 struct CudaIpcMemHandleHash {
   size_t operator()(const cudaIpcMemHandle_t& handle) const {
@@ -97,43 +95,59 @@ struct CudaIpcMemHandleEqual {
   }
 };
 
-static std::unordered_map<cudaIpcMemHandle_t, void*, CudaIpcMemHandleHash, CudaIpcMemHandleEqual>
+static std::unordered_map<cudaIpcMemHandle_t, std::weak_ptr<void>, CudaIpcMemHandleHash, CudaIpcMemHandleEqual>
     openCudaIpcMemHandleMap;
 
 static std::mutex openCudaIpcMemHandleMapMutex;
 
-// Cache open ipc handles to avoid opening multiple times (ROCm may exceed system limit on vm.max_map_count).
-static inline cudaError_t cudaIpcOpenMemHandleWrapper(void** addr, cudaIpcMemHandle_t ipcHandle) {
+static inline cudaError_t cudaIpcOpenMemHandleWrapper(std::shared_ptr<void>& basePtr, cudaIpcMemHandle_t ipcHandle) {
+#if defined(MSCCLPP_DEVICE_HIP)
+  // Cache open ipc handles to avoid opening multiple times (ROCm may exceed system limit on vm.max_map_count).
   std::lock_guard<std::mutex> lock(openCudaIpcMemHandleMapMutex);
   auto it = openCudaIpcMemHandleMap.find(ipcHandle);
   if (it != openCudaIpcMemHandleMap.end()) {
-    *addr = it->second;
-    return cudaSuccess;
+    basePtr = it->second.lock();
+    if (basePtr) {
+      return cudaSuccess;
+    }
+    openCudaIpcMemHandleMap.erase(it);
   }
-  cudaError_t err = cudaIpcOpenMemHandle(addr, ipcHandle, cudaIpcMemLazyEnablePeerAccess);
-  if (err == cudaSuccess) {
-    openCudaIpcMemHandleMap[ipcHandle] = *addr;
+
+  void* rawBasePtr = nullptr;
+  cudaError_t err = cudaIpcOpenMemHandle(&rawBasePtr, ipcHandle, cudaIpcMemLazyEnablePeerAccess);
+  if (err != cudaSuccess) {
+    return err;
   }
-  return err;
-}
 
-static inline cudaError_t cudaIpcCloseMemHandleWrapper(void* addr, cudaIpcMemHandle_t ipcHandle) {
-  std::lock_guard<std::mutex> lock(openCudaIpcMemHandleMapMutex);
-  openCudaIpcMemHandleMap.erase(ipcHandle);
-  return cudaIpcCloseMemHandle(addr);
-}
-
-#else  // !defined(MSCCLPP_DEVICE_HIP)
-
-static inline cudaError_t cudaIpcOpenMemHandleWrapper(void** addr, cudaIpcMemHandle_t ipcHandle) {
-  return cudaIpcOpenMemHandle(addr, ipcHandle, cudaIpcMemLazyEnablePeerAccess);
-}
-
-static inline cudaError_t cudaIpcCloseMemHandleWrapper(void* addr, [[maybe_unused]] cudaIpcMemHandle_t ipcHandle) {
-  return cudaIpcCloseMemHandle(addr);
-}
+  basePtr = std::shared_ptr<void>(rawBasePtr, [ipcHandle](void* ptr) {
+    {
+      std::lock_guard<std::mutex> lock(openCudaIpcMemHandleMapMutex);
+      openCudaIpcMemHandleMap.erase(ipcHandle);
+    }
+    cudaError_t err = cudaIpcCloseMemHandle(ptr);
+    if (err != cudaSuccess) {
+      WARN(GPU, "Failed to close CUDA IPC handle at pointer ", ptr, ": ", cudaGetErrorString(err));
+    }
+  });
+  openCudaIpcMemHandleMap[ipcHandle] = basePtr;
+  return cudaSuccess;
+#else   // !defined(MSCCLPP_DEVICE_HIP)
+  void* rawBasePtr = nullptr;
+  cudaError_t err = cudaIpcOpenMemHandle(&rawBasePtr, ipcHandle, cudaIpcMemLazyEnablePeerAccess);
+  if (err != cudaSuccess) {
+    return err;
+  }
 
+  basePtr = std::shared_ptr<void>(rawBasePtr, [](void* ptr) {
+    cudaError_t err = cudaIpcCloseMemHandle(ptr);
+    if (err != cudaSuccess) {
+      WARN(GPU, "Failed to close CUDA IPC handle at pointer ", ptr, ": ", cudaGetErrorString(err));
+      (void)cudaGetLastError();
+    }
+  });
+  return cudaSuccess;
 #endif  // !defined(MSCCLPP_DEVICE_HIP)
+}
 
 void GpuIpcMemHandle::deleter(GpuIpcMemHandle* handle) {
   if (handle) {
@@ -172,7 +186,8 @@ UniqueGpuIpcMemHandle GpuIpcMemHandle::create(const CUdeviceptr ptr) {
   if (err == cudaSuccess) {
     handle->typeFlags |= GpuIpcMemHandle::Type::RuntimeIpc;
   } else {
-    (void)cudaGetLastError();
+    WARN(GPU, "Failed to create runtime CUDA IPC handle for pointer ", (void*)basePtr,
+         ": error=", static_cast<int>(err), " (", std::string(cudaGetErrorString(err)), ")");
   }
 
 #if !defined(MSCCLPP_DEVICE_HIP)  // Remove when HIP fully supports virtual memory management APIs
@@ -345,17 +360,10 @@ std::shared_ptr<void> GpuIpcMem::map() {
 
   if (type_ == GpuIpcMemHandle::Type::RuntimeIpc) {
     // RuntimeIpc: Open handle and return shared_ptr with cleanup in deleter
-    void* basePtr = nullptr;
-    MSCCLPP_CUDATHROW(cudaIpcOpenMemHandleWrapper(&basePtr, handle_.runtimeIpc.handle));
-    void* dataPtr = static_cast<void*>(static_cast<char*>(basePtr) + handle_.offsetFromBase);
-    cudaIpcMemHandle_t ipcHandle = handle_.runtimeIpc.handle;
-    return std::shared_ptr<void>(dataPtr, [self = shared_from_this(), basePtr, ipcHandle](void*) {
-      cudaError_t err = cudaIpcCloseMemHandleWrapper(basePtr, ipcHandle);
-      if (err != cudaSuccess) {
-        WARN(GPU, "Failed to close CUDA IPC handle at pointer ", basePtr, ": ", cudaGetErrorString(err));
-        (void)cudaGetLastError();
-      }
-    });
+    std::shared_ptr<void> basePtr;
+    MSCCLPP_CUDATHROW(cudaIpcOpenMemHandleWrapper(basePtr, handle_.runtimeIpc.handle));
+    void* rawDataPtr = static_cast<void*>(static_cast<char*>(basePtr.get()) + handle_.offsetFromBase);
+    return std::shared_ptr<void>(basePtr, rawDataPtr);
   }
 
   size_t pageSize = getpagesize();
diff --git a/src/core/registered_memory.cc b/src/core/registered_memory.cc
index f464de2a..c82233ce 100644
--- a/src/core/registered_memory.cc
+++ b/src/core/registered_memory.cc
@@ -97,11 +97,14 @@ MSCCLPP_API_CPP std::vector<char> RegisteredMemory::serialize() const {
   for (auto& entry : pimpl_->transportInfos) {
     detail::serialize(result, entry.transport);
     if (entry.transport == Transport::CudaIpc) {
+      if (entry.gpuIpcMemHandle.typeFlags == GpuIpcMemHandle::Type::None) {
+        THROW(GPU, Error, ErrorCode::InternalError, "GpuIpcMemHandle type is None");
+      }
       detail::serialize(result, entry.gpuIpcMemHandle);
     } else if (AllIBTransports.has(entry.transport)) {
       detail::serialize(result, entry.ibMrInfo);
     } else {
-      throw Error("Unknown transport", ErrorCode::InternalError);
+      THROW(GPU, Error, ErrorCode::InternalError, "Unknown transport");
     }
   }
   return result;
diff --git a/src/ext/collectives/allreduce/allreduce_fullmesh.cu b/src/ext/collectives/allreduce/allreduce_fullmesh.cu
index 24d2a31c..725bdb0d 100644
--- a/src/ext/collectives/allreduce/allreduce_fullmesh.cu
+++ b/src/ext/collectives/allreduce/allreduce_fullmesh.cu
@@ -10,7 +10,7 @@ namespace mscclpp {
 namespace collective {
 
 template <ReduceOp OpType, typename T, typename AccumT = T>
-__global__ void __launch_bounds__(512, 1)
+__global__ void __launch_bounds__(1024, 1)
     allreduceFullmesh(T* buff, T* scratch, T* resultBuff, DeviceHandle<MemoryChannel>* memoryChannels,
                       DeviceHandle<MemoryChannel>* memoryOutChannels, size_t channelOutDataOffset, int rank,
                       int nRanksPerNode, int worldSize, size_t nelems) {
@@ -194,17 +194,6 @@ CommResult AllreduceFullmesh::allreduceKernelFunc(
     MSCCLPP_CUTHROW(cuMemGetAddressRange(&recvBasePtr, &recvBytes, (CUdeviceptr)output));
     channelOutOffset = (char*)output - (char*)recvBasePtr;
   }
-  std::shared_ptr<DeviceHandle<MemoryChannel>> inputChannelHandles;
-  if (this->memoryChannelsMap_.find(input) != this->memoryChannelsMap_.end()) {
-    inputChannelHandles = this->memoryChannelsMap_[input].second;
-  } else {
-    RegisteredMemory localMemory = comm_->registerMemory(const_cast<void*>(input), inputSize, Transport::CudaIpc);
-    std::vector<MemoryChannel> channels =
-        setupMemoryChannels(this->conns_, this->inputScratchSemaphores_, this->remoteScratchMemories_, localMemory,
-                            nChannelsPerConnection_);
-    this->memoryChannelsMap_[input] = std::make_pair(channels, setupMemoryChannelDeviceHandles(channels));
-  }
-  inputChannelHandles = this->memoryChannelsMap_[input].second;
 
   AllreduceFunc allreduce = dispatch<AllreduceAllconnectAdapter>(op, dtype, accumDtype);
   if (!allreduce) {
@@ -220,10 +209,12 @@ CommResult AllreduceFullmesh::allreduceKernelFunc(
   if (numBlocksAndThreads.first == 0 || numBlocksAndThreads.second == 0) {
     numBlocksAndThreads = {35, 512};
   }
+  auto inputChannelDeviceHandles = ctx->memoryChannelDeviceHandles.get();
+  auto outputChannelDeviceHandles = inputChannelDeviceHandles + ctx->memoryChannels.size() / 2;
   cudaError_t error =
-      allreduce(input, this->scratchBuffer_, output, inputChannelHandles.get(), ctx->memoryChannelDeviceHandles.get(),
-                nullptr, nullptr, 0, channelOutOffset, 0, ctx->rank, ctx->nRanksPerNode, ctx->workSize, inputSize,
-                stream, nullptr, 0, 0, numBlocksAndThreads.first, numBlocksAndThreads.second);
+      allreduce(input, this->scratchBuffer_, output, inputChannelDeviceHandles, outputChannelDeviceHandles, nullptr,
+                nullptr, 0, channelOutOffset, 0, ctx->rank, ctx->nRanksPerNode, ctx->workSize, inputSize, stream,
+                nullptr, 0, 0, numBlocksAndThreads.first, numBlocksAndThreads.second);
   if (error != cudaSuccess) {
     WARN("AllreduceAllconnect failed with error: %s", cudaGetErrorString(error));
     return CommResult::CommUnhandledCudaError;
@@ -231,20 +222,21 @@ CommResult AllreduceFullmesh::allreduceKernelFunc(
   return CommResult::CommSuccess;
 }
 
-AlgorithmCtxKey AllreduceFullmesh::generateAllreduceContextKey(const void*, void* output, size_t, DataType,
+AlgorithmCtxKey AllreduceFullmesh::generateAllreduceContextKey(const void* input, void* output, size_t size, DataType,
                                                                bool symmetricMemory) {
-  static int tag = 0;
-  size_t recvBytes;
-  CUdeviceptr recvBasePtr;
-  MSCCLPP_CUTHROW(cuMemGetAddressRange(&recvBasePtr, &recvBytes, (CUdeviceptr)output));
   symmetricMemory_ = symmetricMemory;
   if (!symmetricMemory_) {
-    return AlgorithmCtxKey{nullptr, (void*)recvBasePtr, 0, recvBytes, tag++};
+    return AlgorithmCtxKey{const_cast<void*>(input), output, size, size, 0};
   }
-  return AlgorithmCtxKey{nullptr, (void*)recvBasePtr, 0, recvBytes, 0};
+
+  size_t sendBytes, recvBytes;
+  CUdeviceptr sendBasePtr, recvBasePtr;
+  MSCCLPP_CUTHROW(cuMemGetAddressRange(&sendBasePtr, &sendBytes, (CUdeviceptr)input));
+  MSCCLPP_CUTHROW(cuMemGetAddressRange(&recvBasePtr, &recvBytes, (CUdeviceptr)output));
+  return AlgorithmCtxKey{(void*)sendBasePtr, (void*)recvBasePtr, sendBytes, recvBytes, 0};
 }
 
-std::shared_ptr<void> AllreduceFullmesh::initAllreduceContext(std::shared_ptr<Communicator> comm, const void*,
+std::shared_ptr<void> AllreduceFullmesh::initAllreduceContext(std::shared_ptr<Communicator> comm, const void* input,
                                                               void* output, size_t size, DataType) {
   auto ctx = std::make_shared<AlgorithmCtx>();
   ctx->rank = comm->bootstrap()->getRank();
@@ -263,8 +255,19 @@ std::shared_ptr<void> AllreduceFullmesh::initAllreduceContext(std::shared_ptr<Co
   }
   RegisteredMemory localMemory = comm->registerMemory((void*)recvBasePtr, recvBytes, Transport::CudaIpc);
   ctx->registeredMemories = setupRemoteMemories(comm, ctx->rank, localMemory);
-  ctx->memoryChannels = setupMemoryChannels(this->conns_, ctx->memorySemaphores, ctx->registeredMemories, localMemory,
-                                            nChannelsPerConnection_);
+
+  RegisteredMemory inputMemory = comm->registerMemory(const_cast<void*>(input), size, TransportFlags());
+  std::vector<MemoryChannel> inputMemoryChannels = setupMemoryChannels(
+      this->conns_, this->inputScratchSemaphores_, this->remoteScratchMemories_, inputMemory, nChannelsPerConnection_);
+  std::vector<MemoryChannel> outputMemoryChannels = setupMemoryChannels(
+      this->conns_, this->outputSemaphores_, ctx->registeredMemories, localMemory, nChannelsPerConnection_);
+  ctx->memoryChannels.reserve(inputMemoryChannels.size() + outputMemoryChannels.size());
+  for (auto& channel : inputMemoryChannels) {
+    ctx->memoryChannels.emplace_back(std::move(channel));
+  }
+  for (auto& channel : outputMemoryChannels) {
+    ctx->memoryChannels.emplace_back(std::move(channel));
+  }
   ctx->memoryChannelDeviceHandles = setupMemoryChannelDeviceHandles(ctx->memoryChannels);
   return ctx;
 }
diff --git a/src/ext/collectives/include/allreduce/allreduce_fullmesh.hpp b/src/ext/collectives/include/allreduce/allreduce_fullmesh.hpp
index a54352b3..e0c63a3d 100644
--- a/src/ext/collectives/include/allreduce/allreduce_fullmesh.hpp
+++ b/src/ext/collectives/include/allreduce/allreduce_fullmesh.hpp
@@ -30,8 +30,6 @@ class AllreduceFullmesh : public mscclpp::AlgorithmBuilder {
   std::vector<std::shared_ptr<MemoryDevice2DeviceSemaphore>> inputScratchSemaphores_;
   std::vector<RegisteredMemory> remoteScratchMemories_;
   RegisteredMemory localScratchMemory_;
-  std::unordered_map<const void*, std::pair<std::vector<MemoryChannel>, std::shared_ptr<DeviceHandle<MemoryChannel>>>>
-      memoryChannelsMap_;
   bool symmetricMemory_ = false;
 };
 }  // namespace collective
diff --git a/test/deploy/setup.sh b/test/deploy/setup.sh
index 2a88a310..fbcdc119 100644
--- a/test/deploy/setup.sh
+++ b/test/deploy/setup.sh
@@ -68,7 +68,9 @@ elif [[ "${CUDA_VERSION}" == *"12."* ]]; then
 elif [[ "${CUDA_VERSION}" == *"13."* ]]; then
     pip3 install ".[cuda13,benchmark,test]"
 elif [ "${PLATFORM}" == "rocm" ]; then
-    pip3 install ".[rocm6,benchmark,test]"
+    ROCM_VERSION=$(cat /opt/rocm/.info/version)
+    ROCM_MAJOR="${ROCM_VERSION%%.*}"
+    pip3 install ".[rocm${ROCM_MAJOR},benchmark,test]"
 else
     pip3 install ".[benchmark,test]"
 fi

From 56dc9cba6348f7b2d897b39ac250fa3dc4617d5e Mon Sep 17 00:00:00 2001
From: Changho Hwang <changhohwang@microsoft.com>
Date: Thu, 25 Jun 2026 19:21:51 +0800
Subject: [PATCH 11/16] Update port channel perf tests (#820)

- Make MultiQp bandwidth/flush-stress variants separate perf tests
---
 test/mp_unit/port_channel_tests.cu | 107 ++++++++++++++++++++---------
 1 file changed, 76 insertions(+), 31 deletions(-)

diff --git a/test/mp_unit/port_channel_tests.cu b/test/mp_unit/port_channel_tests.cu
index 47034cdb..eec1760c 100644
--- a/test/mp_unit/port_channel_tests.cu
+++ b/test/mp_unit/port_channel_tests.cu
@@ -629,22 +629,19 @@ PERF_TEST(PortChannelOneToOneTest, BandwidthIbHostNoAtomicMode) {
 static constexpr int kMaxQps = 4;
 __constant__ DeviceHandle<mscclpp::PortChannel> gMultiQpPortChans[kMaxQps];
 
-// Multi-QP bandwidth kernel: barrier on QP 0 only, then putWithSignal on all QPs.
-// Only one signal/wait pair is needed for sync between two GPU kernels.
+// Multi-QP bandwidth kernel: one thread per QP, putWithSignal per QP, parallel waits.
 __global__ void kernelMultiQpBandwidth(int nElemPerChan, int nIters, int numQps) {
-  if (threadIdx.x != 0) return;
+  int q = threadIdx.x;
+  if (q >= numQps) return;
   for (int i = 0; i < nIters; i++) {
-    // Barrier on QP 0 only — syncs both ranks
-    gMultiQpPortChans[0].signal();
-    gMultiQpPortChans[0].wait();
-    // Data transfer: put on all QPs simultaneously
-    for (int q = 0; q < numQps; q++) {
-      gMultiQpPortChans[q].putWithSignal(0, nElemPerChan * sizeof(int));
-    }
-    // Wait for all remote data arrivals
-    for (int q = 0; q < numQps; q++) {
-      gMultiQpPortChans[q].wait();
+    if (q == 0) {
+      gMultiQpPortChans[0].signal();
+      gMultiQpPortChans[0].wait();
     }
+    __syncthreads();
+    gMultiQpPortChans[q].putWithSignal(0, nElemPerChan * sizeof(int));
+    gMultiQpPortChans[q].wait();
+    __syncthreads();
   }
 }
 
@@ -715,15 +712,15 @@ void PortChannelOneToOneTest::testMultiQpBandwidth(IbMode ibMode, int numQps) {
 
   for (int nElemPerChan :
        {256, 16 * 1024, 256 * 1024, 1024 * 1024, 4 * 1024 * 1024, 16 * 1024 * 1024, 32 * 1024 * 1024}) {
-    int nIters = 10000;
+    int nIters = 200;
     // Warm-up
-    kernelMultiQpBandwidth<<<1, 1>>>(nElemPerChan, 10, numQps);
+    kernelMultiQpBandwidth<<<1, numQps>>>(nElemPerChan, 10, numQps);
     MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
     communicator->bootstrap()->barrier();
 
     // Measure
     mscclpp::Timer timer;
-    kernelMultiQpBandwidth<<<1, 1>>>(nElemPerChan, nIters, numQps);
+    kernelMultiQpBandwidth<<<1, numQps>>>(nElemPerChan, nIters, numQps);
     MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
     double elapsedUs = timer.elapsed();
     communicator->bootstrap()->barrier();
@@ -748,20 +745,46 @@ void PortChannelOneToOneTest::testMultiQpBandwidth(IbMode ibMode, int numQps) {
   for (auto& m : remoteMems) registeredMemories.push_back(m);
 }
 
+PERF_TEST(PortChannelOneToOneTest, SingleQpBandwidthIbHostMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::Host);
+  testMultiQpBandwidth(IbMode::Host, /*numQps=*/1);
+}
+
+PERF_TEST(PortChannelOneToOneTest, TwoQpBandwidthIbHostMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::Host);
+  testMultiQpBandwidth(IbMode::Host, /*numQps=*/2);
+}
+
 PERF_TEST(PortChannelOneToOneTest, MultiQpBandwidthIbHostMode) {
   REQUIRE_IBVERBS;
   REQUIRE_GDR_FOR_IB_MODE(IbMode::Host);
-  for (int numQps : {1, 2, 4}) {
-    testMultiQpBandwidth(IbMode::Host, numQps);
-  }
+  testMultiQpBandwidth(IbMode::Host, /*numQps=*/4);
+}
+
+PERF_TEST(PortChannelOneToOneTest, SingleQpBandwidthIbHostNoAtomicMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic);
+  testMultiQpBandwidth(IbMode::HostNoAtomic, /*numQps=*/1);
+}
+
+PERF_TEST(PortChannelOneToOneTest, TwoQpBandwidthIbHostNoAtomicMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic);
+  testMultiQpBandwidth(IbMode::HostNoAtomic, /*numQps=*/2);
+}
+
+PERF_TEST(PortChannelOneToOneTest, ThreeQpBandwidthIbHostNoAtomicMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic);
+  testMultiQpBandwidth(IbMode::HostNoAtomic, /*numQps=*/3);
 }
 
 PERF_TEST(PortChannelOneToOneTest, MultiQpBandwidthIbHostNoAtomicMode) {
   REQUIRE_IBVERBS;
   REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic);
-  for (int numQps : {1, 2, 4}) {
-    testMultiQpBandwidth(IbMode::HostNoAtomic, numQps);
-  }
+  testMultiQpBandwidth(IbMode::HostNoAtomic, /*numQps=*/4);
 }
 
 // Multi-QP flush-stress kernel: one thread per QP, all calling putWithSignalAndFlush
@@ -786,7 +809,7 @@ void PortChannelOneToOneTest::testMultiQpFlushStress(IbMode ibMode, int numQps)
   if (gEnv->rank >= numRanksToUse) return;
 
   const int rank = communicator->bootstrap()->getRank();
-  const int maxElemPerChan = 64 * 1024;
+  const int maxElemPerChan = 8 * 1024 * 1024;
 
   std::vector<std::shared_ptr<int>> sendBuffs;
   std::vector<mscclpp::RegisteredMemory> localMems;
@@ -805,8 +828,8 @@ void PortChannelOneToOneTest::testMultiQpFlushStress(IbMode ibMode, int numQps)
 
   const std::string qpLabel = std::to_string(numQps) + " QP" + (numQps > 1 ? "s" : "");
 
-  for (int nElemPerChan : {256, 4 * 1024, 64 * 1024}) {
-    int nIters = 2000;
+  for (int nElemPerChan : {256, 4 * 1024, 64 * 1024, 256 * 1024, 1024 * 1024, 4 * 1024 * 1024, 8 * 1024 * 1024}) {
+    int nIters = (nElemPerChan >= 256 * 1024) ? 200 : 2000;
     kernelMultiQpFlushStress<<<1, numQps>>>(nElemPerChan, 10, numQps);
     MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
     communicator->bootstrap()->barrier();
@@ -823,8 +846,10 @@ void PortChannelOneToOneTest::testMultiQpFlushStress(IbMode ibMode, int numQps)
       int bytesPerChan = nElemPerChan * (int)sizeof(int);
       std::string sizeLabel = (bytesPerChan >= 1024) ? (std::to_string(bytesPerChan / 1024) + " KB")
                                                      : (std::to_string(bytesPerChan) + " B");
+      double aggGbps = ((double)bytesPerChan * numQps) / usPerIter * 1e-3;  // bytes/us = MB/s × 1e-3 = GB/s
       ::mscclpp::test::reportPerfResult(sizeLabel + " (" + qpLabel + ") per-iter", usPerIter, "us");
       ::mscclpp::test::reportPerfResult(sizeLabel + " (" + qpLabel + ") per-iter/QP", usPerIterPerQp, "us");
+      ::mscclpp::test::reportPerfResult(sizeLabel + " (" + qpLabel + ") aggregate", aggGbps, "GB/s");
     }
   }
 
@@ -834,20 +859,40 @@ void PortChannelOneToOneTest::testMultiQpFlushStress(IbMode ibMode, int numQps)
   for (auto& m : remoteMems) registeredMemories.push_back(m);
 }
 
+PERF_TEST(PortChannelOneToOneTest, SingleQpFlushStressIbHostMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::Host);
+  testMultiQpFlushStress(IbMode::Host, /*numQps=*/1);
+}
+
+PERF_TEST(PortChannelOneToOneTest, TwoQpFlushStressIbHostMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::Host);
+  testMultiQpFlushStress(IbMode::Host, /*numQps=*/2);
+}
+
 PERF_TEST(PortChannelOneToOneTest, MultiQpFlushStressIbHostMode) {
   REQUIRE_IBVERBS;
   REQUIRE_GDR_FOR_IB_MODE(IbMode::Host);
-  for (int numQps : {1, 2, 4}) {
-    testMultiQpFlushStress(IbMode::Host, numQps);
-  }
+  testMultiQpFlushStress(IbMode::Host, /*numQps=*/4);
+}
+
+PERF_TEST(PortChannelOneToOneTest, SingleQpFlushStressIbHostNoAtomicMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic);
+  testMultiQpFlushStress(IbMode::HostNoAtomic, /*numQps=*/1);
+}
+
+PERF_TEST(PortChannelOneToOneTest, TwoQpFlushStressIbHostNoAtomicMode) {
+  REQUIRE_IBVERBS;
+  REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic);
+  testMultiQpFlushStress(IbMode::HostNoAtomic, /*numQps=*/2);
 }
 
 PERF_TEST(PortChannelOneToOneTest, MultiQpFlushStressIbHostNoAtomicMode) {
   REQUIRE_IBVERBS;
   REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic);
-  for (int numQps : {1, 2, 4}) {
-    testMultiQpFlushStress(IbMode::HostNoAtomic, numQps);
-  }
+  testMultiQpFlushStress(IbMode::HostNoAtomic, /*numQps=*/4);
 }
 
 // Same-channel concurrent-flush kernel: N GPU threads on the same PortChannel each call

From 8c3730b495cc889bc8740b182f54b3abb38756c6 Mon Sep 17 00:00:00 2001
From: RJ Souza <rjsouza1995@gmail.com>
Date: Thu, 25 Jun 2026 12:59:58 -0700
Subject: [PATCH 12/16] Migrate to C++20 and drop CUDA 11 support (#822)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR bumps the project's C++ language standard from C++17 to C++20
(host, CUDA, and HIP) and removes support for CUDA 11, which cannot
compile C++20. CI matrices are updated to drop CUDA 11.8 and add CUDA
13.0 alongside 12.9.

 - C++20 enables newer language features across the codebase.
 - C++20 requires CUDA ≥ 12.0, so CUDA 11.8 can no longer be supported.
- Blackwell / sm_100 targets require CUDA ≥ 12.8, so CUDA 13.x is added
to CI coverage.
---
 .azure-pipelines/codecov.yml                |  4 ++++
 .azure-pipelines/integration-test.yml       |  6 ++++--
 .azure-pipelines/multi-nodes-test.yml       |  2 ++
 .azure-pipelines/nccl-api-test.yml          |  4 ++++
 .azure-pipelines/sglang-multi-node-test.yml |  2 ++
 .azure-pipelines/sglang-test.yml            |  2 ++
 .azure-pipelines/ut.yml                     | 16 ++++++++++++----
 .devcontainer/devcontainer.json             |  2 +-
 .devcontainer/devcontainer_amd.json         |  2 +-
 .github/workflows/codeql-analysis.yml       |  2 +-
 .github/workflows/mscclpp-lang.yml          |  2 +-
 CMakeLists.txt                              | 19 ++++++++++++++-----
 docker/build.sh                             |  4 +---
 docs/quickstart.md                          | 13 +++++++++----
 pyproject.toml                              |  4 ----
 python/mscclpp/_core/compiler.py            |  2 +-
 python/mscclpp/utils.py                     |  2 +-
 python/requirements_cuda11.txt              | 11 -----------
 src/core/core.cc                            |  2 +-
 src/core/executor/execution_plan.cc         |  4 ++--
 test/deploy/setup.sh                        |  4 +---
 21 files changed, 64 insertions(+), 45 deletions(-)
 delete mode 100644 python/requirements_cuda11.txt

diff --git a/.azure-pipelines/codecov.yml b/.azure-pipelines/codecov.yml
index 5f3e531e..a806445e 100644
--- a/.azure-pipelines/codecov.yml
+++ b/.azure-pipelines/codecov.yml
@@ -38,6 +38,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
@@ -59,6 +61,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
diff --git a/.azure-pipelines/integration-test.yml b/.azure-pipelines/integration-test.yml
index 45bb1e96..859e2a68 100644
--- a/.azure-pipelines/integration-test.yml
+++ b/.azure-pipelines/integration-test.yml
@@ -30,10 +30,10 @@ jobs:
   displayName: Integration test A100
   strategy:
     matrix:
-      cuda11:
-        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda11.8
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   pool:
     name: msccl-ci
@@ -53,6 +53,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   pool:
     name: msccl-ci-h100
diff --git a/.azure-pipelines/multi-nodes-test.yml b/.azure-pipelines/multi-nodes-test.yml
index ee2766fd..c4e27be4 100644
--- a/.azure-pipelines/multi-nodes-test.yml
+++ b/.azure-pipelines/multi-nodes-test.yml
@@ -31,6 +31,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
   pool:
     name: mscclpp-multi-node
   container:
diff --git a/.azure-pipelines/nccl-api-test.yml b/.azure-pipelines/nccl-api-test.yml
index 85b466ef..a9fe6b25 100644
--- a/.azure-pipelines/nccl-api-test.yml
+++ b/.azure-pipelines/nccl-api-test.yml
@@ -35,6 +35,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
@@ -56,6 +58,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
diff --git a/.azure-pipelines/sglang-multi-node-test.yml b/.azure-pipelines/sglang-multi-node-test.yml
index bf640db2..01fcbb4f 100644
--- a/.azure-pipelines/sglang-multi-node-test.yml
+++ b/.azure-pipelines/sglang-multi-node-test.yml
@@ -67,6 +67,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
   pool:
     name: mscclpp-multi-node
   container:
diff --git a/.azure-pipelines/sglang-test.yml b/.azure-pipelines/sglang-test.yml
index 70e30d35..0b49a0ce 100644
--- a/.azure-pipelines/sglang-test.yml
+++ b/.azure-pipelines/sglang-test.yml
@@ -48,6 +48,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
   pool:
     name: msccl-ci-h100
   container:
diff --git a/.azure-pipelines/ut.yml b/.azure-pipelines/ut.yml
index 0cabdaeb..2a65f361 100644
--- a/.azure-pipelines/ut.yml
+++ b/.azure-pipelines/ut.yml
@@ -34,10 +34,10 @@ jobs:
     name: msccl-ci
   strategy:
     matrix:
-      cuda11:
-        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda11.8
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
@@ -55,10 +55,10 @@ jobs:
     name: msccl-ci
   strategy:
     matrix:
-      cuda11:
-        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda11.8
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
@@ -78,6 +78,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
@@ -97,6 +99,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
@@ -118,6 +122,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
@@ -161,6 +167,8 @@ jobs:
     matrix:
       cuda12:
         containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9
+      cuda13:
+        containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0
 
   container:
     image: $(containerImage)
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
index 6fc7dd32..d0f27675 100644
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@@ -32,7 +32,7 @@
                     "/usr/include"
                 ],
                 "C_Cpp.default.cStandard": "c17",
-                "C_Cpp.default.cppStandard": "c++17"
+                "C_Cpp.default.cppStandard": "c++20"
             }
         }
     },
diff --git a/.devcontainer/devcontainer_amd.json b/.devcontainer/devcontainer_amd.json
index 80d47956..7f1e2c2b 100644
--- a/.devcontainer/devcontainer_amd.json
+++ b/.devcontainer/devcontainer_amd.json
@@ -32,7 +32,7 @@
                     "/usr/include"
                 ],
                 "C_Cpp.default.cStandard": "c17",
-                "C_Cpp.default.cppStandard": "c++17"
+                "C_Cpp.default.cppStandard": "c++20"
             }
         }
     },
diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml
index 982d8633..b9d3a5c1 100644
--- a/.github/workflows/codeql-analysis.yml
+++ b/.github/workflows/codeql-analysis.yml
@@ -40,7 +40,7 @@ jobs:
       fail-fast: false
       matrix:
         language: [ 'cpp', 'python' ]
-        version: [ 'cuda11.8', 'cuda12.9' ]
+        version: [ 'cuda12.9', 'cuda13.0' ]
 
     steps:
     - name: Checkout repository
diff --git a/.github/workflows/mscclpp-lang.yml b/.github/workflows/mscclpp-lang.yml
index a9187e96..6f91eb86 100644
--- a/.github/workflows/mscclpp-lang.yml
+++ b/.github/workflows/mscclpp-lang.yml
@@ -15,7 +15,7 @@ jobs:
     strategy:
         fail-fast: false
         matrix:
-          version: [ 'cuda11.8', 'cuda12.9' ]
+          version: [ 'cuda12.9', 'cuda13.0' ]
 
     steps:
     - uses: actions/checkout@v4
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 92822012..7b8316fb 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -169,8 +169,8 @@ elseif(MSCCLPP_USE_CUDA)
     if(NVIDIA_FOUND)
         set(MSCCLPP_GPU_ARCHS "native")
     else()
-	    if(CUDAToolkit_VERSION VERSION_LESS "11.8")
-	        message(FATAL_ERROR "CUDA 11.8 or higher required, found ${CUDAToolkit_VERSION}")
+	    if(CUDAToolkit_VERSION VERSION_LESS "12.0")
+	        message(FATAL_ERROR "CUDA 12.0 or higher required (C++20 build), found ${CUDAToolkit_VERSION}")
 	    endif()
 	    set(MSCCLPP_GPU_ARCHS 80)
 	    if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.0")
@@ -199,10 +199,11 @@ if(MSCCLPP_USE_ROCM)
 endif()
 
 # Declare project
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra")
 if(MSCCLPP_USE_CUDA)
-    set(CMAKE_CUDA_STANDARD 17)
+    set(CMAKE_CUDA_STANDARD 20)
+    set(CMAKE_CUDA_EXTENSIONS OFF)
     set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wall,-Wextra")
     enable_language(CUDA)
 
@@ -217,7 +218,7 @@ if(MSCCLPP_USE_CUDA)
         set(GPU_LIBRARIES CUDA::cudart CUDA::cuda_driver)
     endif()
 else()
-    set(CMAKE_HIP_STANDARD 17)
+    set(CMAKE_HIP_STANDARD 20)
     set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -Wall -Wextra")
 
     set(CMAKE_HIP_ARCHITECTURES ${MSCCLPP_GPU_ARCHS})
@@ -226,6 +227,14 @@ else()
     set(GPU_INCLUDE_DIRS ${hip_INCLUDE_DIRS})
 endif()
 
+message(STATUS "C++ standard: C++${CMAKE_CXX_STANDARD}")
+if(MSCCLPP_USE_CUDA)
+    message(STATUS "CUDA toolkit version: ${CUDAToolkit_VERSION}")
+    message(STATUS "CUDA language standard: C++${CMAKE_CUDA_STANDARD}")
+else()
+    message(STATUS "HIP language standard: C++${CMAKE_HIP_STANDARD}")
+endif()
+
 if(CMAKE_BUILD_TYPE STREQUAL "Debug")
     add_compile_definitions(DEBUG_BUILD)
 endif()
diff --git a/docker/build.sh b/docker/build.sh
index 73434cb9..a65d0867 100755
--- a/docker/build.sh
+++ b/docker/build.sh
@@ -4,7 +4,6 @@ set -e
 
 declare -A baseImageTable
 baseImageTable=(
-    ["cuda11.8"]="nvidia/cuda:11.8.0-devel-ubuntu22.04"
     ["cuda12.4"]="nvidia/cuda:12.4.1-devel-ubuntu22.04"
     ["cuda12.8"]="nvidia/cuda:12.8.1-devel-ubuntu22.04"
     ["cuda12.9"]="nvidia/cuda:12.9.1-devel-ubuntu24.04"
@@ -21,7 +20,6 @@ extraLdPathTable=(
 
 declare -A ofedVersionTable
 ofedVersionTable=(
-    ["cuda11.8"]="23.07-0.5.1.2"
     ["cuda12.4"]="23.07-0.5.1.2"
     ["cuda12.8"]="24.10-1.1.4.0"
     ["cuda12.9"]="24.10-1.1.4.0"
@@ -34,7 +32,7 @@ TARGET=${1}
 OS_ARCH=$(uname -m)
 
 print_usage() {
-    echo "Usage: $0 [cuda11.8|cuda12.4|cuda12.8|cuda12.9|cuda13.0|rocm6.2|rocm7.2]"
+    echo "Usage: $0 [cuda12.4|cuda12.8|cuda12.9|cuda13.0|rocm6.2|rocm7.2]"
 }
 
 if [[ ! -v "baseImageTable[${TARGET}]" ]]; then
diff --git a/docs/quickstart.md b/docs/quickstart.md
index 6bda2877..42f6a19c 100644
--- a/docs/quickstart.md
+++ b/docs/quickstart.md
@@ -14,10 +14,15 @@
             * [NDm_A100_v4](https://learn.microsoft.com/en-us/azure/virtual-machines/ndm-a100-v4-series)
             * [ND_H100_v5](https://learn.microsoft.com/en-us/azure/virtual-machines/nd-h100-v5-series)
         * Non-Azure Systems
-            * NVIDIA A100 GPUs + CUDA >= 11.8
+            * NVIDIA A100 GPUs + CUDA >= 12.0
             * NVIDIA H100 GPUs + CUDA >= 12.0
             * AMD MI250X GPUs + ROCm >= 5.7
             * AMD MI300X GPUs + ROCm >= 6.0
+* Toolchain
+    * MSCCL++ is built as **C++20** (both host and device code), so a C++20-capable toolchain is required.
+    * [CMake](https://cmake.org/) >= 3.25
+    * A C++20-capable host compiler, e.g., GCC >= 11 or Clang >= 14
+    * On NVIDIA platforms, **CUDA Toolkit >= 12.0** is required. `nvcc` first added `-std=c++20` support in CUDA 12.0, so earlier toolkits (11.x and below) cannot build the project.
 * OS
     * Tested on Ubuntu 20.04 and later
 * Libraries
@@ -109,12 +114,12 @@ $ python -m pip install ".[cuda12]"
 $ CXX=/opt/rocm/bin/hipcc python -m pip install ".[rocm7]"
 ```
 
-> **Note:** A platform extra (`cuda11`, `cuda12`, `cuda13`, `rocm6`, or `rocm7`) is required to install CuPy.
+> **Note:** A platform extra (`cuda12`, `cuda13`, `rocm6`, or `rocm7`) is required to install CuPy.
 > The CUDA extras install pre-built CuPy wheels and CUDA Python bindings. The ROCm extras install CuPy from source
 > and HIP Python for the matching ROCm major version, which require ROCm and may take longer. Running `pip install .` without an extra will not install CuPy.
 
 Optional extras can be installed by specifying them in brackets. Available extras:
-- **`cuda11`**, **`cuda12`**, **`cuda13`**: Install a pre-built CuPy package and CUDA Python bindings for your CUDA version.
+- **`cuda12`**, **`cuda13`**: Install a pre-built CuPy package and CUDA Python bindings for your CUDA version.
 - **`rocm6`**, **`rocm7`**: Install CuPy from source and HIP Python for AMD ROCm platforms.
 - **`benchmark`**: Install benchmark dependencies (mpi4py, prettytable, netifaces, matplotlib).
 - **`test`**: Install test dependencies (pytest, mpi4py, netifaces).
@@ -215,7 +220,7 @@ $ mpirun -np 16 -npernode 8 -hostfile hostfile ./bin/mp_unit_tests -ip_port 10.0
 
 ```bash
 # Install with benchmark dependencies and the appropriate CUDA/ROCm extras.
-# Replace `cuda12` with your platform: cuda11, cuda12, cuda13, rocm6, or rocm7.
+# Replace `cuda12` with your platform: cuda12, cuda13, rocm6, or rocm7.
 $ python3 -m pip install ".[cuda12,benchmark,test]"
 
 ```
diff --git a/pyproject.toml b/pyproject.toml
index 89ee4a73..5d83d7da 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,10 +21,6 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-cuda11 = [
-    "cupy-cuda11x",
-    "cuda-bindings>=11.8,<12",
-]
 cuda12 = [
     "cupy-cuda12x",
     "cuda-bindings>=12,<13",
diff --git a/python/mscclpp/_core/compiler.py b/python/mscclpp/_core/compiler.py
index 3b77ce8e..341f3c84 100644
--- a/python/mscclpp/_core/compiler.py
+++ b/python/mscclpp/_core/compiler.py
@@ -198,7 +198,7 @@ class NativeCodeCompiler:
         self._is_hip = cp.cuda.runtime.is_hip
         self._device_arch = get_device_arch()
         self._compiler = self._get_compiler()
-        self._default_options = ["-std=c++17", "-O3", "--shared"]
+        self._default_options = ["-std=c++20", "-O3", "--shared"]
         python_include = sysconfig.get_path("include")
         pybind11_include = pybind11.get_include()
         self._default_options += [f"-I{python_include}", f"-I{pybind11_include}"]
diff --git a/python/mscclpp/utils.py b/python/mscclpp/utils.py
index 0f0a28d4..99547402 100644
--- a/python/mscclpp/utils.py
+++ b/python/mscclpp/utils.py
@@ -99,7 +99,7 @@ class KernelBuilder:
         self._kernel = Kernel(cubin, kernel_name)
         self.kernel_map[kernel_key] = self._kernel
 
-    def _compile_cuda(self, source_file, output_file, std_version="c++17"):
+    def _compile_cuda(self, source_file, output_file, std_version="c++20"):
         mscclpp_home = os.environ.get("MSCCLPP_HOME", "/usr/local/mscclpp")
         include_dir = os.path.join(mscclpp_home, "include")
         if not cp.cuda.runtime.is_hip:
diff --git a/python/requirements_cuda11.txt b/python/requirements_cuda11.txt
deleted file mode 100644
index 1f575f67..00000000
--- a/python/requirements_cuda11.txt
+++ /dev/null
@@ -1,11 +0,0 @@
-mpi4py
-cupy-cuda11x
-cuda-bindings>=11.8,<12
-prettytable
-netifaces
-pytest
-numpy
-matplotlib
-sortedcontainers
-blake3
-pybind11
\ No newline at end of file
diff --git a/src/core/core.cc b/src/core/core.cc
index 2d67b988..ee630529 100644
--- a/src/core/core.cc
+++ b/src/core/core.cc
@@ -75,7 +75,7 @@ MSCCLPP_API_CPP bool TransportFlags::operator==(TransportFlags other) const {
 }
 
 MSCCLPP_API_CPP bool TransportFlags::operator!=(TransportFlags other) const {
-  return detail::TransportFlagsBase::operator!=(other);
+  return !detail::TransportFlagsBase::operator==(other);
 }
 
 MSCCLPP_API_CPP detail::TransportFlagsBase TransportFlags::toBitset() const { return *this; }
diff --git a/src/core/executor/execution_plan.cc b/src/core/executor/execution_plan.cc
index 98ec3ab6..4c92c26b 100644
--- a/src/core/executor/execution_plan.cc
+++ b/src/core/executor/execution_plan.cc
@@ -232,7 +232,7 @@ void ExecutionPlan::Impl::loadExecutionPlan(size_t inputSize, size_t outputSize,
                                             size_t constDstOffset) {
   std::ifstream file(this->planPath);
   json obj = json::parse(file);
-  if (this->name != obj["name"]) {
+  if (this->name != obj["name"].get<std::string>()) {
     throw Error("Plan name does not match", ErrorCode::ExecutorError);
   }
   this->collective = obj["collective"];
@@ -268,7 +268,7 @@ void ExecutionPlan::Impl::lightLoadExecutionPlan(size_t inputSize, size_t output
                                                  size_t constDstOffset) {
   std::ifstream file(this->planPath);
   json obj = json::parse(file);
-  if (this->name != obj["name"]) {
+  if (this->name != obj["name"].get<std::string>()) {
     throw Error("Plan name does not match", ErrorCode::ExecutorError);
   }
   std::string protocol = obj["protocol"];
diff --git a/test/deploy/setup.sh b/test/deploy/setup.sh
index fbcdc119..0cb6f4b8 100644
--- a/test/deploy/setup.sh
+++ b/test/deploy/setup.sh
@@ -61,9 +61,7 @@ if [ -f "${PIP_CMAKE_ARGS_FILE}" ]; then
 fi
 
 cd /root/mscclpp
-if [[ "${CUDA_VERSION}" == *"11."* ]]; then
-    pip3 install ".[cuda11,benchmark,test]"
-elif [[ "${CUDA_VERSION}" == *"12."* ]]; then
+if [[ "${CUDA_VERSION}" == *"12."* ]]; then
     pip3 install ".[cuda12,benchmark,test]"
 elif [[ "${CUDA_VERSION}" == *"13."* ]]; then
     pip3 install ".[cuda13,benchmark,test]"

From 9c6b28337de52eeb6a95f569563899a6dc7750d1 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Thu, 25 Jun 2026 16:03:13 -0700
Subject: [PATCH 13/16] Fix unit test (#823)

Fix unit test failure.
---
 .azure-pipelines/templates/ut.yml |  4 ++--
 CMakeLists.txt                    | 15 ++++++++++++---
 src/core/gpu_ipc_mem.cc           |  2 ++
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/.azure-pipelines/templates/ut.yml b/.azure-pipelines/templates/ut.yml
index 743c66e6..81a74cf8 100644
--- a/.azure-pipelines/templates/ut.yml
+++ b/.azure-pipelines/templates/ut.yml
@@ -40,8 +40,8 @@ steps:
     name: PyTests
     displayName: Run pytests
     remoteScript: |
-      mpirun --allow-run-as-root -tag-output -x MSCCLPP_HOME=/root/mscclpp -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m pytest ./python/test/test_mscclpp.py -x
-      mpirun --allow-run-as-root -tag-output -x MSCCLPP_HOME=/root/mscclpp -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m pytest ./python/test/test_fp8_accum.py -x
+      mpirun --allow-run-as-root -tag-output -x MSCCLPP_HOME=/root/mscclpp -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m pytest ./python/test/test_mscclpp.py -x
+      mpirun --allow-run-as-root -tag-output -x MSCCLPP_HOME=/root/mscclpp -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m pytest ./python/test/test_fp8_accum.py -x
 
 - template: stop.yml
   parameters:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 7b8316fb..c3fb8fba 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -112,11 +112,20 @@ if(MSCCLPP_ENABLE_COVERAGE)
     if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
         message(STATUS "Code coverage enabled")
         
-        # Add coverage flags to C++ targets only (not CUDA)
-        add_compile_options($<$<COMPILE_LANGUAGE:CXX>:--coverage>)
+        # Add coverage flags to C++ targets only (not CUDA). For ROCm, keep
+        # coverage on the host compilation so HIP device linking does not look
+        # for gcov runtime symbols.
+        if(MSCCLPP_USE_ROCM)
+            add_compile_options($<$<COMPILE_LANGUAGE:CXX>:-Xarch_host>)
+            add_compile_options($<$<COMPILE_LANGUAGE:CXX>:--coverage>)
+            add_link_options($<$<LINK_LANGUAGE:CXX>:-Xarch_host>)
+            add_link_options($<$<LINK_LANGUAGE:CXX>:--coverage>)
+        else()
+            add_compile_options($<$<COMPILE_LANGUAGE:CXX>:--coverage>)
+            add_link_options($<$<LINK_LANGUAGE:CXX>:--coverage>)
+        endif()
         add_compile_options($<$<COMPILE_LANGUAGE:CXX>:-O0>)
         add_compile_options($<$<COMPILE_LANGUAGE:CXX>:-g>)
-        add_link_options($<$<LINK_LANGUAGE:CXX>:--coverage>)
         
         # Find lcov
         find_program(LCOV_PATH lcov)
diff --git a/src/core/gpu_ipc_mem.cc b/src/core/gpu_ipc_mem.cc
index 49dfcf82..b885a1ac 100644
--- a/src/core/gpu_ipc_mem.cc
+++ b/src/core/gpu_ipc_mem.cc
@@ -186,6 +186,8 @@ UniqueGpuIpcMemHandle GpuIpcMemHandle::create(const CUdeviceptr ptr) {
   if (err == cudaSuccess) {
     handle->typeFlags |= GpuIpcMemHandle::Type::RuntimeIpc;
   } else {
+    // The VMM fallback below handles this failure, so do not leak it as CUDA's last error.
+    (void)cudaGetLastError();
     WARN(GPU, "Failed to create runtime CUDA IPC handle for pointer ", (void*)basePtr,
          ": error=", static_cast<int>(err), " (", std::string(cudaGetErrorString(err)), ")");
   }

From 197d62c4b9bd4188c0f1a90aa02ec6324cb56ab0 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Tue, 30 Jun 2026 09:47:14 -0700
Subject: [PATCH 14/16] Add IPC domain rank detection (#824)

Expose the number of ranks in each GPU IPC domain through Bootstrap and
Python bindings, using NVML fabric information on CUDA when available
and falling back to host-local ranks otherwise.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 CMakeLists.txt                      |  4 ++
 include/mscclpp/core.hpp            |  7 +++
 python/csrc/core_py.cpp             |  1 +
 python/mscclpp/_core/comm.py        |  1 +
 src/core/bootstrap/bootstrap.cc     | 43 +++++++++++++++
 src/core/include/utils_internal.hpp |  1 +
 src/core/utils_internal.cc          | 81 +++++++++++++++++++++++++++++
 test/mp_unit/bootstrap_tests.cc     |  8 +++
 test/mp_unit/mp_unit_tests.hpp      |  2 +
 9 files changed, 148 insertions(+)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index c3fb8fba..985c2bd6 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -226,6 +226,10 @@ if(MSCCLPP_USE_CUDA)
     else()
         set(GPU_LIBRARIES CUDA::cudart CUDA::cuda_driver)
     endif()
+    if(NOT TARGET CUDA::nvml)
+        message(FATAL_ERROR "CUDA NVML target CUDA::nvml is required for MSCCLPP CUDA builds.")
+    endif()
+    list(APPEND GPU_LIBRARIES CUDA::nvml)
 else()
     set(CMAKE_HIP_STANDARD 20)
     set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -Wall -Wextra")
diff --git a/include/mscclpp/core.hpp b/include/mscclpp/core.hpp
index 45b56bcc..4c14f1ee 100644
--- a/include/mscclpp/core.hpp
+++ b/include/mscclpp/core.hpp
@@ -46,6 +46,10 @@ class Bootstrap {
   /// @return The total number of ranks per node.
   virtual int getNranksPerNode() const = 0;
 
+  /// Return the number of ranks in this rank's GPU IPC domain.
+  /// @return The number of ranks in the GPU IPC domain.
+  virtual int getNranksPerIpcDomain() const;
+
   /// Send arbitrary data to another process.
   ///
   /// Data sent via `send(senderBuff, size, receiverRank, tag)` can be received via `recv(receiverBuff, size,
@@ -144,6 +148,9 @@ class TcpBootstrap : public Bootstrap {
   /// Return the total number of ranks per node.
   int getNranksPerNode() const override;
 
+  /// Return the number of ranks in this rank's GPU IPC domain.
+  int getNranksPerIpcDomain() const override;
+
   /// Send arbitrary data to another process.
   ///
   /// Data sent via `send(senderBuff, size, receiverRank, tag)` can be received via `recv(receiverBuff, size,
diff --git a/python/csrc/core_py.cpp b/python/csrc/core_py.cpp
index a94f9863..7e9af6c1 100644
--- a/python/csrc/core_py.cpp
+++ b/python/csrc/core_py.cpp
@@ -56,6 +56,7 @@ void register_core(nb::module_& m) {
       .def("get_rank", &Bootstrap::getRank)
       .def("get_n_ranks", &Bootstrap::getNranks)
       .def("get_n_ranks_per_node", &Bootstrap::getNranksPerNode)
+      .def("get_n_ranks_per_ipc_domain", &Bootstrap::getNranksPerIpcDomain)
       .def(
           "send",
           [](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) {
diff --git a/python/mscclpp/_core/comm.py b/python/mscclpp/_core/comm.py
index d42349dd..9e8e6c2c 100644
--- a/python/mscclpp/_core/comm.py
+++ b/python/mscclpp/_core/comm.py
@@ -73,6 +73,7 @@ class CommGroup:
         self.my_rank = self.bootstrap.get_rank()
         self.nranks = self.bootstrap.get_n_ranks()
         self.nranks_per_node = self.bootstrap.get_n_ranks_per_node()
+        self.nranks_per_ipc_domain = self.bootstrap.get_n_ranks_per_ipc_domain()
 
     def barrier(self):
         self.bootstrap.barrier()
diff --git a/src/core/bootstrap/bootstrap.cc b/src/core/bootstrap/bootstrap.cc
index b3032e50..2f18d70c 100644
--- a/src/core/bootstrap/bootstrap.cc
+++ b/src/core/bootstrap/bootstrap.cc
@@ -50,6 +50,8 @@ MSCCLPP_API_CPP void Bootstrap::groupBarrier(const std::vector<int>& ranks) {
   }
 }
 
+MSCCLPP_API_CPP int Bootstrap::getNranksPerIpcDomain() const { return getNranksPerNode(); }
+
 MSCCLPP_API_CPP void Bootstrap::send(const std::vector<char>& data, int peer, int tag) {
   size_t size = data.size();
   send((void*)&size, sizeof(size_t), peer, tag);
@@ -83,6 +85,7 @@ class TcpBootstrap::Impl {
   int getRank();
   int getNranks();
   int getNranksPerNode();
+  int getNranksPerIpcDomain();
   void allGather(void* allData, int size);
   void broadcast(void* data, int size, int root);
   void send(void* data, int size, int peer, int tag);
@@ -95,6 +98,7 @@ class TcpBootstrap::Impl {
   int rank_;
   int nRanks_;
   int nRanksPerNode_;
+  int nRanksPerIpcDomain_;
   bool netInitialized;
   std::unique_ptr<Socket> listenSockRoot_;
   std::unique_ptr<Socket> listenSock_;
@@ -148,6 +152,7 @@ TcpBootstrap::Impl::Impl(int rank, int nRanks)
     : rank_(rank),
       nRanks_(nRanks),
       nRanksPerNode_(0),
+      nRanksPerIpcDomain_(0),
       netInitialized(false),
       peerCommAddresses_(nRanks, SocketAddress()),
       barrierArr_(nRanks, 0),
@@ -451,6 +456,42 @@ int TcpBootstrap::Impl::getNranksPerNode() {
   return nRanksPerNode_;
 }
 
+int TcpBootstrap::Impl::getNranksPerIpcDomain() {
+  if (nRanksPerIpcDomain_ > 0) return nRanksPerIpcDomain_;
+  std::vector<uint64_t> ipcDomainHashes(nRanks_);
+  ipcDomainHashes[rank_] = getIpcDomainHash();
+  allGather(ipcDomainHashes.data(), sizeof(uint64_t));
+
+  std::unordered_map<uint64_t, int> ipcDomainCounts;
+  for (uint64_t ipcDomainHash : ipcDomainHashes) {
+    ipcDomainCounts[ipcDomainHash]++;
+  }
+
+  const int nRanksPerIpcDomain = ipcDomainCounts[ipcDomainHashes[rank_]];
+  const std::string invalidIpcDomainLayout =
+      "IPC domain ranks must have the same size and be arranged in consecutive rank order";
+  if (nRanks_ % nRanksPerIpcDomain != 0) {
+    throw Error(invalidIpcDomainLayout + ": rank count is not divisible by IPC domain size", ErrorCode::InvalidUsage);
+  }
+
+  for (const auto& entry : ipcDomainCounts) {
+    if (entry.second != nRanksPerIpcDomain) {
+      throw Error(invalidIpcDomainLayout + ": IPC domains have different sizes", ErrorCode::InvalidUsage);
+    }
+  }
+  for (int i = 0; i < nRanks_; ++i) {
+    const int ipcDomainFirstRank = i - i % nRanksPerIpcDomain;
+    if (ipcDomainHashes[i] != ipcDomainHashes[ipcDomainFirstRank]) {
+      throw Error(invalidIpcDomainLayout + ": ranks are not grouped by IPC domain", ErrorCode::InvalidUsage);
+    }
+  }
+
+  INFO(MSCCLPP_INIT, "rank %d IPC domain fabric hash 0x%016llx nRanksPerIpcDomain %d", rank_,
+       static_cast<unsigned long long>(ipcDomainHashes[rank_]), nRanksPerIpcDomain);
+  nRanksPerIpcDomain_ = nRanksPerIpcDomain;
+  return nRanksPerIpcDomain_;
+}
+
 void TcpBootstrap::Impl::allGather(void* allData, int size) {
   char* data = static_cast<char*>(allData);
   int rank = rank_;
@@ -592,6 +633,8 @@ MSCCLPP_API_CPP int TcpBootstrap::getNranks() const { return pimpl_->getNranks()
 
 MSCCLPP_API_CPP int TcpBootstrap::getNranksPerNode() const { return pimpl_->getNranksPerNode(); }
 
+MSCCLPP_API_CPP int TcpBootstrap::getNranksPerIpcDomain() const { return pimpl_->getNranksPerIpcDomain(); }
+
 MSCCLPP_API_CPP void TcpBootstrap::send(void* data, int size, int peer, int tag) {
   pimpl_->send(data, size, peer, tag);
 }
diff --git a/src/core/include/utils_internal.hpp b/src/core/include/utils_internal.hpp
index c5c67e26..c6934194 100644
--- a/src/core/include/utils_internal.hpp
+++ b/src/core/include/utils_internal.hpp
@@ -37,6 +37,7 @@ int64_t busIdToInt64(const std::string busId);
 uint64_t getHash(const char* string, int n);
 uint64_t getHostHash();
 uint64_t getPidHash();
+uint64_t getIpcDomainHash();
 void getRandomData(void* buffer, size_t bytes);
 
 struct netIf {
diff --git a/src/core/utils_internal.cc b/src/core/utils_internal.cc
index 8cc55430..d48c4225 100644
--- a/src/core/utils_internal.cc
+++ b/src/core/utils_internal.cc
@@ -6,6 +6,10 @@
 #include <signal.h>
 #include <unistd.h>
 
+#if defined(MSCCLPP_USE_CUDA)
+#include <nvml.h>
+#endif
+
 #include <cstring>
 #include <fstream>
 #include <iostream>
@@ -175,6 +179,83 @@ uint64_t getPidHash(void) {
   return *pidHash;
 }
 
+#if defined(MSCCLPP_USE_CUDA) && defined(NVML_GPU_FABRIC_UUID_LEN)
+namespace {
+
+class NvmlState {
+ public:
+  NvmlState() : initialized_(nvmlInit_v2() == NVML_SUCCESS) {}
+
+  NvmlState(const NvmlState&) = delete;
+  NvmlState& operator=(const NvmlState&) = delete;
+  NvmlState(NvmlState&&) = delete;
+  NvmlState& operator=(NvmlState&&) = delete;
+
+  ~NvmlState() {
+    if (initialized_) {
+      (void)nvmlShutdown();
+    }
+  }
+
+  bool isInitialized() const { return initialized_; }
+
+ private:
+  bool initialized_ = false;
+};
+
+template <typename FabricInfo>
+uint64_t getFabricHash(const FabricInfo& fabricInfo) {
+  char hashData[NVML_GPU_FABRIC_UUID_LEN + sizeof(fabricInfo.cliqueId)];
+  std::memcpy(hashData, fabricInfo.clusterUuid, NVML_GPU_FABRIC_UUID_LEN);
+  std::memcpy(hashData + NVML_GPU_FABRIC_UUID_LEN, &fabricInfo.cliqueId, sizeof(fabricInfo.cliqueId));
+  return getHash(hashData, sizeof(hashData));
+}
+
+bool tryGetNvmlIpcDomainHash(uint64_t& ipcDomainHash) {
+  // Use the current CUDA device; callers must set the rank's device before querying.
+  int deviceId;
+  char pciBusId[] = "00000000:00:00.0";
+  if (cudaGetDevice(&deviceId) != cudaSuccess ||
+      cudaDeviceGetPCIBusId(pciBusId, sizeof(pciBusId), deviceId) != cudaSuccess) {
+    return false;
+  }
+
+  static NvmlState nvml;
+  nvmlDevice_t nvmlDevice;
+  if (!nvml.isInitialized() || nvmlDeviceGetHandleByPciBusId_v2(pciBusId, &nvmlDevice) != NVML_SUCCESS) {
+    return false;
+  }
+
+#if defined(nvmlGpuFabricInfo_v2)
+  nvmlGpuFabricInfoV_t fabricInfo = {};
+  fabricInfo.version = nvmlGpuFabricInfo_v2;
+  nvmlReturn_t result = nvmlDeviceGetGpuFabricInfoV(nvmlDevice, &fabricInfo);
+#else
+  nvmlGpuFabricInfo_t fabricInfo = {};
+  nvmlReturn_t result = nvmlDeviceGetGpuFabricInfo(nvmlDevice, &fabricInfo);
+#endif
+  if (result != NVML_SUCCESS || fabricInfo.state != NVML_GPU_FABRIC_STATE_COMPLETED ||
+      fabricInfo.status != NVML_SUCCESS) {
+    return false;
+  }
+
+  ipcDomainHash = getFabricHash(fabricInfo);
+  return true;
+}
+
+}  // namespace
+#endif
+
+uint64_t getIpcDomainHash(void) {
+#if defined(MSCCLPP_USE_CUDA) && defined(NVML_GPU_FABRIC_UUID_LEN)
+  uint64_t ipcDomainHash;
+  if (tryGetNvmlIpcDomainHash(ipcDomainHash)) {
+    return ipcDomainHash;
+  }
+#endif
+  return getHostHash();
+}
+
 int parseStringList(const char* string, netIf* ifList, int maxList) {
   if (!string) return 0;
 
diff --git a/test/mp_unit/bootstrap_tests.cc b/test/mp_unit/bootstrap_tests.cc
index c28087a4..27ca1f5f 100644
--- a/test/mp_unit/bootstrap_tests.cc
+++ b/test/mp_unit/bootstrap_tests.cc
@@ -42,10 +42,17 @@ void BootstrapTest::bootstrapTestSendRecv(std::shared_ptr<mscclpp::Bootstrap> bo
   }
 }
 
+void BootstrapTest::bootstrapTestIpcDomain(std::shared_ptr<mscclpp::Bootstrap> bootstrap) {
+  int nRanksPerIpcDomain = bootstrap->getNranksPerIpcDomain();
+  EXPECT_GT(nRanksPerIpcDomain, 0);
+  EXPECT_LE(nRanksPerIpcDomain, bootstrap->getNranks());
+}
+
 void BootstrapTest::bootstrapTestAll(std::shared_ptr<mscclpp::Bootstrap> bootstrap) {
   bootstrapTestAllGather(bootstrap);
   bootstrapTestBarrier(bootstrap);
   bootstrapTestSendRecv(bootstrap);
+  bootstrapTestIpcDomain(bootstrap);
 }
 
 TEST(BootstrapTest, WithId) {
@@ -127,6 +134,7 @@ class MPIBootstrap : public mscclpp::Bootstrap {
     MPI_Comm_size(shmcomm, &shmrank);
     return shmrank;
   }
+  int getNranksPerIpcDomain() const override { return getNranksPerNode(); }
   void allGather(void* sendbuf, int size) override {
     MPI_Allgather(MPI_IN_PLACE, 0, MPI_BYTE, sendbuf, size, MPI_BYTE, MPI_COMM_WORLD);
   }
diff --git a/test/mp_unit/mp_unit_tests.hpp b/test/mp_unit/mp_unit_tests.hpp
index eb8b5485..d01857f3 100644
--- a/test/mp_unit/mp_unit_tests.hpp
+++ b/test/mp_unit/mp_unit_tests.hpp
@@ -56,6 +56,8 @@ class BootstrapTest : public MultiProcessTest {
 
   void bootstrapTestSendRecv(std::shared_ptr<mscclpp::Bootstrap> bootstrap);
 
+  void bootstrapTestIpcDomain(std::shared_ptr<mscclpp::Bootstrap> bootstrap);
+
   void bootstrapTestAll(std::shared_ptr<mscclpp::Bootstrap> bootstrap);
 
   // Each test case should finish within 30 seconds.

From 414f6d1a86a14b573d781ef58e61bac66a39803e Mon Sep 17 00:00:00 2001
From: RJ Souza <rjsouza1995@gmail.com>
Date: Thu, 2 Jul 2026 16:24:04 -0700
Subject: [PATCH 15/16] Add support for allgather nvls algorithms (#817)

Add support for testing with NVSL by creating MSCCLPP_FORCE_DISABLE_IB
to disable IB for NVLS runs.

Add needed functions for allgather implementation

Implement initial allgather nvls algorithm.

Results with allgather_nvls_zero_copy.py:

<img width="368" height="209" alt="image"
src="https://github.com/user-attachments/assets/88676976-134d-4ec2-b3e2-d6029c44abfa"
/>
---
 include/mscclpp/npkit/npkit_event.hpp         |  2 +-
 python/mscclpp/language/channel.py            | 68 ++++++++++++++
 .../mscclpp/language/internal/operations.py   | 16 ++--
 python/mscclpp/language/internal/types.py     |  1 +
 .../single_node/allgather_nvls_zero_copy.py   | 90 ++++++++++++++++++
 .../allgather_nvls_zero_copy_pkt.py           | 91 +++++++++++++++++++
 src/core/executor/execution_plan.cc           |  4 +
 src/core/executor/executor.cc                 | 15 +--
 src/core/include/execution_common.hpp         |  2 +
 src/core/include/execution_kernel.hpp         | 63 +++++++++++++
 tools/npkit/npkit_trace_generator.py          |  2 +
 11 files changed, 340 insertions(+), 14 deletions(-)
 create mode 100644 python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy.py
 create mode 100644 python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_pkt.py

diff --git a/include/mscclpp/npkit/npkit_event.hpp b/include/mscclpp/npkit/npkit_event.hpp
index 60429193..5084638e 100644
--- a/include/mscclpp/npkit/npkit_event.hpp
+++ b/include/mscclpp/npkit/npkit_event.hpp
@@ -41,6 +41,6 @@
 #define NPKIT_EVENT_KERNEL_ALLREDUCE_EXIT 0x1C
 
 #define NPKIT_EVENT_EXECUTOR_OP_BASE_ENTRY 0x1D
-#define NPKIT_EVENT_EXECUTOR_OP_BASE_EXIT 0x39
+#define NPKIT_EVENT_EXECUTOR_OP_BASE_EXIT 0x3B
 
 #endif
diff --git a/python/mscclpp/language/channel.py b/python/mscclpp/language/channel.py
index de0f65c5..190e4b25 100644
--- a/python/mscclpp/language/channel.py
+++ b/python/mscclpp/language/channel.py
@@ -959,6 +959,54 @@ class SwitchChannel:
         op = GroupStore(src_chunk, self.buffer_type, buffer_offset, size, tb_channel_ids, self.channel_type)
         get_program().add_operation(self.src_rank, tb, op)
 
+    def broadcast_packets(self, rank, src_chunk: Chunk, buffer_offset, size, tb):
+        """Broadcast data in packet format from the source chunk to all ranks' scratch buffers in the switch channel.
+
+        Performs a specialized broadcast operation that reads data in packet format
+        from the source rank's scratch buffer and broadcasts it to each destination rank's
+        scratch buffer. Both source and destination chunks must be scratch buffers.
+
+        Args:
+            rank (int): The rank that will execute this broadcast operation.
+            src_chunk (Chunk): The source scratch chunk containing packet data to broadcast.
+            buffer_offset (int): The offset in the destination scratch buffer where data will be stored.
+            size (int): The size of data to broadcast.
+            tb (int): The thread block ID that will execute this operation.
+
+        Raises:
+            RuntimeError: If src_chunk rank is not in the rank group, if the source
+                chunk or destination buffer is not a scratch buffer, if chunk size
+                doesn't match the required size, or if buffer size is insufficient.
+
+        Example:
+            >>> channel.broadcast_packets(rank=0, src_chunk=chunk, buffer_offset=0, size=1, tb=0)
+        """
+        self.src_rank = rank
+        if src_chunk.rank not in self.rank_group.ranks:
+            raise RuntimeError(
+                f"Source chunk rank {src_chunk.rank} is not part of the rank group {self.rank_group.ranks}."
+            )
+        if src_chunk.buffer != BufferType.scratch:
+            raise RuntimeError(f"Source chunk must be of type scratch for the packet broadcast.")
+        if self.buffer_type != BufferType.scratch:
+            raise RuntimeError(f"Destination buffer must be of type scratch for the packet broadcast.")
+        if src_chunk.size != size:
+            raise RuntimeError(f"Source chunk size {src_chunk.size} does not match the required size {size}.")
+
+        for rank in self.rank_group.ranks:
+            buffer_size = get_program().gpus[rank].scratch_chunks
+
+            if buffer_size < buffer_offset + size:
+                raise RuntimeError(
+                    f"Buffer size {buffer_size} is smaller than required size {buffer_offset + size} for rank {rank}."
+                )
+
+        tb_channel_ids = get_program().setup_channel(tb, self)
+        op = GroupStore(
+            src_chunk, self.buffer_type, buffer_offset, size, tb_channel_ids, self.channel_type, use_packet=True
+        )
+        get_program().add_operation(self.src_rank, tb, op)
+
     class SwitchChannelRankView:
         """A rank-specific view of a SwitchChannel for performing operations.
 
@@ -1022,3 +1070,23 @@ class SwitchChannel:
                 >>> rank_view.broadcast(src_chunk=chunk, buffer_offset=0, size=1, tb=0)
             """
             return self._channel.broadcast(self._rank, src_chunk, buffer_offset, size, tb)
+
+        def broadcast_packets(self, src_chunk: Chunk, buffer_offset, size, tb):
+            """Perform a packet broadcast operation from this rank's perspective.
+
+            Convenience method that calls the underlying channel's broadcast_packets
+            method with this view's rank automatically provided.
+
+            Args:
+                src_chunk (Chunk): The source chunk containing packet data to broadcast.
+                buffer_offset (int): The offset in the destination buffer where data will be stored.
+                size (int): The size of data to broadcast.
+                tb (int): The thread block ID that will execute this operation.
+
+            Returns:
+                The result of the underlying channel's broadcast_packets operation.
+
+            Example:
+                >>> rank_view.broadcast_packets(src_chunk=chunk, buffer_offset=0, size=1, tb=0)
+            """
+            return self._channel.broadcast_packets(self._rank, src_chunk, buffer_offset, size, tb)
diff --git a/python/mscclpp/language/internal/operations.py b/python/mscclpp/language/internal/operations.py
index 5fb392e3..be6b1f56 100644
--- a/python/mscclpp/language/internal/operations.py
+++ b/python/mscclpp/language/internal/operations.py
@@ -871,6 +871,7 @@ class GroupLoadReduce(BaseOperation):
         fused_operation = None
         if (
             isinstance(other, GroupStore)
+            and other.name == Instruction.group_store
             and self.buffer_type == other.buffer_type
             and self.size == other.size
             and self.dst_chunk == other.src_chunk
@@ -911,8 +912,12 @@ class GroupStore(BaseOperation):
         size: int,
         channel_ids: List[int],
         channel_type: ChannelType = ChannelType.switch,
+        use_packet: bool = False,
     ):
-        super().__init__(Instruction.group_store)
+        if use_packet:
+            super().__init__(Instruction.group_store_packet)
+        else:
+            super().__init__(Instruction.group_store)
         self.src_chunk = src_chunk
         self.buffer_type = buffer_type
         self.buffer_offset = buffer_offset
@@ -926,11 +931,10 @@ class GroupStore(BaseOperation):
 
     def to_dict(self):
         result = {"name": self.name.value}
-        result["src_chunk"] = self.src_chunk.to_dict()
-        result["buffer_type"] = self.buffer_type.value
-        result["buffer_offset"] = self.buffer_offset
-        result["size"] = self.size
-        result["channel_ids"] = self.channel_ids
+        result["src_buff"] = [self.src_chunk.to_dict()]
+        result["dst_buff"] = [
+            {"switch_channel_id": self.channel_ids[0], "index": self.buffer_offset, "size": self.size}
+        ]
         result["channel_type"] = self.channel_type.value
         return result
 
diff --git a/python/mscclpp/language/internal/types.py b/python/mscclpp/language/internal/types.py
index 9bfe1c76..392282f4 100644
--- a/python/mscclpp/language/internal/types.py
+++ b/python/mscclpp/language/internal/types.py
@@ -90,6 +90,7 @@ class Instruction(Enum):
     read_reduce = "rre"
     read_reduce_send = "rres"
     group_store = "gstore"
+    group_store_packet = "gstorepkt"
     group_load_reduce = "glre"
     group_load_reduce_store = "glres"
     pipeline = "pipeline"
diff --git a/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy.py b/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy.py
new file mode 100644
index 00000000..5c138f2a
--- /dev/null
+++ b/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy.py
@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+from mscclpp.language.channel import *
+from mscclpp.language.rank import *
+from mscclpp.language.general import *
+from mscclpp.language.program import *
+from mscclpp.language.collectives import *
+
+
+def allgather_example(name, gpu_size, num_threads_per_block, min_message_size, max_message_size, instances):
+    # Defaults instances=8, num_threads_per_block=256 are tuned for 64-GPU (4x GB200) MNNVL NVLS:
+    # they give the best busbw across 1MB-1GB (instances saturate at 8; tpb=256 beats 512/1024).
+    chunksperloop = 1
+    collective = AllGather(gpu_size, chunksperloop, True)
+    with CollectiveProgram(
+        name,
+        collective,
+        gpu_size,
+        instances=instances,
+        protocol="Simple",
+        num_threads_per_block=num_threads_per_block,
+        use_double_scratch_buffer=False,
+        min_message_size=min_message_size,
+        max_message_size=max_message_size,
+    ):
+        # NVLS multicast channel over the output buffer. For Allgather each
+        # rank stores its own chunk to all ranks' output buffers via the switch.
+        nvls_chan = SwitchChannel(rank_list=[gpu for gpu in range(gpu_size)], buffer_type=BufferType.output)
+        channels = {}
+        for gpu in range(gpu_size):
+            for peer in range(gpu_size):
+                if peer != gpu:
+                    channels[(peer, gpu)] = MemoryChannel(peer, gpu)
+
+        # Synchronization to ensure all the GPUs are ready
+        for gpu in range(gpu_size):
+            src_rank = gpu
+            for peer in range(gpu_size):
+                if peer != src_rank:
+                    dst_rank = peer
+                    channels[(dst_rank, src_rank)].signal(tb=0, relaxed=True)
+            for peer in range(gpu_size):
+                if peer != src_rank:
+                    dst_rank = peer
+                    channels[(dst_rank, src_rank)].wait(tb=0, relaxed=True, data_sync=SyncType.after)
+
+        # Broadcasting each rank's chunk to every rank via NVLS multimem store.
+        # Rank `gpu` owns output chunk `gpu` (its input under in-place AllGather) and
+        # stores it to offset `gpu` across all ranks in the switch group.
+        for gpu in range(gpu_size):
+            rank = Rank(gpu)
+            output_buffer = rank.get_output_buffer()
+            nvls_chan.at_rank(gpu).broadcast(src_chunk=output_buffer[gpu : gpu + 1], buffer_offset=gpu, size=1, tb=0)
+
+        # Synchronization to ensure the GPUs finished
+        for gpu in range(gpu_size):
+            src_rank = gpu
+            for peer in range(gpu_size):
+                if peer != src_rank:
+                    dst_rank = peer
+                    channels[(dst_rank, src_rank)].signal(tb=0, relaxed=True, data_sync=SyncType.before)
+            for peer in range(gpu_size):
+                if peer != src_rank:
+                    dst_rank = peer
+                    channels[(dst_rank, src_rank)].wait(tb=0, relaxed=True)
+
+        print(JSON())
+
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--name", type=str, help="name of the program")
+parser.add_argument("--num_gpus", type=int, help="number of gpus")
+parser.add_argument("--num_threads_per_block", type=int, default=256, help="number of threads per block")
+parser.add_argument("--min_message_size", type=int, default=0, help="minimum message size")
+parser.add_argument("--max_message_size", type=int, default=2**64 - 1, help="maximum message size")
+parser.add_argument("--instances", type=int, default=8, help="number of instances (parallel threadblocks)")
+
+args = parser.parse_args()
+
+allgather_example(
+    args.name,
+    args.num_gpus,
+    args.num_threads_per_block,
+    args.min_message_size,
+    args.max_message_size,
+    args.instances,
+)
diff --git a/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_pkt.py b/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_pkt.py
new file mode 100644
index 00000000..7427e11f
--- /dev/null
+++ b/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_pkt.py
@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+from mscclpp.language.channel import *
+from mscclpp.language.rank import *
+from mscclpp.language.general import *
+from mscclpp.language.program import *
+from mscclpp.language.collectives import *
+
+
+def allgather_example(name, gpu_size, num_threads_per_block, min_message_size, max_message_size, instances):
+    # Packet (LL protocol) NVLS AllGather, tuned for small-message latency.
+    #
+    # Tuned launch defaults (64-GPU GB200 MNNVL, 1K-1M): instances=1, num_threads_per_block=256
+    #
+    # Unlike allgather_nvls_zero_copy.py (Simple protocol + full-mesh barriers around
+    # an NVLS multimem store), this variant carries an LL flag inside every packet, so
+    # the broadcast is self-synchronizing and NO signal/wait barriers are needed. Each
+    # rank packs its own chunk into scratch, multicasts those packets to every rank's
+    # scratch via the switch (gstorepkt / MULTI_STORE_PKT), and unpacks locally.
+    chunksperloop = 1
+    collective = AllGather(gpu_size, chunksperloop, True)
+    with CollectiveProgram(
+        name,
+        collective,
+        gpu_size,
+        instances=instances,
+        protocol="LL",
+        auto_sync=False,
+        num_threads_per_block=num_threads_per_block,
+        use_double_scratch_buffer=True,
+        min_message_size=min_message_size,
+        max_message_size=max_message_size,
+    ):
+        # Scratch holds packet-formatted chunks: gpu_size slots per rank, one per source rank.
+        scratch_buffer = []
+        for gpu in range(gpu_size):
+            scratch_buffer.append(Buffer(gpu, gpu_size))
+
+        # NVLS multicast channel bound to the scratch buffer (the packet staging area).
+        nvls_chan = SwitchChannel(rank_list=[gpu for gpu in range(gpu_size)], buffer_type=BufferType.scratch)
+
+        # Pack each rank's own chunk into its scratch slot `gpu`, then multicast those
+        # packets to slot `gpu` of every rank's scratch via the switch.
+        for gpu in range(gpu_size):
+            rank = Rank(gpu)
+            output_buffer = rank.get_output_buffer()
+            rank.copy_packets(scratch_buffer[gpu][gpu : gpu + 1], output_buffer[gpu : gpu + 1], tb=0)
+            nvls_chan.at_rank(gpu).broadcast_packets(
+                src_chunk=scratch_buffer[gpu][gpu : gpu + 1], buffer_offset=gpu, size=1, tb=0
+            )
+
+        # Unpack every slot from local scratch into the output buffer. Each unpack waits
+        # on the packet flag delivered by the owning rank's multicast (no barrier needed).
+        # Slot j is unpacked on tb=j to parallelize across thread blocks.
+        for gpu in range(gpu_size):
+            rank = Rank(gpu)
+            output_buffer = rank.get_output_buffer()
+            for j in range(gpu_size):
+                rank.unpack_packets(output_buffer[j : j + 1], scratch_buffer[gpu][j : j + 1], tb=j)
+
+        print(JSON())
+
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--name", type=str, help="name of the program")
+parser.add_argument("--num_gpus", type=int, help="number of gpus")
+parser.add_argument("--num_threads_per_block", type=int, default=256, help="number of threads per block")
+parser.add_argument("--min_message_size", type=int, default=1024, help="minimum message size")
+parser.add_argument("--max_message_size", type=int, default=1024 * 1024, help="maximum message size")
+parser.add_argument("--instances", type=int, default=1, help="number of instances (parallel threadblocks)")
+
+args = parser.parse_args()
+
+min_message_size = 1024 * args.instances
+
+if min_message_size > args.min_message_size:
+    raise RuntimeError(
+        f"Minimum message size {args.min_message_size} is too small for the number of instances {args.instances}. The minimum message size must be at least {min_message_size}."
+    )
+
+allgather_example(
+    args.name,
+    args.num_gpus,
+    args.num_threads_per_block,
+    args.min_message_size,
+    args.max_message_size,
+    args.instances,
+)
diff --git a/src/core/executor/execution_plan.cc b/src/core/executor/execution_plan.cc
index 4c92c26b..25ae3156 100644
--- a/src/core/executor/execution_plan.cc
+++ b/src/core/executor/execution_plan.cc
@@ -69,6 +69,10 @@ auto getOpType = [](const std::string& str) {
     return mscclpp::OperationType::REDUCE_COPY_SEND_PACKETS;
   } else if (str == "glres") {
     return mscclpp::OperationType::MULTI_LOAD_REDUCE_STORE;
+  } else if (str == "gstore") {
+    return mscclpp::OperationType::MULTI_STORE;
+  } else if (str == "gstorepkt") {
+    return mscclpp::OperationType::MULTI_STORE_PKT;
   } else if (str == "rlxsignal") {
     return mscclpp::OperationType::RELAXED_SIGNAL;
   } else if (str == "rlxwait") {
diff --git a/src/core/executor/executor.cc b/src/core/executor/executor.cc
index 9ef59bc1..91db3808 100644
--- a/src/core/executor/executor.cc
+++ b/src/core/executor/executor.cc
@@ -94,10 +94,9 @@ struct hash<mscclpp::DeviceExecutionPlanKey> {
 namespace {
 auto hasIBDevices = []() { return mscclpp::getIBDeviceCount() > 0; };
 
-// TODO(binyli): Need to add NVL domain check.
-auto useIB = [](int rank1, int rank2, int nranksPerNode) {
-  bool inSameNode = rank1 / nranksPerNode == rank2 / nranksPerNode;
-  return hasIBDevices() && !inSameNode;
+auto useIB = [](int rank1, int rank2, int nranksPerIpcDomain) {
+  bool inSameIpcDomain = rank1 / nranksPerIpcDomain == rank2 / nranksPerIpcDomain;
+  return hasIBDevices() && !inSameIpcDomain;
 };
 
 static const mscclpp::Transport IBs[] = {mscclpp::Transport::IB0, mscclpp::Transport::IB1, mscclpp::Transport::IB2,
@@ -138,6 +137,7 @@ struct ExecutionContext {
 
 struct Executor::Impl {
   int nranksPerNode;
+  int nranksPerIpcDomain;
   int nranks;
   std::shared_ptr<Communicator> comm;
   const size_t defaultScratchBufferSize = (1 << 27);
@@ -148,6 +148,7 @@ struct Executor::Impl {
   Impl(std::shared_ptr<Communicator> comm, std::shared_ptr<char> defaultScratchBuffer = nullptr)
       : comm(comm), defaultScratchBuffer(defaultScratchBuffer) {
     this->nranksPerNode = comm->bootstrap()->getNranksPerNode();
+    this->nranksPerIpcDomain = comm->bootstrap()->getNranksPerIpcDomain();
     this->nranks = comm->bootstrap()->getNranks();
     this->proxyService = std::make_shared<ProxyService>();
     this->proxyService->startProxy(true);
@@ -217,7 +218,7 @@ struct Executor::Impl {
       if (type == ChannelType::MEMORY) {
         flags |= Transport::CudaIpc;
       } else if (type == ChannelType::PORT) {
-        if (useIB(rank, info.accessRank, this->nranksPerNode)) {
+        if (useIB(rank, info.accessRank, this->nranksPerIpcDomain)) {
           flags |= IBs[rank % this->nranksPerNode];
         } else
           flags |= Transport::CudaIpc;
@@ -275,7 +276,7 @@ struct Executor::Impl {
       std::vector<ChannelInfo> channelInfos = plan.impl_->getChannelInfos(channelType);
       for (const auto& info : channelInfos) {
         for (int peer : info.connectedPeers) {
-          Transport transport = channelType == ChannelType::PORT && useIB(rank, peer, this->nranksPerNode)
+          Transport transport = channelType == ChannelType::PORT && useIB(rank, peer, this->nranksPerIpcDomain)
                                     ? ibTransport
                                     : Transport::CudaIpc;
           connFutures.push_back(this->comm->connect(transport, peer, peerTagCounters[peer]++));
@@ -284,7 +285,7 @@ struct Executor::Impl {
       channelInfos = plan.impl_->getUnpairedChannelInfos(nranks, channelType);
       for (const auto& info : channelInfos) {
         for (int peer : info.connectedPeers) {
-          Transport transport = channelType == ChannelType::PORT && useIB(rank, peer, this->nranksPerNode)
+          Transport transport = channelType == ChannelType::PORT && useIB(rank, peer, this->nranksPerIpcDomain)
                                     ? ibTransport
                                     : Transport::CudaIpc;
           connFutures.push_back(this->comm->connect(transport, peer, peerTagCounters[peer]++));
diff --git a/src/core/include/execution_common.hpp b/src/core/include/execution_common.hpp
index d071ce7d..6ce47d8f 100644
--- a/src/core/include/execution_common.hpp
+++ b/src/core/include/execution_common.hpp
@@ -66,6 +66,8 @@ enum class OperationType : uint8_t {
   PIPELINE,
   SEM_RELEASE,
   SEM_ACQUIRE,
+  MULTI_STORE,
+  MULTI_STORE_PKT,
 };
 
 struct Channels {
diff --git a/src/core/include/execution_kernel.hpp b/src/core/include/execution_kernel.hpp
index 12223181..43ba8449 100644
--- a/src/core/include/execution_kernel.hpp
+++ b/src/core/include/execution_kernel.hpp
@@ -575,6 +575,65 @@ MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(const Operation& op, uint3
     }
   }
 }
+
+template <bool ReuseScratch>
+MSCCLPP_DEVICE_INLINE void handleMultiStore(const Operation& op, void* input, void* output, void* scratch,
+                                            uint32_t offset, uint32_t unitSize) {
+  const uint32_t size = min(op.inputBufferSizes[0] - offset, unitSize);
+  if (size == 0) {
+    return;
+  }
+
+  const uint32_t srcOffset = op.inputOffsets[0] + getOffset<ReuseScratch>(op.inputBufferRefs[0].type, offset);
+  const uint32_t dstOffset = op.outputOffsets[0] + getOffset<ReuseScratch>(op.nvlsOutputBufferType, offset);
+  char* srcBase = static_cast<char*>(getBuffer(input, output, scratch, op.inputBufferRefs[0].type)) + srcOffset;
+  char* dstBase = reinterpret_cast<char*>(nvlsChannels_[op.nvlsOutputIndex].mcPtr) + dstOffset;
+
+  size_t processed = 0;
+  const bool alignedf32x4 = (reinterpret_cast<uintptr_t>(srcBase) % sizeof(f32x4) == 0) &&
+                            (reinterpret_cast<uintptr_t>(dstBase) % sizeof(f32x4) == 0);
+  if (alignedf32x4) {
+    const size_t nf32x4 = size / sizeof(f32x4);
+    f32x4* src16 = reinterpret_cast<f32x4*>(srcBase);
+    f32x4* dst16 = reinterpret_cast<f32x4*>(dstBase);
+    for (size_t idx = threadIdx.x; idx < nf32x4; idx += blockDim.x) {
+      SwitchChannelDeviceHandle::multimemStore(src16[idx], dst16 + idx);
+    }
+    processed = nf32x4 * sizeof(f32x4);
+  }
+
+  // Handle remaining data
+  const size_t startIdx = processed / sizeof(u32x1);
+  const size_t endIdx = size / sizeof(u32x1);
+  u32x1* src4 = reinterpret_cast<u32x1*>(srcBase);
+  u32x1* dst4 = reinterpret_cast<u32x1*>(dstBase);
+  for (size_t idx = threadIdx.x + startIdx; idx < endIdx; idx += blockDim.x) {
+    SwitchChannelDeviceHandle::multimemStore(src4[idx], dst4 + idx);
+  }
+}
+#endif
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+template <typename T, typename PacketType>
+MSCCLPP_DEVICE_INLINE void handleMultiStorePkt(const Operation& op, void* input, void* output, void* scratch) {
+  const uint32_t srcOffset = op.inputOffsets[0];
+  const uint32_t dstOffset = op.outputOffsets[0];
+  const uint32_t size = op.inputBufferSizes[0];
+  uint32_t nPackets = size / sizeof(PacketPayload<PacketType>);
+
+  PacketType* srcPackets =
+      (PacketType*)((char*)getBuffer(input, output, scratch, op.inputBufferRefs[0].type) + (srcOffset << 1));
+  PacketType* multiPkt =
+      (PacketType*)((char*)nvlsChannels_[op.nvlsOutputIndex].mcPtr + scratchOffset_ + (dstOffset << 1));
+
+  static_assert(sizeof(PacketType) == 16 || sizeof(PacketType) == 8, "Unsupported packet size for MULTI_STORE_PKT");
+  using StoreVec = std::conditional_t<sizeof(PacketType) == 16, mscclpp::f32x4, mscclpp::f32x2>;
+  for (size_t idx = threadIdx.x; idx < nPackets; idx += blockDim.x) {
+    PacketPayload<PacketType> data = srcPackets[idx].read(flag_);
+    PacketType pkt(data, flag_);
+    mscclpp::SwitchChannelDeviceHandle::multimemStore(*(StoreVec*)(&pkt), multiPkt + idx);
+  }
+}
 #endif
 
 template <typename T, typename PacketType, bool ReuseScratch>
@@ -708,6 +767,10 @@ MSCCLPP_DEVICE_INLINE void executeDeviceFunction(const Operation& op, T* input,
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
   else if (opType == OperationType::MULTI_LOAD_REDUCE_STORE) {
     handleMultiLoadReduceStore<T, ReuseScratch>(op, offset, unitSize);
+  } else if (opType == OperationType::MULTI_STORE) {
+    handleMultiStore<ReuseScratch>(op, input, output, scratch, offset, unitSize);
+  } else if (opType == OperationType::MULTI_STORE_PKT) {
+    handleMultiStorePkt<T, PacketType>(op, input, output, scratch);
   }
 #endif
   else if (opType == OperationType::PIPELINE) {
diff --git a/tools/npkit/npkit_trace_generator.py b/tools/npkit/npkit_trace_generator.py
index 294516e6..52f899dc 100644
--- a/tools/npkit/npkit_trace_generator.py
+++ b/tools/npkit/npkit_trace_generator.py
@@ -39,6 +39,8 @@ def parse_npkit_event_header(npkit_event_header_path):
         "PIPELINE",
         "SEM_RELEASE",
         "SEM_ACQUIRE",
+        "MULTI_STORE",
+        "MULTI_STORE_PKT",
     ]
     executor_op_to_offset = {}
     for executor_op in executor_ops:

From 9b42ad3062e8e38075e594d852f59721f7996062 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Tue, 7 Jul 2026 08:56:35 -0700
Subject: [PATCH 16/16] Support FP8 NVLS multimem reduction (#826)

## Summary
- add FP16 accumulation support for FP8 NVLS multimem load-reduce
- probe and gate FP8 NVLS algorithms on device support
- keep MNNVL algorithm sizing and buffer-pool changes out of this sub-PR

## Validation
- ./tools/lint.sh
- cmake --build build --target mscclpp_collectives mscclpp_nccl
-j$(nproc)

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 include/mscclpp/switch_channel_device.hpp     | 84 ++++++++++++++-----
 .../allreduce_nvls_block_pipeline.cu          | 38 ++++-----
 .../allreduce/allreduce_nvls_warp_pipeline.cu | 30 +++----
 .../allreduce/allreduce_nvls_zero_copy.cu     | 25 +++---
 ...ollective_utils.cc => collective_utils.cu} | 81 +++++++++++++++++-
 .../allreduce_nvls_block_pipeline.hpp         |  1 +
 .../allreduce_nvls_warp_pipeline.hpp          |  1 +
 .../allreduce/allreduce_nvls_zero_copy.hpp    |  1 +
 .../collectives/include/allreduce/common.hpp  | 19 ++++-
 .../collectives/include/collective_utils.hpp  |  4 +
 src/ext/nccl/CMakeLists.txt                   |  1 +
 src/ext/nccl/algorithm_selector.cc            | 16 +---
 src/ext/nccl/algorithm_selector.hpp           |  1 +
 src/ext/nccl/nccl.cc                          |  4 +
 14 files changed, 222 insertions(+), 84 deletions(-)
 rename src/ext/collectives/{collective_utils.cc => collective_utils.cu} (80%)

diff --git a/include/mscclpp/switch_channel_device.hpp b/include/mscclpp/switch_channel_device.hpp
index b52b6572..7b749f7a 100644
--- a/include/mscclpp/switch_channel_device.hpp
+++ b/include/mscclpp/switch_channel_device.hpp
@@ -37,7 +37,11 @@ struct SwitchChannelDeviceHandle {
     SwitchChannelDeviceHandle::multimemStore(val, reinterpret_cast<T*>(mcPtr) + index);
   }
 
-  template <typename VectorType>
+  /// Vectorized multimem load+reduce. The optional `AccumT` template parameter selects the
+  /// accumulator: when `AccumT == __half` and `VectorType` is an FP8 vector type, the
+  /// `.acc::f16` variant of the instruction is used (faster but lower precision than the
+  /// default FP32 accumulator). For all other types `AccumT` is ignored.
+  template <typename VectorType, typename AccumT = void>
   MSCCLPP_DEVICE_INLINE static VectorType multimemLoadReduce(VectorType* ptr) {
     VectorType val;
     if constexpr (std::is_same_v<VectorType, i32x1>) {
@@ -81,29 +85,71 @@ struct SwitchChannelDeviceHandle {
           : "l"(ptr)
           : "memory");
     } else if constexpr (std::is_same_v<VectorType, f8_e4m3x4>) {
-      asm("multimem.ld_reduce.relaxed.sys.global.add.e4m3x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory");
+      if constexpr (std::is_same_v<AccumT, __half>) {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.e4m3x4 %0, [%1];"
+            : "=r"(val.words[0])
+            : "l"(ptr)
+            : "memory");
+      } else {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.e4m3x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory");
+      }
     } else if constexpr (std::is_same_v<VectorType, f8_e4m3x8>) {
-      asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e4m3x4 {%0,%1}, [%2];"
-          : "=r"(val.words[0]), "=r"(val.words[1])
-          : "l"(ptr)
-          : "memory");
+      if constexpr (std::is_same_v<AccumT, __half>) {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v2.e4m3x4 {%0,%1}, [%2];"
+            : "=r"(val.words[0]), "=r"(val.words[1])
+            : "l"(ptr)
+            : "memory");
+      } else {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e4m3x4 {%0,%1}, [%2];"
+            : "=r"(val.words[0]), "=r"(val.words[1])
+            : "l"(ptr)
+            : "memory");
+      }
     } else if constexpr (std::is_same_v<VectorType, f8_e4m3x16>) {
-      asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e4m3x4 {%0,%1,%2,%3}, [%4];"
-          : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3])
-          : "l"(ptr)
-          : "memory");
+      if constexpr (std::is_same_v<AccumT, __half>) {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v4.e4m3x4 {%0,%1,%2,%3}, [%4];"
+            : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3])
+            : "l"(ptr)
+            : "memory");
+      } else {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e4m3x4 {%0,%1,%2,%3}, [%4];"
+            : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3])
+            : "l"(ptr)
+            : "memory");
+      }
     } else if constexpr (std::is_same_v<VectorType, f8_e5m2x4>) {
-      asm("multimem.ld_reduce.relaxed.sys.global.add.e5m2x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory");
+      if constexpr (std::is_same_v<AccumT, __half>) {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.e5m2x4 %0, [%1];"
+            : "=r"(val.words[0])
+            : "l"(ptr)
+            : "memory");
+      } else {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.e5m2x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory");
+      }
     } else if constexpr (std::is_same_v<VectorType, f8_e5m2x8>) {
-      asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e5m2x4 {%0,%1}, [%2];"
-          : "=r"(val.words[0]), "=r"(val.words[1])
-          : "l"(ptr)
-          : "memory");
+      if constexpr (std::is_same_v<AccumT, __half>) {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v2.e5m2x4 {%0,%1}, [%2];"
+            : "=r"(val.words[0]), "=r"(val.words[1])
+            : "l"(ptr)
+            : "memory");
+      } else {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e5m2x4 {%0,%1}, [%2];"
+            : "=r"(val.words[0]), "=r"(val.words[1])
+            : "l"(ptr)
+            : "memory");
+      }
     } else if constexpr (std::is_same_v<VectorType, f8_e5m2x16>) {
-      asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e5m2x4 {%0,%1,%2,%3}, [%4];"
-          : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3])
-          : "l"(ptr)
-          : "memory");
+      if constexpr (std::is_same_v<AccumT, __half>) {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v4.e5m2x4 {%0,%1,%2,%3}, [%4];"
+            : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3])
+            : "l"(ptr)
+            : "memory");
+      } else {
+        asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e5m2x4 {%0,%1,%2,%3}, [%4];"
+            : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3])
+            : "l"(ptr)
+            : "memory");
+      }
     } else {
       static_assert(dependentFalse<VectorType>, "Not supported type");
     }
diff --git a/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu b/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu
index 2d71cd63..95ec384b 100644
--- a/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu
+++ b/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu
@@ -13,7 +13,7 @@ namespace collective {
 
 __device__ DeviceSemaphore deviceSemaphore[NUM_SEMAPHORES];
 
-template <typename T>
+template <typename T, typename AccumT = T>
 __global__ void __launch_bounds__(1024, 1)
     allreduceNvlsBlockPipeline([[maybe_unused]] const void* src, [[maybe_unused]] void* scratch,
                                [[maybe_unused]] void* dst,
@@ -105,7 +105,7 @@ __global__ void __launch_bounds__(1024, 1)
           deviceSemaphore[oriBid].acquire();
         }
         __syncthreads();
-        handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, reduceIterSize, tid, blockDim.x);
+        handleMultiLoadReduceStore<T, AccumT>(mcBuff, mcBuff, offset, offset, reduceIterSize, tid, blockDim.x);
         __syncthreads();
         if (tid == 0) {
           deviceSemaphore[nBlocksForCopy + bidForReduce * copyReduceRatio + i].release();
@@ -158,24 +158,19 @@ struct NvlsBlockPipelineAdapter {
     } else if constexpr (std::is_same_v<T, __fp8_e4m3b15>) {
       // fp8_e4m3b15 is a software-only type with no hardware NVLS support.
       return cudaErrorNotSupported;
-    } else
-#if defined(__CUDA_ARCH__)  // Skip the __CUDA_ARCH__ < 1000 since FP8 has not been supported for NVLS
-      if constexpr (std::is_same_v<T, __fp8_e4m3> || std::is_same_v<T, __fp8_e5m2>) {
-        return cudaErrorNotSupported;
-      } else
-#endif
-      {
-        using ChannelType = DeviceHandle<BaseMemoryChannel>;
-        allreduceNvlsBlockPipeline<T>
-            <<<nBlocks, nThreadsPerBlock, 0, stream>>>(input, scratch, output, (ChannelType*)memoryChannels,
-                                                       nvlsChannels, inputSize, scratchBufferSize, rank, nRanksPerNode);
-        return cudaGetLastError();
-      }
+    } else {
+      using ChannelType = DeviceHandle<BaseMemoryChannel>;
+      allreduceNvlsBlockPipeline<T, AccumT>
+          <<<nBlocks, nThreadsPerBlock, 0, stream>>>(input, scratch, output, (ChannelType*)memoryChannels, nvlsChannels,
+                                                     inputSize, scratchBufferSize, rank, nRanksPerNode);
+      return cudaGetLastError();
+    }
   }
 };
 
 void AllreduceNvlsBlockPipeline::initialize(std::shared_ptr<Communicator> comm) {
   nSwitchChannels_ = 8;
+  fp8NvlsSupported_ = isFp8NvlsSupported();
   int nBaseChannels = 64;
   this->conns_ = setupConnections(comm);
   // setup semaphores
@@ -187,12 +182,15 @@ void AllreduceNvlsBlockPipeline::initialize(std::shared_ptr<Communicator> comm)
   this->nvlsConnections_ = setupNvlsConnections(comm, nvlsBufferSize_, nSwitchChannels_);
 }
 
-CommResult AllreduceNvlsBlockPipeline::allreduceKernelFunc(const std::shared_ptr<void> ctx_void, const void* input,
-                                                           void* output, size_t inputSize, DataType dtype, ReduceOp op,
-                                                           cudaStream_t stream, int nBlocks, int nThreadsPerBlock,
-                                                           const std::unordered_map<std::string, uintptr_t>& extras,
-                                                           DataType accumDtype) {
+CommResult AllreduceNvlsBlockPipeline::allreduceKernelFunc(
+    const std::shared_ptr<void> ctx_void, const void* input, void* output, size_t inputSize, DataType dtype,
+    ReduceOp op, cudaStream_t stream, int nBlocks, int nThreadsPerBlock,
+    [[maybe_unused]] const std::unordered_map<std::string, uintptr_t>& extras, DataType accumDtype) {
   auto ctx = std::static_pointer_cast<AlgorithmCtx>(ctx_void);
+  if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) {
+    WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction.");
+    return CommResult::CommInvalidArgument;
+  }
   AllreduceFunc allreduce = dispatch<NvlsBlockPipelineAdapter>(op, dtype, accumDtype);
   if (!allreduce) {
     WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast<int>(dtype));
diff --git a/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu b/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu
index 3bb054da..385aebd5 100644
--- a/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu
+++ b/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu
@@ -11,7 +11,7 @@
 namespace mscclpp {
 namespace collective {
 
-template <typename T>
+template <typename T, typename AccumT = T>
 __global__ void __launch_bounds__(1024, 1)
     allreduceNvlsWarpPipeline([[maybe_unused]] const void* src, [[maybe_unused]] void* scratch,
                               [[maybe_unused]] void* dst,
@@ -85,7 +85,8 @@ __global__ void __launch_bounds__(1024, 1)
       asm volatile("bar.sync %0, %1;" ::"r"(1), "r"((NCOPY_WARPS + NREDUCE_WARPS) * WARP_SIZE) : "memory");
       T* mcBuff = (T*)multicastPtr->mcPtr;
       size_t offset = blockScratchOffset + (it * copyPerIter) % scratchSizePerBlock;
-      handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, iterSize, tidInReduce, NREDUCE_WARPS * WARP_SIZE);
+      handleMultiLoadReduceStore<T, AccumT>(mcBuff, mcBuff, offset, offset, iterSize, tidInReduce,
+                                            NREDUCE_WARPS * WARP_SIZE);
       asm volatile("bar.sync %0, %1;" ::"r"(2), "r"((NRECV_COPY_WARPS + NREDUCE_WARPS) * WARP_SIZE) : "memory");
     }
     if (warpId >= startRecvCopyWid && warpId < endRecvCopyWid) {
@@ -121,24 +122,19 @@ struct NvlsWarpPipelineAdapter {
     } else if constexpr (std::is_same_v<T, __fp8_e4m3b15>) {
       // fp8_e4m3b15 is a software-only type with no hardware NVLS support.
       return cudaErrorNotSupported;
-    } else
-#if defined(__CUDA_ARCH__)  // Skip the __CUDA_ARCH__ < 1000 since FP8 has not been supported for NVLS
-      if constexpr (std::is_same_v<T, __fp8_e4m3> || std::is_same_v<T, __fp8_e5m2>) {
-        return cudaErrorNotSupported;
-      } else
-#endif
-      {
-        using ChannelType = DeviceHandle<BaseMemoryChannel>;
-        allreduceNvlsWarpPipeline<T>
-            <<<nBlocks, nThreadsPerBlock, 0, stream>>>(input, scratch, output, (ChannelType*)memoryChannels,
-                                                       nvlsChannels, inputSize, scratchBufferSize, rank, nRanksPerNode);
-        return cudaGetLastError();
-      }
+    } else {
+      using ChannelType = DeviceHandle<BaseMemoryChannel>;
+      allreduceNvlsWarpPipeline<T, AccumT>
+          <<<nBlocks, nThreadsPerBlock, 0, stream>>>(input, scratch, output, (ChannelType*)memoryChannels, nvlsChannels,
+                                                     inputSize, scratchBufferSize, rank, nRanksPerNode);
+      return cudaGetLastError();
+    }
   }
 };
 
 void AllreduceNvlsWarpPipeline::initialize(std::shared_ptr<Communicator> comm) {
   nSwitchChannels_ = 8;
+  fp8NvlsSupported_ = isFp8NvlsSupported();
   int nBaseChannels = 64;
   this->conns_ = setupConnections(comm);
   // setup semaphores
@@ -155,6 +151,10 @@ CommResult AllreduceNvlsWarpPipeline::allreduceKernelFunc(
     ReduceOp op, cudaStream_t stream, int nBlocks, int nThreadsPerBlock,
     [[maybe_unused]] const std::unordered_map<std::string, uintptr_t>& extras, DataType accumDtype) {
   auto ctx = std::static_pointer_cast<AlgorithmCtx>(ctx_void);
+  if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) {
+    WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction.");
+    return CommResult::CommInvalidArgument;
+  }
   AllreduceFunc allreduce = dispatch<NvlsWarpPipelineAdapter>(op, dtype, accumDtype);
   if (!allreduce) {
     WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast<int>(dtype));
diff --git a/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu b/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu
index e7f2028f..48ac67ff 100644
--- a/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu
+++ b/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu
@@ -13,7 +13,7 @@ namespace collective {
 
 constexpr int MAX_NBLOCKS = 32;
 
-template <typename T>
+template <typename T, typename AccumT = T>
 __global__ void __launch_bounds__(1024, 1)
     allreduceNvls([[maybe_unused]] mscclpp::DeviceHandle<mscclpp::BaseMemoryChannel>* memoryChannels,
                   [[maybe_unused]] mscclpp::DeviceHandle<mscclpp::SwitchChannel>* multicast,
@@ -56,8 +56,8 @@ __global__ void __launch_bounds__(1024, 1)
   T* src = (T*)multicastPtr->mcPtr;
   T* dst = (T*)multicastOutPtr->mcPtr;
   if (curBlockSize > 0) {
-    handleMultiLoadReduceStore(src, dst, blockOffset + channelInOffset, blockOffset + channelOutOffset, curBlockSize,
-                               threadIdx.x, blockDim.x);
+    handleMultiLoadReduceStore<T, AccumT>(src, dst, blockOffset + channelInOffset, blockOffset + channelOutOffset,
+                                          curBlockSize, threadIdx.x, blockDim.x);
   }
   __syncthreads();
   if (threadIdx.x < nPeers) {
@@ -80,17 +80,11 @@ struct NvlsAdapter {
     } else if constexpr (std::is_same_v<T, __fp8_e4m3b15>) {
       // fp8_e4m3b15 is a software-only type with no hardware NVLS support.
       return cudaErrorNotSupported;
-    } else
-#if (!defined(__CUDA_ARCH_SPECIFIC__) && !defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) || (__CUDA_ARCH__ < 1000)
-        if constexpr (std::is_same_v<T, __fp8_e4m3> || std::is_same_v<T, __fp8_e5m2>) {
-      return cudaErrorNotSupported;
-    } else
-#endif
-    {
+    } else {
       using ChannelType = DeviceHandle<mscclpp::BaseMemoryChannel>;
-      allreduceNvls<T><<<nBlocks, nThreadsPerBlock, 0, stream>>>((ChannelType*)memoryChannels, nvlsChannels,
-                                                                 nvlsOutChannels, channelInOffset, channelOutOffset,
-                                                                 inputSize, rank, nRanksPerNode);
+      allreduceNvls<T, AccumT>
+          <<<nBlocks, nThreadsPerBlock, 0, stream>>>((ChannelType*)memoryChannels, nvlsChannels, nvlsOutChannels,
+                                                     channelInOffset, channelOutOffset, inputSize, rank, nRanksPerNode);
       return cudaGetLastError();
     }
   }
@@ -102,6 +96,7 @@ void AllreduceNvls::initialize(std::shared_ptr<mscclpp::Communicator> comm) {
   cudaDeviceProp deviceProp;
   MSCCLPP_CUDATHROW(cudaGetDeviceProperties(&deviceProp, device));
   computeCapabilityMajor_ = deviceProp.major;
+  fp8NvlsSupported_ = isFp8NvlsSupported();
   nSwitchChannels_ = 32;
   this->conns_ = setupConnections(comm);
   // setup semaphores
@@ -124,6 +119,10 @@ CommResult AllreduceNvls::allreduceKernelFunc(const std::shared_ptr<void> ctx_vo
     return CommResult::CommInvalidArgument;
   }
   auto ctx = std::static_pointer_cast<AlgorithmCtx>(ctx_void);
+  if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) {
+    WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction.");
+    return CommResult::CommInvalidArgument;
+  }
   AllreduceFunc allreduce = dispatch<NvlsAdapter>(op, dtype, accumDtype);
   if (!allreduce) {
     WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast<int>(dtype));
diff --git a/src/ext/collectives/collective_utils.cc b/src/ext/collectives/collective_utils.cu
similarity index 80%
rename from src/ext/collectives/collective_utils.cc
rename to src/ext/collectives/collective_utils.cu
index 016c4a5c..2868c979 100644
--- a/src/ext/collectives/collective_utils.cc
+++ b/src/ext/collectives/collective_utils.cu
@@ -1,16 +1,93 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-#include "collective_utils.hpp"
-
 #include <algorithm>
 #include <mscclpp/algorithm.hpp>
 #include <mscclpp/core.hpp>
+#include <mscclpp/gpu_utils.hpp>
 #include <mscclpp/memory_channel.hpp>
 #include <mscclpp/switch_channel.hpp>
 
+#include "collective_utils.hpp"
+
 namespace mscclpp {
 namespace collective {
+
+namespace {
+
+#if !defined(MSCCLPP_DEVICE_HIP)
+__global__ void fp8NvlsSupportProbeKernel(int* supported) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
+    (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__))
+  *supported = 1;
+#else
+  *supported = 0;
+#endif
+}
+
+bool detectFp8NvlsSupport() {
+  AvoidCudaGraphCaptureGuard cgcGuard;
+  auto supportedDevice = mscclpp::detail::gpuCallocUnique<int>();
+  int supportedHost = 0;
+  auto stream = gpuStreamPool()->getStream();
+
+  fp8NvlsSupportProbeKernel<<<1, 1, 0, stream>>>(supportedDevice.get());
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    return false;
+  }
+
+  MSCCLPP_CUDATHROW(
+      cudaMemcpyAsync(&supportedHost, supportedDevice.get(), sizeof(supportedHost), cudaMemcpyDeviceToHost, stream));
+  err = cudaStreamSynchronize(stream);
+  if (err != cudaSuccess) {
+    (void)cudaGetLastError();
+    return false;
+  }
+  return supportedHost != 0;
+}
+#endif
+
+}  // namespace
+
+bool isFp8DataType(DataType dtype) {
+  return dtype == DataType::FLOAT8_E4M3FN || dtype == DataType::FLOAT8_E4M3FNUZ || dtype == DataType::FLOAT8_E5M2 ||
+         dtype == DataType::FLOAT8_E5M2FNUZ || dtype == DataType::FLOAT8_E4M3B15;
+}
+
+bool isNativeFp8DataType(DataType dtype) {
+#if defined(__FP8_TYPES_EXIST__)
+#if defined(__FP8_E4M3_IS_FNUZ__)
+  if (dtype == DataType::FLOAT8_E4M3FNUZ) {
+    return true;
+  }
+#else
+  if (dtype == DataType::FLOAT8_E4M3FN) {
+    return true;
+  }
+#endif
+#if defined(__FP8_E5M2_IS_FNUZ__)
+  if (dtype == DataType::FLOAT8_E5M2FNUZ) {
+    return true;
+  }
+#else
+  if (dtype == DataType::FLOAT8_E5M2) {
+    return true;
+  }
+#endif
+#endif
+  return false;
+}
+
+bool isFp8NvlsSupported() {
+#if defined(MSCCLPP_DEVICE_HIP)
+  return false;
+#else
+  static const bool supported = detectFp8NvlsSupport();
+  return supported;
+#endif
+}
+
 std::vector<mscclpp::RegisteredMemory> setupRemoteMemories(std::shared_ptr<mscclpp::Communicator> comm, int rank,
                                                            mscclpp::RegisteredMemory localMemory) {
   std::vector<mscclpp::RegisteredMemory> remoteMemories;
diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp
index 81b74add..b408c64c 100644
--- a/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp
+++ b/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp
@@ -33,6 +33,7 @@ class AllreduceNvlsBlockPipeline : public AlgorithmBuilder {
   std::vector<BaseMemoryChannel> baseChannels_;
   std::vector<Connection> conns_;
   std::vector<std::shared_ptr<NvlsConnection>> nvlsConnections_;
+  bool fp8NvlsSupported_{false};
 };
 }  // namespace collective
 }  // namespace mscclpp
diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp
index 8f02a873..2ce3a4fb 100644
--- a/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp
+++ b/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp
@@ -33,6 +33,7 @@ class AllreduceNvlsWarpPipeline : public AlgorithmBuilder {
   std::vector<BaseMemoryChannel> baseChannels_;
   std::vector<Connection> conns_;
   std::vector<std::shared_ptr<NvlsConnection>> nvlsConnections_;
+  bool fp8NvlsSupported_{false};
 };
 }  // namespace collective
 }  // namespace mscclpp
diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp
index d53ea180..edc6bceb 100644
--- a/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp
+++ b/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp
@@ -36,6 +36,7 @@ class AllreduceNvls : public AlgorithmBuilder {
   std::vector<std::shared_ptr<NvlsConnection>> nvlsConnections_;
   std::vector<std::shared_ptr<NvlsConnection>> nvlsOutConnections_;
   int computeCapabilityMajor_{0};
+  bool fp8NvlsSupported_{false};
 };
 
 }  // namespace collective
diff --git a/src/ext/collectives/include/allreduce/common.hpp b/src/ext/collectives/include/allreduce/common.hpp
index 93b18e26..6f9a3d4c 100644
--- a/src/ext/collectives/include/allreduce/common.hpp
+++ b/src/ext/collectives/include/allreduce/common.hpp
@@ -36,9 +36,21 @@ MSCCLPP_DEVICE_INLINE constexpr std::size_t calcVectorSize() {
   }
 }
 
-template <typename T>
+template <typename T, typename AccumT = T>
 MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t srcOffset, size_t dstOffset, size_t size,
                                                       int tid, int nThreads) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
+    (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__))
+  constexpr bool fp8NvlsArchSupported = true;
+#else
+  constexpr bool fp8NvlsArchSupported = false;
+#endif
+  constexpr bool nativeFp8Type = std::is_same_v<T, __fp8_e4m3> || std::is_same_v<T, __fp8_e5m2>;
+  if constexpr (nativeFp8Type && !fp8NvlsArchSupported) {
+    MSCCLPP_ASSERT_DEVICE(false, "FP8 NVLS multimem reduction is not supported on this architecture");
+    return;
+  }
+
   // nvls can only handle 4 bytes alignment
   MSCCLPP_ASSERT_DEVICE(size % 4 == 0, "size must be 4 bytes aligned");
   constexpr size_t nElem = calcVectorSize<T>();
@@ -54,7 +66,7 @@ MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t src
   vectorType* src4 = (vectorType*)src;
   vectorType* dst4 = (vectorType*)dst;
   for (size_t idx = tid; idx < nVec; idx += nThreads) {
-    auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce(src4 + srcOffset4 + idx);
+    auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce<vectorType, AccumT>(src4 + srcOffset4 + idx);
     mscclpp::SwitchChannelDeviceHandle::multimemStore(val, dst4 + dstOffset4 + idx);
   }
   // handle rest of data
@@ -64,7 +76,8 @@ MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t src
   const size_t startIdx = (srcOffset + processed) / sizeof(restVectorType);
   const size_t endIdx = (srcOffset + size) / sizeof(restVectorType);
   for (size_t idx = tid + startIdx; idx < endIdx; idx += nThreads) {
-    auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce((restVectorType*)src + idx);
+    auto val =
+        mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce<restVectorType, AccumT>((restVectorType*)src + idx);
     mscclpp::SwitchChannelDeviceHandle::multimemStore(val, (restVectorType*)dst + idx);
   }
 }
diff --git a/src/ext/collectives/include/collective_utils.hpp b/src/ext/collectives/include/collective_utils.hpp
index f705a9d1..b7f86473 100644
--- a/src/ext/collectives/include/collective_utils.hpp
+++ b/src/ext/collectives/include/collective_utils.hpp
@@ -32,6 +32,10 @@ constexpr int MAX_NRANKS_PER_NODE = 8;
 
 constexpr int SCRATCH_SIZE = 2 * 1024 * 1024 * 70;  // double buffer * 35 thread-blocks * 8 ranks * 256KB = 70MB
 
+bool isFp8DataType(DataType dtype);
+bool isNativeFp8DataType(DataType dtype);
+bool isFp8NvlsSupported();
+
 std::vector<RegisteredMemory> setupRemoteMemories(std::shared_ptr<Communicator> comm, int rank,
                                                   RegisteredMemory localMemory);
 
diff --git a/src/ext/nccl/CMakeLists.txt b/src/ext/nccl/CMakeLists.txt
index 9767e66f..463b7550 100644
--- a/src/ext/nccl/CMakeLists.txt
+++ b/src/ext/nccl/CMakeLists.txt
@@ -13,6 +13,7 @@ target_include_directories(mscclpp_nccl PRIVATE
     include 
     ${PROJECT_SOURCE_DIR}/include
     ${PROJECT_SOURCE_DIR}/src/core/include
+    ${PROJECT_SOURCE_DIR}/src/ext/collectives/include
     ${GPU_INCLUDE_DIRS}
 )
 target_link_libraries(mscclpp_nccl PUBLIC mscclpp mscclpp_collectives)
diff --git a/src/ext/nccl/algorithm_selector.cc b/src/ext/nccl/algorithm_selector.cc
index c94aab34..1ccac65d 100644
--- a/src/ext/nccl/algorithm_selector.cc
+++ b/src/ext/nccl/algorithm_selector.cc
@@ -6,6 +6,7 @@
 #include <mscclpp/env.hpp>
 #include <mscclpp/utils.hpp>
 
+#include "collective_utils.hpp"
 #include "debug.h"
 
 namespace mscclpp {
@@ -20,24 +21,15 @@ static bool isNvlsSupportedForDataType(const AlgorithmSelectorConfig& config, Da
     return false;
   }
 
-  const bool isFp8 = dtype == DataType::FLOAT8_E4M3FN || dtype == DataType::FLOAT8_E4M3FNUZ ||
-                     dtype == DataType::FLOAT8_E5M2 || dtype == DataType::FLOAT8_E5M2FNUZ;
-
-  if (!isFp8) {
+  if (!collective::isFp8DataType(dtype)) {
     return nvlsSupported;
   }
 
-  // FP8 handling
 #if !defined(__HIP_PLATFORM_AMD__)
-  // NVLS does not support FP8 on devices with compute capability < 10
-  if (config.computeCapability.first < 10) {
+  if (!collective::isNativeFp8DataType(dtype)) {
     return false;
   }
-#if (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__))
-  return true;
-#else
-  return false;
-#endif
+  return nvlsSupported && config.fp8NvlsSupported;
 #else
   return nvlsSupported;
 #endif
diff --git a/src/ext/nccl/algorithm_selector.hpp b/src/ext/nccl/algorithm_selector.hpp
index c8705f8b..2048ea05 100644
--- a/src/ext/nccl/algorithm_selector.hpp
+++ b/src/ext/nccl/algorithm_selector.hpp
@@ -16,6 +16,7 @@ namespace nccl {
 struct AlgorithmSelectorConfig {
   bool symmetricMemory;
   bool nvlsSupported;
+  bool fp8NvlsSupported;
   bool isCuMemMapAllocated;
   bool inCaptureMode;
   std::pair<int, int> computeCapability;
diff --git a/src/ext/nccl/nccl.cc b/src/ext/nccl/nccl.cc
index 2d6c5f9d..8fcc1bb1 100644
--- a/src/ext/nccl/nccl.cc
+++ b/src/ext/nccl/nccl.cc
@@ -20,6 +20,7 @@
 #include <mscclpp/algorithm.hpp>
 
 #include "algorithm_selector.hpp"
+#include "collective_utils.hpp"
 #include "datatype_conversion.hpp"
 
 static constexpr auto MSCCLPP_NCCL = mscclpp::LogSubsys::NCCL;
@@ -239,6 +240,8 @@ static std::shared_ptr<mscclpp::Algorithm> algoSelector(
   static const bool isNvlsSupported = mscclpp::isNvlsSupported();
   static const std::pair<int, int> deviceComputeCapability = getDeviceComputeCapability();
   static const bool ncclSymmetricMemory = mscclpp::env()->ncclSymmetricMemory;
+  const bool fp8NvlsSupported =
+      mscclpp::collective::isNativeFp8DataType(request.dtype) ? mscclpp::collective::isFp8NvlsSupported() : false;
 
   const bool isCuMemMapAllocated = mscclpp::isCuMemMapAllocated(const_cast<void*>(request.inputBuffer)) &&
                                    mscclpp::isCuMemMapAllocated(request.outputBuffer);
@@ -249,6 +252,7 @@ static std::shared_ptr<mscclpp::Algorithm> algoSelector(
 
   mscclpp::nccl::AlgorithmSelectorConfig config{.symmetricMemory = ncclSymmetricMemory,
                                                 .nvlsSupported = isNvlsSupported,
+                                                .fp8NvlsSupported = fp8NvlsSupported,
                                                 .isCuMemMapAllocated = isCuMemMapAllocated,
                                                 .inCaptureMode = inCaptureMode,
                                                 .computeCapability = deviceComputeCapability,