Introduce reduce-scatter algo

include/mscclpp/executor.hpp

@@ -87,6 +87,12 @@ class Executor {
   void execute(int rank, void* sendbuff, void* recvBuff, size_t sendBuffSize, size_t recvBuffSize, DataType dataType,
                const ExecutionPlan& plan, cudaStream_t stream, PacketType packetType = PacketType::LL16);
 
+  /// Release cached execution contexts while retaining the communicator and default scratch buffer.
+  ///
+  /// The caller must ensure that no execution is in flight and that all participating ranks reset collectively before
+  /// launching another execution.
+  void reset();
+
  private:
   struct Impl;
   std::unique_ptr<Impl> impl_;

python/csrc/executor_py.cpp

@@ -35,5 +35,7 @@ void register_executor(nb::module_& m) {
           },
           nb::arg("rank"), nb::arg("send_buff"), nb::arg("recv_buff"), nb::arg("send_buff_size"),
           nb::arg("recv_buff_size"), nb::arg("data_type"), nb::arg("plan"), nb::arg("stream"),
-          nb::arg("packet_type") = PacketType::LL16);
+          nb::arg("packet_type") = PacketType::LL16)
+      .def("reset", &Executor::reset,
+           "Release cached execution contexts while retaining the communicator and default scratch buffer.");
 }

python/mscclpp/default_algos/__init__.py

@@ -3,5 +3,12 @@
 
 from mscclpp.default_algos.allgather_multi_nodes import allgather_multi_nodes
 from mscclpp.default_algos.allreduce_multi_nodes import allreduce_multi_nodes
+from mscclpp.default_algos.reducescatter_multi_nodes import (
+    reducescatter_multi_nodes,
+)
 
-__all__ = ["allgather_multi_nodes", "allreduce_multi_nodes"]
+__all__ = [
+    "allgather_multi_nodes",
+    "allreduce_multi_nodes",
+    "reducescatter_multi_nodes",
+]

python/mscclpp/default_algos/reducescatter_multi_nodes.py

