cutlass/examples/python/CuTeDSL/distributed/all_reduce_two_shot_multimem.py

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