v4.5 tag update (#3202)

* Python DSL examples reorganization.

* v4.5 tag update.

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/README.md

@@ -0,0 +1,139 @@
+# CuTeDSL Distributed Examples
+
+This directory contains distributed examples using CuTeDSL with NVSHMEM for multi-GPU communication. Currently, we do not support to use NVSHMEM for any device side copy/put/get impl, only use the host side setup and allocations.
+
+## NVSHMEM Dependency
+
+These examples require two components:
+
+1. **NVSHMEM4Py** (`nvshmem4py-cu12` / `nvshmem4py-cu13`): A Python package that provides the official Python binding for NVIDIA's NVSHMEM. See the [NVSHMEM4Py Documentation](https://docs.nvidia.com/nvshmem/api/api/language_bindings/python/index.html).
+
+2. **NVSHMEM Library** (`nvidia-nvshmem-cu12` / `nvidia-nvshmem-cu13`): The underlying native library that contains the actual NVSHMEM implementation.
+
+### Overview
+
+**NVSHMEM4Py** (`nvshmem4py-cu12` / `nvshmem4py-cu13`) is a Python binding library that provides a Pythonic interface to NVSHMEM functionality. In these examples, we use it primarily for:
+
+- Allocating tensors that support peer-to-peer (P2P) communication across GPUs
+- Allocating multicast (MC) tensors that can leverage `multimem` instructions for efficient collective operations
+
+**nvidia-nvshmem** (`nvidia-nvshmem-cu12` / `nvidia-nvshmem-cu13`) is the underlying library that wraps NVSHMEM functions into dynamic libraries (`.so` files). NVSHMEM4Py dynamically loads and calls these libraries at runtime.
+
+### Installation
+
+For CUDA 12:
+```bash
+pip install nvshmem4py-cu12 nvidia-nvshmem-cu12
+```
+
+For CUDA 13:
+```bash
+pip install nvshmem4py-cu13 nvidia-nvshmem-cu13
+```
+
+> **Note:** `nvshmem4py` version >= 0.1.3 is recommended.
+
+### Key APIs Used
+
+We primarily use the following APIs from `nvshmem.core`:
+
+| API | Description |
+|-----|-------------|
+| `nvshmem.core.tensor(shape, dtype)` | Allocates a symmetric tensor that supports P2P communication |
+| `nvshmem.core.get_peer_tensor(tensor, pe)` | Returns a tensor handle for accessing the given tensor on a remote PE (processing element) |
+| `nvshmem.core.get_multicast_tensor(tensor)` | Returns a tensor that can be accessed using `multimem` instructions for efficient multicast operations |
+| `nvshmem.core.free_tensor(tensor)` | Explicitly frees the allocated symmetric memory |
+
+### Memory Management
+
+NVSHMEM requires **manual memory management**. Unlike PyTorch tensors that are garbage-collected automatically, NVSHMEM symmetric memory must be explicitly freed using `nvshmem.core.free_tensor()` to avoid memory leaks.
+
+Example:
+```python
+import nvshmem.core
+
+# init the environment
+# refer to the torchrun_uid_init_bcast() in example
+
+# Allocate symmetric tensor
+local_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+
+# Get peer tensors for P2P access
+tensor_list = [nvshmem.core.get_peer_tensor(local_tensor, rank) for rank in range(world_size)]
+
+# ... use tensors ...
+
+# Explicitly free memory when done
+for t in tensor_list:
+    nvshmem.core.free_tensor(t)
+
+# finalize the environment
+# refer to the torchrun_finalize() in example
+
+```
+
+## Multimem Instructions
+
+These examples demonstrate the use of NVIDIA's `multimem` PTX instructions for efficient multi-GPU collective operations. The `multimem` instructions operate on multicast (MC) addresses obtained via `nvshmem.core.get_multicast_tensor()`, enabling hardware-accelerated communication across multiple GPUs.
+
+### Why Multimem is Fast: NVLS (NVLink SHARP)
+
+The `multimem` instructions leverage **NVLS (NVLink SHARP)** technology to perform **in-network computation**. When multiple GPUs map the same symmetric memory region, `multimem` instructions can operate on a multicast address to perform hardware-accelerated reduction or broadcast operations directly in the NVLink/NVSwitch fabric, without requiring data to traverse to GPU memory first.
+
+**Key benefits:**
+- **In-network computation**: Reduction and broadcast operations happen in the NVSwitch hardware, not in GPU compute units
+- **Reduced memory traffic**: Data is processed in-flight within the interconnect, minimizing HBM bandwidth consumption
+- **Lower latency**: Single instruction replaces multiple loads/stores and arithmetic operations
+
+### Instruction Categories
+
+We use three types of `multimem` instructions in these examples:
+
+#### 1. `multimem.ld_reduce` - Reduction
+
+Reads data from a multicast address and returns the **reduced result** (e.g., sum) across all GPUs:
+
+```
+multimem.ld_reduce.sys.relaxed.global.add.v4.f32 {$0, $1, $2, $3}, [$4];
+```
+
+This instruction reads from a multicast address and performs a sum reduction (`.add`) across all GPUs that have mapped this address via NVLS.
+
+**Accumulator Precision**: For lower-precision data types, you can specify a higher accumulator precision to improve numerical accuracy:
+- **FP16 / BF16**: Can use FP32 accumulator (`.acc::f32`)
+- **FP8 (E4M3 / E5M2)**: Can use FP16 accumulator (`.acc::f16`)
+
+Example with FP16 using FP32 accumulator:
+```
+multimem.ld_reduce.sys.relaxed.global.add.acc::f32.v4.f16x2 {$0, $1, $2, $3}, [$4];
+```
+
+#### 2. `multimem.st` - Broadcast via Store
+
+Stores data to a multicast address, which **broadcasts** the data to all participating GPUs:
+
+```
+multimem.st.sys.relaxed.global.v4.f32 [$1], {$2, $3, $4, $5};
+```
+
+This writes data to a multicast address, and the data becomes visible to all GPUs that have mapped this address via NVLS.
+
+#### 3. `multimem.red` - Broadcast via Atomic Reduction
+
+Performs an atomic reduction operation on a multicast address. This is commonly used for **signaling/synchronization** across GPUs:
+
+```
+multimem.red.release.sys.global.add.u32 [$0], 1;
+```
+
+This atomically adds a value to a multicast address. When used with synchronization patterns (e.g., spin locks), it enables efficient inter-GPU barriers where all GPUs can observe the updated value.
+
+## Future Work
+
+The `nvidia-nvshmem-cu12/cu13` packages include LLVM IR bitcode libraries that could potentially be integrated into CuTeDSL in the future. This would enable calling NVSHMEM functions directly from within CuTeDSL kernels, allowing for more fine-grained control over communication patterns at the kernel level.
+
+## References
+
+- [NVSHMEM4Py Documentation](https://docs.nvidia.com/nvshmem/api/api/language_bindings/python/index.html)
+- [NVSHMEM API Reference](https://docs.nvidia.com/nvshmem/api/api/language_bindings/python/index.html)
+- [multimem PTX instruction](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-multimem)

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/all_reduce_one_shot_lamport.py

@@ -0,0 +1,426 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import os
+import torch
+import argparse
+
+import numpy as np
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+import cuda.bindings.driver as cuda
+from cuda.core.experimental import Device
+from cuda.pathfinder import load_nvidia_dynamic_lib
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T
+from cutlass._mlir.dialects import vector
+
+try:
+    import nvshmem.core
+except ImportError as exc:
+    raise ImportError(
+        "nvshmem4py is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvshmem4py-cu12\n"
+        "  For CUDA 13: pip install nvshmem4py-cu13\n"
+        "Note: nvshmem4py version >= 0.1.3 is recommended."
+    ) from None
+
+try:
+    load_nvidia_dynamic_lib("nvshmem_host")
+except RuntimeError as exc:
+    raise ImportError(
+        "nvshmem lib is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvidia-nvshmem-cu12\n"
+        "  For CUDA 13: pip install nvidia-nvshmem-cu13\n"
+    ) from None
+
+"""
+A Distributed One-Shot All-Reduce Example using CuTe DSL and fine-grained memory control. This is a mirrored version of the 
+existing tensorrt_llm kernel:
+https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
+
+In Lamport terminology this is a classic flag-based busy-wait: every participant keeps polling the shared slot until the
+flag changes from the sentinel (negative zero) to real data, which indicates that the Lamport-style logical ordering has
+advanced and the payload is safe to consume.
+
+This example kernel demonstrates a one-shot all-reduce operation using the CuTe DSL with fine-grained memory control.
+It uses dedicated communication buffers for data exchange, and these buffers act as ping-pong buffers. During the 
+process, the kernel uses one buffer for communication and initializes the next buffer to all negative zeros.
+
+In this kernel, each thread is only responsible for 128bits of data. The kernel will write it's local data to every
+buffer at different ranks, then read the data from the local rank buffer. The buffer itself behaves as a barrier, 
+if kernel read negtive 0, then it means data are not ready or not visible yet so that the kernel will read the data again.
+
+If the input tensors from each device are not remotely accessible, this kernel can be used to perform the one-shot all-reduce 
+since it uses communication buffers for data exchange.
+
+The .SYS memory scope and .VOLATILE memory order are used to ensure that the data will be visible at the system scope.
+
+.. code-block:: bash
+
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_one_shot_lamport.py --M 8192 --N 8192
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_one_shot_lamport.py \
+        --M 8192 --N 8192 --benchmark --warmup_iterations 2 --iterations 10
+"""
+
+
+PING_PONG_SIZE = 3
+
+
+class AllReduceOneShotLamportKernel:
+    @cute.jit
+    def __call__(
+        self,
+        rank: cutlass.Constexpr,
+        world_size: cutlass.Constexpr,
+        signal: cutlass.Int32,
+        local_input: cute.Tensor,
+        local_output: cute.Tensor,
+        buffers: list[cute.Tensor],
+        stream: cuda.CUstream,
+    ):
+        copy_bits = 128
+        dtype = local_input.element_type
+        vector_size = copy_bits // dtype.width
+
+        thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+        val_layout = cute.make_ordered_layout((1, vector_size), order=(1, 0))
+        tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+
+        grouped_buffers = [cute.group_modes(buffer, 0, 2) for buffer in buffers]
+        tiled_buffers = [
+            cute.zipped_divide(buffer, (tiler_mn, world_size, PING_PONG_SIZE))
+            for buffer in grouped_buffers
+        ]
+        tiled_input = cute.zipped_divide(local_input, tiler_mn)
+        tiled_output = cute.zipped_divide(local_output, tiler_mn)
+
+        self.kernel(
+            tiled_buffers,
+            tiled_input,
+            tiled_output,
+            thr_layout,
+            val_layout,
+            signal,
+            rank,
+        ).launch(
+            grid=[cute.size(tiled_input, mode=[1]), 1, 1],
+            block=[cute.size(tv_layout, mode=[0]), 1, 1],
+            stream=stream,
+        )
+
+    #  GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        buffers: list[cute.Tensor],
+        local_input: cute.Tensor,
+        local_output: cute.Tensor,
+        thr_layout: cute.Layout,
+        val_layout: cute.Layout,
+        signal: cutlass.Int32,
+        rank: cutlass.Constexpr,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, _, _ = cute.arch.block_idx()
+        ping = signal % 3
+        pong = (signal + 1) % 3
+
+        buffer_local = buffers[rank]
+        cta_coord = ((None, None), bidx)
+        local_tile_in = local_input[cta_coord]
+        local_tile_out = local_output[cta_coord]
+
+        ping_coord = (((None, None), None, ping), bidx)
+        pong_coord = (((None, None), None, pong), bidx)
+
+        read_buffer = buffer_local[ping_coord]
+        clear_buffer = buffer_local[pong_coord]
+
+        write_coord = (((None, None), rank, ping), bidx)
+        write_buffers = [buffer[write_coord] for buffer in buffers]
+
+        # assume all buffers have the same element type with input
+        copy_atom_load = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            buffers[0].element_type,
+            num_bits_per_copy=128,
+            memory_scope=cute.nvgpu.common.MemoryScope.SYS,
+            memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
+        )
+        copy_atom_store = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            buffers[0].element_type,
+            num_bits_per_copy=128,
+            memory_scope=cute.nvgpu.common.MemoryScope.SYS,
+            memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
+        )
+        tiled_copy = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
+        thr_copy = tiled_copy.get_slice(tidx)
+
+        thr_write_buffer_list = [
+            thr_copy.partition_D(tensor) for tensor in write_buffers
+        ]
+        thr_read_buffer = thr_copy.partition_S(read_buffer)
+
+        thr_clear_buffer = thr_copy.partition_D(clear_buffer)
+
+        thr_in = thr_copy.partition_S(local_tile_in)
+        thr_out = thr_copy.partition_D(local_tile_out)
+
+        frg_in = cute.make_fragment_like(thr_in)
+        frg_clear = cute.make_fragment_like(thr_clear_buffer)
+        frg_acc = cute.make_fragment_like(thr_out)
+        frg_acc.fill(0.0)
+
+        # clear a next buffer to be all negtive 0
+        clear_tensor = frg_clear.load()
+        frg_size = cute.size(clear_tensor.shape)
+        neg0_i32_vec = cute.full_like(clear_tensor, 0x80000000, cutlass.Int32)
+        neg0_f32_vec = vector.bitcast(T.vector(frg_size, T.f32()), neg0_i32_vec)
+        neg0_f32_tensor = cute.TensorSSA(
+            neg0_f32_vec, clear_tensor.shape, cutlass.Float32
+        )
+        frg_clear.store(neg0_f32_tensor)
+        cute.copy(copy_atom_store, frg_clear, thr_clear_buffer)
+
+        # read local data to the register
+        cute.copy(copy_atom_load, thr_in, frg_in)
+
+        # write local data to every buffer at different ranks
+        for thr_write_buffer in thr_write_buffer_list:
+            cute.copy(copy_atom_store, frg_in, thr_write_buffer)
+
+        frg_in_vector_neg0_i32 = cute.full_like(
+            frg_in, cutlass.Int32(0x80000000), cutlass.Int32
+        )
+        frg_in_size = cute.size(frg_in.shape)
+
+        # loop over each buffer and accumulate the data
+        for i in cutlass.range_constexpr(len(buffers)):
+            read_coord = (None, 0, 0, i)
+            cute.copy(copy_atom_load, thr_read_buffer[read_coord], frg_in[None, 0, 0])
+            frg_vector = frg_in.load()
+            frg_vector_i32 = cute.TensorSSA(
+                vector.bitcast(T.vector(frg_in_size, T.i32()), frg_vector),
+                frg_in.shape,
+                cutlass.Int32,
+            )
+            isNotNeg0 = cute.all_(
+                cute.TensorSSA(
+                    frg_vector_i32 != frg_in_vector_neg0_i32,
+                    frg_in.shape,
+                    cutlass.Boolean,
+                )
+            )
+            # if the data is negtive 0, it means data are not ready or not visible yet, so we need to read the data again
+            while not isNotNeg0:
+                cute.copy(
+                    copy_atom_load, thr_read_buffer[read_coord], frg_in[None, 0, 0]
+                )
+                frg_vector = frg_in.load()
+                frg_vector_i32 = cute.TensorSSA(
+                    vector.bitcast(T.vector(frg_in_size, T.i32()), frg_vector),
+                    frg_in.shape,
+                    cutlass.Int32,
+                )
+                isNotNeg0 = cute.all_(
+                    cute.TensorSSA(
+                        frg_vector_i32 != frg_in_vector_neg0_i32,
+                        frg_in.shape,
+                        cutlass.Boolean,
+                    )
+                )
+            frg_acc.store(frg_in.load() + frg_acc.load())
+
+        cute.copy(copy_atom_store, frg_acc, thr_out)
+
+
+def run_all_reduce_one_shot(
+    M,
+    N,
+    warmup_iterations=2,
+    iterations=10,
+    skip_ref_check=False,
+    benchmark=True,
+):
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+    if rank == 0:
+        print("\nRunning Elementwise Add test with:")
+        print(f"Tensor dimensions: [{M}, {N}]")
+        print(f"GPU count: {world_size}")
+
+    # init buffer tensors to be neg 0
+    local_buffer_tensor = nvshmem.core.tensor([PING_PONG_SIZE, world_size, M, N,], dtype=torch.float32).neg_()
+    buffer_tensor_list = [nvshmem.core.get_peer_tensor(local_buffer_tensor, rank).permute(2, 3, 1, 0) for rank in range(world_size)]
+    signal = cutlass.Int32(0)
+    input_tensor = torch.randn([M, N], device=f"cuda:{rank}")
+    output_tensor = torch.zeros([M, N], device=f"cuda:{rank}")
+    stream = cutlass.cuda.default_stream()
+    all_reduce_one_shot_lamport_kernel = AllReduceOneShotLamportKernel()
+
+    compiled_func = cute.compile(
+        all_reduce_one_shot_lamport_kernel,
+        rank,
+        world_size,
+        signal,
+        from_dlpack(input_tensor, assumed_align=32),
+        from_dlpack(output_tensor, assumed_align=32),
+        [from_dlpack(t, assumed_align=32) for t in buffer_tensor_list],
+        stream=stream,
+    )
+
+    if not skip_ref_check:
+        compiled_func(
+            signal,
+            from_dlpack(input_tensor, assumed_align=32),
+            from_dlpack(output_tensor, assumed_align=32),
+            [from_dlpack(t, assumed_align=32) for t in buffer_tensor_list],
+            stream,
+        )
+        if rank == 0:
+            print("Verifying results...")
+        dist.all_reduce(input_tensor, op=dist.ReduceOp.SUM)
+        dist.barrier(device_ids=[rank])
+        torch.testing.assert_close(input_tensor.cpu(), output_tensor.cpu())
+        if rank == 0:
+            print("Results verified successfully!")
+
+    for t in buffer_tensor_list:
+        nvshmem.core.free_tensor(t)
+
+    if not benchmark:
+        return
+
+    free_func_and_tensor_pairs = []
+    def add_free_func_and_tensor(free_func, tensor):
+        free_func_and_tensor_pairs.append((free_func, tensor))
+
+    def generate_tensors():
+        local_buffer = nvshmem.core.tensor([PING_PONG_SIZE, world_size, M, N,], dtype=torch.float32).neg_()
+        buffer_tensor_list = [nvshmem.core.get_peer_tensor(local_buffer, rank).permute(2, 3, 1, 0) for rank in range(world_size)]
+        input_tensor = torch.randn([M, N], device=f"cuda:{rank}")
+        output_tensor = torch.zeros([M, N], device=f"cuda:{rank}")
+
+        ja = testing.JitArguments(
+            cutlass.Int32(0),
+            from_dlpack(input_tensor, assumed_align=32),
+            from_dlpack(output_tensor, assumed_align=32),
+            [from_dlpack(t, assumed_align=32) for t in buffer_tensor_list],
+            stream=stream
+        )
+        for tensor in buffer_tensor_list:
+            add_free_func_and_tensor(nvshmem.core.free_tensor, tensor)
+        
+        return ja
+    avg_time_us = testing.benchmark(
+        compiled_func,
+        workspace_generator=generate_tensors,
+        workspace_count=10,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+    )
+
+    # Print execution results
+    if rank == 0:
+        print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
+        print(
+            f"Achieved memory throughput: {((world_size + 1) * output_tensor.numel() * 32 // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s"
+        )
+
+    for free_func, tensor in free_func_and_tensor_pairs:
+        free_func(tensor)
+
+def torchrun_uid_init_bcast():
+    """
+    Initialize NVSHMEM using UniqueID with `torchrun` as the launcher
+
+    It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting
+    """
+    # Set Torch device
+    local_rank = int(os.environ['LOCAL_RANK'])
+    torch.cuda.set_device(local_rank)
+
+    # nvshmem4py requires a cuda.core Device at init time
+    dev = Device(local_rank)
+    dev.set_current()
+    global stream
+    stream = dev.create_stream()
+
+    # Initialize torch.distributed process group
+    dist.init_process_group(
+        backend="cpu:gloo,cuda:nccl",
+    )
+
+    # Extract rank, nranks from process group
+    num_ranks = dist.get_world_size()
+
+    # Create an empty uniqueid for all ranks
+    uid = nvshmem.core.get_unique_id(empty=(local_rank != 0))
+    uid_bytes = uid._data.view(np.uint8).copy()
+    uid_tensor = torch.from_numpy(uid_bytes).cuda()
+    dist.broadcast(uid_tensor, src=0)
+    dist.barrier()
+    uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype)
+
+    nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid")
+
+
+def torchrun_finalize():
+    nvshmem.core.finalize()
+    dist.destroy_process_group()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="example of elementwise add to demonstrate the numpy/pytorch as input for kernels"
+    )
+    parser.add_argument("--M", default=1024, type=int)
+    parser.add_argument("--N", default=1024, type=int)
+    parser.add_argument("--warmup_iterations", default=2, type=int)
+    parser.add_argument("--iterations", default=10, type=int)
+    parser.add_argument("--skip_ref_check", action="store_true")
+    parser.add_argument("--benchmark", action="store_true")
+    args = parser.parse_args()
+
+    torchrun_uid_init_bcast()
+
+    run_all_reduce_one_shot(args.M, args.N, args.warmup_iterations, args.iterations, args.skip_ref_check, args.benchmark)
+
+    torchrun_finalize()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/all_reduce_simple.py