@@ -0,0 +1,194 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Hierarchical multi-node ReduceScatter for the low-latency packet protocol."""
+
+from mscclpp.language.channel import MemoryChannel, PortChannel
+from mscclpp.language.collectives import ReduceScatter
+from mscclpp.language.program import CollectiveProgram
+from mscclpp.language.rank import Buffer, Rank
+from mscclpp.language.thread_block_group import ThreadBlockGroup
+from mscclpp.language.utils import AlgoSpec
+
+
+def reducescatter_multi_nodes(
+    spec: AlgoSpec,
+    thread_block_group_size: int = 1,
+) -> CollectiveProgram:
+    """Build a hierarchical ReduceScatter across nodes and local GPUs."""
+    if not isinstance(spec.collective, ReduceScatter):
+        raise ValueError("reducescatter_multi_nodes requires a ReduceScatter collective")
+    if spec.protocol != "LL":
+        raise ValueError("reducescatter_multi_nodes requires protocol='LL'")
+    if spec.world_size % spec.nranks_per_node != 0:
+        raise ValueError("world_size must be divisible by nranks_per_node")
+    if spec.collective.chunk_factor != 1:
+        raise ValueError("reducescatter_multi_nodes requires chunk_factor=1")
+    if not spec.in_place or not spec.collective.inplace:
+        raise ValueError("reducescatter_multi_nodes requires in-place buffers")
+    if thread_block_group_size <= 0:
+        raise ValueError("thread_block_group_size must be positive")
+
+    num_nodes = spec.world_size // spec.nranks_per_node
+    gpus_per_node = spec.nranks_per_node
+    total_gpus = spec.world_size
+
+    with CollectiveProgram.from_spec(spec) as prog:
+        local_receive_slots = (gpus_per_node - 1) * num_nodes
+        local_send_offset = local_receive_slots
+        remote_receive_offset = local_send_offset + num_nodes
+        local_owner_offset = remote_receive_offset + num_nodes - 1
+        scratch_slots = local_owner_offset + 1
+        scratch_buffers = [Buffer(rank, scratch_slots) for rank in range(total_gpus)]
+        logical_thread_blocks = (gpus_per_node - 1) + num_nodes
+        thread_block_groups = [
+            ThreadBlockGroup(
+                tb_list=[
+                    logical_block * thread_block_group_size + group_offset
+                    for group_offset in range(thread_block_group_size)
+                ]
+            )
+            for logical_block in range(logical_thread_blocks)
+        ]
+
+        intra_node_channels: dict[tuple[int, int], MemoryChannel] = {}
+        for node_id in range(num_nodes):
+            for src_local_rank in range(gpus_per_node):
+                for dst_local_rank in range(gpus_per_node):
+                    if src_local_rank == dst_local_rank:
+                        continue
+                    src_rank = src_local_rank + node_id * gpus_per_node
+                    dst_rank = dst_local_rank + node_id * gpus_per_node
+                    intra_node_channels[(dst_rank, src_rank)] = MemoryChannel(
+                        dst_rank,
+                        src_rank,
+                    )
+
+        inter_node_channels: dict[tuple[int, int], PortChannel] = {}
+        for reducer_local_rank in range(gpus_per_node):
+            for chunk_offset in range(num_nodes):
+                owner_rank = reducer_local_rank * num_nodes + chunk_offset
+                owner_node_id = owner_rank // gpus_per_node
+                for src_node_id in range(num_nodes):
+                    if src_node_id == owner_node_id:
+                        continue
+                    src_rank = reducer_local_rank + src_node_id * gpus_per_node
+                    inter_node_channels[(owner_rank, src_rank)] = PortChannel(
+                        owner_rank,
+                        src_rank,
+                    )
+
+        # Each local GPU reduces one contiguous M / gpus_per_node group. Exchange
+        # those groups with one packet operation per local peer.
+        for node_id in range(num_nodes):
+            for src_local_rank in range(gpus_per_node):
+                src_rank = src_local_rank + node_id * gpus_per_node
+                src_input = Rank(src_rank).get_input_buffer()
+                for dst_local_rank in range(gpus_per_node):
+                    if src_local_rank == dst_local_rank:
+                        continue
+                    dst_rank = dst_local_rank + node_id * gpus_per_node
+                    local_peer_slot = src_local_rank if src_local_rank < dst_local_rank else src_local_rank - 1
+                    dst_peer_slot = dst_local_rank - 1 if src_local_rank < dst_local_rank else dst_local_rank
+                    chunk_index = dst_local_rank * num_nodes
+                    scratch_slot = local_peer_slot * num_nodes
+                    intra_node_channels[(dst_rank, src_rank)].put_packets(
+                        scratch_buffers[dst_rank][scratch_slot : scratch_slot + num_nodes],
+                        src_input[chunk_index : chunk_index + num_nodes],
+                        tb_group=thread_block_groups[dst_peer_slot],
+                    )
+
+        # Reduce each contiguous group locally. Remote nodes send their partials
+        # directly to the standard owner while the owner node transfers its one
+        # local partial over NVLink.
+        local_reduce_offset = gpus_per_node - 1
+        for src_node_id in range(num_nodes):
+            for reducer_local_rank in range(gpus_per_node):
+                src_rank = reducer_local_rank + src_node_id * gpus_per_node
+                rank = Rank(src_rank)
+                input_buffer = rank.get_input_buffer()
+
+                for chunk_offset in range(num_nodes):
+                    thread_block_group = thread_block_groups[local_reduce_offset + chunk_offset]
+                    chunk_index = reducer_local_rank * num_nodes + chunk_offset
+                    owner_rank = chunk_index
+                    owner_node_id = owner_rank // gpus_per_node
+                    local_packets = []
+                    for peer_local_rank in range(gpus_per_node):
+                        if peer_local_rank == reducer_local_rank:
+                            continue
+                        local_peer_slot = (
+                            peer_local_rank if peer_local_rank < reducer_local_rank else peer_local_rank - 1
+                        )
+                        scratch_slot = local_peer_slot * num_nodes + chunk_offset
+                        local_packets.append(scratch_buffers[src_rank][scratch_slot : scratch_slot + 1])
+
+                    local_reduced_chunk = input_buffer[chunk_index : chunk_index + 1]
+                    if local_packets:
+                        rank.reduce(
+                            local_reduced_chunk,
+                            local_packets,
+                            tb_group=thread_block_group,
+                            packet=True,
+                        )
+
+                    if src_node_id == owner_node_id:
+                        if src_rank != owner_rank:
+                            intra_node_channels[(owner_rank, src_rank)].put_packets(
+                                scratch_buffers[owner_rank][local_owner_offset : local_owner_offset + 1],
+                                local_reduced_chunk,
+                                tb_group=thread_block_group,
+                            )
+                        continue
+
+                    local_packet_slot = local_send_offset + chunk_offset
+                    rank.copy_packets(
+                        scratch_buffers[src_rank][local_packet_slot : local_packet_slot + 1],
+                        local_reduced_chunk,
+                        tb_group=thread_block_group,
+                    )
+                    remote_node_slot = src_node_id if src_node_id < owner_node_id else src_node_id - 1
+                    inter_node_channels[(owner_rank, src_rank)].read_put_packets(
+                        scratch_buffers[owner_rank][
+                            remote_receive_offset + remote_node_slot : remote_receive_offset + remote_node_slot + 1
+                        ],
+                        scratch_buffers[src_rank][local_packet_slot : local_packet_slot + 1],
+                        tb_group=thread_block_group,
+                    )
+
+        if num_nodes == 1:
+            return prog
+
+        # Every rank receives one standard shard. The owner-node handoff and
+        # direct IB transfers from the other nodes can progress concurrently.
+        for owner_rank in range(total_gpus):
+            owner = Rank(owner_rank)
+            owner_input = owner.get_input_buffer()
+            owner_node_id = owner_rank // gpus_per_node
+            reducer_local_rank = owner_rank // num_nodes
+            chunk_offset = owner_rank % num_nodes
+            reducer_rank = reducer_local_rank + owner_node_id * gpus_per_node
+            thread_block_group = thread_block_groups[local_reduce_offset + chunk_offset]
+            owner_chunk = owner_input[owner_rank : owner_rank + 1]
+
+            if reducer_rank != owner_rank:
+                owner.unpack_packets(
+                    owner_chunk,
+                    scratch_buffers[owner_rank][local_owner_offset : local_owner_offset + 1],
+                    tb_group=thread_block_group,
+                )
+
+            remote_packets = [
+                scratch_buffers[owner_rank][
+                    remote_receive_offset + remote_node_slot : remote_receive_offset + remote_node_slot + 1
+                ]
+                for remote_node_slot in range(num_nodes - 1)
+            ]
+            owner.reduce(
+                owner_chunk,
+                remote_packets,
+                tb_group=thread_block_group,
+                packet=True,
+            )
+
+    return prog

python/mscclpp/language/channel.py

@@ -737,7 +737,13 @@ def put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int):
 
         get_program().add_operation(self.src_rank, tb, op)
 
-    def read_put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int):
+    def read_put_packets(
+        self,
+        dst_chunk: Chunk,
+        src_chunk: Chunk,
+        tb: int = None,
+        tb_group: ThreadBlockGroup = None,
+    ):
         """Transfer data in packet format from local to remote scratch buffer.
 
         Performs a specialized put operation that transfers data in packet format
