From d690896bc5b93a39815bf4da834352474dd081ef Mon Sep 17 00:00:00 2001
From: Caio Rocha <caiorocha@microsft.com>
Date: Thu, 30 Jul 2026 22:05:21 +0000
Subject: [PATCH 01/12] update for nvml

---
 src/core/utils_internal.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/core/utils_internal.cc b/src/core/utils_internal.cc
index d48c4225..9703e36e 100644
--- a/src/core/utils_internal.cc
+++ b/src/core/utils_internal.cc
@@ -238,6 +238,10 @@ bool tryGetNvmlIpcDomainHash(uint64_t& ipcDomainHash) {
       fabricInfo.status != NVML_SUCCESS) {
     return false;
   }
+  const char emptyClusterUuid[NVML_GPU_FABRIC_UUID_LEN] = {};
+  if (std::memcmp(fabricInfo.clusterUuid, emptyClusterUuid, sizeof(emptyClusterUuid)) == 0) {
+    return false;
+  }
 
   ipcDomainHash = getFabricHash(fabricInfo);
   return true;

From 421f02ca174a18bc7e14684e7ae5d3aeb8a3ea1d Mon Sep 17 00:00:00 2001
From: Caio Rocha <caiorocha@microsft.com>
Date: Wed, 5 Aug 2026 05:19:23 +0000
Subject: [PATCH 02/12] allgather support

---
 python/mscclpp/default_algos/__init__.py      |   3 +-
 .../default_algos/allgather_multi_nodes.py    | 300 ++++++++++++++++++
 python/test/test_default_algos.py             |  91 ++++++
 src/core/bootstrap/socket.cc                  |   6 +-
 4 files changed, 397 insertions(+), 3 deletions(-)
 create mode 100644 python/mscclpp/default_algos/allgather_multi_nodes.py
 create mode 100644 python/test/test_default_algos.py

diff --git a/python/mscclpp/default_algos/__init__.py b/python/mscclpp/default_algos/__init__.py
index 1767aab6..32de1ac7 100644
--- a/python/mscclpp/default_algos/__init__.py
+++ b/python/mscclpp/default_algos/__init__.py
@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+from mscclpp.default_algos.allgather_multi_nodes import allgather_multi_nodes
 from mscclpp.default_algos.allreduce_multi_nodes import allreduce_multi_nodes
 
