merge main

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

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

python/mscclpp/__init__.py

@@ -100,6 +100,7 @@ __all__ = [
     "AlgorithmCollection",
     "CommGroup",
     "GpuBuffer",
+    "GpuBufferGranularity",
 ]
 
 

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)

python/mscclpp/_core/comm.py

@@ -80,6 +80,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()

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

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.
@@ -953,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.
 
@@ -1016,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)

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

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
 

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"

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

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

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)

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

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

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:

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

python/mscclpp_benchmark/bench_collective.py

@@ -0,0 +1,648 @@
+# 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, runtime_name, version
+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)
+    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 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)
+            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()

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

python/mscclpp_benchmark/correctness.py

@@ -0,0 +1,401 @@
+# 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)
+        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

python/mscclpp_benchmark/gpu.py

@@ -0,0 +1,203 @@
+# 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"),
+    "runtime_get_version": ("hipRuntimeGetVersion", "cudaRuntimeGetVersion"),
+}
+
+
+@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 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)
+    _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}")

python/mscclpp_benchmark/tuner.py

@@ -0,0 +1,83 @@
+# 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
+                    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

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

python/requirements_cuda12.txt

@@ -1,5 +1,6 @@
 mpi4py
 cupy-cuda12x
+cuda-bindings>=12,<13
 prettytable
 netifaces
 pytest

python/requirements_cuda13.txt

@@ -1,5 +1,6 @@
 mpi4py
 cupy-cuda13x
+cuda-bindings>=13,<14
 prettytable
 netifaces
 pytest

python/requirements_rocm6.txt

@@ -7,4 +7,5 @@ numpy
 matplotlib
 sortedcontainers
 blake3
-pybind11
+pybind11
+hip-python>=6,<7

python/requirements_rocm7.txt

@@ -1,5 +1,5 @@
 mpi4py
-cupy-cuda11x
+cupy
 prettytable
 netifaces
 pytest
@@ -7,4 +7,5 @@ numpy
 matplotlib
 sortedcontainers
 blake3
-pybind11
+pybind11
+hip-python>=7,<8

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

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

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)