@@ -747,11 +753,14 @@ def read_put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int):
         Args:
             dst_chunk (Chunk): The destination scratch chunk on the destination rank.
             src_chunk (Chunk): The source scratch chunk on the source rank.
-            tb (int): The thread block ID that will execute this operation.
+            tb (int, optional): The thread block ID that will execute this operation.
+            tb_group (ThreadBlockGroup, optional): The thread block group that will
+                partition and execute this operation.
 
         Raises:
             RuntimeError: If chunk ranks don't match channel configuration, if chunks
                 are not scratch buffers, or if chunk sizes don't match.
+            ValueError: If neither a thread block ID nor a thread block group is provided.
 
         Example:
             >>> channel.read_put_packet(dst_chunk, src_chunk, tb=0)
@@ -773,20 +782,59 @@ def read_put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int):
                 f"Destination chunk size {dst_chunk.size} does not match source chunk size {src_chunk.size}."
             )
 
-        remote_chunk = RemoteBuffer(src_chunk.rank, dst_chunk.rank, dst_chunk.buffer, self.channel_type)
-        tb_chunk_id = get_program().setup_remote_chunk(self.src_rank, tb, remote_chunk, self.channel_type)
-        tb_channel_ids = get_program().setup_channel(tb, self)
+        if tb is not None:
+            tb_list = [tb]
+        elif tb_group is not None:
+            tb_list = tb_group.tb_list
+        else:
+            raise ValueError(
+                "Either 'tb' (thread block ID) or 'tb_group' " "(ThreadBlockGroup) must be provided, but both are None."
+            )
 