@@ -0,0 +1,312 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+import os
+import time
+import importlib
+import argparse
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from cuda.core.experimental import Device
+from cuda.pathfinder import load_nvidia_dynamic_lib
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+from cutlass.cute.runtime import from_dlpack
+
+try:
+    import nvshmem.core
+except ImportError as exc:
+    raise ImportError(
+        "nvshmem4py is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvshmem4py-cu12\n"
+        "  For CUDA 13: pip install nvshmem4py-cu13\n"
+        "Note: nvshmem4py version >= 0.1.3 is recommended."
+    ) from None
+
+try:
+    load_nvidia_dynamic_lib("nvshmem_host")
+except RuntimeError as exc:
+    raise ImportError(
+        "nvshmem lib is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvidia-nvshmem-cu12\n"
+        "  For CUDA 13: pip install nvidia-nvshmem-cu13\n"
+    ) from None
+
+"""
+A Distributed All-Reduce Addition Example using CuTe DSL and PyTorch Symmetric Memory.
+
+This example kernel demonstrates distributed all-reduce across multiple GPUs using the SIMT copy 
+of CuTe DSL and PyTorch's symmetric memory feature. Basic CuTe layout calculation is derived
+from the elementwise_add.py example.
+
+This kernel is a simple version of all-reduce. It will directly copy data from remote memory to 
+registers, then accumulate the data and finally store the accumulated data back to local global memory.
+If the input tensors from each device are remotely accessible, then this kernel can be used to perform the all-reduce.
+
+On the host side, we use `torch.distributed._symmetric_memory` to manage the symmetric memory. We use `symm_mem.empty`
+and `symm_mem.rendezvous` to create a symmetric tensor. Then we use `get_buffer` to get tensors that are accessible from all devices.
+In this way, we can hide the details of CUDA driver API calls to enable access to remote memory.
+
+.. code-block:: python
+
+    t = symm_mem.empty((M, N), device=torch.device(f"cuda:{rank}"))
+    hdl = symm_mem.rendezvous(t, dist.group.WORLD)
+    # get tensors from other devices from the symmetric memory
+    tensor_list = [hdl.get_buffer(rank, t.shape, t.dtype) for rank in range(world_size)]
+
+To run this example:
+
+.. code-block:: bash
+
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_simple.py --M 1024 --N 512
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_simple.py \
+        --M 1024 --N 1024 --benchmark --warmup_iterations 2 --iterations 100
+"""
+
+
+@cute.kernel
+def all_reduce_simple_kernel(
+    inputs: list[cute.Tensor],
+    gOut: cute.Tensor,
+    thr_layout: cute.Layout,
+    val_layout: cute.Layout,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    # slice for CTAs
+    # logical id -> address
+    blk_coord = ((None, None), bidx)
+    local_tile_out = gOut[blk_coord]
+    local_tile_list = [t[blk_coord] for t in inputs]
+
+    assert all(t.element_type == inputs[0].element_type for t in inputs)
+
+    copy_atom_load = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        inputs[0].element_type,
+    )
+    copy_atom_store = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        inputs[0].element_type,
+    )
+    tiled_copy = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
+    thr_copy = tiled_copy.get_slice(tidx)
+
+    thr_tensor_list = [thr_copy.partition_S(tensor) for tensor in local_tile_list]
+    thr_out = thr_copy.partition_D(local_tile_out)
+    frg_tensor_list = [cute.make_fragment_like(tensor) for tensor in thr_tensor_list]
+    frg_acc = cute.make_fragment_like(thr_out)
+    frg_acc.fill(0.0)
+
+    # load the frg at the same offset from all devices and accumulate the result in frg_acc
+    for thr, frg in zip(thr_tensor_list, frg_tensor_list):
+        cute.copy(copy_atom_load, thr, frg)
+        tmp = frg.load() + frg_acc.load()
+        frg_acc.store(tmp)
+
+    # copy from register memory to global memory
+    cute.copy(copy_atom_store, frg_acc, thr_out)
+
+
+@cute.jit
+def all_reduce_simple(
+    inputs: list[cute.Tensor], output: cute.Tensor, copy_bits: cutlass.Constexpr = 128
+):
+    dtype = inputs[0].element_type
+    vector_size = copy_bits // dtype.width
+
+    thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+    val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))
+    tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+
+    divided_inputs = [cute.zipped_divide(tensor, tiler_mn) for tensor in inputs]
+    gOut = cute.zipped_divide(output, tiler_mn)  # ((Tile),(Rest))
+    all_reduce_simple_kernel(
+        divided_inputs,
+        gOut,
+        thr_layout,
+        val_layout,
+    ).launch(
+        grid=[cute.size(gOut, mode=[1]), 1, 1],
+        block=[cute.size(tv_layout, mode=[0]), 1, 1],
+    )
+
+
+def run_all_reduce_simple(
+    M,
+    N,
+    warmup_iterations=2,
+    iterations=10,
+    skip_ref_check=False,
+    benchmark=True,
+):
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+    if rank == 0:
+        print("\nRunning Elementwise Add test with:")
+        print(f"Tensor dimensions: [{M}, {N}]")
+        print(f"GPU count: {world_size}")
+
+    local_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+    local_tensor.random_(0, 100)
+    tensor_list = [nvshmem.core.get_peer_tensor(local_tensor, rank) for rank in range(world_size)]
+    output = torch.zeros((M, N), device=f"cuda:{rank}")
+    
+    if rank == 0:
+        print("Compiling kernel with cute.compile ...")
+    start_time = time.time()
+    compiled_func = cute.compile(all_reduce_simple, [from_dlpack(t) for t in tensor_list], from_dlpack(output))
+    compilation_time = time.time() - start_time
+    if rank == 0:
+        print(f"Compilation time: {compilation_time:.4f} seconds")
+        print("Executing vector add kernel...")
+
+    if not skip_ref_check:
+        dist.barrier(device_ids=[rank])
+        compiled_func([from_dlpack(t) for t in tensor_list], from_dlpack(output))
+        if rank == 0:
+            print("Verifying results...")
+        dist.barrier(device_ids=[rank])
+        torch.testing.assert_close(sum([t.cpu() for t in tensor_list]), output.cpu())
+        if rank == 0:
+            print("Results verified successfully!")
+    
+    for t in tensor_list:
+        nvshmem.core.free_tensor(t)
+
+    if not benchmark:
+        return
+
+    free_func_and_tensor_pairs = []
+    def add_free_func_and_tensor(free_func, tensor):
+        free_func_and_tensor_pairs.append((free_func, tensor))
+
+    def generate_tensors():
+        local_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+        local_tensor.random_(0, 100)
+        tensor_list = [nvshmem.core.get_peer_tensor(local_tensor, rank) for rank in range(world_size)]
+        output = torch.zeros((M, N), device=f"cuda:{rank}")
+
+        ja = testing.JitArguments(
+            [from_dlpack(t) for t in tensor_list],
+            from_dlpack(output),
+        )
+        for tensor in tensor_list:
+            add_free_func_and_tensor(nvshmem.core.free_tensor, tensor)
+        return ja
+   
+    avg_time_us = testing.benchmark(
+        compiled_func,
+        workspace_generator=generate_tensors,
+        workspace_count=10,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+    )
+
+    # Print execution results
+    if rank == 0:
+        print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
+        print(
+            f"Achieved memory throughput: {((world_size + 1) * output.numel() * 32 // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s"
+        )
+        print(f"First few elements of result: \n{output[:3, :3]}")
+
+    for free_func, tensor in free_func_and_tensor_pairs:
+        free_func(tensor)
+
+
+def torchrun_uid_init_bcast():
+    """
+    Initialize NVSHMEM using UniqueID with `torchrun` as the launcher
+
+    It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting
+    """
+    # Set Torch device
+    local_rank = int(os.environ['LOCAL_RANK'])
+    torch.cuda.set_device(local_rank)
+
+    # nvshmem4py requires a cuda.core Device at init time
+    dev = Device(local_rank)
+    dev.set_current()
+    global stream
+    stream = dev.create_stream()
+
+    # Initialize torch.distributed process group
+    dist.init_process_group(
+        backend="cpu:gloo,cuda:nccl",
+    )
+
+    # Extract rank, nranks from process group
+    num_ranks = dist.get_world_size()
+
+    # Create an empty uniqueid for all ranks
+    uid = nvshmem.core.get_unique_id(empty=(local_rank != 0))
+    uid_bytes = uid._data.view(np.uint8).copy()
+    uid_tensor = torch.from_numpy(uid_bytes).cuda()
+    dist.broadcast(uid_tensor, src=0)
+    dist.barrier()
+    uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype)
+
+    nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid")
+
+
+def torchrun_finalize():
+    nvshmem.core.finalize()
+    dist.destroy_process_group()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="example of elementwise add to demonstrate the numpy/pytorch as input for kernels"
+    )
+    parser.add_argument("--M", default=1024, type=int)
+    parser.add_argument("--N", default=1024, type=int)
+    parser.add_argument("--warmup_iterations", default=2, type=int)
+    parser.add_argument("--iterations", default=10, type=int)
+    parser.add_argument("--skip_ref_check", action="store_true")
+    parser.add_argument("--benchmark", action="store_true")
+
+    args = parser.parse_args()
+
+    torchrun_uid_init_bcast()
+
+    run_all_reduce_simple(args.M, args.N, args.warmup_iterations, args.iterations, args.skip_ref_check, args.benchmark)
+
+    torchrun_finalize()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/all_reduce_tma.py