-__all__ = ["allreduce_multi_nodes"]
+__all__ = ["allgather_multi_nodes", "allreduce_multi_nodes"]
diff --git a/python/mscclpp/default_algos/allgather_multi_nodes.py b/python/mscclpp/default_algos/allgather_multi_nodes.py
new file mode 100644
index 00000000..5193b955
--- /dev/null
+++ b/python/mscclpp/default_algos/allgather_multi_nodes.py
@@ -0,0 +1,300 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Hierarchical multi-node AllGather for the low-latency packet protocol."""
+
+from mscclpp.language.channel import MemoryChannel, PortChannel
+from mscclpp.language.collectives import AllGather
+from mscclpp.language.program import CollectiveProgram
+from mscclpp.language.rank import Buffer, Rank
+from mscclpp.language.utils import AlgoSpec
+
+
+def allgather_multi_nodes(
+    spec: AlgoSpec,
+    thread_block_group_size: int,
+) -> CollectiveProgram:
+    """Build a hierarchical AllGather across nodes and local GPUs."""
+    if not isinstance(spec.collective, AllGather):
+        raise ValueError("allgather_multi_nodes requires an AllGather collective")
+    if spec.protocol != "LL":
+        raise ValueError("allgather_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("allgather_multi_nodes requires chunk_factor=1")
+    if spec.in_place != spec.collective.inplace:
+        raise ValueError(
+            "spec.in_place must match spec.collective.inplace"
+        )
+    if thread_block_group_size != 1:
+        raise ValueError(
+            "allgather_multi_nodes currently requires thread_block_group_size=1"
+        )
+
+    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:
+        scratch_slots = (
+            (gpus_per_node - 1)
+            + 2 * (num_nodes - 1)
+            + (gpus_per_node - 1) * (num_nodes - 1)
+        )
+        scratch_buffers = [Buffer(rank, scratch_slots) for rank in range(total_gpus)]
+
+        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 local_rank in range(gpus_per_node):
+            for src_node_id in range(num_nodes):
+                for dst_node_id in range(num_nodes):
+                    if src_node_id == dst_node_id:
+                        continue
+                    src_rank = local_rank + src_node_id * gpus_per_node
+                    dst_rank = local_rank + dst_node_id * gpus_per_node
+                    inter_node_channels[(dst_rank, src_rank)] = PortChannel(
+                        dst_rank,
+                        src_rank,
+                    )
+
+        thread_block_offset = 1
+        local_sources = []
+        for rank_id in range(total_gpus):
+            rank = Rank(rank_id)
+            output_chunk = rank.get_output_buffer()[rank_id : rank_id + 1]
+            if spec.collective.inplace:
+                local_sources.append(output_chunk)
+            else:
+                input_chunk = rank.get_input_buffer()[0:1]
+                rank.copy(output_chunk, input_chunk, tb=0)
+                local_sources.append(input_chunk)
+
+        # Phase 0: exchange contributions within each node.
+        phase_0_send_offset = thread_block_offset
+        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
+                    scratch_slot = (
+                        src_local_rank
+                        if src_local_rank < dst_local_rank
+                        else src_local_rank - 1
+                    )
+                    thread_block = (
+                        dst_local_rank - 1
+                        if src_local_rank < dst_local_rank
+                        else dst_local_rank
+                    )
+                    intra_node_channels[(dst_rank, src_rank)].put_packets(
+                        scratch_buffers[dst_rank][scratch_slot : scratch_slot + 1],
+                        local_sources[src_rank],
+                        tb=phase_0_send_offset + thread_block,
+                    )
+
+        phase_0_unpack_offset = phase_0_send_offset + gpus_per_node - 1
+        for node_id in range(num_nodes):
+            for dst_local_rank in range(gpus_per_node):
+                dst_rank = dst_local_rank + node_id * gpus_per_node
+                rank = Rank(dst_rank)
+                for src_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
+                    scratch_slot = (
+                        src_local_rank - 1
+                        if dst_local_rank < src_local_rank
+                        else src_local_rank
+                    )
+                    rank.unpack_packets(
+                        rank.get_output_buffer()[src_rank : src_rank + 1],
+                        scratch_buffers[dst_rank][scratch_slot : scratch_slot + 1],
+                        tb=phase_0_unpack_offset + scratch_slot,
+                    )
+
+        # Phase 1: exchange same-local-rank contributions across nodes.
+        phase_1_send_offset = phase_0_unpack_offset + gpus_per_node - 1
+        remote_receive_offset = gpus_per_node - 1
+        local_packet_offset = remote_receive_offset + num_nodes - 1
+        for local_rank in range(gpus_per_node):
+            for src_node_id in range(num_nodes):
+                src_rank = local_rank + src_node_id * gpus_per_node
+                rank = Rank(src_rank)
+                for dst_node_id in range(num_nodes):
+                    if src_node_id == dst_node_id:
+                        continue
+                    thread_block = (
+                        dst_node_id - 1 if src_node_id < dst_node_id else dst_node_id
+                    )
+                    local_packet_slot = local_packet_offset + thread_block
+                    rank.copy_packets(
+                        scratch_buffers[src_rank][
+                            local_packet_slot : local_packet_slot + 1
+                        ],
+                        local_sources[src_rank],
+                        tb=phase_1_send_offset + thread_block,
+                    )
+
+                    dst_rank = local_rank + dst_node_id * gpus_per_node
+                    remote_scratch_slot = remote_receive_offset + (
+                        src_node_id if src_node_id < dst_node_id else src_node_id - 1
+                    )
+                    inter_node_channels[(dst_rank, src_rank)].read_put_packets(
+                        scratch_buffers[dst_rank][
+                            remote_scratch_slot : remote_scratch_slot + 1
+                        ],
+                        scratch_buffers[src_rank][
+                            local_packet_slot : local_packet_slot + 1
+                        ],
+                        tb=phase_1_send_offset + thread_block,
+                    )
+
+        phase_1_unpack_offset = phase_1_send_offset + num_nodes - 1
+        for local_rank in range(gpus_per_node):
+            for dst_node_id in range(num_nodes):
+                dst_rank = local_rank + dst_node_id * gpus_per_node
+                rank = Rank(dst_rank)
+                for src_node_id in range(num_nodes):
+                    if src_node_id == dst_node_id:
+                        continue
+                    src_rank = local_rank + src_node_id * gpus_per_node
+                    remote_node_slot = (
+                        src_node_id - 1 if dst_node_id < src_node_id else src_node_id
+                    )
+                    scratch_slot = remote_receive_offset + remote_node_slot
+                    rank.unpack_packets(
+                        rank.get_output_buffer()[src_rank : src_rank + 1],
+                        scratch_buffers[dst_rank][scratch_slot : scratch_slot + 1],
+                        tb=phase_1_unpack_offset + remote_node_slot,
+                    )
+
+        # Phase 2: fan out remote-node contributions within each node.
+        phase_2_send_offset = phase_1_unpack_offset + num_nodes - 1
+        local_fanout_offset = local_packet_offset + num_nodes - 1
+        for dst_node_id in range(num_nodes):
+            for src_node_id in range(num_nodes):
+                if src_node_id == dst_node_id:
+                    continue
+                remote_node_slot = (
+                    src_node_id - 1 if dst_node_id < src_node_id else src_node_id
+                )
+                for src_local_rank in range(gpus_per_node):
+                    src_rank = src_local_rank + dst_node_id * gpus_per_node
+                    remote_scratch_slot = remote_receive_offset + remote_node_slot
+                    for dst_local_rank in range(gpus_per_node):
+                        if src_local_rank == dst_local_rank:
+                            continue
+                        dst_rank = dst_local_rank + dst_node_id * gpus_per_node
+                        local_peer_slot = (
+                            src_local_rank
+                            if src_local_rank < dst_local_rank
+                            else src_local_rank - 1
+                        )
+                        fanout_slot = (
+                            local_fanout_offset
+                            + local_peer_slot
+                            + (gpus_per_node - 1) * remote_node_slot
+                        )
+                        thread_block = (
+                            dst_local_rank - 1
+                            if src_local_rank < dst_local_rank
+                            else dst_local_rank
+                        ) + (gpus_per_node - 1) * remote_node_slot
+                        intra_node_channels[(dst_rank, src_rank)].read_put_packets(
+                            scratch_buffers[dst_rank][fanout_slot : fanout_slot + 1],
+                            scratch_buffers[src_rank][
+                                remote_scratch_slot : remote_scratch_slot + 1
+                            ],
+                            tb=phase_2_send_offset + thread_block,
+                        )
+
+        phase_2_unpack_offset = phase_2_send_offset + (num_nodes - 1) * (
+            gpus_per_node - 1
+        )
+        for dst_node_id in range(num_nodes):
+            for src_node_id in range(num_nodes):
+                if src_node_id == dst_node_id:
+                    continue
+                remote_node_slot = (
+                    src_node_id - 1 if dst_node_id < src_node_id else src_node_id
+                )
+                for dst_local_rank in range(gpus_per_node):
+                    dst_rank = dst_local_rank + dst_node_id * gpus_per_node
+                    rank = Rank(dst_rank)
+                    for src_local_rank in range(gpus_per_node):
+                        if src_local_rank == dst_local_rank:
+                            continue
+                        src_rank = src_local_rank + src_node_id * gpus_per_node
+                        local_peer_slot = (
+                            src_local_rank - 1
+                            if dst_local_rank < src_local_rank
+                            else src_local_rank
+                        )
+                        fanout_slot = (
+                            local_fanout_offset
+                            + local_peer_slot
+                            + (gpus_per_node - 1) * remote_node_slot
+                        )
+                        thread_block = (
+                            local_peer_slot + (gpus_per_node - 1) * remote_node_slot
+                        )
+                        rank.unpack_packets(
+                            rank.get_output_buffer()[src_rank : src_rank + 1],
+                            scratch_buffers[dst_rank][fanout_slot : fanout_slot + 1],
+                            tb=phase_2_unpack_offset + thread_block,
+                        )
+
+    return prog
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--name", type=str, required=True)
+    parser.add_argument("--num_gpus", type=int, required=True)
+    parser.add_argument("--gpus_per_node", type=int, required=True)
+    parser.add_argument("--tbg", type=int, default=1)
+    parser.add_argument("--num_threads_per_block", type=int, default=1024)
+    parser.add_argument("--min_message_size", type=int, default=1 << 10)
+    parser.add_argument("--max_message_size", type=int, default=8 << 20)
+    parser.add_argument(
+        "--in_place",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    args = parser.parse_args()
+
+    algo_spec = AlgoSpec(
+        name=args.name,
+        collective=AllGather(args.num_gpus, 1, args.in_place),
+        nranks_per_node=args.gpus_per_node,
+        world_size=args.num_gpus,
+        in_place=args.in_place,
+        instances=1,
+        protocol="LL",
+        auto_sync=False,
+        num_threads_per_block=args.num_threads_per_block,
+        reuse_resources=True,
+        use_double_scratch_buffer=True,
+        min_message_size=args.min_message_size,
+        max_message_size=args.max_message_size,
+        tags={"default": 1},
+    )
+    program = allgather_multi_nodes(algo_spec, args.tbg)
+    print(program.to_json())
diff --git a/python/test/test_default_algos.py b/python/test/test_default_algos.py
new file mode 100644
index 00000000..b9dc9406
--- /dev/null
+++ b/python/test/test_default_algos.py
@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import json
+
+import pytest
+
+from mscclpp.default_algos import allgather_multi_nodes
+from mscclpp.language.collectives import AllGather, AllReduce
+from mscclpp.language.utils import AlgoSpec
+
+
+def _allgather_spec(*, in_place: bool = False) -> AlgoSpec:
+    return AlgoSpec(
+        name="test_allgather_multi_nodes",
+        collective=AllGather(16, 1, in_place),
+        nranks_per_node=8,
+        world_size=16,
+        in_place=in_place,
+        instances=1,
+        protocol="LL",
+        auto_sync=False,
+        num_threads_per_block=1024,
+        reuse_resources=True,
+        use_double_scratch_buffer=True,
+        min_message_size=1 << 10,
+        max_message_size=8 << 20,
+        tags={"default": 1},
+    )
+
+
+@pytest.mark.parametrize("in_place", [False, True])
+def test_allgather_multi_nodes_builds_serializable_program(
+    in_place: bool,
+) -> None:
+    program = allgather_multi_nodes(_allgather_spec(in_place=in_place), 1)
+    program.post_process_operations()
+    payload = json.loads(program.to_json())
+
+    assert payload["name"] == "test_allgather_multi_nodes"
+    assert payload["collective"] == "allgather"
+    assert len(payload["gpus"]) == 16
+    assert all(gpu["input_chunks"] == 1 for gpu in payload["gpus"])
+    assert all(gpu["output_chunks"] == 16 for gpu in payload["gpus"])
+    assert all(gpu["scratch_chunks"] == 16 for gpu in payload["gpus"])
+
+    if not in_place:
+        for rank, gpu in enumerate(payload["gpus"]):
+            local_copy = gpu["threadblocks"][0]["ops"][0]
+            assert local_copy["name"] == "copy"
+            assert local_copy["src_buff"] == [{"type": "i", "index": 0, "size": 1}]
+            assert local_copy["dst_buff"] == [{"type": "o", "index": rank, "size": 1}]
+
+
+def test_allgather_multi_nodes_rejects_invalid_spec() -> None:
+    spec = _allgather_spec()
+    invalid_collective = AlgoSpec(
+        **{
+            **spec.__dict__,
+            "collective": AllReduce(16, 1, False),
+        }
+    )
+
+    with pytest.raises(ValueError, match="AllGather"):
+        allgather_multi_nodes(invalid_collective, 1)
+    with pytest.raises(ValueError, match="thread_block_group_size=1"):
+        allgather_multi_nodes(spec, 2)
+    invalid_protocol = AlgoSpec(
+        **{
+            **spec.__dict__,
+            "protocol": "Simple",
+        }
+    )
+    with pytest.raises(ValueError, match="protocol='LL'"):
+        allgather_multi_nodes(invalid_protocol, 1)
+    invalid_chunk_factor = AlgoSpec(
+        **{
+            **spec.__dict__,
+            "collective": AllGather(16, 2, False),
+        }
+    )
+    with pytest.raises(ValueError, match="chunk_factor=1"):
+        allgather_multi_nodes(invalid_chunk_factor, 1)
+    inconsistent_in_place = AlgoSpec(
+        **{
+            **spec.__dict__,
+            "in_place": True,
+        }
+    )
+    with pytest.raises(ValueError, match="must match"):
+        allgather_multi_nodes(inconsistent_in_place, 1)
diff --git a/src/core/bootstrap/socket.cc b/src/core/bootstrap/socket.cc
index a9d9e20d..f0f8db33 100644
--- a/src/core/bootstrap/socket.cc
+++ b/src/core/bootstrap/socket.cc
@@ -674,10 +674,12 @@ void Socket::pollConnect() {
   if (ret == -1) throw SysError("poll failed", errno);
   if (ret == 0) return;
 
-  /* check socket status */
-  if ((ret == 1 && (pfd.revents & POLLOUT)) == 0) {
+  if (pfd.revents & POLLNVAL) {
     throw Error("poll failed", ErrorCode::InternalError);
   }
+  if ((pfd.revents & (POLLOUT | POLLERR | POLLHUP)) == 0) return;
+
+  /* Check SO_ERROR for both successful and failed nonblocking connects. */
   if (getsockopt(fd_, SOL_SOCKET, SO_ERROR, (void*)&ret, &rlen) == -1) {
     throw SysError("getsockopt failed", errno);
   }

From 6df33574adeb3f79c405176768c5507f0ef2ba79 Mon Sep 17 00:00:00 2001
From: Caio Rocha <caiorocha@microsft.com>
Date: Wed, 5 Aug 2026 19:02:15 +0000
Subject: [PATCH 03/12] WIP

---
 .../default_algos/allgather_multi_nodes.py    | 17 ++-----
 python/test/test_default_algos.py             | 48 +++++++++++++++----
 2 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/python/mscclpp/default_algos/allgather_multi_nodes.py b/python/mscclpp/default_algos/allgather_multi_nodes.py
index 5193b955..474571d8 100644
--- a/python/mscclpp/default_algos/allgather_multi_nodes.py
+++ b/python/mscclpp/default_algos/allgather_multi_nodes.py
@@ -10,10 +10,7 @@
 from mscclpp.language.utils import AlgoSpec
 
 
-def allgather_multi_nodes(
-    spec: AlgoSpec,
-    thread_block_group_size: int,
-) -> CollectiveProgram:
+def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
     """Build a hierarchical AllGather across nodes and local GPUs."""
     if not isinstance(spec.collective, AllGather):
         raise ValueError("allgather_multi_nodes requires an AllGather collective")
@@ -24,14 +21,7 @@ def allgather_multi_nodes(
     if spec.collective.chunk_factor != 1:
         raise ValueError("allgather_multi_nodes requires chunk_factor=1")
     if spec.in_place != spec.collective.inplace:
-        raise ValueError(
-            "spec.in_place must match spec.collective.inplace"
-        )
-    if thread_block_group_size != 1:
-        raise ValueError(
-            "allgather_multi_nodes currently requires thread_block_group_size=1"
-        )
-
+        raise ValueError("spec.in_place must match spec.collective.inplace")
     num_nodes = spec.world_size // spec.nranks_per_node
     gpus_per_node = spec.nranks_per_node
     total_gpus = spec.world_size
@@ -269,7 +259,6 @@ def allgather_multi_nodes(
     parser.add_argument("--name", type=str, required=True)
     parser.add_argument("--num_gpus", type=int, required=True)
     parser.add_argument("--gpus_per_node", type=int, required=True)
-    parser.add_argument("--tbg", type=int, default=1)
     parser.add_argument("--num_threads_per_block", type=int, default=1024)
     parser.add_argument("--min_message_size", type=int, default=1 << 10)
     parser.add_argument("--max_message_size", type=int, default=8 << 20)
@@ -296,5 +285,5 @@ def allgather_multi_nodes(
         max_message_size=args.max_message_size,
         tags={"default": 1},
     )
-    program = allgather_multi_nodes(algo_spec, args.tbg)
+    program = allgather_multi_nodes(algo_spec)
     print(program.to_json())
diff --git a/python/test/test_default_algos.py b/python/test/test_default_algos.py
index b9dc9406..ab49cee8 100644
--- a/python/test/test_default_algos.py
+++ b/python/test/test_default_algos.py
@@ -5,7 +5,7 @@
 
 import pytest
 
-from mscclpp.default_algos import allgather_multi_nodes
+from mscclpp.default_algos import allgather_multi_nodes, allreduce_multi_nodes
 from mscclpp.language.collectives import AllGather, AllReduce
 from mscclpp.language.utils import AlgoSpec
 
@@ -29,11 +29,45 @@ def _allgather_spec(*, in_place: bool = False) -> AlgoSpec:
     )
 
 
+def _allreduce_spec() -> AlgoSpec:
+    return AlgoSpec(
+        name="test_allreduce_multi_nodes",
+        collective=AllReduce(16, 1, True),
+        nranks_per_node=8,
+        world_size=16,
+        in_place=True,
+        instances=1,
+        protocol="LL",
+        auto_sync=False,
+        num_threads_per_block=1024,
+        reuse_resources=True,
+        use_double_scratch_buffer=True,
+        min_message_size=1 << 10,
+        max_message_size=8 << 20,
+        tags={"default": 1},
+    )
+
+
+def test_allreduce_multi_nodes_builds_serializable_program() -> None:
+    program = allreduce_multi_nodes(_allreduce_spec(), 1)
+    program.post_process_operations()
+    payload = json.loads(program.to_json())
+
+    assert payload["name"] == "test_allreduce_multi_nodes"
+    assert payload["collective"] == "allreduce"
+    assert payload["inplace"] is True
+    assert len(payload["gpus"]) == 16
+    assert all(gpu["input_chunks"] == 16 for gpu in payload["gpus"])
+    assert all(gpu["output_chunks"] == 16 for gpu in payload["gpus"])
+    assert all(gpu["scratch_chunks"] == 36 for gpu in payload["gpus"])
+    assert all(gpu["threadblocks"] for gpu in payload["gpus"])
+
+
 @pytest.mark.parametrize("in_place", [False, True])
 def test_allgather_multi_nodes_builds_serializable_program(
     in_place: bool,
 ) -> None:
-    program = allgather_multi_nodes(_allgather_spec(in_place=in_place), 1)
+    program = allgather_multi_nodes(_allgather_spec(in_place=in_place))
     program.post_process_operations()
     payload = json.loads(program.to_json())
 
@@ -62,9 +96,7 @@ def test_allgather_multi_nodes_rejects_invalid_spec() -> None:
     )
 
     with pytest.raises(ValueError, match="AllGather"):
-        allgather_multi_nodes(invalid_collective, 1)
-    with pytest.raises(ValueError, match="thread_block_group_size=1"):
-        allgather_multi_nodes(spec, 2)
+        allgather_multi_nodes(invalid_collective)
     invalid_protocol = AlgoSpec(
         **{
             **spec.__dict__,
@@ -72,7 +104,7 @@ def test_allgather_multi_nodes_rejects_invalid_spec() -> None:
         }
     )
     with pytest.raises(ValueError, match="protocol='LL'"):
-        allgather_multi_nodes(invalid_protocol, 1)
+        allgather_multi_nodes(invalid_protocol)
     invalid_chunk_factor = AlgoSpec(
         **{
             **spec.__dict__,
@@ -80,7 +112,7 @@ def test_allgather_multi_nodes_rejects_invalid_spec() -> None:
         }
     )
     with pytest.raises(ValueError, match="chunk_factor=1"):
-        allgather_multi_nodes(invalid_chunk_factor, 1)
+        allgather_multi_nodes(invalid_chunk_factor)
     inconsistent_in_place = AlgoSpec(
         **{
             **spec.__dict__,
@@ -88,4 +120,4 @@ def test_allgather_multi_nodes_rejects_invalid_spec() -> None:
         }
     )
     with pytest.raises(ValueError, match="must match"):
-        allgather_multi_nodes(inconsistent_in_place, 1)
+        allgather_multi_nodes(inconsistent_in_place)

From ff803bd242ef471686f4964a76ab6609fdd2c265 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 5 Aug 2026 19:07:10 +0000
Subject: [PATCH 04/12] Format multinode allgather for Black

Co-authored-by: Binyang2014 <9415966+Binyang2014@users.noreply.github.com>
---
 .../default_algos/allgather_multi_nodes.py    | 96 ++++---------------
 1 file changed, 21 insertions(+), 75 deletions(-)

diff --git a/python/mscclpp/default_algos/allgather_multi_nodes.py b/python/mscclpp/default_algos/allgather_multi_nodes.py
index 474571d8..1e8d8532 100644
--- a/python/mscclpp/default_algos/allgather_multi_nodes.py
+++ b/python/mscclpp/default_algos/allgather_multi_nodes.py
@@ -27,11 +27,7 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
     total_gpus = spec.world_size
 
     with CollectiveProgram.from_spec(spec) as prog:
-        scratch_slots = (
-            (gpus_per_node - 1)
-            + 2 * (num_nodes - 1)
-            + (gpus_per_node - 1) * (num_nodes - 1)
-        )
+        scratch_slots = (gpus_per_node - 1) + 2 * (num_nodes - 1) + (gpus_per_node - 1) * (num_nodes - 1)
         scratch_buffers = [Buffer(rank, scratch_slots) for rank in range(total_gpus)]
 
         intra_node_channels: dict[tuple[int, int], MemoryChannel] = {}
@@ -81,16 +77,8 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
                         continue
                     src_rank = src_local_rank + node_id * gpus_per_node
                     dst_rank = dst_local_rank + node_id * gpus_per_node
-                    scratch_slot = (
-                        src_local_rank
-                        if src_local_rank < dst_local_rank
-                        else src_local_rank - 1
-                    )
-                    thread_block = (
-                        dst_local_rank - 1
-                        if src_local_rank < dst_local_rank
-                        else dst_local_rank
-                    )
+                    scratch_slot = src_local_rank if src_local_rank < dst_local_rank else src_local_rank - 1
+                    thread_block = dst_local_rank - 1 if src_local_rank < dst_local_rank else dst_local_rank
                     intra_node_channels[(dst_rank, src_rank)].put_packets(
                         scratch_buffers[dst_rank][scratch_slot : scratch_slot + 1],
                         local_sources[src_rank],
@@ -106,11 +94,7 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
                     if src_local_rank == dst_local_rank:
                         continue
                     src_rank = src_local_rank + node_id * gpus_per_node
-                    scratch_slot = (
-                        src_local_rank - 1
-                        if dst_local_rank < src_local_rank
-                        else src_local_rank
-                    )
+                    scratch_slot = src_local_rank - 1 if dst_local_rank < src_local_rank else src_local_rank
                     rank.unpack_packets(
                         rank.get_output_buffer()[src_rank : src_rank + 1],
                         scratch_buffers[dst_rank][scratch_slot : scratch_slot + 1],
@@ -128,14 +112,10 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
                 for dst_node_id in range(num_nodes):
                     if src_node_id == dst_node_id:
                         continue
-                    thread_block = (
-                        dst_node_id - 1 if src_node_id < dst_node_id else dst_node_id
-                    )
+                    thread_block = dst_node_id - 1 if src_node_id < dst_node_id else dst_node_id
                     local_packet_slot = local_packet_offset + thread_block
                     rank.copy_packets(
-                        scratch_buffers[src_rank][
-                            local_packet_slot : local_packet_slot + 1
-                        ],
+                        scratch_buffers[src_rank][local_packet_slot : local_packet_slot + 1],
                         local_sources[src_rank],
                         tb=phase_1_send_offset + thread_block,
                     )
@@ -145,12 +125,8 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
                         src_node_id if src_node_id < dst_node_id else src_node_id - 1
                     )
                     inter_node_channels[(dst_rank, src_rank)].read_put_packets(
-                        scratch_buffers[dst_rank][
-                            remote_scratch_slot : remote_scratch_slot + 1
-                        ],
-                        scratch_buffers[src_rank][
-                            local_packet_slot : local_packet_slot + 1
-                        ],
+                        scratch_buffers[dst_rank][remote_scratch_slot : remote_scratch_slot + 1],
+                        scratch_buffers[src_rank][local_packet_slot : local_packet_slot + 1],
                         tb=phase_1_send_offset + thread_block,
                     )
 
@@ -163,9 +139,7 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
                     if src_node_id == dst_node_id:
                         continue
                     src_rank = local_rank + src_node_id * gpus_per_node
-                    remote_node_slot = (
-                        src_node_id - 1 if dst_node_id < src_node_id else src_node_id
-                    )
+                    remote_node_slot = src_node_id - 1 if dst_node_id < src_node_id else src_node_id
                     scratch_slot = remote_receive_offset + remote_node_slot
                     rank.unpack_packets(
                         rank.get_output_buffer()[src_rank : src_rank + 1],
@@ -180,9 +154,7 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
             for src_node_id in range(num_nodes):
                 if src_node_id == dst_node_id:
                     continue
-                remote_node_slot = (
-                    src_node_id - 1 if dst_node_id < src_node_id else src_node_id
-                )
+                remote_node_slot = src_node_id - 1 if dst_node_id < src_node_id else src_node_id
                 for src_local_rank in range(gpus_per_node):
                     src_rank = src_local_rank + dst_node_id * gpus_per_node
                     remote_scratch_slot = remote_receive_offset + remote_node_slot
@@ -190,39 +162,23 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
                         if src_local_rank == dst_local_rank:
                             continue
                         dst_rank = dst_local_rank + dst_node_id * gpus_per_node
-                        local_peer_slot = (
-                            src_local_rank
-                            if src_local_rank < dst_local_rank
-                            else src_local_rank - 1
-                        )
-                        fanout_slot = (
-                            local_fanout_offset
-                            + local_peer_slot
-                            + (gpus_per_node - 1) * remote_node_slot
-                        )
-                        thread_block = (
-                            dst_local_rank - 1
-                            if src_local_rank < dst_local_rank
-                            else dst_local_rank
-                        ) + (gpus_per_node - 1) * remote_node_slot
+                        local_peer_slot = src_local_rank if src_local_rank < dst_local_rank else src_local_rank - 1
+                        fanout_slot = local_fanout_offset + local_peer_slot + (gpus_per_node - 1) * remote_node_slot
+                        thread_block = (dst_local_rank - 1 if src_local_rank < dst_local_rank else dst_local_rank) + (
+                            gpus_per_node - 1
+                        ) * remote_node_slot
                         intra_node_channels[(dst_rank, src_rank)].read_put_packets(
                             scratch_buffers[dst_rank][fanout_slot : fanout_slot + 1],
-                            scratch_buffers[src_rank][
-                                remote_scratch_slot : remote_scratch_slot + 1
-                            ],
+                            scratch_buffers[src_rank][remote_scratch_slot : remote_scratch_slot + 1],
                             tb=phase_2_send_offset + thread_block,
                         )
 
-        phase_2_unpack_offset = phase_2_send_offset + (num_nodes - 1) * (
-            gpus_per_node - 1
-        )
+        phase_2_unpack_offset = phase_2_send_offset + (num_nodes - 1) * (gpus_per_node - 1)
         for dst_node_id in range(num_nodes):
             for src_node_id in range(num_nodes):
                 if src_node_id == dst_node_id:
                     continue
-                remote_node_slot = (
-                    src_node_id - 1 if dst_node_id < src_node_id else src_node_id
-                )
+                remote_node_slot = src_node_id - 1 if dst_node_id < src_node_id else src_node_id
                 for dst_local_rank in range(gpus_per_node):
                     dst_rank = dst_local_rank + dst_node_id * gpus_per_node
                     rank = Rank(dst_rank)
@@ -230,19 +186,9 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
                         if src_local_rank == dst_local_rank:
                             continue
                         src_rank = src_local_rank + src_node_id * gpus_per_node
-                        local_peer_slot = (
-                            src_local_rank - 1
-                            if dst_local_rank < src_local_rank
-                            else src_local_rank
-                        )
-                        fanout_slot = (
-                            local_fanout_offset
-                            + local_peer_slot
-                            + (gpus_per_node - 1) * remote_node_slot
-                        )
-                        thread_block = (
-                            local_peer_slot + (gpus_per_node - 1) * remote_node_slot
-                        )
+                        local_peer_slot = src_local_rank - 1 if dst_local_rank < src_local_rank else src_local_rank
+                        fanout_slot = local_fanout_offset + local_peer_slot + (gpus_per_node - 1) * remote_node_slot
+                        thread_block = local_peer_slot + (gpus_per_node - 1) * remote_node_slot
                         rank.unpack_packets(
                             rank.get_output_buffer()[src_rank : src_rank + 1],
                             scratch_buffers[dst_rank][fanout_slot : fanout_slot + 1],

From 8e25731f90d9234d0935ce0eab08499969664512 Mon Sep 17 00:00:00 2001
From: Caio Rocha <caiorocha@microsft.com>
Date: Wed, 5 Aug 2026 22:09:43 +0000
Subject: [PATCH 05/12] address the comments

---
 .../default_algos/allgather_multi_nodes.py    |   2 +-
 python/test/test_default_algos.py             | 123 ------------------
 src/core/bootstrap/socket.cc                  |   2 +-
 3 files changed, 2 insertions(+), 125 deletions(-)
 delete mode 100644 python/test/test_default_algos.py

diff --git a/python/mscclpp/default_algos/allgather_multi_nodes.py b/python/mscclpp/default_algos/allgather_multi_nodes.py
index 1e8d8532..8406b765 100644
--- a/python/mscclpp/default_algos/allgather_multi_nodes.py
+++ b/python/mscclpp/default_algos/allgather_multi_nodes.py
@@ -56,7 +56,7 @@ def allgather_multi_nodes(spec: AlgoSpec) -> CollectiveProgram:
                         src_rank,
                     )
 
-        thread_block_offset = 1
+        thread_block_offset = 0
         local_sources = []
         for rank_id in range(total_gpus):
             rank = Rank(rank_id)
diff --git a/python/test/test_default_algos.py b/python/test/test_default_algos.py
deleted file mode 100644
index ab49cee8..00000000
--- a/python/test/test_default_algos.py
+++ /dev/null
@@ -1,123 +0,0 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
-import json
-
-import pytest
-
-from mscclpp.default_algos import allgather_multi_nodes, allreduce_multi_nodes
-from mscclpp.language.collectives import AllGather, AllReduce
-from mscclpp.language.utils import AlgoSpec
-
-
-def _allgather_spec(*, in_place: bool = False) -> AlgoSpec:
-    return AlgoSpec(
-        name="test_allgather_multi_nodes",
-        collective=AllGather(16, 1, in_place),
-        nranks_per_node=8,
-        world_size=16,
-        in_place=in_place,
-        instances=1,
-        protocol="LL",
-        auto_sync=False,
-        num_threads_per_block=1024,
-        reuse_resources=True,
-        use_double_scratch_buffer=True,
-        min_message_size=1 << 10,
-        max_message_size=8 << 20,
-        tags={"default": 1},
-    )
-
-
-def _allreduce_spec() -> AlgoSpec:
-    return AlgoSpec(
-        name="test_allreduce_multi_nodes",
-        collective=AllReduce(16, 1, True),
-        nranks_per_node=8,
-        world_size=16,
-        in_place=True,
-        instances=1,
-        protocol="LL",
-        auto_sync=False,
-        num_threads_per_block=1024,
-        reuse_resources=True,
-        use_double_scratch_buffer=True,
-        min_message_size=1 << 10,
-        max_message_size=8 << 20,
-        tags={"default": 1},
-    )
-
-
-def test_allreduce_multi_nodes_builds_serializable_program() -> None:
-    program = allreduce_multi_nodes(_allreduce_spec(), 1)
-    program.post_process_operations()
-    payload = json.loads(program.to_json())
-
-    assert payload["name"] == "test_allreduce_multi_nodes"
-    assert payload["collective"] == "allreduce"
-    assert payload["inplace"] is True
-    assert len(payload["gpus"]) == 16
-    assert all(gpu["input_chunks"] == 16 for gpu in payload["gpus"])
-    assert all(gpu["output_chunks"] == 16 for gpu in payload["gpus"])
-    assert all(gpu["scratch_chunks"] == 36 for gpu in payload["gpus"])
-    assert all(gpu["threadblocks"] for gpu in payload["gpus"])
-
-
-@pytest.mark.parametrize("in_place", [False, True])
-def test_allgather_multi_nodes_builds_serializable_program(
-    in_place: bool,
-) -> None:
-    program = allgather_multi_nodes(_allgather_spec(in_place=in_place))
-    program.post_process_operations()
-    payload = json.loads(program.to_json())
-
-    assert payload["name"] == "test_allgather_multi_nodes"
-    assert payload["collective"] == "allgather"
-    assert len(payload["gpus"]) == 16
-    assert all(gpu["input_chunks"] == 1 for gpu in payload["gpus"])
-    assert all(gpu["output_chunks"] == 16 for gpu in payload["gpus"])
-    assert all(gpu["scratch_chunks"] == 16 for gpu in payload["gpus"])
-
-    if not in_place:
-        for rank, gpu in enumerate(payload["gpus"]):
-            local_copy = gpu["threadblocks"][0]["ops"][0]
-            assert local_copy["name"] == "copy"
-            assert local_copy["src_buff"] == [{"type": "i", "index": 0, "size": 1}]
-            assert local_copy["dst_buff"] == [{"type": "o", "index": rank, "size": 1}]
-
-
-def test_allgather_multi_nodes_rejects_invalid_spec() -> None:
-    spec = _allgather_spec()
-    invalid_collective = AlgoSpec(
-        **{
-            **spec.__dict__,
-            "collective": AllReduce(16, 1, False),
-        }
-    )
-
-    with pytest.raises(ValueError, match="AllGather"):
-        allgather_multi_nodes(invalid_collective)
-    invalid_protocol = AlgoSpec(
-        **{
-            **spec.__dict__,
-            "protocol": "Simple",
-        }
-    )
-    with pytest.raises(ValueError, match="protocol='LL'"):
-        allgather_multi_nodes(invalid_protocol)
-    invalid_chunk_factor = AlgoSpec(
-        **{
-            **spec.__dict__,
-            "collective": AllGather(16, 2, False),
-        }
-    )
-    with pytest.raises(ValueError, match="chunk_factor=1"):
-        allgather_multi_nodes(invalid_chunk_factor)
-    inconsistent_in_place = AlgoSpec(
-        **{
-            **spec.__dict__,
-            "in_place": True,
-        }
-    )
-    with pytest.raises(ValueError, match="must match"):
-        allgather_multi_nodes(inconsistent_in_place)
diff --git a/src/core/bootstrap/socket.cc b/src/core/bootstrap/socket.cc
index f0f8db33..312aaa6e 100644
--- a/src/core/bootstrap/socket.cc
+++ b/src/core/bootstrap/socket.cc
@@ -688,7 +688,7 @@ void Socket::pollConnect() {
     state_ = SocketStateConnected;
   } else if (ret == ECONNREFUSED || ret == ETIMEDOUT) {
     if (++connectRetries_ % 1000 == 0) {
-      INFO(MSCCLPP_ALL, "Call to connect returned %s, retrying", strerror(errno));
+      INFO(MSCCLPP_ALL, "Call to connect returned %s, retrying", strerror(ret));
     }
     usleep(SLEEP_INT);
 

From a6a78ae408671930a68cb29471fe1033c383f90d Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Wed, 5 Aug 2026 22:47:06 +0000
Subject: [PATCH 06/12] WIP

---
 src/core/ib.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/core/ib.cc b/src/core/ib.cc
index bd34321c..55ac21c4 100644
--- a/src/core/ib.cc
+++ b/src/core/ib.cc
@@ -404,7 +404,7 @@ void IbQp::postSend() {
   numPostedSignaledSend_ += numStagedSignaledSend_;
   numStagedSignaledSend_ = 0;
   if (numPostedSignaledSend_ + 4 > maxSendCqSize_) {
-    WARN(NET, "IB: CQ is almost full (", numPostedSignaledSend_, " / ", maxSendCqSize_,
+    INFO(NET, "IB: CQ is almost full (", numPostedSignaledSend_, " / ", maxSendCqSize_,
          "). The connection needs to be flushed to prevent timeout errors.");
   }
 }

From c5df28c16f4107d508001299a6011ff1a16aebc9 Mon Sep 17 00:00:00 2001
From: Changho Hwang <changhohwang@microsoft.com>
Date: Fri, 7 Aug 2026 15:54:00 +0000
Subject: [PATCH 07/12] Proxy FIFO latency improvement

---
 include/mscclpp/fifo.hpp                |  14 ++-
 include/mscclpp/fifo_device.hpp         |  42 +++++----
 include/mscclpp/port_channel_device.hpp |   4 +-
 src/core/fifo.cc                        |  48 +++++++---
 src/core/proxy.cc                       |   6 +-
 test/unit/fifo_perf_tests.cu            |   6 +-
 test/unit/fifo_tests.cu                 | 112 +++++++++++++++++++++---
 7 files changed, 177 insertions(+), 55 deletions(-)

diff --git a/include/mscclpp/fifo.hpp b/include/mscclpp/fifo.hpp
index 6aae03b5..4094036d 100644
--- a/include/mscclpp/fifo.hpp
+++ b/include/mscclpp/fifo.hpp
@@ -16,15 +16,21 @@ constexpr size_t DEFAULT_FIFO_SIZE = 512;
 class Fifo {
  public:
   /// Constructor.
-  /// @param size Number of entries (default: DEFAULT_FIFO_SIZE).
+  /// @param size Number of entries. Must be a power of two (default: DEFAULT_FIFO_SIZE).
+  /// @throws Error with ErrorCode::InvalidUsage if size is not a positive power of two.
   Fifo(int size = DEFAULT_FIFO_SIZE);
 
   /// Destructor.
   ~Fifo();
 
-  /// Poll and get the trigger at the head.
-  /// @return ProxyTrigger at the head of the FIFO.
-  ProxyTrigger poll();
+  /// Poll for the trigger at the head.
+  ///
+  /// A trigger carries no reserved payload value, so readiness is reported separately rather than
+  /// encoded in the trigger itself.
+  ///
+  /// @param trigger Set to the trigger at the head if one is ready. Untouched otherwise.
+  /// @return True if a trigger was ready and written to @p trigger.
+  bool poll(ProxyTrigger& trigger);
 
   /// Remove the head trigger.
   void pop();
diff --git a/include/mscclpp/fifo_device.hpp b/include/mscclpp/fifo_device.hpp
index d5ae75f6..294d5498 100644
--- a/include/mscclpp/fifo_device.hpp
+++ b/include/mscclpp/fifo_device.hpp
@@ -25,6 +25,8 @@ constexpr unsigned int TriggerBitsOffset = 32;
 constexpr unsigned int TriggerBitsMemoryId = 9;
 constexpr unsigned int TriggerBitsType = 3;
 constexpr unsigned int TriggerBitsSemaphoreId = 10;
+// The FIFO uses the reserved bit to mark a slot as written, so a trigger must not carry data
+// there. See FifoDeviceHandle::push().
 constexpr unsigned int TriggerBitsFifoReserved = 1;
 
 /// Pair of 64-bit unsigned integers used as a trigger for the proxy.
@@ -71,7 +73,6 @@ union alignas(16) ProxyTrigger {
     MSCCLPP_ASSERT_DEVICE(dstOffset < (1ULL << TriggerBitsOffset), "dstOffset is too large");
     MSCCLPP_ASSERT_DEVICE(srcId < (1ULL << TriggerBitsMemoryId), "srcId is too large");
     MSCCLPP_ASSERT_DEVICE(srcOffset < (1ULL << TriggerBitsOffset), "srcOffset is too large");
-    MSCCLPP_ASSERT_DEVICE(bytes != 0, "bytes must not be zero");
     MSCCLPP_ASSERT_DEVICE(bytes < (1ULL << TriggerBitsSize), "bytes is too large");
     MSCCLPP_ASSERT_DEVICE(semaphoreId < (1ULL << TriggerBitsSemaphoreId), "semaphoreId is too large");
     constexpr uint64_t maskSize = (1ULL << TriggerBitsSize) - 1;
@@ -106,29 +107,32 @@ struct FifoDeviceHandle {
   MSCCLPP_DEVICE_INLINE uint64_t push(ProxyTrigger trigger, int64_t maxSpinCount = 1000000) {
     uint64_t prevHead = atomicFetchAdd<uint64_t, scopeDevice>(head, 1, memoryOrderRelaxed);
 
-    // Flip the last bit for safe polling; host will revert.
-    constexpr uint64_t flipMask = uint64_t{1} << uint64_t{63};
-    trigger.snd ^= flipMask;
-
-    // Wait until the trigger is freed by the host.
+    // Wait until the slot's previous occupant has been consumed. Lap parity identifies a stale
+    // trigger but does not prevent overwriting a live one; this does.
     if (prevHead >= size + *tailCache) {
       sync(prevHead - size, maxSpinCount);
     }
 
-    ProxyTrigger* triggerPtr = &(triggers[prevHead % size]);
+    // Commit bit: the parity of the lap this slot is on, so that the value left by the previous
+    // lap reads as stale. Lap 0 writes 1, which makes a zero-initialized buffer read as empty.
+    trigger.fields.reserved = ((prevHead >> sizeShift) & 1ULL) ^ 1ULL;
+
+    ProxyTrigger* triggerPtr = &(triggers[prevHead & sizeMask]);
 
+    // snd is the commit word: a consumer that observes this lap's parity must also observe the
+    // payload that goes with it.
 #if defined(MSCCLPP_DEVICE_CUDA)
-#if __CUDA_ARCH__ == 800
-    // This is faster than release for A100.
-    __threadfence_system();
-    asm volatile("st.global.relaxed.sys.v2.u64 [%0], {%1,%2};" ::"l"(triggerPtr), "l"(trigger.fst), "l"(trigger.snd));
-#else
+    // One 128-bit store publishes both words together, so no ordering is needed between them.
+    // The proxy already relies on this: a torn store would let it read a new fst against the
+    // previous lap's snd, and dispatch on a stale semaphoreId.
+    //
+    // sm_80 used __threadfence_system() plus a relaxed store here, which was once faster. On
+    // A100 with CUDA 12.9 it is no longer, according to `FifoTest.Fifo`.
     asm volatile("st.global.release.sys.v2.u64 [%0], {%1,%2};" ::"l"(triggerPtr), "l"(trigger.fst), "l"(trigger.snd));
-#endif
 #else   // !defined(MSCCLPP_DEVICE_CUDA)
-    // Store snd no later than fst.
-    atomicStore(&(triggerPtr->snd), trigger.snd, memoryOrderRelaxed);
-    atomicStore(&(triggerPtr->fst), trigger.fst, memoryOrderRelease);
+    // No vector store here, so order the payload ahead of the commit explicitly.
+    atomicStore(&(triggerPtr->fst), trigger.fst, memoryOrderRelaxed);
+    atomicStore(&(triggerPtr->snd), trigger.snd, memoryOrderRelease);
 #endif  // !defined(MSCCLPP_DEVICE_CUDA)
 
     return prevHead;
@@ -168,8 +172,12 @@ struct FifoDeviceHandle {
   uint64_t* tail;
   /// Cached tail value.
   uint64_t* tailCache;
-  /// FIFO size.
+  /// FIFO size. Always a power of two.
   int size;
+  /// size - 1, for mapping a position to a slot.
+  uint64_t sizeMask;
+  /// log2(size), for extracting the lap from a position.
+  uint64_t sizeShift;
 };
 
 }  // namespace mscclpp
diff --git a/include/mscclpp/port_channel_device.hpp b/include/mscclpp/port_channel_device.hpp
index 74fa3d89..fd575b4c 100644
--- a/include/mscclpp/port_channel_device.hpp
+++ b/include/mscclpp/port_channel_device.hpp
@@ -72,7 +72,7 @@ struct BasePortChannelDeviceHandle {
   }
 
   /// Push a TriggerFlag to the FIFO.
-  MSCCLPP_DEVICE_INLINE void signal() { fifo_.push({TriggerFlag, 0, 0, 0, 0, 1, semaphoreId_}); }
+  MSCCLPP_DEVICE_INLINE void signal() { fifo_.push({TriggerFlag, 0, 0, 0, 0, 0, semaphoreId_}); }
 
   /// Push a TriggerData and a TriggerFlag at the same time to the FIFO.
   /// @param dstId The ID of destination memory region.
@@ -122,7 +122,7 @@ struct BasePortChannelDeviceHandle {
   /// Push a TriggerSync to the FIFO.
   /// @param maxSpinCount The maximum number of spin counts before asserting. Never assert if negative.
   MSCCLPP_DEVICE_INLINE void flush(int64_t maxSpinCount = 1000000) {
-    uint64_t pos = fifo_.push({TriggerSync, 0, 0, 0, 0, 1, semaphoreId_});
+    uint64_t pos = fifo_.push({TriggerSync, 0, 0, 0, 0, 0, semaphoreId_});
     detail::waitFlush(flushDonePos_, pos, maxSpinCount);
   }
 
diff --git a/src/core/fifo.cc b/src/core/fifo.cc
index b11775d8..b44cfecc 100644
--- a/src/core/fifo.cc
+++ b/src/core/fifo.cc
@@ -8,25 +8,42 @@
 
 #include "api.h"
 #include "atomic.hpp"
+#include "logger.hpp"
 
 namespace mscclpp {
 
+namespace {
+// size is validated to be a positive power of two, so this is exact.
+uint64_t shiftOf(int size) {
+  uint64_t shift = 0;
+  while ((1 << shift) < size) shift++;
+  return shift;
+}
+}  // namespace
+
 struct Fifo::Impl {
   detail::UniqueGpuHostPtr<ProxyTrigger> triggers;
   detail::UniqueGpuPtr<uint64_t> head;
   detail::UniqueGpuHostPtr<uint64_t> tail;
   detail::UniqueGpuPtr<uint64_t> tailCache;
   const int size;
+  const uint64_t sizeMask;
+  const uint64_t sizeShift;
 
   Impl(int size)
       : triggers(detail::gpuCallocHostUnique<ProxyTrigger>(size)),
         head(detail::gpuCallocUnique<uint64_t>()),
         tail(detail::gpuCallocHostUnique<uint64_t>()),
         tailCache(detail::gpuCallocUnique<uint64_t>()),
-        size(size) {}
+        size(size),
+        sizeMask(uint64_t(size) - 1),
+        sizeShift(shiftOf(size)) {}
 };
 
 MSCCLPP_API_CPP Fifo::Fifo(int size) {
+  if (size <= 0 || (size & (size - 1)) != 0) {
+    THROW(GPU, Error, ErrorCode::InvalidUsage, "FIFO size must be a positive power of two, got ", size);
+  }
   int device;
   MSCCLPP_CUDATHROW(cudaGetDevice(&device));
   int numaNode = getDeviceNumaNode(device);
@@ -38,19 +55,26 @@ MSCCLPP_API_CPP Fifo::Fifo(int size) {
 
 MSCCLPP_API_CPP Fifo::~Fifo() = default;
 
-MSCCLPP_API_CPP ProxyTrigger Fifo::poll() {
-  ProxyTrigger trigger;
-  ProxyTrigger* ptr = &pimpl_->triggers.get()[*(pimpl_->tail) % pimpl_->size];
-  // we are loading fst first. if fst is non-zero then snd is also valid
-  trigger.fst = atomicLoad(&(ptr->fst), memoryOrderAcquire);
-  trigger.snd = ptr->snd;
-  return trigger;
+MSCCLPP_API_CPP bool Fifo::poll(ProxyTrigger& trigger) {
+  const uint64_t curTail = *(pimpl_->tail);
+  ProxyTrigger* ptr = &pimpl_->triggers.get()[curTail & pimpl_->sizeMask];
+
+  // snd is the commit word: the producer writes it last, with release, carrying the parity of the
+  // lap this slot is on. A match means the payload written before it is visible too.
+  const uint64_t expectedParity = ((curTail >> pimpl_->sizeShift) & 1ULL) ^ 1ULL;
+  ProxyTrigger candidate;
+  candidate.snd = atomicLoad(&(ptr->snd), memoryOrderAcquire);
+  if (candidate.fields.reserved != expectedParity) return false;
+
+  candidate.fields.reserved = 0;
+  candidate.fst = atomicLoad(&(ptr->fst), memoryOrderRelaxed);
+  trigger = candidate;
+  return true;
 }
 
 MSCCLPP_API_CPP void Fifo::pop() {
-  uint64_t curTail = *(pimpl_->tail);
-  pimpl_->triggers.get()[curTail % pimpl_->size].fst = 0;
-  atomicStore(pimpl_->tail.get(), curTail + 1, memoryOrderRelease);
+  // The slot is not cleared: the next lap's parity makes what it holds stale.
+  atomicStore(pimpl_->tail.get(), *(pimpl_->tail) + 1, memoryOrderRelease);
 }
 
 MSCCLPP_API_CPP uint64_t Fifo::tail() const { return *(pimpl_->tail); }
@@ -64,6 +88,8 @@ MSCCLPP_API_CPP FifoDeviceHandle Fifo::deviceHandle() const {
   deviceHandle.tail = pimpl_->tail.get();
   deviceHandle.tailCache = pimpl_->tailCache.get();
   deviceHandle.size = pimpl_->size;
+  deviceHandle.sizeMask = pimpl_->sizeMask;
+  deviceHandle.sizeShift = pimpl_->sizeShift;
   return deviceHandle;
 }
 
diff --git a/src/core/proxy.cc b/src/core/proxy.cc
index 554336e8..d0982148 100644
--- a/src/core/proxy.cc
+++ b/src/core/proxy.cc
@@ -71,11 +71,9 @@ MSCCLPP_API_CPP void Proxy::start(bool blocking) {
       if (progressHandler) progressHandler();
 
       // Poll to see if we are ready to send anything
-      trigger = fifo->poll();
-      if (trigger.fst == 0 || trigger.snd == 0) {  // TODO: this check is a potential pitfall for custom triggers
-        continue;                                  // there is one in progress
+      if (!fifo->poll(trigger)) {
+        continue;  // no trigger has been committed to this slot yet
       }
-      trigger.snd ^= (uint64_t{1} << uint64_t{63});  // this is where the last bit of snd is reverted.
 
       ProxyHandlerResult result = handler(trigger);
 
diff --git a/test/unit/fifo_perf_tests.cu b/test/unit/fifo_perf_tests.cu
index 34b5d6bc..ffdb4452 100644
--- a/test/unit/fifo_perf_tests.cu
+++ b/test/unit/fifo_perf_tests.cu
@@ -36,14 +36,12 @@ static bool consumePerfTriggers(std::unique_ptr<mscclpp::Fifo>& hostFifo, int nu
   for (int i = 0; i < totalTriggers; ++i) {
     mscclpp::ProxyTrigger trigger;
     uint64_t spin = 0;
-    do {
-      trigger = hostFifo->poll();
+    while (!hostFifo->poll(trigger)) {
       if (spin++ > TIMEOUT_SPINS) {
         return false;
       }
-    } while (trigger.fst == 0 || trigger.snd == 0);
+    }
 
-    trigger.snd ^= ((uint64_t)1 << (uint64_t)63);
     trigger.snd = trigger.snd ^ trigger.fst;
     if (triggerCounts[trigger.snd] + 1 != trigger.fst) {
       return false;  // Validation failed
diff --git a/test/unit/fifo_tests.cu b/test/unit/fifo_tests.cu
index 8d30ca5e..45e3f151 100644
--- a/test/unit/fifo_tests.cu
+++ b/test/unit/fifo_tests.cu
@@ -17,7 +17,8 @@ __global__ void kernelFifoTest() {
 
   mscclpp::FifoDeviceHandle& fifo = gFifoTestFifoDeviceHandle;
   mscclpp::ProxyTrigger trigger;
-  for (uint64_t i = 1; i < ITER + 1; ++i) {
+  // Payloads start at 0: no trigger value is reserved by the FIFO any more.
+  for (uint64_t i = 0; i < ITER; ++i) {
     trigger.fst = i;
     trigger.snd = i;
     uint64_t curFifoHead = fifo.push(trigger);
@@ -45,24 +46,16 @@ TEST(FifoTest, Fifo) {
   MSCCLPP_CUDATHROW(cudaGetLastError());
 
   mscclpp::ProxyTrigger trigger;
-  trigger.fst = 0;
-  trigger.snd = 0;
-
   uint64_t spin = 0;
   mscclpp::Timer timer(3);
   for (uint64_t i = 0; i < ITER; ++i) {
-    trigger = hostFifo.poll();
-    while (trigger.fst == 0 || trigger.snd == 0) {
-      trigger = hostFifo.poll();
-
+    while (!hostFifo.poll(trigger)) {
       if (spin++ > 1000000) {
-        FAIL() << "Polling is stuck.";
+        FAIL() << "Polling is stuck at trigger " << i;
       }
     }
-    // see `src/proxy.cc` for the reason of this line
-    trigger.snd ^= ((uint64_t)1 << (uint64_t)63);
-    ASSERT_TRUE(trigger.fst == (i + 1));
-    ASSERT_TRUE(trigger.snd == (i + 1));
+    ASSERT_TRUE(trigger.fst == i);
+    ASSERT_TRUE(trigger.snd == i);
     hostFifo.pop();
     spin = 0;
   }
@@ -73,3 +66,96 @@ TEST(FifoTest, Fifo) {
 
   MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
 }
+
+__constant__ mscclpp::FifoDeviceHandle gFifoZeroTestHandle;
+
+// A trigger whose words are both zero must round-trip. Under the old protocol a zero first word
+// meant "not yet written", so this trigger was invisible and the FIFO stalled on its slot.
+__global__ void kernelFifoZeroTrigger(int count) {
+  if (threadIdx.x + blockIdx.x * blockDim.x != 0) return;
+  mscclpp::FifoDeviceHandle& fifo = gFifoZeroTestHandle;
+  for (int i = 0; i < count; ++i) {
+    mscclpp::ProxyTrigger trigger;
+    trigger.fst = 0;
+    trigger.snd = 0;
+    fifo.push(trigger);
+  }
+}
+
+TEST(FifoTest, ZeroTrigger) {
+  const int count = 32;
+  mscclpp::Fifo hostFifo;
+  mscclpp::FifoDeviceHandle devFifo = hostFifo.deviceHandle();
+  MSCCLPP_CUDATHROW(cudaMemcpyToSymbol(gFifoZeroTestHandle, &devFifo, sizeof(devFifo)));
+
+  kernelFifoZeroTrigger<<<1, 1>>>(count);
+  MSCCLPP_CUDATHROW(cudaGetLastError());
+
+  mscclpp::ProxyTrigger trigger;
+  for (int i = 0; i < count; ++i) {
+    uint64_t spin = 0;
+    while (!hostFifo.poll(trigger)) {
+      if (spin++ > 1000000) {
+        FAIL() << "Polling is stuck on a zero-valued trigger " << i;
+      }
+    }
+    ASSERT_TRUE(trigger.fst == 0);
+    ASSERT_TRUE(trigger.snd == 0);
+    hostFifo.pop();
+  }
+  MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
+}
+
+__constant__ mscclpp::FifoDeviceHandle gFifoWrapTestHandle;
+
+// Push exactly a whole number of laps. A parity polarity error shows up here: the consumer either
+// stalls at a lap boundary or accepts the previous lap's trigger a second time.
+__global__ void kernelFifoWrap(int laps, int fifoSize) {
+  if (threadIdx.x + blockIdx.x * blockDim.x != 0) return;
+  mscclpp::FifoDeviceHandle& fifo = gFifoWrapTestHandle;
+  for (int i = 0; i < laps * fifoSize; ++i) {
+    mscclpp::ProxyTrigger trigger;
+    trigger.fst = uint64_t(i);
+    trigger.snd = ~uint64_t(i);
+    trigger.fields.reserved = 0;  // the FIFO owns this bit
+    uint64_t head = fifo.push(trigger);
+    if ((i + 1) % fifoSize == 0) fifo.sync(head);
+  }
+}
+
+TEST(FifoTest, WrapAtLapBoundary) {
+  const int laps = 4;
+  mscclpp::Fifo hostFifo;
+  const int fifoSize = hostFifo.size();
+  mscclpp::FifoDeviceHandle devFifo = hostFifo.deviceHandle();
+  MSCCLPP_CUDATHROW(cudaMemcpyToSymbol(gFifoWrapTestHandle, &devFifo, sizeof(devFifo)));
+
+  kernelFifoWrap<<<1, 1>>>(laps, fifoSize);
+  MSCCLPP_CUDATHROW(cudaGetLastError());
+
+  mscclpp::ProxyTrigger trigger;
+  for (int i = 0; i < laps * fifoSize; ++i) {
+    uint64_t spin = 0;
+    while (!hostFifo.poll(trigger)) {
+      if (spin++ > 10000000) {
+        FAIL() << "Polling is stuck at position " << i << " (lap " << i / fifoSize << ")";
+      }
+    }
+    ASSERT_TRUE(trigger.fst == uint64_t(i));
+    mscclpp::ProxyTrigger expected;
+    expected.snd = ~uint64_t(i);
+    expected.fields.reserved = 0;
+    ASSERT_TRUE(trigger.snd == expected.snd);
+    hostFifo.pop();
+  }
+  MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
+}
+
+TEST(FifoTest, RejectsNonPowerOfTwoSize) {
+  try {
+    mscclpp::Fifo fifo(500);
+    FAIL() << "Expected a non-power-of-two FIFO size to throw";
+  } catch (const mscclpp::Error& e) {
+    EXPECT_TRUE(e.getErrorCode() == mscclpp::ErrorCode::InvalidUsage);
+  }
+}

From bdf9fd8762e826873d5b182a95a01493f5141daf Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 7 Aug 2026 18:26:47 +0000
Subject: [PATCH 08/12] update

---
 python/mscclpp/default_algos/__init__.py      |   9 +-
 .../reducescatter_multi_nodes.py              | 193 ++++++++++++++++++
 python/mscclpp/language/channel.py            |  75 +++++--
 src/core/executor/executor.cc                 |  75 ++++---
 4 files changed, 311 insertions(+), 41 deletions(-)
 create mode 100644 python/mscclpp/default_algos/reducescatter_multi_nodes.py

diff --git a/python/mscclpp/default_algos/__init__.py b/python/mscclpp/default_algos/__init__.py
index 32de1ac7..f613d0c7 100644
--- a/python/mscclpp/default_algos/__init__.py
+++ b/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",
+]
diff --git a/python/mscclpp/default_algos/reducescatter_multi_nodes.py b/python/mscclpp/default_algos/reducescatter_multi_nodes.py
new file mode 100644
index 00000000..9c245b30
--- /dev/null
+++ b/python/mscclpp/default_algos/reducescatter_multi_nodes.py
@@ -0,0 +1,193 @@
+# 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]
+                    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
+        # seven direct IB transfers can progress concurrently before this reduce.
+        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
diff --git a/python/mscclpp/language/channel.py b/python/mscclpp/language/channel.py
index 190e4b25..80c9b0a3 100644
--- a/python/mscclpp/language/channel.py
+++ b/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,7 +753,9 @@ 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
@@ -773,20 +781,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
diff --git a/src/core/executor/executor.cc b/src/core/executor/executor.cc
index c272b17a..b4268743 100644
--- a/src/core/executor/executor.cc
+++ b/src/core/executor/executor.cc
@@ -155,10 +155,10 @@ struct Executor::Impl {
   }
   ~Impl() = default;
 
-  ExecutionContext setupExecutionContext(int rank, void* sendbuff, void* recvbuff, size_t inputMessageSize,
-                                         size_t outputMessageSize, size_t constSrcOffset, size_t constDstOffset,
-                                         size_t sendMemRange, size_t recvMemRange, const ExecutionPlan& plan,
-                                         std::shared_ptr<ProxyService> proxyService) {
+  const ExecutionContext& setupExecutionContext(int rank, void* sendbuff, void* recvbuff, size_t inputMessageSize,
+                                                size_t outputMessageSize, size_t constSrcOffset, size_t constDstOffset,
+                                                size_t sendMemRange, size_t recvMemRange, const ExecutionPlan& plan,
+                                                std::shared_ptr<ProxyService> proxyService) {
     ExecutionContextKey key = {sendbuff, recvbuff, sendMemRange, recvMemRange, plan.impl_->name};
     DeviceExecutionPlanKey devicePlanKey = {inputMessageSize, outputMessageSize, constSrcOffset, constDstOffset};
 
@@ -166,24 +166,25 @@ struct Executor::Impl {
     if (plan.impl_->reuseResources) {
       key = {nullptr, nullptr, 0, 0, plan.impl_->name};
     }
-    if (this->contexts.find(key) != this->contexts.end()) {
-      auto& devicePlans = this->contexts[key].deviceExecutionPlans;
-      if (this->contexts[key].currentDevicePlan == devicePlanKey) {
-        return this->contexts[key];
+    auto contextIt = this->contexts.find(key);
+    if (contextIt != this->contexts.end()) {
+      auto& context = contextIt->second;
+      auto& devicePlans = context.deviceExecutionPlans;
+      if (context.currentDevicePlan == devicePlanKey) {
+        return context;
       } else if (devicePlans.find(devicePlanKey) != devicePlans.end()) {
-        this->contexts[key].currentDevicePlan = devicePlanKey;
-        return this->contexts[key];
+        context.currentDevicePlan = devicePlanKey;
+        return context;
       }
       plan.impl_->operationsReset();
       plan.impl_->lightLoadExecutionPlan(inputMessageSize, outputMessageSize, constSrcOffset, constDstOffset);
-      this->setupDeviceExecutionPlan(this->contexts[key], devicePlanKey, plan);
-      this->contexts[key].deviceExecutionPlansBuffers[devicePlanKey] =
+      this->setupDeviceExecutionPlan(context, devicePlanKey, plan);
+      context.deviceExecutionPlansBuffers[devicePlanKey] =
           GpuBuffer(devicePlans[devicePlanKey].size() * sizeof(DeviceExecutionPlan)).memory();
-      gpuMemcpy(this->contexts[key].deviceExecutionPlansBuffers[devicePlanKey].get(),
-                (char*)devicePlans[devicePlanKey].data(),
+      gpuMemcpy(context.deviceExecutionPlansBuffers[devicePlanKey].get(), (char*)devicePlans[devicePlanKey].data(),
                 devicePlans[devicePlanKey].size() * sizeof(DeviceExecutionPlan), cudaMemcpyHostToDevice);
-      this->contexts[key].currentDevicePlan = devicePlanKey;
-      return this->contexts[key];
+      context.currentDevicePlan = devicePlanKey;
+      return context;
     }
 
     plan.impl_->reset();
@@ -208,8 +209,11 @@ struct Executor::Impl {
               (char*)context.deviceExecutionPlans[devicePlanKey].data(),
               context.deviceExecutionPlans[devicePlanKey].size() * sizeof(DeviceExecutionPlan), cudaMemcpyHostToDevice);
     context.currentDevicePlan = devicePlanKey;
-    this->contexts.insert({key, context});
-    return context;
+    auto [insertedIt, inserted] = this->contexts.emplace(key, std::move(context));
+    if (!inserted) {
+      throw Error("Execution context insertion failed", ErrorCode::ExecutorError);
+    }
+    return insertedIt->second;
   }
 
   TransportFlags getTransportFlags(const BufferInfo& info, int rank) {
@@ -227,6 +231,24 @@ struct Executor::Impl {
     return flags;
   };
 
+  bool usesIbTransport(int rank, const ExecutionPlan& plan) {
+    auto hasIbPeer = [&](const std::vector<ChannelInfo>& channelInfos) {
+      for (const auto& info : channelInfos) {
+        for (int peer : info.connectedPeers) {
+          if (useIB(rank, peer, this->nranksPerIpcDomain)) {
+            return true;
+          }
+        }
+      }
+      return false;
+    };
+
+    if (hasIbPeer(plan.impl_->getChannelInfos(ChannelType::PORT))) {
+      return true;
+    }
+    return hasIbPeer(plan.impl_->getUnpairedChannelInfos(this->nranks, ChannelType::PORT));
+  }
+
   void setupScratchBuffer(ExecutionContext& context, size_t sendBuffSize, size_t recvBuffSize,
                           const ExecutionPlan& plan) {
     size_t scratchBufferSize = plan.impl_->calScratchBufferSize(std::min(sendBuffSize, plan.impl_->maxMessageSize),
@@ -309,9 +331,10 @@ struct Executor::Impl {
                                size_t recvBufferSize, int rank, const ExecutionPlan& plan) {
     // Add local src,dst and scratch to registeredMemoryIds
     context.localMemoryIdBegin = context.proxyService->nextMemoryId(3);
+    bool registerIb = this->usesIbTransport(rank, plan);
     for (auto& bufferType : {BufferType::INPUT, BufferType::OUTPUT, BufferType::SCRATCH}) {
       TransportFlags flags = Transport::CudaIpc;
-      if (hasIBDevices()) flags |= IBs[rank % this->nranksPerNode];
+      if (registerIb) flags |= IBs[rank % this->nranksPerNode];
       RegisteredMemory localMemory;
       auto bufferInfo = getBufferInfo(bufferType, sendbuff, recvbuff, context.scratchBuffer.get(), sendBufferSize,
                                       recvBufferSize, context.scratchBufferSize);
@@ -469,10 +492,10 @@ struct Executor::Impl {
   }
 
   template <typename PacketType>
-  void launchKernelHelper(ExecutionContext& context, int rank, void* sendbuff, void* recvbuff, DataType dataType,
+  void launchKernelHelper(const ExecutionContext& context, int rank, void* sendbuff, void* recvbuff, DataType dataType,
                           cudaStream_t stream, uint32_t sharedMemSize, const uint32_t& flag) {
     DeviceExecutionPlanKey key = context.currentDevicePlan;
-    int nthreadblocks = context.deviceExecutionPlans[key].size();
+    int nthreadblocks = context.deviceExecutionPlans.at(key).size();
     void* scratchBuffer = context.scratchBuffer.get();
     size_t scratchOffset = 0;
     if (context.doubleScratchBuff && (flag & 0x1) == 0) {
@@ -481,23 +504,23 @@ struct Executor::Impl {
     if (context.reuseResources) {
       ExecutionKernel::launchKernel<PacketType, true>(
           rank, nthreadblocks, context.nthreadsPerBlock, sendbuff, recvbuff, scratchBuffer, scratchOffset,
-          context.scratchChunkSize, dataType, (DeviceExecutionPlan*)context.deviceExecutionPlansBuffers[key].get(),
+          context.scratchChunkSize, dataType, (DeviceExecutionPlan*)context.deviceExecutionPlansBuffers.at(key).get(),
           (DeviceSemaphore*)context.smemaphores.get(), context.localMemoryIdBegin, sharedMemSize, stream, flag);
     } else {
       ExecutionKernel::launchKernel<PacketType, false>(
           rank, nthreadblocks, context.nthreadsPerBlock, sendbuff, recvbuff, scratchBuffer, scratchOffset,
-          context.scratchChunkSize, dataType, (DeviceExecutionPlan*)context.deviceExecutionPlansBuffers[key].get(),
+          context.scratchChunkSize, dataType, (DeviceExecutionPlan*)context.deviceExecutionPlansBuffers.at(key).get(),
           (DeviceSemaphore*)context.smemaphores.get(), context.localMemoryIdBegin, sharedMemSize, stream, flag);
     }
   }
 
-  void launchKernel(ExecutionContext& context, int rank, void* sendbuff, void* recvbuff, DataType dataType,
+  void launchKernel(const ExecutionContext& context, int rank, void* sendbuff, void* recvbuff, DataType dataType,
                     cudaStream_t stream, PacketType packetType) {
     static uint32_t flag = 0;
 #if defined(ENABLE_NPKIT)
 #if defined(MSCCLPP_USE_ROCM)
     DeviceExecutionPlanKey key = context.currentDevicePlan;
-    int nthreadblocks = context.deviceExecutionPlans[key].size();
+    int nthreadblocks = context.deviceExecutionPlans.at(key).size();
     if (nthreadblocks > NPKIT_MAX_NUM_GPU_THREADBLOCKS) {
       throw Error("Executor plan launching " + std::to_string(nthreadblocks) +
                       " thread blocks, exceeding NPKit support (" + std::to_string(NPKIT_MAX_NUM_GPU_THREADBLOCKS) +
@@ -536,7 +559,7 @@ void Executor::execute(int rank, void* sendbuff, void* recvbuff, size_t sendBuff
   size_t offsetIn = (char*)sendbuff - (char*)sendBasePtr;
   size_t offsetOut = (char*)recvbuff - (char*)recvBasePtr;
 
-  ExecutionContext context = this->impl_->setupExecutionContext(
+  const ExecutionContext& context = this->impl_->setupExecutionContext(
       rank, (void*)sendBasePtr, (void*)recvBasePtr, sendBuffSize, recvBuffSize, offsetIn, offsetOut, sendMemRange,
       recvMemRange, plan, this->impl_->proxyService);
   this->impl_->launchKernel(context, rank, sendbuff, recvbuff, dataType, stream, packetType);

From 2a843faf06f641cbde08458402e80cec3d2b6ed0 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Fri, 7 Aug 2026 19:00:07 +0000
Subject: [PATCH 09/12] Revert enhanced proxy FIFO changes

Restore the previous FIFO protocol because the parity-based commit scheme stalls PortChannel proxy traffic on TP64.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 60d476fd-fd9a-46ef-a7f8-c4d7fd05ef1d
---
 include/mscclpp/fifo.hpp                |  14 +--
 include/mscclpp/fifo_device.hpp         |  42 ++++-----
 include/mscclpp/port_channel_device.hpp |   4 +-
 src/core/fifo.cc                        |  48 +++-------
 src/core/proxy.cc                       |   6 +-
 test/unit/fifo_perf_tests.cu            |   6 +-
 test/unit/fifo_tests.cu                 | 112 +++---------------------
 7 files changed, 55 insertions(+), 177 deletions(-)

diff --git a/include/mscclpp/fifo.hpp b/include/mscclpp/fifo.hpp
index 4094036d..6aae03b5 100644
--- a/include/mscclpp/fifo.hpp
+++ b/include/mscclpp/fifo.hpp
@@ -16,21 +16,15 @@ constexpr size_t DEFAULT_FIFO_SIZE = 512;
 class Fifo {
  public:
   /// Constructor.
-  /// @param size Number of entries. Must be a power of two (default: DEFAULT_FIFO_SIZE).
-  /// @throws Error with ErrorCode::InvalidUsage if size is not a positive power of two.
+  /// @param size Number of entries (default: DEFAULT_FIFO_SIZE).
   Fifo(int size = DEFAULT_FIFO_SIZE);
 
   /// Destructor.
   ~Fifo();
 
-  /// Poll for the trigger at the head.
-  ///
-  /// A trigger carries no reserved payload value, so readiness is reported separately rather than
-  /// encoded in the trigger itself.
-  ///
-  /// @param trigger Set to the trigger at the head if one is ready. Untouched otherwise.
-  /// @return True if a trigger was ready and written to @p trigger.
-  bool poll(ProxyTrigger& trigger);
+  /// Poll and get the trigger at the head.
+  /// @return ProxyTrigger at the head of the FIFO.
+  ProxyTrigger poll();
 
   /// Remove the head trigger.
   void pop();
diff --git a/include/mscclpp/fifo_device.hpp b/include/mscclpp/fifo_device.hpp
index 294d5498..d5ae75f6 100644
--- a/include/mscclpp/fifo_device.hpp
+++ b/include/mscclpp/fifo_device.hpp
@@ -25,8 +25,6 @@ constexpr unsigned int TriggerBitsOffset = 32;
 constexpr unsigned int TriggerBitsMemoryId = 9;
 constexpr unsigned int TriggerBitsType = 3;
 constexpr unsigned int TriggerBitsSemaphoreId = 10;
-// The FIFO uses the reserved bit to mark a slot as written, so a trigger must not carry data
-// there. See FifoDeviceHandle::push().
 constexpr unsigned int TriggerBitsFifoReserved = 1;
 
 /// Pair of 64-bit unsigned integers used as a trigger for the proxy.
@@ -73,6 +71,7 @@ union alignas(16) ProxyTrigger {
     MSCCLPP_ASSERT_DEVICE(dstOffset < (1ULL << TriggerBitsOffset), "dstOffset is too large");
     MSCCLPP_ASSERT_DEVICE(srcId < (1ULL << TriggerBitsMemoryId), "srcId is too large");
     MSCCLPP_ASSERT_DEVICE(srcOffset < (1ULL << TriggerBitsOffset), "srcOffset is too large");
+    MSCCLPP_ASSERT_DEVICE(bytes != 0, "bytes must not be zero");
     MSCCLPP_ASSERT_DEVICE(bytes < (1ULL << TriggerBitsSize), "bytes is too large");
     MSCCLPP_ASSERT_DEVICE(semaphoreId < (1ULL << TriggerBitsSemaphoreId), "semaphoreId is too large");
     constexpr uint64_t maskSize = (1ULL << TriggerBitsSize) - 1;
@@ -107,32 +106,29 @@ struct FifoDeviceHandle {
   MSCCLPP_DEVICE_INLINE uint64_t push(ProxyTrigger trigger, int64_t maxSpinCount = 1000000) {
     uint64_t prevHead = atomicFetchAdd<uint64_t, scopeDevice>(head, 1, memoryOrderRelaxed);
 
-    // Wait until the slot's previous occupant has been consumed. Lap parity identifies a stale
-    // trigger but does not prevent overwriting a live one; this does.
+    // Flip the last bit for safe polling; host will revert.
+    constexpr uint64_t flipMask = uint64_t{1} << uint64_t{63};
+    trigger.snd ^= flipMask;
+
+    // Wait until the trigger is freed by the host.
     if (prevHead >= size + *tailCache) {
       sync(prevHead - size, maxSpinCount);
     }
 
-    // Commit bit: the parity of the lap this slot is on, so that the value left by the previous
-    // lap reads as stale. Lap 0 writes 1, which makes a zero-initialized buffer read as empty.
-    trigger.fields.reserved = ((prevHead >> sizeShift) & 1ULL) ^ 1ULL;
-
-    ProxyTrigger* triggerPtr = &(triggers[prevHead & sizeMask]);
+    ProxyTrigger* triggerPtr = &(triggers[prevHead % size]);
 
-    // snd is the commit word: a consumer that observes this lap's parity must also observe the
-    // payload that goes with it.
 #if defined(MSCCLPP_DEVICE_CUDA)
-    // One 128-bit store publishes both words together, so no ordering is needed between them.
-    // The proxy already relies on this: a torn store would let it read a new fst against the
-    // previous lap's snd, and dispatch on a stale semaphoreId.
-    //
-    // sm_80 used __threadfence_system() plus a relaxed store here, which was once faster. On
-    // A100 with CUDA 12.9 it is no longer, according to `FifoTest.Fifo`.
+#if __CUDA_ARCH__ == 800
+    // This is faster than release for A100.
+    __threadfence_system();
+    asm volatile("st.global.relaxed.sys.v2.u64 [%0], {%1,%2};" ::"l"(triggerPtr), "l"(trigger.fst), "l"(trigger.snd));
+#else
     asm volatile("st.global.release.sys.v2.u64 [%0], {%1,%2};" ::"l"(triggerPtr), "l"(trigger.fst), "l"(trigger.snd));
+#endif
 #else   // !defined(MSCCLPP_DEVICE_CUDA)
-    // No vector store here, so order the payload ahead of the commit explicitly.
-    atomicStore(&(triggerPtr->fst), trigger.fst, memoryOrderRelaxed);
-    atomicStore(&(triggerPtr->snd), trigger.snd, memoryOrderRelease);
+    // Store snd no later than fst.
+    atomicStore(&(triggerPtr->snd), trigger.snd, memoryOrderRelaxed);
+    atomicStore(&(triggerPtr->fst), trigger.fst, memoryOrderRelease);
 #endif  // !defined(MSCCLPP_DEVICE_CUDA)
 
     return prevHead;
@@ -172,12 +168,8 @@ struct FifoDeviceHandle {
   uint64_t* tail;
   /// Cached tail value.
   uint64_t* tailCache;
-  /// FIFO size. Always a power of two.
+  /// FIFO size.
   int size;
-  /// size - 1, for mapping a position to a slot.
-  uint64_t sizeMask;
-  /// log2(size), for extracting the lap from a position.
-  uint64_t sizeShift;
 };
 
 }  // namespace mscclpp
diff --git a/include/mscclpp/port_channel_device.hpp b/include/mscclpp/port_channel_device.hpp
index fd575b4c..74fa3d89 100644
--- a/include/mscclpp/port_channel_device.hpp
+++ b/include/mscclpp/port_channel_device.hpp
@@ -72,7 +72,7 @@ struct BasePortChannelDeviceHandle {
   }
 
   /// Push a TriggerFlag to the FIFO.
-  MSCCLPP_DEVICE_INLINE void signal() { fifo_.push({TriggerFlag, 0, 0, 0, 0, 0, semaphoreId_}); }
+  MSCCLPP_DEVICE_INLINE void signal() { fifo_.push({TriggerFlag, 0, 0, 0, 0, 1, semaphoreId_}); }
 
   /// Push a TriggerData and a TriggerFlag at the same time to the FIFO.
   /// @param dstId The ID of destination memory region.
@@ -122,7 +122,7 @@ struct BasePortChannelDeviceHandle {
   /// Push a TriggerSync to the FIFO.
   /// @param maxSpinCount The maximum number of spin counts before asserting. Never assert if negative.
   MSCCLPP_DEVICE_INLINE void flush(int64_t maxSpinCount = 1000000) {
-    uint64_t pos = fifo_.push({TriggerSync, 0, 0, 0, 0, 0, semaphoreId_});
+    uint64_t pos = fifo_.push({TriggerSync, 0, 0, 0, 0, 1, semaphoreId_});
     detail::waitFlush(flushDonePos_, pos, maxSpinCount);
   }
 
diff --git a/src/core/fifo.cc b/src/core/fifo.cc
index b44cfecc..b11775d8 100644
--- a/src/core/fifo.cc
+++ b/src/core/fifo.cc
@@ -8,42 +8,25 @@
 
 #include "api.h"
 #include "atomic.hpp"
-#include "logger.hpp"
 
 namespace mscclpp {
 
-namespace {
-// size is validated to be a positive power of two, so this is exact.
-uint64_t shiftOf(int size) {
-  uint64_t shift = 0;
-  while ((1 << shift) < size) shift++;
-  return shift;
-}
-}  // namespace
-
 struct Fifo::Impl {
   detail::UniqueGpuHostPtr<ProxyTrigger> triggers;
   detail::UniqueGpuPtr<uint64_t> head;
   detail::UniqueGpuHostPtr<uint64_t> tail;
   detail::UniqueGpuPtr<uint64_t> tailCache;
   const int size;
-  const uint64_t sizeMask;
-  const uint64_t sizeShift;
 
   Impl(int size)
       : triggers(detail::gpuCallocHostUnique<ProxyTrigger>(size)),
         head(detail::gpuCallocUnique<uint64_t>()),
         tail(detail::gpuCallocHostUnique<uint64_t>()),
         tailCache(detail::gpuCallocUnique<uint64_t>()),
-        size(size),
-        sizeMask(uint64_t(size) - 1),
-        sizeShift(shiftOf(size)) {}
+        size(size) {}
 };
 
 MSCCLPP_API_CPP Fifo::Fifo(int size) {
-  if (size <= 0 || (size & (size - 1)) != 0) {
-    THROW(GPU, Error, ErrorCode::InvalidUsage, "FIFO size must be a positive power of two, got ", size);
-  }
   int device;
   MSCCLPP_CUDATHROW(cudaGetDevice(&device));
   int numaNode = getDeviceNumaNode(device);
@@ -55,26 +38,19 @@ MSCCLPP_API_CPP Fifo::Fifo(int size) {
 
 MSCCLPP_API_CPP Fifo::~Fifo() = default;
 
-MSCCLPP_API_CPP bool Fifo::poll(ProxyTrigger& trigger) {
-  const uint64_t curTail = *(pimpl_->tail);
-  ProxyTrigger* ptr = &pimpl_->triggers.get()[curTail & pimpl_->sizeMask];
-
-  // snd is the commit word: the producer writes it last, with release, carrying the parity of the
-  // lap this slot is on. A match means the payload written before it is visible too.
-  const uint64_t expectedParity = ((curTail >> pimpl_->sizeShift) & 1ULL) ^ 1ULL;
-  ProxyTrigger candidate;
-  candidate.snd = atomicLoad(&(ptr->snd), memoryOrderAcquire);
-  if (candidate.fields.reserved != expectedParity) return false;
-
-  candidate.fields.reserved = 0;
-  candidate.fst = atomicLoad(&(ptr->fst), memoryOrderRelaxed);
-  trigger = candidate;
-  return true;
+MSCCLPP_API_CPP ProxyTrigger Fifo::poll() {
+  ProxyTrigger trigger;
+  ProxyTrigger* ptr = &pimpl_->triggers.get()[*(pimpl_->tail) % pimpl_->size];
+  // we are loading fst first. if fst is non-zero then snd is also valid
+  trigger.fst = atomicLoad(&(ptr->fst), memoryOrderAcquire);
+  trigger.snd = ptr->snd;
+  return trigger;
 }
 
 MSCCLPP_API_CPP void Fifo::pop() {
-  // The slot is not cleared: the next lap's parity makes what it holds stale.
-  atomicStore(pimpl_->tail.get(), *(pimpl_->tail) + 1, memoryOrderRelease);
+  uint64_t curTail = *(pimpl_->tail);
+  pimpl_->triggers.get()[curTail % pimpl_->size].fst = 0;
+  atomicStore(pimpl_->tail.get(), curTail + 1, memoryOrderRelease);
 }
 
 MSCCLPP_API_CPP uint64_t Fifo::tail() const { return *(pimpl_->tail); }
@@ -88,8 +64,6 @@ MSCCLPP_API_CPP FifoDeviceHandle Fifo::deviceHandle() const {
   deviceHandle.tail = pimpl_->tail.get();
   deviceHandle.tailCache = pimpl_->tailCache.get();
   deviceHandle.size = pimpl_->size;
-  deviceHandle.sizeMask = pimpl_->sizeMask;
-  deviceHandle.sizeShift = pimpl_->sizeShift;
   return deviceHandle;
 }
 
diff --git a/src/core/proxy.cc b/src/core/proxy.cc
index d0982148..554336e8 100644
--- a/src/core/proxy.cc
+++ b/src/core/proxy.cc
@@ -71,9 +71,11 @@ MSCCLPP_API_CPP void Proxy::start(bool blocking) {
       if (progressHandler) progressHandler();
 
       // Poll to see if we are ready to send anything
-      if (!fifo->poll(trigger)) {
-        continue;  // no trigger has been committed to this slot yet
+      trigger = fifo->poll();
+      if (trigger.fst == 0 || trigger.snd == 0) {  // TODO: this check is a potential pitfall for custom triggers
+        continue;                                  // there is one in progress
       }
+      trigger.snd ^= (uint64_t{1} << uint64_t{63});  // this is where the last bit of snd is reverted.
 
       ProxyHandlerResult result = handler(trigger);
 
diff --git a/test/unit/fifo_perf_tests.cu b/test/unit/fifo_perf_tests.cu
index ffdb4452..34b5d6bc 100644
--- a/test/unit/fifo_perf_tests.cu
+++ b/test/unit/fifo_perf_tests.cu
@@ -36,12 +36,14 @@ static bool consumePerfTriggers(std::unique_ptr<mscclpp::Fifo>& hostFifo, int nu
   for (int i = 0; i < totalTriggers; ++i) {
     mscclpp::ProxyTrigger trigger;
     uint64_t spin = 0;
-    while (!hostFifo->poll(trigger)) {
+    do {
+      trigger = hostFifo->poll();
       if (spin++ > TIMEOUT_SPINS) {
         return false;
       }
-    }
+    } while (trigger.fst == 0 || trigger.snd == 0);
 
+    trigger.snd ^= ((uint64_t)1 << (uint64_t)63);
     trigger.snd = trigger.snd ^ trigger.fst;
     if (triggerCounts[trigger.snd] + 1 != trigger.fst) {
       return false;  // Validation failed
diff --git a/test/unit/fifo_tests.cu b/test/unit/fifo_tests.cu
index 45e3f151..8d30ca5e 100644
--- a/test/unit/fifo_tests.cu
+++ b/test/unit/fifo_tests.cu
@@ -17,8 +17,7 @@ __global__ void kernelFifoTest() {
 
   mscclpp::FifoDeviceHandle& fifo = gFifoTestFifoDeviceHandle;
   mscclpp::ProxyTrigger trigger;
-  // Payloads start at 0: no trigger value is reserved by the FIFO any more.
-  for (uint64_t i = 0; i < ITER; ++i) {
+  for (uint64_t i = 1; i < ITER + 1; ++i) {
     trigger.fst = i;
     trigger.snd = i;
     uint64_t curFifoHead = fifo.push(trigger);
@@ -46,16 +45,24 @@ TEST(FifoTest, Fifo) {
   MSCCLPP_CUDATHROW(cudaGetLastError());
 
   mscclpp::ProxyTrigger trigger;
+  trigger.fst = 0;
+  trigger.snd = 0;
+
   uint64_t spin = 0;
   mscclpp::Timer timer(3);
   for (uint64_t i = 0; i < ITER; ++i) {
-    while (!hostFifo.poll(trigger)) {
+    trigger = hostFifo.poll();
+    while (trigger.fst == 0 || trigger.snd == 0) {
+      trigger = hostFifo.poll();
+
       if (spin++ > 1000000) {
-        FAIL() << "Polling is stuck at trigger " << i;
+        FAIL() << "Polling is stuck.";
       }
     }
-    ASSERT_TRUE(trigger.fst == i);
-    ASSERT_TRUE(trigger.snd == i);
+    // see `src/proxy.cc` for the reason of this line
+    trigger.snd ^= ((uint64_t)1 << (uint64_t)63);
+    ASSERT_TRUE(trigger.fst == (i + 1));
+    ASSERT_TRUE(trigger.snd == (i + 1));
     hostFifo.pop();
     spin = 0;
   }
@@ -66,96 +73,3 @@ TEST(FifoTest, Fifo) {
 
   MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
 }
-
-__constant__ mscclpp::FifoDeviceHandle gFifoZeroTestHandle;
-
-// A trigger whose words are both zero must round-trip. Under the old protocol a zero first word
-// meant "not yet written", so this trigger was invisible and the FIFO stalled on its slot.
-__global__ void kernelFifoZeroTrigger(int count) {
-  if (threadIdx.x + blockIdx.x * blockDim.x != 0) return;
-  mscclpp::FifoDeviceHandle& fifo = gFifoZeroTestHandle;
-  for (int i = 0; i < count; ++i) {
-    mscclpp::ProxyTrigger trigger;
-    trigger.fst = 0;
-    trigger.snd = 0;
-    fifo.push(trigger);
-  }
-}
-
-TEST(FifoTest, ZeroTrigger) {
-  const int count = 32;
-  mscclpp::Fifo hostFifo;
-  mscclpp::FifoDeviceHandle devFifo = hostFifo.deviceHandle();
-  MSCCLPP_CUDATHROW(cudaMemcpyToSymbol(gFifoZeroTestHandle, &devFifo, sizeof(devFifo)));
-
-  kernelFifoZeroTrigger<<<1, 1>>>(count);
-  MSCCLPP_CUDATHROW(cudaGetLastError());
-
-  mscclpp::ProxyTrigger trigger;
-  for (int i = 0; i < count; ++i) {
-    uint64_t spin = 0;
-    while (!hostFifo.poll(trigger)) {
-      if (spin++ > 1000000) {
-        FAIL() << "Polling is stuck on a zero-valued trigger " << i;
-      }
-    }
-    ASSERT_TRUE(trigger.fst == 0);
-    ASSERT_TRUE(trigger.snd == 0);
-    hostFifo.pop();
-  }
-  MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
-}
-
-__constant__ mscclpp::FifoDeviceHandle gFifoWrapTestHandle;
-
-// Push exactly a whole number of laps. A parity polarity error shows up here: the consumer either
-// stalls at a lap boundary or accepts the previous lap's trigger a second time.
-__global__ void kernelFifoWrap(int laps, int fifoSize) {
-  if (threadIdx.x + blockIdx.x * blockDim.x != 0) return;
-  mscclpp::FifoDeviceHandle& fifo = gFifoWrapTestHandle;
-  for (int i = 0; i < laps * fifoSize; ++i) {
-    mscclpp::ProxyTrigger trigger;
-    trigger.fst = uint64_t(i);
-    trigger.snd = ~uint64_t(i);
-    trigger.fields.reserved = 0;  // the FIFO owns this bit
-    uint64_t head = fifo.push(trigger);
-    if ((i + 1) % fifoSize == 0) fifo.sync(head);
-  }
-}
-
-TEST(FifoTest, WrapAtLapBoundary) {
-  const int laps = 4;
-  mscclpp::Fifo hostFifo;
-  const int fifoSize = hostFifo.size();
-  mscclpp::FifoDeviceHandle devFifo = hostFifo.deviceHandle();
-  MSCCLPP_CUDATHROW(cudaMemcpyToSymbol(gFifoWrapTestHandle, &devFifo, sizeof(devFifo)));
-
-  kernelFifoWrap<<<1, 1>>>(laps, fifoSize);
-  MSCCLPP_CUDATHROW(cudaGetLastError());
-
-  mscclpp::ProxyTrigger trigger;
-  for (int i = 0; i < laps * fifoSize; ++i) {
-    uint64_t spin = 0;
-    while (!hostFifo.poll(trigger)) {
-      if (spin++ > 10000000) {
-        FAIL() << "Polling is stuck at position " << i << " (lap " << i / fifoSize << ")";
-      }
-    }
-    ASSERT_TRUE(trigger.fst == uint64_t(i));
-    mscclpp::ProxyTrigger expected;
-    expected.snd = ~uint64_t(i);
-    expected.fields.reserved = 0;
-    ASSERT_TRUE(trigger.snd == expected.snd);
-    hostFifo.pop();
-  }
-  MSCCLPP_CUDATHROW(cudaDeviceSynchronize());
-}
-
-TEST(FifoTest, RejectsNonPowerOfTwoSize) {
-  try {
-    mscclpp::Fifo fifo(500);
-    FAIL() << "Expected a non-power-of-two FIFO size to throw";
-  } catch (const mscclpp::Error& e) {
-    EXPECT_TRUE(e.getErrorCode() == mscclpp::ErrorCode::InvalidUsage);
-  }
-}

From cefe801b3783cfc642c6e34f90a2ccde07d564bd Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Tue, 11 Aug 2026 18:35:02 +0000
Subject: [PATCH 10/12] update review comment

---
 .../default_algos/reducescatter_multi_nodes.py    | 15 ++++++++-------
 python/mscclpp/language/channel.py                |  1 +
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/python/mscclpp/default_algos/reducescatter_multi_nodes.py b/python/mscclpp/default_algos/reducescatter_multi_nodes.py
index 9c245b30..9a13a9d6 100644
--- a/python/mscclpp/default_algos/reducescatter_multi_nodes.py
+++ b/python/mscclpp/default_algos/reducescatter_multi_nodes.py
@@ -124,12 +124,13 @@ def reducescatter_multi_nodes(
                         local_packets.append(scratch_buffers[src_rank][scratch_slot : scratch_slot + 1])
 
                     local_reduced_chunk = input_buffer[chunk_index : chunk_index + 1]
-                    rank.reduce(
-                        local_reduced_chunk,
-                        local_packets,
-                        tb_group=thread_block_group,
-                        packet=True,
-                    )
+                    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:
@@ -159,7 +160,7 @@ def reducescatter_multi_nodes(
             return prog
 
         # Every rank receives one standard shard. The owner-node handoff and
-        # seven direct IB transfers can progress concurrently before this reduce.
+        # 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()
diff --git a/python/mscclpp/language/channel.py b/python/mscclpp/language/channel.py
index 80c9b0a3..fc148dde 100644
--- a/python/mscclpp/language/channel.py
+++ b/python/mscclpp/language/channel.py
@@ -760,6 +760,7 @@ def read_put_packets(
         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)

From 752715660744dfd0a2086e1e2d4186d652e94ede Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Tue, 11 Aug 2026 18:55:13 +0000
Subject: [PATCH 11/12] WIP

---
 src/core/executor/executor.cc | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/core/executor/executor.cc b/src/core/executor/executor.cc
index b4268743..8b3a0220 100644
--- a/src/core/executor/executor.cc
+++ b/src/core/executor/executor.cc
@@ -209,10 +209,7 @@ struct Executor::Impl {
               (char*)context.deviceExecutionPlans[devicePlanKey].data(),
               context.deviceExecutionPlans[devicePlanKey].size() * sizeof(DeviceExecutionPlan), cudaMemcpyHostToDevice);
     context.currentDevicePlan = devicePlanKey;
-    auto [insertedIt, inserted] = this->contexts.emplace(key, std::move(context));
-    if (!inserted) {
-      throw Error("Execution context insertion failed", ErrorCode::ExecutorError);
-    }
+    auto insertedIt = this->contexts.insert_or_assign(key, std::move(context)).first;
     return insertedIt->second;
   }
 

From 9ff8ad3c883b4f3fe949c52ae8ddad7abca69485 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Wed, 12 Aug 2026 04:32:38 +0000
Subject: [PATCH 12/12] add reset API for executor

---
 include/mscclpp/executor.hpp  |  6 ++++++
 python/csrc/executor_py.cpp   |  4 +++-
 python/test/test_mscclpp.py   | 19 +++++++++++++++++++
 src/core/executor/executor.cc |  7 +++++++
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/include/mscclpp/executor.hpp b/include/mscclpp/executor.hpp
index 4109844f..b2c6b55b 100644
--- a/include/mscclpp/executor.hpp
+++ b/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_;
diff --git a/python/csrc/executor_py.cpp b/python/csrc/executor_py.cpp
index 350a1e7a..43113b51 100644
--- a/python/csrc/executor_py.cpp
+++ b/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.");
 }
diff --git a/python/test/test_mscclpp.py b/python/test/test_mscclpp.py
index 6b3119cb..aa297354 100644
--- a/python/test/test_mscclpp.py
+++ b/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()
diff --git a/src/core/executor/executor.cc b/src/core/executor/executor.cc
index 8b3a0220..d7a81df7 100644
--- a/src/core/executor/executor.cc
+++ b/src/core/executor/executor.cc
@@ -562,6 +562,13 @@ void Executor::execute(int rank, void* sendbuff, void* recvbuff, size_t sendBuff
   this->impl_->launchKernel(context, rank, sendbuff, recvbuff, dataType, stream, packetType);
 }
 
+void Executor::reset() {
+  this->impl_->proxyService->stopProxy();
+  this->impl_->contexts.clear();
+  this->impl_->proxyService = std::make_shared<ProxyService>();
+  this->impl_->proxyService->startProxy(true);
+}
+
 Executor::~Executor() = default;
 
 }  // namespace mscclpp