-        op = PutOperation(
-            src_buff=[LocalChunk(src_chunk.buffer, src_chunk.index, src_chunk.size)],
-            dst_buff=[RemoteChunk(dst_chunk.buffer, dst_chunk.index, dst_chunk.size, tb_chunk_id)],
-            channel_ids=tb_channel_ids,
-            channel_type=self.channel_type,
-            from_packet=True,
-            to_packet=True,
+        remote_chunk = RemoteBuffer(
+            src_chunk.rank,
+            dst_chunk.rank,
+            dst_chunk.buffer,
+            self.channel_type,
         )
-
-        get_program().add_operation(self.src_rank, tb, op)
+        for tb_id in tb_list:
+            tb_chunk_id = get_program().setup_remote_chunk(
+                self.src_rank,
+                tb_id,
+                remote_chunk,
+                self.channel_type,
+            )
+            tb_channel_ids = get_program().setup_channel(tb_id, self)
+            op = PutOperation(
+                src_buff=[
+                    LocalChunk(
+                        src_chunk.buffer,
+                        src_chunk.index,
+                        src_chunk.size,
+                    )
+                ],
+                dst_buff=[
+                    RemoteChunk(
+                        dst_chunk.buffer,
+                        dst_chunk.index,
+                        dst_chunk.size,
+                        tb_chunk_id,
+                    )
+                ],
+                channel_ids=tb_channel_ids,
+                channel_type=self.channel_type,
+                tbg_info=(
+                    ThreadBlockGroupInfo(
+                        tb_group.get_internal_id(tb_id),
+                        tb_group.numtb(),
+                    )
+                    if tb_group is not None
+                    else None
+                ),
+                from_packet=True,
+                to_packet=True,
+            )
+            get_program().add_operation(self.src_rank, tb_id, op)
 
 
 @dataclass

python/test/test_mscclpp.py

@@ -694,6 +694,25 @@ def test_executor(mpi_group: MpiGroup, filename: str):
     )
     stream.synchronize()
     assert cp.allclose(sendbuf, expected, atol=1e-3 * mpi_group.comm.size)
+
+    mscclpp_group.barrier()
+    executor.reset()
+    mscclpp_group.barrier()
+    for i in range(nelems_per_rank):
+        sendbuf[i] = sub_arrays[mpi_group.comm.rank][i]
+    executor.execute(
+        mpi_group.comm.rank,
+        sendbuf.data.ptr,
+        sendbuf.data.ptr,
+        sendbuf.nbytes,
+        sendbuf.nbytes,
+        DataType.float16,
+        execution_plan,
+        stream.ptr,
+    )
+    stream.synchronize()
+    assert cp.allclose(sendbuf, expected, atol=1e-3 * mpi_group.comm.size)
+
     if npkit_dump_dir is not None:
         npkit.dump(npkit_dump_dir)
         npkit.shutdown()