# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: # 1. Redistributions of source code must retain the above copyright notice, this # list of conditions and the following disclaimer. # 2. Redistributions in binary form must reproduce the above copyright notice, # this list of conditions and the following disclaimer in the documentation # and/or other materials provided with the distribution. # 3. Neither the name of the copyright holder nor the names of its # contributors may be used to endorse or promote products derived from # this software without specific prior written permission. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import traceback import typing import cuda.bindings.driver as cuda # Required for pre-Python 3.12 instead of typing.override. from typing_extensions import override import torch import cutlass import cutlass.cute.testing as testing import cutlass.torch as cutlass_torch from common_dense_gemm_efc import DenseGemmEFC """ A high-performance persistent batched dense GEMM with custom epilogue fusion for the NVIDIA Blackwell SM100 architecture using CUTE DSL and Epilogue Fusion Configuration (EFC). This example demonstrates a GEMM with a custom fused epilogue that performs: Y = A * B (accumulator stored to Y) D = (A * B) * alpha + C * beta + X * x_factor Tensor dimensions: - 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 (read-only input), C can be row-major("N") or column-major("M") - Matrix D is MxNxL (output), D can be row-major("N") or column-major("M") - Matrix X is MxNxL (read-only input), same layout as C/D - Matrix Y is MxNxL (output), same layout as C/D - alpha, beta, and x_factor are scalar scale factors 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 - Supports persistent tile scheduling to better overlap memory load/store with mma between tiles - Supports warp specialization to avoid explicit pipelining between mainloop load and mma - Uses Epilogue Fusion Configuration (EFC) to define custom epilogue operations 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 (defined via EFC): - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. - Load C and X matrices from global memory (GMEM) to shared memory (SMEM) using TMA, then to registers (RMEM). - Compute Y = accumulator (copy of A*B result) - Compute D = accumulator * alpha + C * beta + X * x_factor - Type convert D and Y matrices to output types. - Store D and Y matrices from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations 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. Example usage: .. code-block:: bash python custom_epilogue_dense_gemm.py \ --ab_dtype Float16 --c_dtype Float16 --d_dtype Float16 --acc_dtype Float32 --epi_dtype Float32 \ --x_dtype Float16 --y_dtype Float16 \ --mma_tiler_mn 128,128 --cluster_shape_mn 2,1 \ --mnkl 8192,8192,8192,1 \ --use_2cta_instrs --alpha 2.0 --beta 1.0 --x_factor 3.0 To collect performance with NCU profiler: .. code-block:: bash ncu python custom_epilogue_dense_gemm.py \ --ab_dtype Float16 --c_dtype Float16 --d_dtype Float16 --acc_dtype Float32 --epi_dtype Float32 \ --x_dtype Float16 --y_dtype Float16 \ --mma_tiler_mn 128,128 --cluster_shape_mn 2,1 \ --mnkl 8192,8192,8192,1 \ --use_2cta_instrs --alpha 2.0 --beta 1.0 --x_factor 3.0 \ --warmup_iterations 1 --iterations 10 --skip_ref_check Constraints: * Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2) * A/B tensors must have the same data type * C/D/X/Y tensors must have the same major order * 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 all 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 DenseGemmAlphaBeta(DenseGemmEFC): """Implements batched GEMM with custom epilogue fusion using EFC. This class extends DenseGemmEFC to provide a fused epilogue that: - Reads from input tensors C and X - Writes to output tensors D and Y - Performs: Y = A*B and D = (A*B) * alpha + C * beta + X * x_factor The class provides CLI argument parsing and tensor creation for the specific epilogue configuration with C, D, X, Y tensors and alpha, beta, x_factor scalar parameters. """ class CLIParser(DenseGemmEFC.CLIParser): @override def more_parsing(self): self.parser.add_argument( "--alpha", type=float, default=1.0, help="alpha scale factor" ) self.parser.add_argument( "--beta", type=float, default=0.0, help="beta scale factor" ) self.parser.add_argument( "--c_dtype", type=cutlass.dtype, default=cutlass.Float32 ) self.parser.add_argument( "--d_dtype", type=cutlass.dtype, default=cutlass.Float32 ) self.parser.add_argument( "--x_dtype", type=cutlass.dtype, default=cutlass.Float32 ) self.parser.add_argument( "--x_factor", type=float, default=3.0, help="x_factor scale factor" ) self.parser.add_argument( "--y_dtype", type=cutlass.dtype, default=cutlass.Float32 ) @override def create_arguments( self, l, m, n, k, a_major, b_major, cd_major, ab_dtype, # For the supplemental tensors. c_dtype, d_dtype, x_dtype, y_dtype, ): """Create arguments for GEMM operations with epilogue tensors. Creates tensors for A, B (from parent class) and epilogue-specific tensors C, D, X, Y with appropriate data types and layouts. :return: Tuple of (a_tensor, b_tensor, a_torch_cpu, b_torch_cpu, c_tensor, c_torch_cpu, c_torch_gpu, d_tensor, d_torch_cpu, d_torch_gpu, x_tensor, x_torch_cpu, x_torch_gpu, y_tensor, y_torch_cpu, y_torch_gpu) """ # Get standard arguments from parent class std_args = super().create_arguments( l, m, n, k, a_major, b_major, cd_major, ab_dtype ) # Add the auxiliary accumulator tensors c_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", c_dtype) c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 ) d_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", d_dtype) d_tensor, d_torch_gpu = cutlass_torch.cute_tensor_like( d_torch_cpu, d_dtype, is_dynamic_layout=True, assumed_align=16 ) x_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", x_dtype) x_tensor, x_torch_gpu = cutlass_torch.cute_tensor_like( x_torch_cpu, x_dtype, is_dynamic_layout=True, assumed_align=16 ) y_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", y_dtype) y_tensor, y_torch_gpu = cutlass_torch.cute_tensor_like( y_torch_cpu, y_dtype, is_dynamic_layout=True, assumed_align=16 ) return ( *std_args, c_tensor, c_torch_cpu, c_torch_gpu, d_tensor, d_torch_cpu, d_torch_gpu, x_tensor, x_torch_cpu, x_torch_gpu, y_tensor, y_torch_cpu, y_torch_gpu, ) def compare( self, a_torch_cpu, b_torch_cpu, epi_dtype, tolerance, # For the tensor check. c_torch_gpu, d_torch_gpu, x_torch_gpu, y_torch_gpu, # The EFC epilogue arguments. c_torch_cpu, d_torch_cpu, alpha, beta, x_torch_cpu, x_factor, y_torch_cpu, ): """Compare GPU results against CPU reference implementation. Evaluates the epilogue computation on CPU and validates that: - Output tensor D matches CPU computation - Output tensor Y matches CPU computation - Input tensors C and X remain unchanged (read-only) :param a_torch_cpu: Input matrix A on CPU :param b_torch_cpu: Input matrix B on CPU :param epi_dtype: Data type for epilogue computation :param tolerance: Tolerance for numerical comparison :param c_torch_gpu: Input matrix C on GPU (to verify unchanged) :param d_torch_gpu: Output matrix D on GPU (to compare) :param x_torch_gpu: Input matrix X on GPU (to verify unchanged) :param y_torch_gpu: Output matrix Y on GPU (to compare) :param c_torch_cpu: Input matrix C on CPU :param d_torch_cpu: Output matrix D on CPU (reference) :param alpha: Scale factor for accumulator :param beta: Scale factor for C :param x_torch_cpu: Input matrix X on CPU :param x_factor: Scale factor for X :param y_torch_cpu: Output matrix Y on CPU (reference) """ # Compute reference result self.evaluate_on_cpu( a_torch_cpu, b_torch_cpu, epi_dtype, c_torch_cpu, d_torch_cpu, alpha, beta, x_torch_cpu, x_factor, y_torch_cpu, ) # Assert close results. torch.testing.assert_close( d_torch_gpu.cpu(), d_torch_cpu, atol=tolerance, rtol=1e-05 ) torch.testing.assert_close( y_torch_gpu.cpu(), y_torch_cpu, atol=tolerance, rtol=1e-05 ) # Assert that the read tensors has not been changed. torch.testing.assert_close( c_torch_gpu.cpu(), c_torch_cpu, atol=tolerance, rtol=1e-05 ) torch.testing.assert_close( x_torch_gpu.cpu(), x_torch_cpu, atol=tolerance, rtol=1e-05 ) @staticmethod def format_as_cli_args( mnkl: typing.Tuple[int, int, int, int], ab_dtype: typing.Type[cutlass.Numeric], acc_dtype: typing.Type[cutlass.Numeric], epi_dtype: typing.Type[cutlass.Numeric], a_major: str, b_major: str, cd_major: str, mma_tiler_mn: typing.Tuple[int, int], cluster_shape_mn: typing.Tuple[int, int], use_2cta_instrs: bool, c_dtype: typing.Type[cutlass.Numeric], d_dtype: typing.Type[cutlass.Numeric], alpha: float, beta: float, x_dtype: typing.Type[cutlass.Numeric], x_factor: float, y_dtype: typing.Type[cutlass.Numeric], tolerance: float, ) -> str: """Format test parameters as CLI arguments for custom_epilogue_dense_gemm.py Formats all test parameters into a CLI command that can be directly copy-pasted to reproduce the test case. Includes base parameters from DenseGemmEFC and epilogue-specific parameters (c_dtype, d_dtype, x_dtype, y_dtype, x_factor). :return: Formatted CLI command string """ # Get base command from parent class base_cmd = DenseGemmEFC.format_as_cli_args( "custom_epilogue_dense_gemm.py", mnkl, ab_dtype, acc_dtype, epi_dtype, a_major, b_major, cd_major, mma_tiler_mn, cluster_shape_mn, use_2cta_instrs, tolerance, ) # Add epilogue-specific parameters specific_args = ( f" --alpha {alpha} " f"--beta {beta} " f"--c_dtype {DenseGemmEFC.dtype_name(c_dtype)} " f"--d_dtype {DenseGemmEFC.dtype_name(d_dtype)} " f"--x_dtype {DenseGemmEFC.dtype_name(x_dtype)} " f"--y_dtype {DenseGemmEFC.dtype_name(y_dtype)} " f"--x_factor {x_factor}" ) return base_cmd + specific_args def run( mnkl: typing.Tuple[int, int, int, int], ab_dtype: typing.Type[cutlass.Numeric], acc_dtype: typing.Type[cutlass.Numeric], epi_dtype: typing.Type[cutlass.Numeric], a_major: str, b_major: str, cd_major: str, # Epilogue EFC arguments. c_dtype: typing.Type[cutlass.Numeric], d_dtype: typing.Type[cutlass.Numeric], alpha: float, beta: float, x_dtype: typing.Type[cutlass.Numeric], x_factor: float, y_dtype: typing.Type[cutlass.Numeric], # Common arguments. mma_tiler_mn: typing.Tuple[int, int], cluster_shape_mn: typing.Tuple[int, int], use_2cta_instrs: bool, tolerance: float, warmup_iterations: int = 0, iterations: int = 1, skip_ref_check: bool = False, ): """Execute batched GEMM with custom epilogue fusion. Performs: Y = A * B D = (A * B) * alpha + C * beta + X * x_factor :param mnkl: Matrix dimensions (M, N, K, L) where L is batch dimension :param ab_dtype: Data type for input matrices A and B :param acc_dtype: Data type for accumulator :param epi_dtype: Data type for epilogue computation :param a_major: Major order for A matrix ('k' or 'm') :param b_major: Major order for B matrix ('k' or 'n') :param cd_major: Major order for C/D/X/Y matrices ('n' or 'm') :param c_dtype: Data type for input matrix C :param d_dtype: Data type for output matrix D :param alpha: Scale factor for accumulator in D computation :param beta: Scale factor for C in D computation :param x_dtype: Data type for input matrix X :param x_factor: Scale factor for X in D computation :param y_dtype: Data type for output matrix Y :param mma_tiler_mn: MMA tiler dimensions (M, N) :param cluster_shape_mn: Cluster shape (M, N) :param use_2cta_instrs: Whether to use 2CTA instructions :param tolerance: Tolerance for validation :param warmup_iterations: Number of warmup iterations :param iterations: Number of iterations to run :param skip_ref_check: Skip reference checking """ print("Running Blackwell Persistent Dense GEMM test with:") print(f"mnkl: {mnkl}") print(f"AB dtype: {ab_dtype}, Acc dtype: {acc_dtype}, Epi dtype: {epi_dtype}") print( f"Matrix majors - A: {a_major}, B: {b_major}, loaded: {cd_major}, stored: {cd_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"Tolerance: {tolerance}") print(f"Warmup iterations: {warmup_iterations}") print(f"Iterations: {iterations}") print(f"Skip reference checking: {skip_ref_check}") print("Epilogue:") print(f"\t{c_dtype = !s}, {d_dtype = !s}") print(f"\t{alpha = }, {beta = }") print(f"\t{x_dtype = !s}, {x_factor = }") print(f"\t{y_dtype = !s}") # 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) # The order of the parameters here is defining the one to be used in all the # other API calls. The epilogue function does not return anything and at # least one tensor .store() is required to have a useful computation. # efc_config exposes many features to the programmer, like activation # functions or accessing some implementation details. See EFC.Configuration. def epilogue(efc_config, C, D, alpha, beta, X, x_factor, Y): # All the .load() happen before any .store(). # Store the accumulator to Y Y.store(efc_config.accum()) # Compute the result with alpha, beta scaling and X factor result = ( efc_config.relu(efc_config.accum() * alpha + C.load() * beta) + X.load() * x_factor ) D.store(result) # Build GEMM object with EFC configuration: # TODO: generalize acc_dtype and epi_dtype gemm = DenseGemmAlphaBeta( acc_dtype, epi_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, epilogue, ) ( a_tensor, b_tensor, a_torch_cpu, b_torch_cpu, # The supplemental tensors. c_tensor, c_torch_cpu, c_torch_gpu, d_tensor, d_torch_cpu, d_torch_gpu, x_tensor, x_torch_cpu, x_torch_gpu, y_tensor, y_torch_cpu, y_torch_gpu, ) = gemm.create_arguments( l, m, n, k, a_major, b_major, cd_major, ab_dtype, # For the supplemental tensors. c_dtype, d_dtype, x_dtype, y_dtype, ) # Check if the configuration can be implemented. Raise a ValueError # otherwise. gemm.check_implementable(a_tensor, b_tensor, d_tensor) max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) compiled_gemm = gemm.compile( a_tensor, b_tensor, max_active_clusters, current_stream, # Here are the supplemental arguments in the same order as for the # epilogue configuration function. c_tensor, d_tensor, alpha, beta, x_tensor, x_factor, y_tensor, # Not really useful here but this is an example of how to pass CuTe # DSL compilation options. options="--opt-level=3 --enable-assertions --generate-line-info", ) compiled_gemm( a_tensor, b_tensor, current_stream, # Here are the supplemental arguments in the same order as for the # epilogue configuration function. c_tensor, d_tensor, alpha, beta, x_tensor, x_factor, y_tensor, ) # TODO: unify with modern way to do benchmarking. exec_time = testing.benchmark( compiled_gemm, kernel_arguments=testing.JitArguments( a_tensor, b_tensor, current_stream, # Here are the supplemental arguments in the same order as for the # epilogue configuration function. c_tensor, d_tensor, alpha, beta, x_tensor, x_factor, y_tensor, ), stream=current_stream, warmup_iterations=warmup_iterations, iterations=iterations, ) print(f"Execution time: {exec_time} us") # Compute reference result if not skip_ref_check: gemm.compare( # The usual arguments. a_torch_cpu, b_torch_cpu, epi_dtype, tolerance, # For the tensor check. c_torch_gpu, d_torch_gpu, x_torch_gpu, y_torch_gpu, # The EFC epilogue arguments. c_torch_cpu, d_torch_cpu, alpha, beta, x_torch_cpu, x_factor, y_torch_cpu, ) if __name__ == "__main__": args = DenseGemmAlphaBeta.CLIParser().parse() try: run( args.mnkl, args.ab_dtype, args.acc_dtype, args.epi_dtype, args.a_major, args.b_major, args.cd_major, args.c_dtype, args.d_dtype, args.alpha, args.beta, args.x_dtype, args.x_factor, args.y_dtype, args.mma_tiler_mn, args.cluster_shape_mn, args.use_2cta_instrs, args.tolerance, args.warmup_iterations, args.iterations, args.skip_ref_check, ) print("PASS") except Exception as exc: traceback.print_exception(exc) raise