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