@@ -0,0 +1,688 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""
+A Distributed All-Reduce Example using TMA (Tensor Memory Accelerator).
+
+This example demonstrates distributed all-reduce across multiple GPUs using TMA
+for data movement. It serves as a tutorial for TMA-based distributed operations,
+not as a performance-optimized implementation.
+
+Tensor Semantics:
+    - Input:  Logical shape (world_size, S), where S is the per-rank tensor size
+    - Output: Logical shape (world_size, S), each rank gets the sum of all inputs
+
+Kernel Parameters:
+    - input:  List of world_size tensors, each with shape S (accessible via NVSHMEM)
+    - output: Single tensor with shape S, using multicast address for broadcast
+
+Algorithm (Two-Shot):
+    1. Each CTA loads data from all ranks at its assigned tile position (TMA Load)
+    2. Accumulates the data locally in registers
+    3. Stores the result via TMA multicast (broadcasts to all ranks)
+    4. Cross-GPU barrier ensures completion before kernel exit
+
+Tile Assignment:
+    - Total tiles = ceil(S / elems_per_cta)
+    - Each rank processes ceil(total_tiles / world_size) CTAs
+    - CTA i on rank r processes global_tile_id = r * ctas_per_rank + i
+
+TMA Usage Notes (for tutorial purposes, not perf-optimal):
+    - Uses 1D TMA load to load from remote GPU memory via NVSHMEM addresses
+    - Uses 1D TMA load to store to multicast address for broadcasting to all ranks
+    - Supports any input shape by flattening to 1D and tiling linearly
+    - Pipeline with 2 stages overlaps TMA loads across ranks
+
+To run this example:
+
+.. code-block:: bash
+
+    torchrun --nproc-per-node 8 examples/distributed/all_reduce_tma.py --shape 1024,1024
+    torchrun --nproc-per-node 8 examples/distributed/all_reduce_tma.py --shape 4,6,8,10,12
+"""
+
+import cutlass
+import cutlass.utils as utils
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync
+
+
+class AllReduceTmaKernel:
+    """
+    TMA-based distributed All-Reduce kernel.
+
+    This kernel performs an all-reduce operation across multiple GPUs using TMA
+    (Tensor Memory Accelerator) for efficient data movement.
+
+    Algorithm (Two-Shot):
+        1. Each CTA loads data from all ranks at its assigned tile position
+        2. Accumulates the data locally in registers
+        3. Stores the result via TMA multicast (broadcasts to all ranks)
+        4. Cross-GPU barrier ensures completion before kernel exit
+
+    The input/output tensors can be of any rank, as long as:
+        - All input tensors and output tensor share the same layout
+        - The layout is compact (no holes in memory)
+
+    We traverse the tensors linearly in codomain (physical offset) order,
+    which guarantees consistent logical coordinate access across all tensors.
+    """
+
+    _elems_per_cta: int = 128 * 128  # Elements processed per CTA
+    _tma_threads: int = 32
+    _consumer_threads: int = 128
+    _threads_per_cta: int = _tma_threads + _consumer_threads
+    _num_stages: int = 2
+
+    def __init__(self, dtype):
+        self.dtype = dtype
+
+        # SMEM layout shape (will be converted to Layout in JIT context)
+        self.smem_layout_shape = (self._elems_per_cta,)
+        self.tiler = (self._elems_per_cta,)
+
+        # TMA transaction bytes (computed from dtype size)
+        # dtype.width is in bits, divide by 8 to get bytes
+        self.tma_bytes = (dtype.width // 8) * self._elems_per_cta
+
+        # Dynamically create SharedStorage type based on dtype
+        elems = self._elems_per_cta
+        stages = self._num_stages
+
+        @cute.struct
+        class SharedStorage:
+            mbar_array: cute.struct.MemRange[cutlass.Int64, stages * 2]
+            smem_buffer: cute.struct.Align[
+                cute.struct.MemRange[dtype, elems * stages],  # stages 个 tile
+                128,
+            ]
+
+        self._SharedStorage = SharedStorage
+
+    @cute.jit
+    def __call__(
+        self,
+        input_tensors: list[cute.Tensor],
+        output_tensor_mc: cute.Tensor,
+        flag: cute.Tensor,
+        flag_mc: cute.Tensor,
+        local_rank: cutlass.Constexpr,
+        world_size: cutlass.Constexpr,
+    ):
+        """
+        Host-side JIT function: creates TMA descriptors and launches kernel.
+
+        Args:
+            input_tensors: List of input tensors from each rank (world_size tensors)
+            output_tensor_mc: Output tensor with multicast address
+            flag: Synchronization flag (local view)
+            flag_mc: Synchronization flag (multicast view)
+            local_rank: This rank's ID
+            world_size: Total number of ranks
+        """
+        # ======================================================================
+        # Layout validation
+        # ======================================================================
+        ref_layout = input_tensors[0].layout
+        ref_size = cute.size(ref_layout)
+        ref_cosize = cute.cosize(ref_layout)
+
+        # Check compact: size == cosize (no holes in memory)
+        assert ref_size == ref_cosize, (
+            f"Input tensor must be compact: size={ref_size}, cosize={ref_cosize}"
+        )
+        assert self.tma_bytes % 16 == 0, f"Not aligned to 16B, TMA should not be used."
+
+        # Check all input tensors have the same layout
+        for i in cutlass.range_constexpr(world_size):
+            assert input_tensors[i].layout == ref_layout, (
+                f"All input tensors must have the same layout. "
+                f"input_tensors[0].layout={ref_layout}, "
+                f"input_tensors[{i}].layout={input_tensors[i].layout}"
+            )
+
+        # Check output tensor has the same layout
+        assert output_tensor_mc.layout == ref_layout, (
+            f"Output tensor must have the same layout as input tensors. "
+            f"input layout={ref_layout}, output layout={output_tensor_mc.layout}"
+        )
+
+        # ======================================================================
+        # Extract tensor info
+        # ======================================================================
+        # Verify dtype matches
+        assert input_tensors[0].element_type == self.dtype, (
+            f"Input tensor dtype mismatch: expected {self.dtype}, "
+            f"got {input_tensors[0].element_type}"
+        )
+
+        total_elems = ref_size
+
+        # Flatten layout: treat tensor as 1D in codomain order
+        flat_layout = cute.make_layout((total_elems,))
+
+        # SMEM layout (created in JIT context)
+        smem_layout = cute.make_layout(self.smem_layout_shape)
+
+        # Create TMA load descriptors (one per rank)
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp()
+        tma_load_atoms = []
+        tma_load_tensors = []
+
+        for i in cutlass.range_constexpr(world_size):
+            flat_input = cute.make_tensor(input_tensors[i].iterator, flat_layout)
+            tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(
+                tma_load_op,
+                flat_input,
+                smem_layout,
+                self.tiler,
+            )
+            tma_load_atoms.append(tma_atom)
+            tma_load_tensors.append(tma_tensor)
+
+        # Create TMA store descriptor
+        tma_store_op = cpasync.CopyBulkTensorTileS2GOp()
+        flat_output = cute.make_tensor(output_tensor_mc.iterator, flat_layout)
+        tma_store_atom, tma_store_tensor = cpasync.make_tiled_tma_atom(
+            tma_store_op,
+            flat_output,
+            smem_layout,
+            self.tiler,
+        )
+
+        # Grid calculation
+        num_tiles_total = cute.ceil_div(total_elems, self._elems_per_cta)
+        ctas_per_rank = cute.ceil_div(num_tiles_total, world_size)
+
+        # SMEM size from SharedStorage
+        smem_bytes = self._SharedStorage.size_in_bytes()
+
+        # Launch kernel
+        self.kernel(
+            tma_load_atoms,
+            tma_load_tensors,
+            tma_store_atom,
+            tma_store_tensor,
+            flag,
+            flag_mc,
+            local_rank,
+            world_size,
+            num_tiles_total,
+            ctas_per_rank,
+        ).launch(
+            grid=[ctas_per_rank, 1, 1],
+            block=[self._threads_per_cta, 1, 1],
+            smem=smem_bytes,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        # TMA atoms and tensors for loading from each rank
+        tma_load_atoms: list[cute.CopyAtom],
+        tma_load_tensors: list[cute.Tensor],
+        # TMA atom and tensor for storing to multicast address
+        tma_store_atom: cute.CopyAtom,
+        tma_store_tensor: cute.Tensor,
+        # Synchronization flags
+        flag: cute.Tensor,
+        flag_mc: cute.Tensor,
+        # Rank info
+        local_rank: cutlass.Constexpr,
+        world_size: cutlass.Constexpr,
+        # Grid info for tile calculation
+        num_tiles_total: cutlass.Constexpr,
+        ctas_per_rank: cutlass.Constexpr,
+    ):
+        # ======================================================================
+        # Thread/Block indexing
+        # ======================================================================
+        tidx = cute.arch.thread_idx()[0]
+        bidx = cute.arch.block_idx()[0]
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+        # ======================================================================
+        # SMEM allocation
+        # ======================================================================
+        staged_smem_layout = cute.make_layout((self._elems_per_cta, self._num_stages))
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self._SharedStorage)
+        mbar_ptr = storage.mbar_array.data_ptr()
+        staged_smem_tensor = storage.smem_buffer.get_tensor(staged_smem_layout)
+
+        # ======================================================================
+        # TMA Pipeline setup
+        # ======================================================================
+        producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+        consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self._consumer_threads
+        )
+
+        tma_pipeline = pipeline.PipelineTmaAsync.create(
+            barrier_storage=mbar_ptr,
+            num_stages=self._num_stages,
+            producer_group=producer_group,
+            consumer_group=consumer_group,
+            tx_count=self.tma_bytes,
+            cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)),
+        )
+
+        global_tile_id = local_rank * ctas_per_rank + bidx
+
+        if global_tile_id < num_tiles_total:
+            # ======================================================================
+            # Warp 0: Producer - TMA Load from all ranks
+            # ======================================================================
+            if warp_idx == 0:
+                producer_state = pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Producer, self._num_stages
+                )
+
+                for rank_i in cutlass.range_constexpr(world_size):
+                    tma_pipeline.producer_acquire(producer_state)
+
+                    stage_idx = producer_state.index
+                    smem_tile = cute.slice_(staged_smem_tensor, (None, stage_idx))
+
+                    g_tensor_tiled = cute.zipped_divide(
+                        tma_load_tensors[rank_i], self.tiler
+                    )
+                    g_tile = g_tensor_tiled[(None,), global_tile_id]
+
+                    g_tile_flat = cute.group_modes(g_tile, 0, cute.rank(g_tile))
+                    s_tile_flat = cute.group_modes(smem_tile, 0, cute.rank(smem_tile))
+
+                    s_part, g_part = cute.nvgpu.cpasync.tma_partition(
+                        tma_load_atoms[rank_i],
+                        0,
+                        cute.make_layout(1),
+                        s_tile_flat,
+                        g_tile_flat,
+                    )
+
+                    cute.copy(
+                        tma_load_atoms[rank_i],
+                        g_part,
+                        s_part,
+                        tma_bar_ptr=tma_pipeline.producer_get_barrier(producer_state),
+                    )
+
+                    tma_pipeline.producer_commit(producer_state)
+                    producer_state.advance()
+
+            # ======================================================================
+            # Warp 1-4: Consumer - Load from smem, ADD, Store to smem
+            # ======================================================================
+            else:
+                consumer_tid = tidx - self._tma_threads
+
+                vec_size = 4
+                chunk_size = vec_size * self._consumer_threads
+
+                # ------------------------------------------------------------------
+                # Initialize accumulator using stage 0's layout
+                # ------------------------------------------------------------------
+                # (elems, stages) -> (elems,)
+                smem_tensor_wo_stage = cute.slice_(staged_smem_tensor, (None, 0))
+                # (elems,) -> ((thr_vec,), (num_chunks,))
+                smem_tensor_tiled_by_thr_vec = cute.zipped_divide(
+                    smem_tensor_wo_stage, (chunk_size,)
+                )
+                # ((thr_vec,), (num_chunks,)) -> (((vec, threads),), (num_chunks,))
+                smem_tensor_tiled_by_thr_vec_tiled_by_vec = cute.logical_divide(
+                    smem_tensor_tiled_by_thr_vec, (vec_size,)
+                )
+                # (((vec, threads),), (num_chunks,)) -> ((vec,), (num_chunks,))
+                per_thread_smem_tensor = cute.slice_(
+                    smem_tensor_tiled_by_thr_vec_tiled_by_vec,
+                    ((None, consumer_tid), None),
+                )
+
+                accum = cute.make_rmem_tensor(per_thread_smem_tensor.layout, self.dtype)
+                accum.fill(self.dtype(0.0))
+
+                # ------------------------------------------------------------------
+                # Main loop: load from SMEM and accumulate
+                # ------------------------------------------------------------------
+                consumer_state = pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Consumer, self._num_stages
+                )
+
+                for rank_i in cutlass.range_constexpr(world_size):
+                    tma_pipeline.consumer_wait(consumer_state)
+
+                    stage_idx = consumer_state.index
+                    smem_tile = cute.slice_(staged_smem_tensor, (None, stage_idx))
+
+                    # (elems,) -> ((thr_vec,), (num_chunks,))
+                    smem_tiled_by_thr_vec = cute.zipped_divide(smem_tile, (chunk_size,))
+                    # ((thr_vec,), (num_chunks,)) -> (((vec, threads),), (num_chunks,))
+                    smem_tiled_by_thr_vec_tiled_by_vec = cute.logical_divide(
+                        smem_tiled_by_thr_vec, (vec_size,)
+                    )
+                    # (((vec, threads),), (num_chunks,)) -> ((vec,), (num_chunks,))
+                    per_thread_smem_view = cute.slice_(
+                        smem_tiled_by_thr_vec_tiled_by_vec,
+                        ((None, consumer_tid), None),
+                    )
+
+                    fragment = per_thread_smem_view.load()
+                    accum.store(accum.load() + fragment)
+
+                    tma_pipeline.sync_object_empty.arrive(
+                        consumer_state.index, tma_pipeline.consumer_mask
+                    )
+                    consumer_state.advance()
+
+                # Store accumulated result back to SMEM (stage 0)
+                per_thread_smem_tensor.store(accum.load())
+
+            # ======================================================================
+            # Sync point: all warps meet here
+            # ======================================================================
+            cute.arch.sync_threads()
+
+            # ======================================================================
+            # Warp 0: TMA Store to multicast output
+            # ======================================================================
+            if warp_idx == 0:
+                # Fence to ensure SMEM writes are visible
+                cute.arch.fence_proxy("async.shared", space="cta")
+
+                smem_tile_out = cute.slice_(staged_smem_tensor, (None, 0))
+
+                g_output_tiled = cute.zipped_divide(tma_store_tensor, self.tiler)
+                g_output_tile = g_output_tiled[(None,), global_tile_id]
+
+                g_out_flat = cute.group_modes(
+                    g_output_tile, 0, cute.rank(g_output_tile)
+                )
+                s_out_flat = cute.group_modes(
+                    smem_tile_out, 0, cute.rank(smem_tile_out)
+                )
+
+                s_part, g_part = cute.nvgpu.cpasync.tma_partition(
+                    tma_store_atom,
+                    0,
+                    cute.make_layout(1),
+                    s_out_flat,
+                    g_out_flat,
+                )
+
+                cute.copy(tma_store_atom, s_part, g_part)
+                cute.arch.cp_async_bulk_commit_group()
+                cute.arch.cp_async_bulk_wait_group(0)
+
+        # ==================================================================
+        # Cross-GPU barrier synchronization (thread 0 only)
+        # ==================================================================
+        if tidx == 0:
+            sm_id_linear = (
+                cute.arch.block_idx()[0]
+                + cute.arch.block_idx()[1] * cute.arch.grid_dim()[0]
+                + cute.arch.block_idx()[2]
+                * cute.arch.grid_dim()[0]
+                * cute.arch.grid_dim()[1]
+            )
+
+            # Signal completion to all ranks
+            utils.distributed.multimem_red_add1(
+                flag_mc.iterator + sm_id_linear,
+                scope="sys",
+                order="release",
+            )
+
+            # The same idx ctas wait until all peer ranks' ctas complete
+            utils.distributed.spin_lock_atom_cas_relaxed_wait(
+                flag.iterator + sm_id_linear,
+                expected_val=world_size,
+                reset_val=0,
+                scope="sys",
+            )
+
+
+# =============================================================================
+# HOST-SIDE DRIVER CODE
+# =============================================================================
+
+import os
+import argparse
+import math
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from cuda.core.experimental import Device
+from cuda.pathfinder import load_nvidia_dynamic_lib
+
+from cutlass.cute.runtime import from_dlpack
+
+try:
+    import nvshmem.core
+except ImportError as exc:
+    raise ImportError(
+        "nvshmem4py is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvshmem4py-cu12\n"
+        "  For CUDA 13: pip install nvshmem4py-cu13\n"
+        "Note: nvshmem4py version >= 0.1.3 is recommended."
+    ) from None
+
+try:
+    load_nvidia_dynamic_lib("nvshmem_host")
+except RuntimeError as exc:
+    raise ImportError(
+        "nvshmem lib is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvidia-nvshmem-cu12\n"
+        "  For CUDA 13: pip install nvidia-nvshmem-cu13\n"
+    ) from None
+
+
+def torchrun_uid_init_bcast():
+    """Initialize NVSHMEM using UniqueID with torchrun as launcher."""
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+
+    dev = Device(local_rank)
+    dev.set_current()
+    global stream
+    stream = dev.create_stream()
+
+    dist.init_process_group(backend="cpu:gloo,cuda:nccl")
+    num_ranks = dist.get_world_size()
+
+    uid = nvshmem.core.get_unique_id(empty=(local_rank != 0))
+    uid_bytes = uid._data.view(np.uint8).copy()
+    uid_tensor = torch.from_numpy(uid_bytes).cuda()
+    dist.broadcast(uid_tensor, src=0)
+    dist.barrier()
+    uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype)
+
+    nvshmem.core.init(
+        device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid"
+    )
+
+
+def torchrun_finalize():
+    """Finalize NVSHMEM and destroy process group."""
+    nvshmem.core.finalize()
+    dist.destroy_process_group()
+
+
+def run_all_reduce_tma(
+    shape: tuple,
+    skip_ref_check: bool = False,
+):
+    """
+    Run the TMA-based All-Reduce kernel.
+
+    Args:
+        shape: Tensor shape tuple, e.g., (4, 6, 8, 10)
+        skip_ref_check: If True, skip reference result verification
+    """
+    local_rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    # Calculate total elements
+    total_elems = math.prod(shape)
+
+    if local_rank == 0:
+        print("\nRunning TMA All-Reduce test with:")
+        print(f"  Tensor shape: {shape}")
+        print(f"  Total elements: {total_elems}")
+        print(f"  GPU count: {world_size}")
+
+    # Allocate input tensor (symmetric memory, accessible from all ranks)
+    local_input_tensor = nvshmem.core.tensor(shape, dtype=torch.float32)
+    local_input_tensor.random_(0, 100)
+
+    # Get peer tensors (views into each rank's input)
+    peer_input_tensors = [
+        nvshmem.core.get_peer_tensor(local_input_tensor, r) for r in range(world_size)
+    ]
+
+    if local_rank == 0:
+        print(f"  Input tensor ptr: {local_input_tensor.data_ptr():#x}")
+
+    # Allocate output tensor with multicast address
+    local_output_tensor = nvshmem.core.tensor(shape, dtype=torch.float32)
+    local_output_tensor.fill_(0)
+    output_tensor_mc = nvshmem.core.get_multicast_tensor(
+        nvshmem.core.Teams.TEAM_NODE, local_output_tensor
+    )
+
+    # Allocate synchronization flags
+    # Flag size = ctas_per_rank (matches kernel's bidx indexing)
+    elems_per_cta = AllReduceTmaKernel._elems_per_cta
+    num_tiles = (total_elems + elems_per_cta - 1) // elems_per_cta
+    ctas_per_rank = (num_tiles + world_size - 1) // world_size
+    local_flag = nvshmem.core.tensor((ctas_per_rank,), dtype=torch.int32)
+    local_flag.fill_(0)
+    flag_mc = nvshmem.core.get_multicast_tensor(
+        nvshmem.core.Teams.TEAM_NODE, local_flag
+    )
+
+    if local_rank == 0:
+        print(f"  Number of tiles: {num_tiles}")
+        print(f"  CTAs per rank: {ctas_per_rank}")
+        print("Compiling kernel...")
+
+    # Create kernel instance and compile
+    kernel = AllReduceTmaKernel(cutlass.Float32)
+
+    compiled_func = cute.compile(
+        kernel,
+        [from_dlpack(t) for t in peer_input_tensors],
+        from_dlpack(output_tensor_mc),
+        from_dlpack(local_flag),
+        from_dlpack(flag_mc),
+        local_rank,
+        world_size,
+    )
+
+    if local_rank == 0:
+        print("Compilation successful!")
+
+    if not skip_ref_check:
+        if local_rank == 0:
+            print("Executing kernel...")
+
+        dist.barrier(device_ids=[local_rank])
+        compiled_func(
+            [from_dlpack(t) for t in peer_input_tensors],
+            from_dlpack(output_tensor_mc),
+            from_dlpack(local_flag),
+            from_dlpack(flag_mc),
+        )
+        dist.barrier(device_ids=[local_rank])
+
+        if local_rank == 0:
+            print("Verifying results...")
+
+        # Compute expected result: sum of all inputs
+        expected = sum([t.cpu() for t in peer_input_tensors])
+
+        # Compare with actual output
+        torch.testing.assert_close(expected, local_output_tensor.cpu())
+
+        if local_rank == 0:
+            print("Results verified successfully!")
+
+    # Cleanup
+    for i in range(world_size):
+        if i != local_rank:
+            nvshmem.core.free_tensor(peer_input_tensors[i])
+
+    nvshmem.core.free_tensor(output_tensor_mc)
+    nvshmem.core.free_tensor(flag_mc)
+    nvshmem.core.free_tensor(local_input_tensor)
+    nvshmem.core.free_tensor(local_output_tensor)
+    nvshmem.core.free_tensor(local_flag)
+
+
+def parse_shape(shape_str: str) -> tuple:
+    """
+    Parse shape string into tuple.
+    Examples:
+        "1024,1024" -> (1024, 1024)
+        "2,3,4,5,6,7,8" -> (2, 3, 4, 5, 6, 7, 8)
+    """
+    return tuple(int(x.strip()) for x in shape_str.split(","))
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="TMA-based distributed all-reduce example"
+    )
+    parser.add_argument(
+        "--shape",
+        default="1024,1024",
+        type=str,
+        help="Tensor shape as comma-separated values, e.g., '1024,1024' or 4,6,8,10,12'",
+    )
+    parser.add_argument(
+        "--skip_ref_check",
+        action="store_true",
+        help="Skip reference result verification",
+    )
+
+    args = parser.parse_args()
+    shape = parse_shape(args.shape)
+
+    torchrun_uid_init_bcast()
+    run_all_reduce_tma(
+        shape=shape,
+        skip_ref_check=args.skip_ref_check,
+    )
+    torchrun_finalize()
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/all_reduce_two_shot_multimem.py

@@ -0,0 +1,393 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+import os
+import time
+import argparse
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from cuda.core.experimental import Device
+from cuda.pathfinder import load_nvidia_dynamic_lib
+
+import cutlass
+import cutlass.utils as utils
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+import cutlass.torch as cutlass_torch
+from cutlass.cute.runtime import from_dlpack
+
+try:
+    import nvshmem.core
+except ImportError as exc:
+    raise ImportError(
+        "nvshmem4py is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvshmem4py-cu12\n"
+        "  For CUDA 13: pip install nvshmem4py-cu13\n"
+        "Note: nvshmem4py version >= 0.1.3 is recommended."
+    ) from None
+
+try:
+    load_nvidia_dynamic_lib("nvshmem_host")
+except RuntimeError as exc:
+    raise ImportError(
+        "nvshmem lib is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvidia-nvshmem-cu12\n"
+        "  For CUDA 13: pip install nvidia-nvshmem-cu13\n"
+    ) from None
+
+
+"""
+A Distributed Two-Shot All-Reduce Example using CuTe DSL and PyTorch Symmetric Memory.
+
+This example kernel demonstrates how to leverage the multimem feature to do a two-shot all-reduce.
+The multimem instruction is operated on symmetric memory, it can offload the broadcast and reduce 
+to the Nvlink Switch so that the nvlink traffic will be reduced.
+
+When calling a 'multimem.ld_reduce addrA', the corresponding data from each remote device will be sent to the NVLS 
+and return the reduced data as result. And for 'multimem.st dataA addrA', the data will be sent to the NVLS once and 
+the data will be broadcast to each remote device. So the memory traffic and instruction count is reduced by 8 times 
+with multimem.
+
+In this example, we are using two-shot styled all-reduce which means each device computes a portion
+of data and stores them to each device. Compared to the one-shot styled all-reduce, the two-shot one can 
+maximize the performance of throughput. The input and output are symmetric memory so we don't need extra 
+communication buffers here. We use the `sm_wise_inter_gpu_multimem_barrier` to synchronize the data 
+between each device. It is to make sure that each device has done the data transfer.
+
+To run this example:
+
+.. code-block:: bash
+
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_two_shot_multimem.py --M 1024 --N 512
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_two_shot_multimem.py \
+        --M 1024 --N 1024 --benchmark --warmup_iterations 2 --iterations 100
+"""
+
+
+@cute.kernel
+def all_reduce_multimem_kernel(
+    gIn: cute.Tensor,
+    gOut: cute.Tensor,
+    flag: cute.Tensor,
+    flag_mc: cute.Tensor,
+    thr_layout: cute.Layout,
+    val_layout: cute.Layout,
+    local_rank: cutlass.Constexpr,
+    world_size: cutlass.Constexpr,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    # slice for CTAs
+    # logical id -> address
+
+    num_ctas = cute.size(gIn, mode=[1])
+    chunk_size = num_ctas // world_size
+    blk_idx = local_rank * chunk_size + bidx
+
+    blk_coord = ((None, None), blk_idx)
+    local_tile_out = gOut[blk_coord]
+    local_tile_in = gIn[blk_coord]
+
+    assert gIn.element_type == gOut.element_type
+
+    copy_atom_load = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        gIn.element_type,
+        num_bits_per_copy=128,
+    )
+    tiled_copy = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
+    thr_copy = tiled_copy.get_slice(tidx)
+
+    thr_in = thr_copy.partition_S(local_tile_in)
+    thr_out = thr_copy.partition_D(local_tile_out)
+
+    (_, rest_m), _, _ = thr_in.shape
+    (_, rest_m_stride), _, _ = thr_in.stride
+
+    for i in cutlass.range_constexpr(rest_m):
+        x, y, z, w = utils.distributed.multimem_ld_reduce_4xf32(
+            thr_in[(None, i), 0, 0].iterator
+        )
+        utils.distributed.multimem_st_4xb32(
+            thr_out[(None, i), 0, 0].iterator, x, y, z, w
+        )
+
+    # Ensure all threads in cta have finish issue multimem.ld_reduce and multimem.st instructions
+    cute.arch.sync_threads()
+
+    if tidx == 0:
+        # Linear id of current SM.
+        sm_id_linear = (
+            cute.arch.block_idx()[0]
+            + cute.arch.block_idx()[1] * cute.arch.grid_dim()[0]
+            + cute.arch.block_idx()[2]
+            * cute.arch.grid_dim()[0]
+            * cute.arch.grid_dim()[1]
+        )
+        # Release flag with sys scope
+        utils.distributed.multimem_red_add1(
+            flag_mc.iterator + sm_id_linear,
+            scope="sys",
+            order="release",
+        )
+        # Relaxed spin-lock wait flag with sys scope
+        utils.distributed.spin_lock_atom_cas_relaxed_wait(
+            flag.iterator + sm_id_linear,
+            expected_val=world_size,
+            reset_val=0,
+            scope="sys",
+        )
+
+@cute.jit
+def all_reduce_multimem(
+    mIn: cute.Tensor,
+    mOut: cute.Tensor,
+    flag: cute.Tensor,
+    flag_mc: cute.Tensor,
+    local_rank: cutlass.Constexpr,
+    world_size: cutlass.Constexpr,
+    copy_bits: cutlass.Constexpr = 128,
+):
+    dtype = mIn.element_type
+    vector_size = copy_bits // dtype.width
+
+    # we choose a 128x128 tile for a CTA
+    thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+    val_layout = cute.make_ordered_layout((32, vector_size), order=(1, 0))
+    tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+
+    gIn = cute.zipped_divide(mIn, tiler_mn)
+    gOut = cute.zipped_divide(mOut, tiler_mn)
+
+    all_reduce_multimem_kernel(
+        gIn,
+        gOut,
+        flag,
+        flag_mc,
+        thr_layout,
+        val_layout,
+        local_rank,
+        world_size,
+    ).launch(
+        grid=[cute.size(gOut, mode=[1]) // world_size, 1, 1],
+        block=[cute.size(tv_layout, mode=[0]), 1, 1],
+    )
+
+
+def run_all_reduce_multimem(
+    M,
+    N,
+    warmup_iterations=2,
+    iterations=10,
+    skip_ref_check=False,
+    benchmark=True,
+):
+    local_rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    tile_m = 128
+    tile_n = 128
+
+    if local_rank == 0:
+        print("\nRunning Elementwise Add test with:")
+        print(f"Tensor dimensions: [{M}, {N}]")
+        print(f"GPU count: {world_size}")
+
+    local_input_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+    input_tensor = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_input_tensor)
+
+    local_output_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+    output_tensor = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_output_tensor)
+
+    local_flag = nvshmem.core.tensor((M*N//(tile_m*tile_n)), dtype=torch.int32)
+    flag_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_flag)
+    
+    if local_rank == 0:
+        print("Compiling kernel with cute.compile ...")
+    start_time = time.time()
+    compiled_func = cute.compile(
+        all_reduce_multimem,
+        from_dlpack(input_tensor),
+        from_dlpack(output_tensor),
+        from_dlpack(local_flag),
+        from_dlpack(flag_mc),
+        local_rank,
+        world_size,
+    )
+    compilation_time = time.time() - start_time
+    if local_rank == 0:
+        print(f"Compilation time: {compilation_time:.4f} seconds")
+        print("Executing all-reduce two shot multimem kernel...")
+
+    if not skip_ref_check:
+        dist.barrier(device_ids=[local_rank])
+        compiled_func(
+            from_dlpack(input_tensor),
+            from_dlpack(output_tensor),
+            from_dlpack(local_flag),
+            from_dlpack(flag_mc),
+        )
+        dist.barrier(device_ids=[local_rank])
+        if local_rank == 0:
+            print("Verifying results...")
+
+        local_buffers = [nvshmem.core.get_peer_tensor(local_input_tensor, local_rank) for local_rank in range(world_size)]
+        torch.testing.assert_close(sum([buffer.cpu() for buffer in local_buffers]), local_output_tensor.cpu())
+        if local_rank == 0:
+            print("Results verified successfully!")
+        for i in range(world_size):
+            if i != local_rank:
+                nvshmem.core.free_tensor(local_buffers[i])
+
+    # always free the multicast tensors first
+    nvshmem.core.free_tensor(input_tensor)
+    nvshmem.core.free_tensor(output_tensor)
+    nvshmem.core.free_tensor(flag_mc)
+    nvshmem.core.free_tensor(local_input_tensor)
+    nvshmem.core.free_tensor(local_output_tensor)
+    nvshmem.core.free_tensor(local_flag)
+
+    if not benchmark:
+        return
+    
+    free_func_and_tensor_pairs = []
+    def add_free_func_and_tensor(free_func, tensor):
+        free_func_and_tensor_pairs.append((free_func, tensor))
+
+    def generate_tensors():
+        local_input_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+        input_tensor_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_input_tensor)
+
+        local_output_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+        output_tensor_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_output_tensor)
+
+        local_flag = nvshmem.core.tensor((M*N//(tile_m*tile_n)), dtype=torch.int32)
+        flag_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_flag)
+
+        ja = testing.JitArguments(
+            from_dlpack(input_tensor_mc),
+            from_dlpack(output_tensor_mc),
+            from_dlpack(local_flag),
+            from_dlpack(flag_mc),
+        )
+        tensors_to_free = [input_tensor_mc, output_tensor_mc, flag_mc, local_input_tensor, local_output_tensor, local_flag]
+        for tensor in tensors_to_free:
+            add_free_func_and_tensor(nvshmem.core.free_tensor, tensor)
+        return ja
+
+    dist.barrier(device_ids=[local_rank])
+    avg_time_us = testing.benchmark(
+        compiled_func,
+        workspace_generator=generate_tensors,
+        workspace_count=10,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+    )
+    dist.barrier(device_ids=[local_rank])
+    torch.cuda.synchronize()
+
+    # Print execution results
+    if local_rank == 0:
+        print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
+        print(
+            f"Achieved memory throughput: {((world_size + 1) * output_tensor.numel() * 32 // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s"
+        )
+    
+    for free_func, tensor in free_func_and_tensor_pairs:
+        free_func(tensor)
+    return
+
+
+def torchrun_uid_init_bcast():
+    """
+    Initialize NVSHMEM using UniqueID with `torchrun` as the launcher
+
+    It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting
+    """
+    # Set Torch device
+    local_rank = int(os.environ['LOCAL_RANK'])
+    torch.cuda.set_device(local_rank)
+
+    # nvshmem4py requires a cuda.core Device at init time
+    dev = Device(local_rank)
+    dev.set_current()
+    global stream
+    stream = dev.create_stream()
+
+    # Initialize torch.distributed process group
+    dist.init_process_group(
+        backend="cpu:gloo,cuda:nccl",
+    )
+
+    # Extract rank, nranks from process group
+    num_ranks = dist.get_world_size()
+
+    # Create an empty uniqueid for all ranks
+    uid = nvshmem.core.get_unique_id(empty=(local_rank != 0))
+    uid_bytes = uid._data.view(np.uint8).copy()
+    uid_tensor = torch.from_numpy(uid_bytes).cuda()
+    dist.broadcast(uid_tensor, src=0)
+    dist.barrier()
+    uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype)
+
+    nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid")
+
+
+def torchrun_finalize():
+    nvshmem.core.finalize()
+    dist.destroy_process_group()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="example of elementwise add to demonstrate the numpy/pytorch as input for kernels"
+    )
+    parser.add_argument("--M", default=1024, type=int)
+    parser.add_argument("--N", default=1024, type=int)
+    parser.add_argument("--warmup_iterations", default=2, type=int)
+    parser.add_argument("--iterations", default=10, type=int)
+    parser.add_argument("--skip_ref_check", action="store_true")
+    parser.add_argument("--benchmark", action="store_true")
+
+    args = parser.parse_args()
+
+    torchrun_uid_init_bcast()
+
+    run_all_reduce_multimem(args.M, args.N, args.warmup_iterations, args.iterations, args.skip_ref_check, args.benchmark)
+
+    torchrun_finalize()
+
+    return
+
+
+if __name__ == "__main__":
+    main()