From 52ae719eda738e665e95cf4fa99409a7a650dd5c Mon Sep 17 00:00:00 2001 From: bangyu shen <94283495+shubaoyu2@users.noreply.github.com> Date: Tue, 2 Dec 2025 11:25:40 +0800 Subject: [PATCH] [examples][CuTeDSL] init commit for distirbuted examples (#2806) * init commit for distirbuted examples * better OOB protection * and try import to nvshmem for better error message and a READMME.md to introduce nvshmem and multimem instructions * add some lamport explanation * enhance f8 output and warn that f8 output can have nan in it * tell user why we need complicate data conversions in ref check part * tell user we don't support nvshmem device function --------- Co-authored-by: bangyus --- examples/python/CuTeDSL/distributed/README.md | 139 + .../all_reduce_one_shot_lamport.py | 426 +++ .../CuTeDSL/distributed/all_reduce_simple.py | 312 ++ .../all_reduce_two_shot_multimem.py | 393 +++ .../distributed_all_gather_gemm_blackwell.py | 2396 +++++++++++++++ .../distributed_gemm_all_reduce_blackwell.py | 2455 +++++++++++++++ ...stributed_gemm_reduce_scatter_blackwell.py | 2624 +++++++++++++++++ 7 files changed, 8745 insertions(+) create mode 100644 examples/python/CuTeDSL/distributed/README.md create mode 100644 examples/python/CuTeDSL/distributed/all_reduce_one_shot_lamport.py create mode 100644 examples/python/CuTeDSL/distributed/all_reduce_simple.py create mode 100644 examples/python/CuTeDSL/distributed/all_reduce_two_shot_multimem.py create mode 100644 examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py create mode 100644 examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py create mode 100644 examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py diff --git a/examples/python/CuTeDSL/distributed/README.md b/examples/python/CuTeDSL/distributed/README.md new file mode 100644 index 00000000..eaa519ef --- /dev/null +++ b/examples/python/CuTeDSL/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) diff --git a/examples/python/CuTeDSL/distributed/all_reduce_one_shot_lamport.py b/examples/python/CuTeDSL/distributed/all_reduce_one_shot_lamport.py new file mode 100644 index 00000000..80a13887 --- /dev/null +++ b/examples/python/CuTeDSL/distributed/all_reduce_one_shot_lamport.py @@ -0,0 +1,426 @@ +# Copyright (c) 2025 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() diff --git a/examples/python/CuTeDSL/distributed/all_reduce_simple.py b/examples/python/CuTeDSL/distributed/all_reduce_simple.py new file mode 100644 index 00000000..eb75db98 --- /dev/null +++ b/examples/python/CuTeDSL/distributed/all_reduce_simple.py @@ -0,0 +1,312 @@ +# Copyright (c) 2025 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() diff --git a/examples/python/CuTeDSL/distributed/all_reduce_two_shot_multimem.py b/examples/python/CuTeDSL/distributed/all_reduce_two_shot_multimem.py new file mode 100644 index 00000000..d6f3b653 --- /dev/null +++ b/examples/python/CuTeDSL/distributed/all_reduce_two_shot_multimem.py @@ -0,0 +1,393 @@ +# Copyright (c) 2025 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() diff --git a/examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py b/examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py new file mode 100644 index 00000000..8ea68dfd --- /dev/null +++ b/examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py @@ -0,0 +1,2396 @@ +# Copyright (c) 2025 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 argparse +import os +from typing import Optional, Tuple, Type, Union + +import numpy as np +import torch +import torch.distributed._symmetric_memory as symm_mem +import torch.distributed as dist +import cuda.bindings.driver as cuda +from cuda.bindings import driver +from cuda.core.experimental import Device +from cuda.pathfinder import load_nvidia_dynamic_lib + +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.distributed as dist_helpers +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +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 high-performance All gather + dense GEMM example for the NVIDIA Blackwell SM100 architecture +using CUTE DSL. +The example assumed to be run on multiple GPUs, the MNKL are the per-GPU problem size +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Support persistent tile scheduling to better overlap memory load/store with mma between tiles + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Input arguments to this example is same as dense_gemm_persistent.py. + +.. code-block:: bash + + torchrun --nproc_per_node=8 examples/distributed/distributed_all_gather_gemm_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu torchrun --nproc_per_node=8 examples/distributed/distributed_all_gather_gemm_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + + +Constraints are same as dense_gemm_persistent.py: +* A and C must be row-major +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), + see detailed valid dtype combinations in below PersistentDenseGemmKernel class documentation +* A/B tensor must have the same data type +* Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* Mma tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled +""" + + +class SyncNvlDevices: + """A class for synchronizing multiple NVIDIA devices using NVL (NVLink) barriers. + + This class implements a barrier synchronization mechanism for distributed computing + across multiple GPUs, ensuring all participating devices reach a synchronization point + before proceeding to the next phase of computation. + + :param num_of_parallelism: Total number of participating devices (NP) + :type num_of_parallelism: int + """ + + def __init__(self, num_of_parallelism: int): + """Initialize the synchronization class. + + :param num_of_parallelism: Total number of participating devices (NP) + :type num_of_parallelism: int + """ + self.num_of_parallelism = num_of_parallelism + + @cute.kernel + def kernel( + self, + device_idx: cutlass.Int32, + device_arrival_counters: cute.Tensor, # Shape: (num_of_parallelism,), dtype: int32 + iteration_flags: cute.Tensor, # Shape: (num_of_parallelism, 1), dtype: int32 (e.g., input_ready_flags) + ): + """ + Synchronizes participating devices and resets iteration flags. + + Args: + device_idx: The logical ID of the current device. + device_arrival_counters: Shared counters for barrier synchronization. + iteration_flags: Flags used by the subsequent kernel (will be reset to 0). + """ + tidx, _, _ = cute.arch.thread_idx() + + # Constants for barrier logic + val_one = cutlass.Int32(1) + val_zero = cutlass.Int32(0) + max_arrivals = cutlass.Int32(self.num_of_parallelism - 1) + + # --- Phase 1: Reset Iteration Flags (only thread 0) --- + # Assumes iteration_flags has shape (num_of_parallelism, 1) or similar + # If iteration_flags represents multiple sets of flags (like Iterations in C++), + # this loop needs adjustment. + for i in range(self.num_of_parallelism): + iteration_flags[i] = val_zero # Reset flags for the next kernel + + # --- Phase 2: Signal Arrival to Peers (only thread 0) --- + for peer_rank in range(self.num_of_parallelism): + # Signal arrival to all *other* devices + if peer_rank != device_idx: + counter_ptr = cute.make_ptr( + cutlass.Int32, + device_arrival_counters[peer_rank], + cute.AddressSpace.gmem, + assumed_align=4, + ) + _ = utils.distributed.atomicAdd(counter_ptr, val_one) + + # --- Phase 3: Wait for Arrivals (only thread 0) --- + # Poll local arrival counter + local_counter_ptr = cute.make_ptr( + cutlass.Int32, + device_arrival_counters[device_idx], + cute.AddressSpace.gmem, + assumed_align=4, + ) + local_counter_tensor = cute.make_tensor( + local_counter_ptr, cute.make_layout(shape=(1,), stride=(1,)) + ) + current_arrivals = utils.distributed.ld_bypass(local_counter_tensor)[0] + # Wait until NP-1 peers have signaled + # Add a small sleep/yield if supported and necessary to prevent excessive polling load + while current_arrivals < max_arrivals: + # Consider adding a small delay here if possible (__nanosleep equivalent) + # For now, just poll: + current_arrivals = utils.distributed.ld_bypass(local_counter_tensor)[0] + + # --- Phase 4: Reset Local Counter (only thread 0) --- + # Atomically subtract (NP-1) from the local counter + reset_val = cutlass.Int32(-(self.num_of_parallelism - 1)) + _ = utils.distributed.atomicAdd(local_counter_ptr, reset_val) + + @cute.jit + def __call__( + self, + device_idx: cutlass.Int32, + device_arrival_counters: cute.Tensor, + iteration_flags: cute.Tensor, + stream: cuda.CUstream, + ): + """Execute the synchronization kernel. + + :param device_idx: The logical ID of the current device + :type device_idx: cutlass.Int32 + :param device_arrival_counters: Shared counters for barrier synchronization + :type device_arrival_counters: cute.Tensor + :param iteration_flags: Flags used by the subsequent kernel (will be reset to 0) + :type iteration_flags: cute.Tensor + :param stream: CUDA stream for kernel execution + :type stream: cuda.CUstream + """ + grid = [1, 1, 1] + self.kernel(device_idx, device_arrival_counters, iteration_flags).launch( + grid=grid, block=[1, 1, 1], cluster=(1, 1, 1), smem=0, stream=stream + ) + +def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + c_smem_layout: Union[cute.Layout, None], +) -> Tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + :param c_smem_layout: Layout of C operand in shared memory, or None if not using TMA store. + :type c_smem_layout: Union[cute.Layout, None] + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C with 1-stage + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1 + ) + + ab_bytes_per_stage = cute.size_in_bytes( + a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes + # Add remaining unused smem to epilogue + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage + + +class PersistentDenseGemmKernel: + """This class implements batched matrix multiplication (C = A x B) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Whether to use Tensor Memory Access (TMA) for storing results + :type use_tma_store: bool + + :note: In current version, A and B tensor must have the same data type + - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + + :note: Supported A/B data types: + - TFloat32 + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported C data types: + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + + Example: + >>> gemm = PersistentDenseGemmKernel( + ... acc_dtype=cutlass.Float32, + ... use_2cta_instrs=True, + ... mma_tiler_mn=(128, 128), + ... cluster_shape_mn=(2, 2) + ... ) + >>> gemm(a_tensor, b_tensor, c_tensor, max_active_clusters, stream) + """ + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_tma_store: bool, + gated_a_load: bool, + ): + """Initializes the configuration for a Blackwell dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant + with cta_group=2 should be used. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + 3. Output C tensor store mode: + - use_tma_store: Boolean indicating whether to use Tensor Memory Access (TMA) for storing results. + + 4. Gated A load: + - gated_a_load: Boolean indicating whether to gating A load in the kernel. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Use Tensor Memory Access (TMA) or normal store for output C tensor. + :type use_tma_store: bool + """ + + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.use_tma_store = use_tma_store + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilog_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_cta = 32 * len( + (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id) + ) + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.gated_a_load = gated_a_load + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C stage counts in shared memory + - Computing A/B/C shared memory layout + - Computing tensor memory allocation columns + """ + # Configure tiled mma + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + # Compute epilogue subtile + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + c_smem_layout = None + if cutlass.const_expr(self.use_tma_store): + c_smem_layout = sm100_utils.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.smem_capacity, + self.occupancy, + self.use_tma_store, + c_smem_layout, + ) + + # Compute A/B/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = None + if self.use_tma_store: + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + # Compute the number of tensor memory allocation columns + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + flag_offset: int, + gate_a_flags: cute.Tensor, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + :raises AssertionError: If OOB (Out-Of-Bounds) tiles are present when TMA store is disabled. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + + # Specific requirements from ring based all gather + a_layout_enum = utils.LayoutEnum.from_tensor(a) + c_layout_enum = utils.LayoutEnum.from_tensor(c) + + if cutlass.const_expr(a_layout_enum != utils.LayoutEnum.ROW_MAJOR): + raise TypeError(f"A must be row-major: {a_layout_enum}") + if cutlass.const_expr(c_layout_enum != utils.LayoutEnum.ROW_MAJOR): + raise TypeError(f"C must be row-major: {c_layout_enum}") + self.a_major_mode = a_layout_enum.mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + # Setup TMA store for C + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + + # Launch the kernel synchronously + gate_a_flag_tensor = cute.make_tensor( + gate_a_flags.iterator + flag_offset, + cute.make_layout((1,), stride=(1,)), + ) + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + gate_a_flag_tensor, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: Optional[cute.CopyAtom], + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + gate_a_flag: cute.Tensor, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.use_tma_store): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + lane_idx = cute.arch.lane_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + tmem_dealloc_barrier = None + if cutlass.const_expr(not self.use_tma_store): + tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_dealloc_sync_bar_id, + num_threads=32 * len(self.epilog_warp_id), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + # + # Setup smem tensor A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ######################################################### + if self.gated_a_load: + # wait_on_flag + if lane_idx == 0: + # Wait until input_ready_flags[iteration_i] is set + ready_flag = utils.distributed.ld_bypass(gate_a_flag)[0] + # Need to use volatile load to prevent compiler optimizing away the polling loop + while ready_flag == 0: + # Keep polling until ready using volatile load + ready_flag = utils.distributed.ld_bypass(gate_a_flag)[0] + cute.arch.sync_warp() + ######################################################### + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + # + # Tma load loop + # + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_producer.tail() + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + + # tCtAcc += tCrA * tCrB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + handle.release() + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + sC = None + if cutlass.const_expr(self.use_tma_store): + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # Persistent tile scheduling loop for epilogue + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + if cutlass.const_expr(self.use_tma_store): + # {$nv-internal-release begin} + # (EPI_TILE_M, EPI_TILE_N, STAGE) + # sC = storage.sC.get_tensor( + # c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + # ) + # {$nv-internal-release end} + + assert tma_atom_c is not None and sC is not None + self.epilogue_tma_store( + tidx, + warp_idx, + acc_pipeline, + tiled_mma, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + tile_sched, + epilogue_op, + ) + else: + self.epilogue( + tidx, + acc_pipeline, + tiled_mma, + tCtAcc_base, + tCgC, + epi_tile, + tile_sched, + epilogue_op, + tmem_dealloc_barrier, + ) + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + @cute.jit + def epilogue_tma_store( + self, + epi_tidx: cutlass.Int32, + warp_idx: cutlass.Int32, + acc_pipeline: pipeline.PipelineAsync, + tiled_mma: cute.TiledMma, + tma_atom_c: cute.CopyAtom, + # Input of epilogue + tCtAcc_base: cute.Tensor, + # Staging of epilogue + sC: cute.Tensor, + # Output of epilogue + tCgC: cute.Tensor, + epi_tile: cute.Tile, + tile_sched: utils.StaticPersistentTileScheduler, + epilogue_op: cutlass.Constexpr, + ) -> None: + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + tCgC_epi = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], epi_tile + ) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=32 * len(self.epilog_warp_id), + ) + + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilog_sync_barrier.arrive_and_wait() + + epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Wait for C store complete + c_pipeline.producer_tail() + + @cute.jit + def epilogue( + self, + epi_tidx: cutlass.Int32, + acc_pipeline: pipeline.PipelineAsync, + tiled_mma: cute.TiledMma, + tCtAcc_base: cute.Tensor, + tCgC: cute.Tensor, + epi_tile: cute.Tile, + tile_sched: utils.StaticPersistentTileScheduler, + epilogue_op: cutlass.Constexpr, + tmem_dealloc_barrier: pipeline.NamedBarrier, + ) -> None: + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs + ) + + gC_epi = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tTR_gC_partitioned = thr_copy_t2r.partition_D(gC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_rmem_tensor( + tTR_gC_partitioned[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype + ) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_gC = tTR_gC_partitioned[ + (None, None, None, None, None, *mma_tile_coord_mnl) + ] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tTR_rAcc.load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tTR_rC.store(acc_vec) + + # + # Store C to global memory + # + cute.copy(simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)]) + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def _compute_num_tmem_alloc_cols( + tiled_mma: cute.TiledMma, + mma_tiler: Tuple[int, int, int], + num_acc_stage: int, + ) -> int: + """ + Compute the number of tensor memory allocation columns. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler: The shape (M, N, K) of the MMA tile. + :type mma_tiler: tuple[int, int, int] + :param num_acc_stage: The stage of the accumulator tensor. + :type num_acc_stage: int + + :return: The number of tensor memory allocation columns. + :rtype: int + """ + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + + return num_tmem_alloc_cols + + def is_valid_dtypes( + self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] + ) -> bool: + """ + Check if the dtypes are valid + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes are valid, False otherwise + :rtype: bool + """ + valid_ab_dtypes = { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Uint8, + cutlass.Int8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + if ab_dtype not in valid_ab_dtypes: + return False + + if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: + return False + + # Define compatibility mapping between accumulator type and AB type + acc_ab_compatibility = { + cutlass.Float32: { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, # Float32 accumulator supports floating point AB types only + cutlass.Float16: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, + } + # Check compatibility between accumulator type and AB type + if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: + return False + + # Define compatibility mapping between accumulator type and C type + acc_c_compatibility = { + cutlass.Float32: { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + cutlass.Float16: { + cutlass.BFloat16, + cutlass.Float16, + }, + cutlass.Int32: { + cutlass.BFloat16, + cutlass.Float16, + cutlass.Float32, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + } + # Check compatibility between accumulator type and C type + if c_dtype not in acc_c_compatibility[self.acc_dtype]: + return False + + return True + + def is_valid_mma_tiler_and_cluster_shape(self) -> bool: + """Check if the mma tiler and cluster shape are valid. + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if not ( + (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) + or (self.use_2cta_instrs and self.mma_tiler_mn[0] in [128, 256]) + ): + is_valid = False + if self.mma_tiler_mn[1] not in range(32, 257, 32): + is_valid = False + # Skip illegal cluster shape + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16 + or self.cluster_shape_mn[0] <= 0 + or self.cluster_shape_mn[1] <= 0 + or not is_power_of_2(self.cluster_shape_mn[0]) + or not is_power_of_2(self.cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + def is_valid_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + # TODO: move to utils + def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid + + def is_valid_epilog_store_option(self, m: int, n: int) -> bool: + """ + Check if the epilogue store option is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + + :return: True if the epilogue store option is valid, False otherwise + :rtype: bool + """ + + is_valid = True + # None TMA store version does not have predication, can not support OOB tiles + cta_tile_shape_mn = ( + self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler_mn[1], + ) + if not self.use_tma_store: + if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): + is_valid = False + return is_valid + + def can_implement(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor) -> bool: + """Check if the given tensors can be implemented by this kernel. + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + + :return: True if the gemm supports the given config, False otherwise + :rtype: bool + """ + m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] + + # infer a_major, b_major, c_major + is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a() + is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b() + is_m_major_c = utils.LayoutEnum.from_tensor(c).is_m_major_c() + a_major = "m" if is_m_major_a else "k" + b_major = "n" if is_n_major_b else "k" + c_major = "m" if is_m_major_c else "n" + + can_implement = True + # Skip unsupported types + if not self.is_valid_dtypes(a.element_type, c.element_type): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not self.is_valid_mma_tiler_and_cluster_shape(): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not self.is_valid_tensor_alignment( + m, n, k, l, a.element_type, c.element_type, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid epilogue store option + if not self.is_valid_epilog_store_option(m, n): + can_implement = False + + return can_implement + + +def create_tensors( + l, + m_per_iteration, + n, + k, + a_major, + b_major, + c_major, + ab_dtype, + c_dtype, + num_steps, + local_rank, +): + torch.manual_seed(1111) + a_torch_cpu = cutlass_torch.matrix(l, m_per_iteration, k, a_major == "m", ab_dtype) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype) + c_torch_cpu = cutlass_torch.matrix(l, m_per_iteration, n, c_major == "m", c_dtype) + + # create local buffers + a_local_list = [] + a_local_torch_list = [] + c_local_list = [] + c_local_torch_list = [] + for m_dim_offset in range(num_steps): + a_temp = a_torch_cpu.clone() + a_temp.zero_() + if m_dim_offset == local_rank: + a_temp.copy_(a_torch_cpu) + a_local, a_local_torch = cutlass_torch.cute_tensor_like( + a_temp, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + a_local_list.append(a_local) + a_local_torch_list.append(a_local_torch) + c_temp = c_torch_cpu.clone().zero_() + c_local, c_local_torch = cutlass_torch.cute_tensor_like( + c_temp, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_local_list.append(c_local) + c_local_torch_list.append(c_local_torch) + a_torch_unique = torch.empty( + a_torch_cpu.shape, device="cuda", dtype=a_torch_cpu.dtype + ) + a_torch_unique.copy_(a_torch_cpu) + a_tensor_symm = nvshmem.core.tensor(a_torch_cpu.shape, dtype=a_torch_cpu.dtype) + a_tensor_symm.copy_(a_local_torch_list[local_rank]) + a_tensor_peers = [nvshmem.core.get_peer_tensor(a_tensor_symm, rank) for rank in range(world_size)] + + b_tensor, b_torch = cutlass_torch.cute_tensor_like( + b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + gate_a_flags = torch.zeros(num_steps, device="cuda", dtype=torch.int32) + + device_arrival_counters = torch.tensor([0], dtype=torch.int32, device="cuda") + device_arrival_counters_symm = nvshmem.core.tensor(1, dtype=torch.int32) + device_arrival_counters_symm.copy_(device_arrival_counters) + device_arrival_counters_peers = [nvshmem.core.get_peer_tensor(device_arrival_counters_symm, rank) for rank in range(world_size)] + + return ( + a_local_list, + a_local_torch_list, + a_torch_unique, + a_tensor_symm, + a_tensor_peers, + b_tensor, + b_torch, + c_local_list, + c_local_torch_list, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + gate_a_flags, + device_arrival_counters_symm, + device_arrival_counters_peers, + ) + + +def compare( + a_torch_unique, + b_torch_cpu, + c_torch_local_list, + c_dtype, + tolerance, + rank, + world_size, +): + # Copy gpu result back + kernel_result = torch.cat(c_torch_local_list, dim=0).cpu() + + all_gather_a = [torch.zeros_like(a_torch_unique) for _ in range(world_size)] + dist.all_gather(all_gather_a, a_torch_unique) + a = torch.cat(all_gather_a, dim=0) + + # Convert GPU tensors to CPU before reference computation + # This ensures the same computation path as dense_gemm_persistent.py + a_cpu = a.cpu() + + # Compute reference result using CPU tensors (like dense_gemm_persistent.py) + num_rows = a_cpu.shape[0] + ref = torch.einsum( + "mkl,nkl->mnl", + a_cpu.to(dtype=torch.float32), + b_torch_cpu.to(dtype=torch.float32), + ) + + # Convert ref to c_dtype + _, ref_torch_gpu = cutlass_torch.cute_tensor_like( + ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + ref_result = ref_torch_gpu.cpu() + # Assert close results + torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05) + + +def run( + rank, + world_size, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """Execute a patthen of All gather + persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + + This function prepares input tensors and flags needed for sync between communicaton and gemm, configures and launches the persistent GEMM kernel and Peer mem copys and memset, + and construt them into a cuda graph. + optionally performs reference validation, and benchmarks the execution performance. + + :param mnkl: Problem size (M, N, K, L) + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C + :type c_dtype: Type[cutlass.Numeric] + :param acc_dtype: Data type for accumulation during matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param a_major/b_major/c_major: Memory layout of tensor A/B/C + :type a_major/b_major/c_major: str + :param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the + default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters. + :type mma_tiler_mn: Tuple[int, int], optional + :param cluster_shape_mn: Cluster shape. If not specified in the decorator parameters, the autotuner will use the + default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters. + :type cluster_shape_mn: Tuple[int, int], optional + :param use_2cta_instrs: Whether to use 2CTA instructions. If not specified in the decorator parameters, the autotuner + will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. + :type use_2cta_instrs: bool, optional + :param use_tma_store: Whether to use TMA store. If not specified in the decorator parameters, the autotuner will use + the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. + :type use_tma_store: bool, optional + :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 1 + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False + :type use_cold_l2: bool, optional + :raises RuntimeError: If CUDA GPU is not available + :raises ValueError: If the configuration is invalid or unsupported by the kernel + :return: Execution time of the GEMM kernel + :rtype: float + """ + print("Running Blackwell Persistent Dense GEMM test with:") + print(f"per GPU mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + + # Unpack parameters + m, n, k, l = mnkl + m_per_step = m // world_size + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + ( + a_local_list, + a_local_torch_list, + a_torch_unique, + a_tensor_symm, + a_tensor_peers, + b_tensor, + b_torch, + c_local_list, + c_local_torch_list, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + gate_a_flags, + device_arrival_counters_symm, + device_arrival_counters_peers, + ) = create_tensors( + l, + m_per_step, + n, + k, + a_major, + b_major, + c_major, + ab_dtype, + c_dtype, + world_size, + rank, + ) + + # Build GEMM object + gemm = PersistentDenseGemmKernel( + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + gated_a_load=False, + ) + gemm_with_gated_a_load = PersistentDenseGemmKernel( + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + gated_a_load=True, + ) + + # Check if configuration can be implemented + can_implement = gemm.can_implement(a_local_list[0], b_tensor, c_local_list[0]) + if not can_implement: + raise ValueError( + f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " + f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " + f"use_tma_store = {use_tma_store}" + ) + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + compiled_gemm = cute.compile( + gemm, + a_local_list[0], + b_tensor, + c_local_list[0], + max_active_clusters, + current_stream, + flag_offset=0, + gate_a_flags=from_dlpack(gate_a_flags), + ) + compiled_gemm_with_gated_a_load = cute.compile( + gemm_with_gated_a_load, + a_local_list[0], + b_tensor, + c_local_list[0], + max_active_clusters, + current_stream, + flag_offset=0, + gate_a_flags=from_dlpack(gate_a_flags), + ) + + device_arrival_counters_ptrs = torch.tensor( + [tensor.data_ptr() for tensor in device_arrival_counters_peers], device="cuda", dtype=torch.int64 + ) + sync_nvl_devices_instance = SyncNvlDevices(world_size) + compiled_sync_nvl_devices = cute.compile( + sync_nvl_devices_instance, + device_idx=rank, + device_arrival_counters=from_dlpack(device_arrival_counters_ptrs), + iteration_flags=from_dlpack(gate_a_flags), + stream=current_stream, + ) + + g = torch.cuda.CUDAGraph() + capture_stream = torch.cuda.Stream(local_rank) + copy_stream = torch.cuda.Stream(local_rank) + gemm_stream = torch.cuda.Stream(local_rank) + + num_steps = world_size + with torch.cuda.graph(g, stream=capture_stream): + compiled_sync_nvl_devices( + rank, + from_dlpack(device_arrival_counters_ptrs), + from_dlpack(gate_a_flags), + cuda.CUstream(capture_stream.cuda_stream), + ) + gemm_stream.wait_stream(capture_stream) + for j in range(num_steps): + m_dim_offset = (local_rank + j) % num_steps + if j == 0: # gemm process local data directly + compiled_gemm( + a_local_list[m_dim_offset], + b_tensor, + c_local_list[m_dim_offset], + cuda.CUstream(gemm_stream.cuda_stream), + flag_offset=m_dim_offset, + gate_a_flags=from_dlpack(gate_a_flags), + ) + else: # gemm process remote data with gated a load + compiled_gemm_with_gated_a_load( + a_local_list[m_dim_offset], + b_tensor, + c_local_list[m_dim_offset], + cuda.CUstream(gemm_stream.cuda_stream), + flag_offset=m_dim_offset, + gate_a_flags=from_dlpack(gate_a_flags), + ) + + copy_stream.wait_stream(capture_stream) + for j in range(num_steps): + m_dim_offset = (local_rank + j) % num_steps + if m_dim_offset != local_rank: + copy_bytes = ( + a_local_torch_list[m_dim_offset].element_size() + * a_local_torch_list[m_dim_offset].numel() + ) + dst_address = a_local_torch_list[m_dim_offset].data_ptr() + src_address = a_tensor_peers[m_dim_offset].data_ptr() + driver.cuMemcpyDtoDAsync( + dst_address, + src_address, + copy_bytes, + copy_stream.cuda_stream, + ) + memset_address = ( + gate_a_flags.data_ptr() + m_dim_offset * gate_a_flags.element_size() + ) + driver.cuMemsetD32Async( + memset_address, + 1, # value_to_set + 1, # size + copy_stream.cuda_stream, + ) + + capture_stream.wait_stream(copy_stream) + capture_stream.wait_stream(gemm_stream) + + replay_stream = torch.cuda.Stream(local_rank) + # warmup + with torch.cuda.stream(replay_stream): + for i in range(warmup_iterations): + g.replay() + torch.cuda.synchronize() + # benchmark + with torch.cuda.stream(replay_stream): + # INSERT_YOUR_CODE + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for i in range(iterations): + g.replay() + end_event.record() + torch.cuda.synchronize() + exec_time = start_event.elapsed_time(end_event) + # Calculate average time per iteration + avg_time_per_iteration = (exec_time * 1000) / iterations + + # Collect times from all GPUs and find maximum + # Convert to tensor for distributed communication + time_tensor = torch.tensor([avg_time_per_iteration], device="cuda") + + # Gather all times to rank 0 + gathered_times = [torch.zeros_like(time_tensor) for _ in range(world_size)] + dist.all_gather(gathered_times, time_tensor) + + # Find maximum time across all GPUs + exec_time = max(tensor.cpu().item() for tensor in gathered_times) + + # reference check + if not skip_ref_check: + # compiled_gemm(a_tensor, b_tensor, c_tensor, current_stream) + compare( + a_torch_unique, + b_torch_cpu, + c_local_torch_list, + c_dtype, + tolerance, + rank, + world_size, + ) + + for i in range(world_size): + if i != dist.get_rank(): + nvshmem.core.free_tensor(a_tensor_peers[i]) + nvshmem.core.free_tensor(device_arrival_counters_peers[i]) + nvshmem.core.free_tensor(a_tensor_symm) + nvshmem.core.free_tensor(device_arrival_counters_symm) + + return exec_time # Return execution time in microseconds + + +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() + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser( + description="Example of All Gather + Dense Persistent GEMM on Blackwell." + ) + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.TFloat32) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--use_tma_store", action="store_true", help="Use tma store or not" + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + torchrun_uid_init_bcast() + local_rank = dist.get_rank() + world_size = dist.get_world_size() + print(f"World size: {world_size}, local rank: {local_rank}") + torch.cuda.set_device(local_rank) + + run( + local_rank, + world_size, + args.mnkl, + args.ab_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.use_tma_store, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + torchrun_finalize() + print("PASS") diff --git a/examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py b/examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py new file mode 100644 index 00000000..4ff104d6 --- /dev/null +++ b/examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py @@ -0,0 +1,2455 @@ +# Copyright (c) 2025 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 argparse +from typing import Optional, Tuple, Type, Union + +import numpy as np +import torch +import torch.distributed as dist +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 +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Int32, Float16, BFloat16, Float32, Float8E5M2, Float8E4M3FN +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +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 high-performance persistent batched dense GEMM AR example for the NVIDIA Blackwell SM100 architecture +using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") +- Matrix C_mc is a multicast C matrix that changes can be broadcasted to all GPUs by multimem instructions. + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Support persistent tile scheduling to better overlap memory load/store with mma between tiles + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + - Support all-reduce epilogue with multimem instructions to distribute the workload to all GPUs + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) +4. All reduce epilogue: + - Load and reduce the 128bit data from all ranks by multimem instructions. + - Broadcast the reduced data to all ranks by multimem instructions. + - current implementation only supports LDMCxSTMC all-reduce which means each rank only computes a portion of + the output tensor and broadcast the result to all ranks. + - the all-reduce epilogue is only supported when use_tma_store is True. + - the all-reduce epilogue is only supported when c_dtype is Float16, Float32, BFloat16, Float8E4M3FN, Float8E5M2. + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Input arguments to this example is same as dense_gemm.py. + +.. code-block:: bash + + torchrun --nproc-per-node 8 examples/distributed/distributed_dense_gemm_persistent_all_reduce.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 1024,1000,1024,1 --warmup_iterations 3 --iterations 10 \ + --use_tma_store --use_2cta_instrs --all_reduce LDMCxSTMC + +To collect performance with NSYS profiler: + +.. code-block:: bash + + nsys profile --gpu-metrics-devices=cuda-visible \ + torchrun --nproc-per-node 8 examples/distributed/distributed_dense_gemm_persistent_all_reduce.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs --warmup_iterations 3 --iterations 10 \ + --skip_ref_check --all_reduce LDMCxSTMC + + +Constraints are same as dense_gemm_persistent.py: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), + see detailed valid dtype combinations in below PersistentDenseGemmKernel class documentation +* A/B tensor must have the same data type +* Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* Mma tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled +* when all_reduce is "two_shot", M and N must be multiple of 128, world_size must be 8 + +""" + + +def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + c_smem_layout: Union[cute.Layout, None], +) -> Tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + :param c_smem_layout: Layout of C operand in shared memory, or None if not using TMA store. + :type c_smem_layout: Union[cute.Layout, None] + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C with 1-stage + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1 + ) + + ab_bytes_per_stage = cute.size_in_bytes( + a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes + # Add remaining unused smem to epilogue + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage + + +class PersistentDenseGemmKernel: + """This class implements batched matrix multiplication (C = A x B) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Whether to use Tensor Memory Access (TMA) for storing results + :type use_tma_store: bool + :param all_reduce: All-reduce mode, can be "none", "LDMCxSTMC" + :type all_reduce: str + + :note: only support two_shot all-reduce mode in a single node right now + + :note: In current version, A and B tensor must have the same data type + - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + + :note: Supported A/B data types: + - TFloat32 + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported C data types: + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + + Example: + >>> gemm = PersistentDenseGemmKernel( + ... acc_dtype=cutlass.Float32, + ... use_2cta_instrs=True, + ... mma_tiler_mn=(128, 128), + ... cluster_shape_mn=(2, 2) + ... ) + >>> gemm(a_tensor, b_tensor, c_tensor, max_active_clusters, stream) + """ + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_tma_store: bool, + rank_id: int, + world_size: int, + all_reduce="two_shot", + ): + """Initializes the configuration for a Blackwell dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant + with cta_group=2 should be used. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + 3. Output C tensor store mode: + - use_tma_store: Boolean indicating whether to use Tensor Memory Access (TMA) for storing results. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Use Tensor Memory Access (TMA) or normal store for output C tensor. + :type use_tma_store: bool + :param all_reduce: All-reduce mode, can be "none", "LDMCxSTMC" + :type all_reduce: str + """ + + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.use_tma_store = use_tma_store + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.all_reduce_warp_id = (6, 7, 8, 9) + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + *self.epilogue_warp_id, + *self.all_reduce_warp_id, + ) + ) + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.epilogue_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.all_reduce_sync_bar_id = 4 + self.all_reduce_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.all_reduce_sync_bar_id, + num_threads=32 * len(self.all_reduce_warp_id), + ) + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.all_reduce = all_reduce + self.rank_id = rank_id + self.num_ranks = world_size + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C stage counts in shared memory + - Computing A/B/C shared memory layout + - Computing tensor memory allocation columns + """ + # Configure tiled mma + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + # Compute epilogue subtile + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + c_smem_layout = None + if cutlass.const_expr(self.use_tma_store): + c_smem_layout = sm100_utils.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.smem_capacity, + self.occupancy, + self.use_tma_store, + c_smem_layout, + ) + + # Compute A/B/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = None + if self.use_tma_store: + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + # Compute the number of tensor memory allocation columns + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + c_mc: cute.Tensor, + barrier_flag: cute.Tensor, + barrier_flag_mc: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + :param c_mc: Output symmetric tensor C_mc, any write or read to a multicast tensor will be broadcasted to all GPUs + :type c_mc: cute.Tensor + :param barrier_flag: flag tensor acting as a barrier for GPUs to synchronize + :type barrier_flag: cute.Tensor + :param barrier_flag_mc: the multicast tensor of barrier_flag + :type barrier_flag_mc: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + :raises AssertionError: If OOB (Out-Of-Bounds) tiles are present when TMA store is disabled. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + # Setup TMA store for C + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + c_mc, + barrier_flag, + barrier_flag_mc, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: Optional[cute.CopyAtom], + mC_mnl: cute.Tensor, + c_mc: cute.Tensor, + barrier_flag: cute.Tensor, + barrier_flag_mc: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.use_tma_store): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem_dealloc_barrier = None + if cutlass.const_expr(not self.use_tma_store): + tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_dealloc_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + # + # Setup smem tensor A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_producer.tail() + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + # tCtAcc += tCrA * tCrB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + handle.release() + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + sC = None + if cutlass.const_expr(self.use_tma_store): + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop for epilogue + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + + if cutlass.const_expr(self.use_tma_store): + # {$nv-internal-release begin} + # (EPI_TILE_M, EPI_TILE_N, STAGE) + # sC = storage.sC.get_tensor( + # c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + # ) + # {$nv-internal-release end} + + assert tma_atom_c is not None and sC is not None + self.epilogue_tma_store_release_flag( + tidx, + warp_idx, + acc_pipeline, + tiled_mma, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + tile_sched, + epilogue_op, + flag_base=barrier_flag_mc, + flag_mem_scope="gpu", + ) + else: + self.epilogue_release_flag( + tidx, + acc_pipeline, + tiled_mma, + tCtAcc_base, + tCgC, + epi_tile, + tile_sched, + epilogue_op, + tmem_dealloc_barrier, + flag_base=barrier_flag_mc, + flag_mem_scope="gpu", + ) + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Allreduce warps + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.all_reduce == "LDMCxSTMC"): + if warp_idx >= self.all_reduce_warp_id[0]: + # /////////////////////////////////////////////////////////////////////////////// + # Add persistent tile loop + # /////////////////////////////////////////////////////////////////////////////// + + rank_id = self.rank_id + num_ranks = Int32(self.num_ranks) + lane_id = cute.arch.lane_idx() + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # we want 128bit ld/st for better performance + atom_val = 128 // c_mc.element_type.width + atom_thr_n = self.mma_tiler[1] // atom_val + atom_thr_m = len(self.all_reduce_warp_id) * cute.arch.WARP_SIZE // atom_thr_n + thr_layout = cute.make_layout( + (atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1) + ) + val_layout = cute.make_layout((1, atom_val), stride=(atom_val, 1)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), c_mc.element_type + ) + tiled_copy_fake = cute.make_tiled_copy_tv( + copy_atom_load, thr_layout, val_layout + ) + thr_copy_fake = tiled_copy_fake.get_slice( + tidx - self.all_reduce_warp_id[0] * 32 + ) + # predicate tensor + idC = cute.make_identity_tensor(c_mc.shape) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + tile_id = Int32( + tile_sched._current_work_linear_idx + * cute.size(self.cluster_shape_mn) + + cute.arch.block_idx_in_cluster() + ) + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # System barrier to make sure that data from each GPU is in memory before allreduce + if warp_idx == self.all_reduce_warp_id[0]: + with cute.arch.elect_one(): + flag = barrier_flag.iterator + tile_id + utils.distributed.spin_lock_atom_cas_relaxed_wait( + lock_ptr=flag, + expected_val=num_ranks, + reset_val=0, + scope="gpu", + ) + + self.all_reduce_sync_barrier.arrive_and_wait() + # partition and slice at tile level + gC_mc = cute.local_tile( + c_mc, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + cC = cute.local_tile( + idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + + tCgC_mc = thr_mma.partition_C(gC_mc) + tCpC = thr_mma.partition_C(cC) + tCgC_mc_slice = tCgC_mc[((None, None), 0, 0, *mma_tile_coord_mnl)] + tCpC_slice = tCpC[((None, None), 0, 0, *mma_tile_coord_mnl)] + + # partition based on the number of GPUs + cta_mma_tile_m = self.mma_tiler[0] // cute.size( + tiled_mma.thr_id.shape + ) + m_local_rank = int(cta_mma_tile_m / self.num_ranks) + tCgC_mc_slice_partitioned = cute.zipped_divide( + tCgC_mc_slice, (m_local_rank, self.mma_tiler[1]) + ) + tCpC_slice_partitioned = cute.zipped_divide(tCpC_slice, (m_local_rank, self.mma_tiler[1])) + tCgC_mc_local_rank = cute.slice_( + tCgC_mc_slice_partitioned, ((None, None), (rank_id, 0)) + ) + tCpC_local_rank = cute.slice_(tCpC_slice_partitioned, ((None, None), (rank_id, 0))) + + # partition at thread level + frgC_mc = thr_copy_fake.partition_S(tCgC_mc_local_rank) + frpC = thr_copy_fake.partition_S(tCpC_local_rank) + atom, loop_m, loop_n = frgC_mc.shape + for i in cutlass.range_constexpr(loop_m): + for j in cutlass.range_constexpr(loop_n): + if cute.elem_less(frpC[0, i, j], c_mc.shape): + mc_ptr = frgC_mc[None, i, j].iterator + x, y, z, w = 0, 0, 0, 0 + if cutlass.const_expr(self.c_dtype == Float16): + x, y, z, w = utils.distributed.multimem_ld_reduce_8xf16( + mc_ptr + ) + elif cutlass.const_expr(self.c_dtype == Float32): + x, y, z, w = utils.distributed.multimem_ld_reduce_4xf32( + mc_ptr + ) + elif cutlass.const_expr(self.c_dtype == BFloat16): + x, y, z, w = ( + utils.distributed.multimem_ld_reduce_8xbf16(mc_ptr) + ) + elif cutlass.const_expr(self.c_dtype == Float8E4M3FN): + x, y, z, w = ( + utils.distributed.multimem_ld_reduce_16xe4m3(mc_ptr) + ) + elif cutlass.const_expr(self.c_dtype == Float8E5M2): + x, y, z, w = ( + utils.distributed.multimem_ld_reduce_16xe5m2(mc_ptr) + ) + utils.distributed.multimem_st_4xb32(mc_ptr, x, y, z, w) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + self.all_reduce_sync_barrier.arrive_and_wait() + + # + # Set Per SM Flag with Release + # + # This ensure + # 1. no rank early exit while other ranks are still issuing multimem.ld_reduce + # 2. each rank's prior multimem.st have become visiable to all other ranks in the system (w/ .SYS scope) + if warp_idx == self.all_reduce_warp_id[0]: + with cute.arch.elect_one(): + # Offset to last tile flag idx + last_tile_id_linear = cute.size( + tile_sched.params.problem_layout_ncluster_mnl + ) * cute.size(self.cluster_shape_mn) + # 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( + lock_ptr=barrier_flag_mc.iterator + + last_tile_id_linear + + sm_id_linear, + scope="sys", + order="release", + ) + # Relaxed spin-lock wait flag with sys scope + utils.distributed.spin_lock_atom_cas_relaxed_wait( + lock_ptr=barrier_flag.iterator + + last_tile_id_linear + + sm_id_linear, + expected_val=num_ranks, + reset_val=0, + scope="sys", + ) + + @cute.jit + def epilogue_tma_store_release_flag( + self, + epi_tidx: cutlass.Int32, + warp_idx: cutlass.Int32, + acc_pipeline: pipeline.PipelineAsync, + tiled_mma: cute.TiledMma, + tma_atom_c: cute.CopyAtom, + # Input of epilogue + tCtAcc_base: cute.Tensor, + # Staging of epilogue + sC: cute.Tensor, + # Output of epilogue + tCgC: cute.Tensor, + epi_tile: cute.Tile, + tile_sched: utils.StaticPersistentTileScheduler, + epilogue_op: cutlass.Constexpr, + flag_base: cute.Tensor, + flag_mem_scope: str, + ) -> None: + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilogue_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilogue_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + tCgC_epi = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], epi_tile + ) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + + epilogue_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilogue_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + epilogue_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilogue_sync_barrier.arrive_and_wait() + + epilogue_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Set Per Output Tile Flag with Release + # + # 1D linear index of current output tile + tile_id_linear = Int32( + tile_sched._current_work_linear_idx * cute.size(self.cluster_shape_mn) + + cute.arch.block_idx_in_cluster() + ) + # Wait for C store complete + # Unlike regular epilogue where we only wait C store complete once at end of each kernel. + # Here we need to wait for C store complete for each output tile before we set the release flag. + c_pipeline.producer_tail() + # Update flag with release semantic with GPU scope + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile = flag_base.iterator + tile_id_linear + utils.distributed.multimem_red_add1( + lock_ptr=flag_curr_tile, + scope=flag_mem_scope, + order="release", + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + @cute.jit + def epilogue_release_flag( + self, + epi_tidx: cutlass.Int32, + acc_pipeline: pipeline.PipelineAsync, + tiled_mma: cute.TiledMma, + tCtAcc_base: cute.Tensor, + tCgC: cute.Tensor, + epi_tile: cute.Tile, + tile_sched: utils.StaticPersistentTileScheduler, + epilogue_op: cutlass.Constexpr, + tmem_dealloc_barrier: pipeline.NamedBarrier, + flag_base: cute.Tensor, + flag_mem_scope: str, + ) -> None: + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilogue_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs + ) + + gC_epi = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tTR_gC_partitioned = thr_copy_t2r.partition_D(gC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_rmem_tensor( + tTR_gC_partitioned[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype + ) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_gC = tTR_gC_partitioned[ + (None, None, None, None, None, *mma_tile_coord_mnl) + ] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tTR_rAcc.load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tTR_rC.store(acc_vec) + + # + # Store C to global memory + # + cute.copy(simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)]) + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Set Per Output Tile Flag with Release + # + # 1D linear index of current output tile + tile_id_linear = Int32( + tile_sched._current_work_linear_idx * cute.size(self.cluster_shape_mn) + + cute.arch.block_idx_in_cluster() + ) + # Wait for C store complete + # Unlike regular epilogue where we only wait C store complete once at end of each kernel. + # Here we need to wait for C store complete for each output tile before we set the release flag. + c_pipeline.producer_tail() + # Update flag with release semantic with GPU scope + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile = flag_base.iterator + tile_id_linear + utils.distributed.multimem_red_add1( + lock_ptr=flag_curr_tile, + scope=flag_mem_scope, + order="release", + ) + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() + + def epilogue_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilogue_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def _compute_num_tmem_alloc_cols( + tiled_mma: cute.TiledMma, + mma_tiler: Tuple[int, int, int], + num_acc_stage: int, + ) -> int: + """ + Compute the number of tensor memory allocation columns. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler: The shape (M, N, K) of the MMA tile. + :type mma_tiler: tuple[int, int, int] + :param num_acc_stage: The stage of the accumulator tensor. + :type num_acc_stage: int + + :return: The number of tensor memory allocation columns. + :rtype: int + """ + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + + return num_tmem_alloc_cols + + def is_valid_dtypes( + self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] + ) -> bool: + """ + Check if the dtypes are valid + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes are valid, False otherwise + :rtype: bool + """ + valid_ab_dtypes = { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Uint8, + cutlass.Int8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + if ab_dtype not in valid_ab_dtypes: + return False + + if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: + return False + + # Define compatibility mapping between accumulator type and AB type + acc_ab_compatibility = { + cutlass.Float32: { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, # Float32 accumulator supports floating point AB types only + cutlass.Float16: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, + } + # Check compatibility between accumulator type and AB type + if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: + return False + + # Define compatibility mapping between accumulator type and C type + acc_c_compatibility = { + cutlass.Float32: { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + cutlass.Float16: { + cutlass.BFloat16, + cutlass.Float16, + }, + cutlass.Int32: { + cutlass.BFloat16, + cutlass.Float16, + cutlass.Float32, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + } + # Check compatibility between accumulator type and C type + if c_dtype not in acc_c_compatibility[self.acc_dtype]: + return False + + # check if c_dtype is supported by multimem all-reduce + if cutlass.const_expr( + c_dtype not in {cutlass.Float16, cutlass.Float32, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2} + ): + return False + + return True + + def is_valid_mma_tiler_and_cluster_shape(self) -> bool: + """Check if the mma tiler and cluster shape are valid. + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if not ( + (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) + or (self.use_2cta_instrs and self.mma_tiler_mn[0] in [128, 256]) + ): + is_valid = False + if self.mma_tiler_mn[1] not in range(32, 257, 32): + is_valid = False + # Skip illegal cluster shape + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16 + or self.cluster_shape_mn[0] <= 0 + or self.cluster_shape_mn[1] <= 0 + or not is_power_of_2(self.cluster_shape_mn[0]) + or not is_power_of_2(self.cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + def is_valid_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + + return is_valid + + def is_valid_epilogue_store_option(self, m: int, n: int) -> bool: + """ + Check if the epilogue store option is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + + :return: True if the epilogue store option is valid, False otherwise + :rtype: bool + """ + + is_valid = True + # None TMA store version does not have predication, can not support OOB tiles + cta_tile_shape_mn = ( + self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler_mn[1], + ) + if not self.use_tma_store: + if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): + is_valid = False + return is_valid + + def can_implement(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor) -> bool: + """Check if the given tensors can be implemented by this kernel. + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + + :return: True if the gemm supports the given config, False otherwise + :rtype: bool + """ + m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] + + # infer a_major, b_major, c_major + is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a() + is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b() + is_m_major_c = utils.LayoutEnum.from_tensor(c).is_m_major_c() + a_major = "m" if is_m_major_a else "k" + b_major = "n" if is_n_major_b else "k" + c_major = "m" if is_m_major_c else "n" + + can_implement = True + # Skip unsupported types + if not self.is_valid_dtypes(a.element_type, c.element_type): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not self.is_valid_mma_tiler_and_cluster_shape(): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not self.is_valid_tensor_alignment( + m, n, k, l, a.element_type, c.element_type, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid epilogue store option + if not self.is_valid_epilogue_store_option(m, n): + can_implement = False + # currently only support world_size = 2, 4, 8 + if dist.get_world_size() not in [2, 4, 8]: + can_implement = False + return can_implement + + +def create_mc_tensor(torch_tensor_cpu, dtype, leading_dim, is_dynamic_layout=True): + torch_tensor_gpu_local = nvshmem.core.tensor(torch_tensor_cpu.shape, dtype=torch_tensor_cpu.dtype) + torch_tensor_gpu_local.copy_(torch_tensor_cpu) + torch_tensor_gpu_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, torch_tensor_gpu_local) + + # create MC tensor memref + cute_tensor_mc = from_dlpack(torch_tensor_gpu_mc, assumed_align=16) + cute_tensor = from_dlpack(torch_tensor_gpu_local, assumed_align=16) + + cute_tensor.element_type = dtype + cute_tensor_mc.element_type = dtype + + if is_dynamic_layout: + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim) + + return cute_tensor, cute_tensor_mc, torch_tensor_gpu_local, torch_tensor_gpu_mc + + +def create_tensors( + l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype, mma_tiler_mn, use_2cta_instrs +): + torch.manual_seed(1111) + cta_tile_shape_mn = ( + mma_tiler_mn[0] // (2 if use_2cta_instrs else 1), + mma_tiler_mn[1], + ) + problem_shape_ntile_mn = (m // cta_tile_shape_mn[0], n // cta_tile_shape_mn[1]) + num_tiles = problem_shape_ntile_mn[0] * problem_shape_ntile_mn[1] + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + + a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype).to(torch.float32).normal_().round_().to(dtype=cutlass_torch.dtype(ab_dtype)) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype).to(torch.float32).normal_().round_().to(dtype=cutlass_torch.dtype(ab_dtype)) + c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype) + barrier_flag_torch_cpu = torch.zeros(num_tiles+num_sms, device="cuda", dtype=torch.int32) + + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_tensor_mc, c_torch_gpu, c_torch_gpu_mc = create_mc_tensor( + c_torch_cpu, c_dtype, (1 if c_major == "n" else 0), is_dynamic_layout=True + ) + barrier_flag, barrier_flag_mc, barrier_flag_torch_gpu, barrier_flag_torch_gpu_mc = create_mc_tensor( + barrier_flag_torch_cpu, cutlass.Int32, 0, is_dynamic_layout=True + ) + + return ( + a_tensor, + b_tensor, + c_tensor, + c_tensor_mc, + barrier_flag, + barrier_flag_mc, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + barrier_flag_torch_cpu, + c_torch_gpu, + c_torch_gpu_mc, + barrier_flag_torch_gpu, + barrier_flag_torch_gpu_mc, + ) + + +def compare( + a_torch_cpu, b_torch_cpu, c_torch_gpu, ab_dtype, c_dtype, tolerance +): + + # Compute reference result + ref = torch.einsum( + "mkl,nkl->mnl", + a_torch_cpu.to(dtype=torch.float32), + b_torch_cpu.to(dtype=torch.float32), + ) + + acc_type = cutlass.Float32 + if c_dtype.width == 16: + acc_dtype = cutlass.Float32 + elif c_dtype.width == 8: + acc_dtype = cutlass.Float16 + + torch_acc_dtype = cutlass_torch.dtype(acc_dtype) + torch_c_dtype = cutlass_torch.dtype(c_dtype) + + # note that in this example, f8 reduce sum is done by f16 acc precision and f16 by f32 + ref = ref.to(torch_c_dtype).to(torch_acc_dtype).cuda() + torch.distributed.all_reduce(ref, op=torch.distributed.ReduceOp.SUM) + ref = ref.to(torch_c_dtype) + + # convert to higher precision to satisfy assert_close function + ref_result = ref.cpu().to(torch_acc_dtype) + kernel_result = c_torch_gpu.view(torch_c_dtype).cpu().to(torch_acc_dtype) + + max_val = torch.finfo(kernel_result.dtype).max + min_val = torch.finfo(kernel_result.dtype).min + kernel_result = torch.nan_to_num(kernel_result, nan=max_val, posinf=max_val, neginf=min_val) + ref_result = torch.nan_to_num(ref_result, nan=max_val, posinf=max_val, neginf=min_val) + + # Assert close results + if dist.get_rank() == 0: + print(f"kernel_result: {kernel_result[:8,:8,0]}") + print(f"ref_result: {ref_result[:8,:8,0]}") + torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05) + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + all_reduce: str = "two_shot", + **kwargs, +): + """Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + + This function prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks the execution performance. + + :param mnkl: Problem size (M, N, K, L) + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C + :type c_dtype: Type[cutlass.Numeric] + :param acc_dtype: Data type for accumulation during matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param a_major/b_major/c_major: Memory layout of tensor A/B/C + :type a_major/b_major/c_major: str + :param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the + default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters. + :type mma_tiler_mn: Tuple[int, int], optional + :param cluster_shape_mn: Cluster shape. If not specified in the decorator parameters, the autotuner will use the + default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters. + :type cluster_shape_mn: Tuple[int, int], optional + :param use_2cta_instrs: Whether to use 2CTA instructions. If not specified in the decorator parameters, the autotuner + will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. + :type use_2cta_instrs: bool, optional + :param use_tma_store: Whether to use TMA store. If not specified in the decorator parameters, the autotuner will use + the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. + :type use_tma_store: bool, optional + :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 1 + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False + :type use_cold_l2: bool, optional + :param all_reduce: All-reduce mode, can be "none", "LDMCxSTMC" + :type all_reduce: str, optional + :raises RuntimeError: If CUDA GPU is not available + :raises ValueError: If the configuration is invalid or unsupported by the kernel + :return: Execution time of the GEMM kernel + :rtype: float + """ + print("Running Blackwell Persistent Dense GEMM AR test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + print(f"Fused AllReduce Op: {all_reduce}") + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + ( + a_tensor, + b_tensor, + c_tensor, + c_tensor_mc, + barrier_flag, + barrier_flag_mc, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + barrier_flag_torch_cpu, + c_torch_gpu, + c_torch_gpu_mc, + barrier_flag_torch_gpu, + barrier_flag_torch_gpu_mc, + ) = create_tensors( + l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype, mma_tiler_mn, use_2cta_instrs + ) + + # Build GEMM object + gemm = PersistentDenseGemmKernel( + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + dist.get_rank(), + dist.get_world_size(), + all_reduce=all_reduce, + ) + + # Check if configuration can be implemented + can_implement = gemm.can_implement(a_tensor, b_tensor, c_tensor) + if not can_implement: + raise ValueError( + f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " + f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " + f"use_tma_store = {use_tma_store}" + ) + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + c_tensor, + c_tensor_mc, + barrier_flag, + barrier_flag_mc, + max_active_clusters, + current_stream, + ) + + if not skip_ref_check: + compiled_gemm( + a_tensor, + b_tensor, + c_tensor, + c_tensor_mc, + barrier_flag, + barrier_flag_mc, + current_stream, + ) + compare( + a_torch_cpu, + b_torch_cpu, + c_torch_gpu, + ab_dtype, + c_dtype, + tolerance, + ) + + # always free the multicast tensors first + nvshmem.core.free_tensor(c_torch_gpu_mc) + nvshmem.core.free_tensor(c_torch_gpu) + nvshmem.core.free_tensor(barrier_flag_torch_gpu_mc) + nvshmem.core.free_tensor(barrier_flag_torch_gpu) + + 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(): + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_tensor_mc, c_torch_gpu, c_torch_gpu_mc = create_mc_tensor( + c_torch_cpu, c_dtype, (1 if c_major == "n" else 0), is_dynamic_layout=True + ) + barrier_flag, barrier_flag_mc, barrier_flag_torch_gpu, barrier_flag_torch_gpu_mc = create_mc_tensor( + barrier_flag_torch_cpu, cutlass.Int32, 0, is_dynamic_layout=True + ) + + ja = testing.JitArguments( + a_tensor, + b_tensor, + c_tensor, + c_tensor_mc, + barrier_flag, + barrier_flag_mc, + current_stream, + ) + tensors_to_free = [c_torch_gpu_mc, c_torch_gpu, barrier_flag_torch_gpu_mc, barrier_flag_torch_gpu] + for tensor in tensors_to_free: + add_free_func_and_tensor(nvshmem.core.free_tensor, tensor) + return ja + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch_cpu.numel() * a_torch_cpu.element_size() + + b_torch_cpu.numel() * b_torch_cpu.element_size() + + c_torch_cpu.numel() * c_torch_cpu.element_size() + + barrier_flag_torch_cpu.numel() * barrier_flag_torch_cpu.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + if dist.get_rank() == 0: + print(f"exec_time: {exec_time}") + + for free_func, tensor in free_func_and_tensor_pairs: + free_func(tensor) + + return exec_time # Return execution time in microseconds + + +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() + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser( + description="Example of Dense Persistent GEMM on Blackwell." + ) + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.TFloat32) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--use_tma_store", action="store_true", help="Use tma store or not" + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + parser.add_argument( + "--all_reduce", + choices=["none", "LDMCxSTMC"], + type=str, + default="two_shot", + help="Allreduce algorithm to fuse with gemm", + ) + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + if args.c_dtype.width == 8: + import warnings + warnings.warn("f8 output is easy to overflow and can be nan here") + + torchrun_uid_init_bcast() + + run( + args.mnkl, + args.ab_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.use_tma_store, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + all_reduce=args.all_reduce, + ) + + torchrun_finalize() diff --git a/examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py b/examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py new file mode 100644 index 00000000..2cee2b31 --- /dev/null +++ b/examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py @@ -0,0 +1,2624 @@ +# Copyright (c) 2025 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 argparse +import ctypes +import os +from math import prod +from functools import partial +from typing import Optional, Tuple, Type, Union + +import numpy as np +import torch +import torch.distributed as dist +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 +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Pointer, Int32, Float16, BFloat16, Float32, Uint32, Float8E4M3FN, Float8E5M2 +from cutlass.cutlass_dsl import dsl_user_op, T +from cutlass._mlir.dialects import llvm, nvvm +from cutlass._mlir.dialects.nvvm import ( + MemOrderKind, + MemScopeKind, + AtomicOpKind, +) + +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 high-performance distributed persistent batched dense GEMM example for the NVIDIA Blackwell SM100 architecture +using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") +- Matrix C_mc is a multicast C matrix that changes can be broadcasted to all GPUs by multimem instructions. + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Support persistent tile scheduling to better overlap memory load/store with mma between tiles + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + - Support all-reduce epilogue with multimem instructions to distribute the workload to all GPUs + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) +4. Reduce scatter epilogue: + - Load and reduce the 128bit data from all ranks by multimem instructions. + - Broadcast the reduced data to all ranks by multimem instructions. + - current implementation only supports two_shot all-reduce which means each rank only computes a portion of + the output tensor and broadcast the result to all ranks. + - the all-reduce epilogue is only supported when use_tma_store is True. + - the all-reduce epilogue is only supported when c_dtype is Float16, Float32, BFloat16, Float8E4M3FN, Float8E5M2. + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Input arguments to this example is same as dense_gemm.py. + +.. code-block:: bash + + torchrun --nproc-per-node 8 examples/distributed/distributed_gemm_reduce_scatter_blackwell.py \ + --ab_dtype Float8E4M3FN --c_dtype Float16 --acc_dtype Float32 \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 16384,4080,4096,1 --warmup_iterations 3 --iterations 10 \ + --use_tma_store --use_2cta_instrs --reduce_scatter two_shot + +To collect performance with NSYS profiler: + +.. code-block:: bash + + nsys profile --gpu-metrics-devices=cuda-visible \ + torchrun --nproc-per-node 8 examples/distributed/distributed_gemm_reduce_scatter_blackwell.py \ + --ab_dtype Float8E5M2 --c_dtype Float16 --acc_dtype Float32 \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 16384,4096,4096,1 \ + --use_tma_store --use_2cta_instrs --warmup_iterations 3 --iterations 10 \ + --reduce_scatter two_shot + + +Constraints are same as dense_gemm_persistent.py: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), + see detailed valid dtype combinations in below PersistentDenseGemmKernel class documentation +* A/B tensor must have the same data type +* Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* Mma tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled +* when reduce_scatter, M must be multiple of 128, world_size must be 2, 4, 8 + +""" + + +class PersistentDenseGemmKernel: + """This class implements batched matrix multiplication (C = A x B) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Whether to use Tensor Memory Access (TMA) for storing results + :type use_tma_store: bool + :param reduce_scatter: reduce scatter mode, can be "two_shot" + :type reduce_scatter: str + + :note: In current version, A and B tensor must have the same data type + - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + + :note: Supported A/B data types: + - TFloat32 + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported C data types: + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + + Example: + >>> gemm = PersistentDenseGemmKernel( + ... acc_dtype=cutlass.Float32, + ... use_2cta_instrs=True, + ... mma_tiler_mn=(128, 128), + ... cluster_shape_mn=(2, 2) + ... ) + >>> gemm(a_tensor, b_tensor, c_tensor, max_active_clusters, stream) + """ + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_tma_store: bool, + reduce_scatter="two_shot", + ): + """Initializes the configuration for a Blackwell dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant + with cta_group=2 should be used. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + 3. Output C tensor store mode: + - use_tma_store: Boolean indicating whether to use Tensor Memory Access (TMA) for storing results. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Use Tensor Memory Access (TMA) or normal store for output C tensor. + :type use_tma_store: bool + :param reduce_scatter: reduce scatter mode, can be "two_shot" + :type reduce_scatter: str + """ + + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.use_tma_store = use_tma_store + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.reduce_scatter = reduce_scatter + + self.occupancy = 1 + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.reduce_scatter = reduce_scatter + self.reduce_scatter_warp_id = (6, 7, 8, 9) + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + *self.epilog_warp_id, + *self.reduce_scatter_warp_id, + ) + ) + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.cta_sync_bar_id = 0 + self.epilog_sync_bar_id = 1 + self.tmem_ptr_sync_bar_id = 2 + self.reduce_scatter_sync_bar_id = 3 + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + + self.num_ranks = 1 + self.rank_id = 0 + self.num_ranks = torch.distributed.get_world_size() + self.rank_id = torch.distributed.get_rank() + + def is_valid(self): + mma_m, mma_n = self.mma_tile_shape_mn + if (mma_m // (2 if self.use_2cta_instrs else 1)) not in [64, 128]: + return False + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + return False + if self.cluster_shape_mn[0] == 4 and self.cluster_shape_mn[1] == 4: + return False + return True + + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C stage counts in shared memory + - Computing A/B/C shared memory layout + - Computing tensor memory allocation columns + """ + # Configure tiled mma + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + # Compute epilogue subtile + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.smem_capacity, + self.occupancy, + self.use_tma_store, + ) + + # Compute A/B/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.c_smem_layout_staged = ( + sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + if self.use_tma_store + else None + ) + + # Compute the number of tensor memory allocation columns + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + c_mc: cute.Tensor = None, + c_peer_tensors: list = None, + barrier_flag: cute.Tensor = None, + barrier_flag_mc: cute.Tensor = None, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + :param c_mc: Output symmetric tensor C_mc, any write or read to a multicast tensor will be broadcasted to all GPUs + :type c_mc: cute.Tensor + :param c_peer_tensors: List of peer tensors for all-reduce + :type c_peer_tensors: List[cute.Tensor] + :param barrier_flag: Barrier flag for all-reduce + :type barrier_flag: cute.Tensor + :param barrier_flag_mc: Barrier flag for multicast tensor + :type barrier_flag_mc: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + :raises AssertionError: If OOB (Out-Of-Bounds) tiles are present when TMA store is disabled. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + # Setup TMA store for C + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + + self.buffer_align_bytes = 1024 + + c_smem_size = ( + cute.cosize(self.c_smem_layout_staged.outer) if self.use_tma_store else 0 + ) + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + c_smem_size, + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + c_mc, + c_peer_tensors, + barrier_flag, + barrier_flag_mc, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: Optional[cute.CopyAtom], + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + c_mc: cute.Tensor, + c_peer_tensors: list, + barrier_flag: cute.Tensor, + barrier_flag_mc: cute.Tensor, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.use_tma_store): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr + tmem_holding_buf = storage.tmem_holding_buf + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Tensor memory dealloc barrier init + if use_2cta_instrs: + if warp_idx == self.tma_warp_id: + num_tmem_dealloc_threads = 32 + with cute.arch.elect_one(): + cute.arch.mbarrier_init( + tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads + ) + cute.arch.mbarrier_init_fence() + + # Cluster arrive after barrier init + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + # + # Setup smem tensor A/B/C + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = ( + storage.sC.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + if self.use_tma_store + else None + ) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # + # Cluster wait before tensor memory alloc + # + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + cute.arch.barrier( + barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta + ) + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id)) + cute.arch.barrier( + barrier_id=self.tmem_ptr_sync_bar_id, + number_of_threads=tmem_ptr_read_threads, + ) + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = cute.arch.retrieve_tmem_ptr( + self.acc_dtype, + alignment=16, + ptr_to_buffer_holding_addr=tmem_holding_buf, + ) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait( + ab_consumer_state, peek_ab_full_status + ) + + # tCtAcc += tCrA * tCrB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = ( + None, + None, + kblock_idx, + ab_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord], + tCrB[kblock_coord], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + if warp_idx == self.epilog_warp_id[0]: + cute.arch.alloc_tmem( + self.num_tmem_alloc_cols, + tmem_holding_buf, + is_two_cta=use_2cta_instrs, + ) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id)) + cute.arch.barrier( + barrier_id=self.tmem_ptr_sync_bar_id, + number_of_threads=tmem_ptr_read_threads, + ) + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = cute.arch.retrieve_tmem_ptr( + self.acc_dtype, + alignment=16, + ptr_to_buffer_holding_addr=tmem_holding_buf, + ) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue + # + epi_tidx = tidx + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs + ) + + tTR_rC = None + tiled_copy_r2s = None + simt_atom = None + tRS_rC = None + tRS_sC = None + bSG_sC = None + bSG_gC_partitioned = None + tTR_gC_partitioned = None + if cutlass.const_expr(self.use_tma_store): + tTR_rC = cute.make_fragment(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + else: + ( + simt_atom, + tTR_rC, + tTR_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tiled_copy_t2r, tCgC, epi_tile, sC + ) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + c_pipeline = None + if cutlass.const_expr(self.use_tma_store): + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id) + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + bSG_gC = None + tTR_gC = None + if cutlass.const_expr(self.use_tma_store): + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + else: + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + if cutlass.const_expr(self.use_tma_store): + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + else: + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + if cutlass.const_expr(self.use_tma_store): + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + epilog_threads = 32 * len(self.epilog_warp_id) + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, + number_of_threads=epilog_threads, + ) + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, + number_of_threads=epilog_threads, + ) + else: + # + # Convert to C type + # + acc_vec = tTR_rAcc.load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tTR_rC.store(acc_vec) + + # + # Store C to global memory + # + cute.copy( + simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)] + ) + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # Reduce Scatter + if cutlass.const_expr(self.reduce_scatter == "two_shot"): + tile_id = Int32( + tile_sched._current_work_linear_idx + * cute.size(self.cluster_shape_mn) + cute.arch.block_idx_in_cluster() + ) + if warp_idx == self.epilog_warp_id[0]: + cute.arch.cp_async_bulk_wait_group(0, read=False) + # System barrier to make sure that data from each GPU is in memory before reduce scatter + with cute.arch.elect_one(): + flag = barrier_flag_mc.iterator + tile_id + utils.distributed.multimem_red_add1( + lock_ptr=flag, + scope="gpu", + order="release", + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Dealloc the tensor memory buffer + # + if warp_idx == self.epilog_warp_id[0]: + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=use_2cta_instrs) + epilog_threads = 32 * len(self.epilog_warp_id) + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads + ) + if warp_idx == self.epilog_warp_id[0]: + if use_2cta_instrs: + cute.arch.mbarrier_arrive( + tmem_dealloc_mbar_ptr, cta_rank_in_cluster ^ 1 + ) + cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0) + cute.arch.dealloc_tmem( + tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=use_2cta_instrs + ) + # + # Wait for C store complete + # + if cutlass.const_expr(self.use_tma_store): + c_pipeline.producer_tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Reduce Scatter warps + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.reduce_scatter == "two_shot"): + if warp_idx >= self.reduce_scatter_warp_id[0]: + # /////////////////////////////////////////////////////////////////////////////// + # Add persistent tile loop + # /////////////////////////////////////////////////////////////////////////////// + + rank_id = self.rank_id + num_ranks = Int32(self.num_ranks) + lane_id = cute.arch.lane_idx() + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # we want 128bit ld/st for better performance + atom_val = 128 // c_mc.element_type.width + atom_thr_n = self.mma_tiler[1] // atom_val + atom_thr_m = len(self.reduce_scatter_warp_id) * cute.arch.WARP_SIZE // atom_thr_n + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1)) + val_layout = cute.make_layout((1, atom_val), stride=(atom_val, 1)) + + copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_mc.element_type) + tiled_copy_fake = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + thr_copy_fake = tiled_copy_fake.get_slice(tidx-self.reduce_scatter_warp_id[0]*32) + + # predicate tensor + idC = cute.make_identity_tensor(c_mc.shape) + + # partition and slice at tile level + gC_mc = cute.local_tile( + c_mc, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + cC = cute.local_tile( + idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + + m_tiles_in_total = gC_mc.shape[2] + n_tiles_in_total = gC_mc.shape[3] + m_tiles_per_rank = m_tiles_in_total // self.num_ranks + + while work_tile.is_valid_tile: + + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + ((cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape)) % m_tiles_per_rank) + self.rank_id * m_tiles_per_rank, + cur_tile_coord[1], + cur_tile_coord[2], + ) + + chunk_id = ((cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape)) // m_tiles_per_rank) + tile_id = mma_tile_coord_mnl[0] + mma_tile_coord_mnl[1] * m_tiles_in_total + tile_id = tile_id * cute.size(tiled_mma.thr_id.shape) + + if not is_leader_cta: + tile_id = tile_id + 1 + + # System barrier to make sure that data from each GPU is in memory before Reduce Scatter + flag = barrier_flag.iterator + tile_id + + if warp_idx == self.reduce_scatter_warp_id[0]: + if lane_id == 0: + res = 0 + while res < self.num_ranks: + res = nvvm.load_ext(T.i32(), flag.llvm_ptr, order=MemOrderKind.RELAXED, scope=MemScopeKind.GPU) + cute.arch.barrier( + barrier_id=self.reduce_scatter_sync_bar_id, + number_of_threads=32 * len(self.reduce_scatter_warp_id), + ) + if warp_idx == self.reduce_scatter_warp_id[0]: + if lane_id == 0: + res = nvvm.atomicrmw( + T.i32(), + AtomicOpKind.ADD, + flag.llvm_ptr, + Int32(1).ir_value(), + mem_order=MemOrderKind.RELAXED, + syncscope=MemScopeKind.SYS, + ) + res = nvvm.load_ext(T.i32(), flag.llvm_ptr, order=MemOrderKind.RELAXED, scope=MemScopeKind.SYS) + if res == self.num_ranks*2: + nvvm.store_ext(Int32(0), flag.llvm_ptr, order=MemOrderKind.RELAXED, scope=MemScopeKind.SYS) + tCgC_mc = thr_mma.partition_C(gC_mc) + tCpC = thr_mma.partition_C(cC) + + tCgC_mc_slice = tCgC_mc[((None, None), 0, 0, *mma_tile_coord_mnl)] + tCpC_slice = tCpC[((None, None), 0, 0, *mma_tile_coord_mnl)] + + cta_mma_tile_m = self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) + m_local_rank = int(cta_mma_tile_m / self.num_ranks) + tCgC_mc_slice_partitioned = cute.zipped_divide(tCgC_mc_slice, (m_local_rank, self.mma_tiler[1])) + tCpC_slice_partitioned = cute.zipped_divide(tCpC_slice, (m_local_rank, self.mma_tiler[1])) + + tCgC_mc_local_rank = cute.slice_(tCgC_mc_slice_partitioned, ((None, None), (chunk_id, 0))) + tCpC_local_rank = cute.slice_(tCpC_slice_partitioned, ((None, None), (chunk_id, 0))) + frgC_mc = thr_copy_fake.partition_S(tCgC_mc_local_rank) + frpC = thr_copy_fake.partition_S(tCpC_local_rank) + + m_idx = gC_mc.shape[0] * mma_tile_coord_mnl[0] + m_per_rank = c_mc.shape[0] // self.num_ranks + dst_rank_this_tile = m_idx // m_per_rank + + peer = c_peer_tensors[rank_id] + + gC_peer = cute.local_tile( + peer, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + tCgC_peer = thr_mma.partition_C(gC_peer) + tCgC_peer_slice = tCgC_peer[((None, None), 0, 0, *mma_tile_coord_mnl)] + cta_mma_tile_m = self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) + m_local_rank = int(cta_mma_tile_m / self.num_ranks) + tCgC_peer_slice_partitioned = cute.zipped_divide(tCgC_peer_slice, (m_local_rank, self.mma_tiler[1])) + tCgC_peer_local_rank = cute.slice_(tCgC_peer_slice_partitioned, ((None, None), (chunk_id, 0))) + frgC_peer = thr_copy_fake.partition_S(tCgC_peer_local_rank) + + atom, loop_m, loop_n = frgC_mc.shape + tmp_results = cute.make_rmem_tensor((4, loop_m, loop_n), cutlass.Int32) + local_chunk_lower_bound = (rank_id * m_per_rank, c_mc.shape[1], c_mc.shape[2]) + local_chunk_upper_bound = ((rank_id + 1) * m_per_rank, c_mc.shape[1], c_mc.shape[2]) + + for i in cutlass.range_constexpr(loop_m): + for j in cutlass.range_constexpr(loop_n): + if cute.elem_less(frpC[0, i, j], local_chunk_upper_bound) and not cute.elem_less(frpC[0, i, j], local_chunk_lower_bound): + mc_ptr = frgC_mc[None, i, j].iterator + if cutlass.const_expr(self.c_dtype == Float16): + x, y, z, w = utils.distributed.multimem_ld_reduce_8xf16(mc_ptr) + elif cutlass.const_expr(self.c_dtype == Float32): + x, y, z, w = utils.distributed.multimem_ld_reduce_4xf32(mc_ptr) + elif cutlass.const_expr(self.c_dtype == BFloat16): + x, y, z, w = utils.distributed.multimem_ld_reduce_8xbf16(mc_ptr) + elif cutlass.const_expr(self.c_dtype == Float8E4M3FN): + x, y, z, w = utils.distributed.multimem_ld_reduce_16xe4m3(mc_ptr) + elif cutlass.const_expr(self.c_dtype == Float8E5M2): + x, y, z, w = utils.distributed.multimem_ld_reduce_16xe5m2(mc_ptr) + tmp_results[0, i, j] = x + tmp_results[1, i, j] = y + tmp_results[2, i, j] = z + tmp_results[3, i, j] = w + + for i in cutlass.range_constexpr(loop_m): + for j in cutlass.range_constexpr(loop_n): + if cute.elem_less(frpC[0, i, j], local_chunk_upper_bound) and not cute.elem_less(frpC[0, i, j], local_chunk_lower_bound): + ptr_int = frgC_peer[None, i, j].iterator.toint().ir_value() + x, y, z, w = tmp_results[0, i, j].ir_value(), tmp_results[1, i, j].ir_value(), tmp_results[2, i, j].ir_value(), tmp_results[3, i, j].ir_value() + llvm.inline_asm( + T.i32(), + [ptr_int, x, y, z, w], + "st.global.sys.relaxed.v4.f32 [$1], {$2, $3, $4, $5};", + "=r,l,r,r,r,r", + has_side_effects=True, + asm_dialect=0 + ) + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + cute.arch.barrier( + barrier_id=self.reduce_scatter_sync_bar_id, + number_of_threads=32 * len(self.reduce_scatter_warp_id), + ) + # System barrier to make sure all the peer memory transfers are completed. + last_flag_idx = cute.size( + tile_sched.params.problem_layout_ncluster_mnl + ) * cute.size(self.cluster_shape_mn) + if warp_idx == self.reduce_scatter_warp_id[0]: + with cute.arch.elect_one(): + # Offset to last tile flag idx + last_tile_id_linear = cute.size( + tile_sched.params.problem_layout_ncluster_mnl + ) * cute.size(self.cluster_shape_mn) + # 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( + lock_ptr=barrier_flag_mc.iterator + + last_tile_id_linear + + sm_id_linear, + scope="sys", + order="release", + ) + # Relaxed spin-lock wait flag with sys scope + utils.distributed.spin_lock_atom_cas_relaxed_wait( + lock_ptr=barrier_flag.iterator + + last_tile_id_linear + + sm_id_linear, + expected_val=self.num_ranks, + reset_val=0, + scope="sys", + ) + + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_fragment( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + - partition register array (source) and global memory (destination) for none TMA store version; + - partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing either: + - For TMA store: (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + - For non-TMA store: (simt_atom, tTR_rC, tTR_gC) where: + - simt_atom: The SIMT copy atom + - tTR_rC: The register tensor C + - tTR_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + if cutlass.const_expr(self.use_tma_store): + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + return tma_atom_c, bSG_sC, bSG_gC + else: + tiled_copy_t2r = atom + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_gC = thr_copy_t2r.partition_D(gC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_fragment( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype + ) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + return simt_atom, tTR_rC, tTR_gC + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + ) -> Tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + c_smem_layout_staged_one = ( + sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + if use_tma_store + else None + ) + ab_bytes_per_stage = cute.size_in_bytes( + a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + c_bytes_per_stage = ( + cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + if use_tma_store + else 0 + ) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes + # Add remaining unused smem to epilogue + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def _compute_num_tmem_alloc_cols( + tiled_mma: cute.TiledMma, + mma_tiler: Tuple[int, int, int], + num_acc_stage: int, + ) -> int: + """ + Compute the number of tensor memory allocation columns. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler: The shape (M, N, K) of the MMA tile. + :type mma_tiler: tuple[int, int, int] + :param num_acc_stage: The stage of the accumulator tensor. + :type num_acc_stage: int + + :return: The number of tensor memory allocation columns. + :rtype: int + """ + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + + return num_tmem_alloc_cols + + def is_valid_dtypes( + self, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + ) -> bool: + """ + Check if the dtypes are valid + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes are valid, False otherwise + :rtype: bool + """ + valid_ab_dtypes = { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Uint8, + cutlass.Int8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + if ab_dtype not in valid_ab_dtypes: + return False + + if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: + return False + + # Define compatibility mapping between accumulator type and AB type + acc_ab_compatibility = { + cutlass.Float32: { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, # Float32 accumulator supports floating point AB types only + cutlass.Float16: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, + } + # Check compatibility between accumulator type and AB type + if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: + return False + + # Define compatibility mapping between accumulator type and C type + acc_c_compatibility = { + cutlass.Float32: { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + cutlass.Float16: { + cutlass.BFloat16, + cutlass.Float16, + }, + cutlass.Int32: { + cutlass.BFloat16, + cutlass.Float16, + cutlass.Float32, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + } + # Check compatibility between accumulator type and C type + if c_dtype not in acc_c_compatibility[self.acc_dtype]: + return False + + # check if c_dtype is supported by multimem all-reduce + if cutlass.const_expr(c_dtype not in {cutlass.Float16, cutlass.Float32, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2}): + return False + + return True + + def is_valid_mma_tiler_and_cluster_shape(self) -> bool: + """Check if the mma tiler and cluster shape are valid. + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if not ( + (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) + or (self.use_2cta_instrs and self.mma_tiler_mn[0] in [128, 256]) + ): + is_valid = False + if self.mma_tiler_mn[1] not in range(32, 257, 32): + is_valid = False + # Skip illegal cluster shape + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16 + or self.cluster_shape_mn[0] <= 0 + or self.cluster_shape_mn[1] <= 0 + or not is_power_of_2(self.cluster_shape_mn[0]) + or not is_power_of_2(self.cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + def is_valid_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + + return is_valid + + def is_valid_epilog_store_option( + self, + m: int, + n: int, + ) -> bool: + """ + Check if the epilogue store option is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + + :return: True if the epilogue store option is valid, False otherwise + :rtype: bool + """ + + is_valid = True + # None TMA store version does not have predication, can not support OOB tiles + cta_tile_shape_mn = ( + self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler_mn[1], + ) + if not self.use_tma_store: + if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): + is_valid = False + return is_valid + + def can_implement(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor) -> bool: + """Check if the given tensors can be implemented by this kernel. + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + + :return: True if the gemm supports the given config, False otherwise + :rtype: bool + """ + m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] + + # infer a_major, b_major, c_major + is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a() + is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b() + is_m_major_c = utils.LayoutEnum.from_tensor(c).is_m_major_c() + a_major = "m" if is_m_major_a else "k" + b_major = "n" if is_n_major_b else "k" + c_major = "m" if is_m_major_c else "n" + + can_implement = True + # Skip unsupported types + if not self.is_valid_dtypes(a.element_type, c.element_type): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not self.is_valid_mma_tiler_and_cluster_shape(): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not self.is_valid_tensor_alignment( + m, n, k, l, a.element_type, c.element_type, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid epilogue store option + if not self.is_valid_epilog_store_option(m, n): + can_implement = False + if dist.get_world_size() not in [2, 4, 8]: + can_implement = False + if m % 128 != 0: + # cannot support OOB tiles when m is not divisible by 128 + can_implement = False + + return can_implement + + +def create_mc_tensor(torch_tensor_cpu, dtype, leading_dim, is_dynamic_layout=True): + torch_tensor_gpu_local = nvshmem.core.tensor(torch_tensor_cpu.shape, dtype=torch_tensor_cpu.dtype) + torch_tensor_gpu_local.copy_(torch_tensor_cpu) + torch_tensor_gpu_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, torch_tensor_gpu_local) + cute_tensor_mc = from_dlpack( + torch_tensor_gpu_mc, + assumed_align=16, + ) + + cute_tensor_c_peer_torch_tensors = [nvshmem.core.get_peer_tensor(torch_tensor_gpu_local, rank) for rank in range(dist.get_world_size())] + cute_tensor_c_peer_tensors = [from_dlpack(t) for t in cute_tensor_c_peer_torch_tensors] + + if is_dynamic_layout: + cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim) + + cute_tensor = from_dlpack(torch_tensor_gpu_local, assumed_align=16) + cute_tensor.element_type = dtype + if is_dynamic_layout: + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + cute_tensor = cutlass_torch.convert_cute_tensor( + torch_tensor_gpu_local, + cute_tensor, + dtype, + is_dynamic_layout=is_dynamic_layout, + ) + return cute_tensor, cute_tensor_mc, torch_tensor_gpu_local, torch_tensor_gpu_mc, cute_tensor_c_peer_torch_tensors, cute_tensor_c_peer_tensors + + +def create_tensors( + l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype +): + torch.manual_seed(1111) + + a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype).to(torch.float32).normal_().round_().to(dtype=cutlass_torch.dtype(ab_dtype)) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype).to(torch.float32).normal_().round_().to(dtype=cutlass_torch.dtype(ab_dtype)) + c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype) + + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( + c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + c_tensor, c_tensor_mc, c_torch_gpu, c_torch_gpu_mc, c_peer_torch_tensors, c_peer_tensors = create_mc_tensor( + c_torch_cpu, c_dtype, (1 if c_major == "n" else 0), is_dynamic_layout=True + ) + + return ( + a_tensor, + b_tensor, + c_tensor, + c_tensor_mc, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + c_torch_gpu, + c_torch_gpu_mc, + c_peer_torch_tensors, + c_peer_tensors, + ) + + +def compare( + a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance +): + # Copy gpu result back + kernel_result = c_torch_gpu.cpu() + + # Compute reference result + ref = torch.einsum( + "mkl,nkl->mnl", + a_torch_cpu.to(dtype=torch.float32), + b_torch_cpu.to(dtype=torch.float32), + ) + + acc_type = cutlass.Float32 + if c_dtype.width == 16: + acc_dtype = cutlass.Float32 + elif c_dtype.width == 8: + acc_dtype = cutlass.Float16 + + torch_acc_dtype = cutlass_torch.dtype(acc_dtype) + torch_c_dtype = cutlass_torch.dtype(c_dtype) + + # note that in this example, f8 reduce sum is done by f16 acc precision and f16 by f32 + ref = ref.to(torch_c_dtype).to(torch_acc_dtype).cuda() + torch.distributed.all_reduce(ref, op=torch.distributed.ReduceOp.SUM) + ref = ref.to(torch_c_dtype) + + # convert to higher precision to satisfy assert_close function + ref_result = ref.cpu().to(torch_acc_dtype) + kernel_result = c_torch_gpu.view(torch_c_dtype).cpu().to(torch_acc_dtype) + + max_val = torch.finfo(kernel_result.dtype).max + min_val = torch.finfo(kernel_result.dtype).min + kernel_result = torch.nan_to_num(kernel_result, nan=max_val, posinf=max_val, neginf=min_val) + ref_result = torch.nan_to_num(ref_result, nan=max_val, posinf=max_val, neginf=min_val) + + rank_id = dist.get_rank() + chunk_per_rank = kernel_result.shape[0] // dist.get_world_size() + start_idx = rank_id * chunk_per_rank + end_idx = start_idx + chunk_per_rank + + # Assert close results + torch.testing.assert_close(kernel_result[start_idx:end_idx,:], ref_result[start_idx:end_idx,:], atol=tolerance, rtol=1e-05) + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + reduce_scatter: str = "two_shot", + **kwargs, +): + """Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + + This function prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks the execution performance. + + :param mnkl: Problem size (M, N, K, L) + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C + :type c_dtype: Type[cutlass.Numeric] + :param acc_dtype: Data type for accumulation during matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param a_major/b_major/c_major: Memory layout of tensor A/B/C + :type a_major/b_major/c_major: str + :param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the + default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters. + :type mma_tiler_mn: Tuple[int, int], optional + :param cluster_shape_mn: Cluster shape. If not specified in the decorator parameters, the autotuner will use the + default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters. + :type cluster_shape_mn: Tuple[int, int], optional + :param use_2cta_instrs: Whether to use 2CTA instructions. If not specified in the decorator parameters, the autotuner + will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. + :type use_2cta_instrs: bool, optional + :param use_tma_store: Whether to use TMA store. If not specified in the decorator parameters, the autotuner will use + the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. + :type use_tma_store: bool, optional + :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 1 + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False + :type use_cold_l2: bool, optional + :param reduce_scatter: reduce scatter mode, can be "two_shot" + :type reduce_scatter: str, optional + :raises RuntimeError: If CUDA GPU is not available + :raises ValueError: If the configuration is invalid or unsupported by the kernel + :return: Execution time of the GEMM kernel + :rtype: float + """ + print(f"Running Blackwell Persistent Dense GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + print(f"Fused Reduce Scatter Op: {reduce_scatter}") + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + ( + a_tensor, + b_tensor, + c_tensor, + c_tensor_mc, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + c_torch_gpu, + c_torch_gpu_mc, + c_peer_torch_tensors, + c_peer_tensors, + ) = create_tensors( + l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype + ) + + # Build GEMM object + gemm = PersistentDenseGemmKernel( + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + reduce_scatter=reduce_scatter, + ) + + # Check if configuration can be implemented + can_implement = gemm.can_implement(a_tensor, b_tensor, c_tensor) + if not can_implement: + raise ValueError( + f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " + f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " + f"use_tma_store = {use_tma_store}" + ) + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + def create_barrier_flags(): + cta_tile_shape_mn = ( + mma_tiler_mn[0] // (2 if use_2cta_instrs else 1), + mma_tiler_mn[1], + ) + problem_shape_ntile_mn = (m // cta_tile_shape_mn[0], n // cta_tile_shape_mn[1]) + num_tiles = problem_shape_ntile_mn[0] * problem_shape_ntile_mn[1] + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + + barrier_flag_torch = nvshmem.core.tensor( + (num_tiles + num_sms,), dtype=torch.int32 + ) + barrier_flag_torch.fill_(0) + barrier_flag_torch_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, barrier_flag_torch) + + barrier_flag = from_dlpack(barrier_flag_torch) + barrier_flag = barrier_flag.mark_layout_dynamic() + barrier_flag_mc = from_dlpack(barrier_flag_torch_mc) + barrier_flag_mc = barrier_flag_mc.mark_layout_dynamic() + + return barrier_flag_torch, barrier_flag_torch_mc, barrier_flag, barrier_flag_mc + + barrier_flag_torch, barrier_flag_torch_mc, barrier_flag, barrier_flag_mc = create_barrier_flags() + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + c_tensor, + max_active_clusters, + current_stream, + c_mc=c_tensor_mc, + c_peer_tensors=c_peer_tensors, + barrier_flag=barrier_flag, + barrier_flag_mc=barrier_flag_mc, + ) + + if not skip_ref_check: + compiled_gemm( + a_tensor, + b_tensor, + c_tensor, + current_stream, + c_mc=c_tensor_mc, + c_peer_tensors=c_peer_tensors, + barrier_flag=barrier_flag, + barrier_flag_mc=barrier_flag_mc, + ) + compare( + a_torch_cpu, + b_torch_cpu, + c_torch_gpu, + c_dtype, + tolerance, + ) + + nvshmem.core.free_tensor(c_torch_gpu_mc) + nvshmem.core.free_tensor(c_torch_gpu) + for i in range(len(c_peer_tensors)): + if i != dist.get_rank(): + nvshmem.core.free_tensor(c_peer_torch_tensors[i]) + nvshmem.core.free_tensor(barrier_flag_torch_mc) + nvshmem.core.free_tensor(barrier_flag_torch) + + 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(): + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, _ = cutlass_torch.cute_tensor_like( + c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_tensor_mc, c_torch_gpu, c_torch_gpu_mc, c_peer_torch_tensors, c_peer_tensors = create_mc_tensor( + c_torch_cpu, + c_dtype, + (1 if c_major == "n" else 0), + is_dynamic_layout=True, + ) + barrier_flag_torch, barrier_flag_torch_mc, barrier_flag, barrier_flag_mc = create_barrier_flags() + add_free_func_and_tensor(nvshmem.core.free_tensor, c_torch_gpu_mc) + add_free_func_and_tensor(nvshmem.core.free_tensor, c_torch_gpu) + for i in range(len(c_peer_torch_tensors)): + if i != dist.get_rank(): + add_free_func_and_tensor(nvshmem.core.free_tensor, c_peer_torch_tensors[i]) + add_free_func_and_tensor(nvshmem.core.free_tensor, barrier_flag_torch_mc) + add_free_func_and_tensor(nvshmem.core.free_tensor, barrier_flag_torch) + return testing.JitArguments( + a_tensor, + b_tensor, + c_tensor, + current_stream, + c_mc=c_tensor_mc, + c_peer_tensors=c_peer_tensors, + barrier_flag=barrier_flag, + barrier_flag_mc=barrier_flag_mc, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch_cpu.numel() * a_torch_cpu.element_size() + + b_torch_cpu.numel() * b_torch_cpu.element_size() + + c_torch_cpu.numel() * c_torch_cpu.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + for free_func, tensor in free_func_and_tensor_pairs: + free_func(tensor) + print(f"exec_time: {exec_time}") + + return exec_time # Return execution time in microseconds + + +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() + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser( + description="Example of Dense Persistent GEMM on Blackwell." + ) + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.TFloat32) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--use_tma_store", action="store_true", help="Use tma store or not" + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + parser.add_argument( + "--reduce_scatter", + choices=["two_shot"], + type=str, + default="two_shot", + help="Reduce Scatter algorithm to fuse with gemm", + ) + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + if args.c_dtype.width == 8: + import warnings + warnings.warn("f8 output is easy to overflow and can be nan here") + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + torchrun_uid_init_bcast() + + run( + args.mnkl, + args.ab_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.use_tma_store, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + reduce_scatter=args.reduce_scatter, + ) + + torchrun_finalize() + +