diff --git a/CHANGELOG.md b/CHANGELOG.md index 21814c68..17048ffa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,10 +2,23 @@ # CUTLASS 4.x -## [4.4.0](https://github.com/NVIDIA/cutlass/tree/main) (2026-01-23) +## [4.4.0](https://github.com/NVIDIA/cutlass/releases/tag/v4.4.0) (2026-02-14) ### CuTe DSL * New features + - CuTe DSL now supports CUDA toolkit 13.1! + + Set up with cutlass/python/CuTeDSL/setup.sh --cu13 + + Refer to https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/quick_start.html for more details + - GB300 is now supported in CuTe DSL with CTK 13.1 + + Refer to [SM103 batched 3xFP4 blockscaled GEMM kernel](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py) for example kernel + - cute.experimental: introduce a higher-level, composable layer on top of existing CuTe DSL APIs (not a separate abstraction), which can be mixed with existing Cute DSL building blocks. + + Fragment-free programming model: copy/dot APIs take memrefs directly instead of descriptors/fragments. + + Automatic TMA descriptor generation and update insertion. + + Automatic vectorization and predication for SIMT copies. + + New pipeline abstraction with convenience wrappers + + New Partition ops to simplify partitioning logic. + + Device-side TMA descriptor allocation, initialization, and management + + These examples can be found here https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/experimental - Ahead of Time (AoT) compilation is now available! + Refer to files under https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/cute/export for example usage - JAX support - you can now use CuTeDSL along with JAX @@ -14,18 +27,54 @@ + cutlass.__version__ for a string representation of DSL version + cutlass.CUDA_VERSION for a version class to tell the CUDA version used for DSL - Added CopyDsmemStoreOp to store data to distributed shared memory with explicit synchronization. + - Grouped GEMM example now supports device-only problem shapes. + - We allow grid carve-out without problem shapes being available on host. + - Tma+LdMatrix features for loading+unpacking narrow-width types (refer to mixed_input_fmha_decode.py for example usage). + - It is possible now to have customized epilogue fusion for persistent dense GEMM through a Python Epilogue Fusion Configuration (EFC) function, somewhat similar to CUTLASS C++ EVT. It also provides a PyTorch evaluator to compare the results. + +* More examples of authorizing peak-performance kernels + - [SM103 batched 3xFP4 blockscaled GEMM kernel](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py) + - Mixed input FMHA decode example with support for int4 KV (int8 KV supported in 4.3) + - New acc_scale grouped mixed input gemm kernel variant is introduced to deliver better performance for decoding cases. + - All mixed_input_gemm examples are moved into a separate folder `mixed_input_gemm`. Common utility functions are also extracted into mixed_input_host_utils.py under the same folder. * Bug fixing and improvements + - Fixed an issue that both branches of if are executed - Fixed `cute.printf` with f-string - - Fixed an issue that cutlass.cuda.initialize_cuda_context() silently kills python + - Fixed an indexing issue of scalar tensor + - Fixed small K reference check error for cta_tile_n = 256 case with overlapping accumulator optimization in [Blackwell SM100 persistent dense blockscaled GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py). * API changes - Deprecate get_num_tmem_alloc_cols from blackwell_helpers.py. Use the one from tmem_allocator.py instead. - Deprecate SM100_TMEM_CAPACITY_COLUMNS and SM100_TMEM_MIN_ALLOC_COLUMNS. - LdMatrix16x16x8bOp and StMatrix16x8x8bOp now require explicit transpose=True when calling __init__, to avoid ambiguity in data transposition. - LdMatrix16x16x8bOp copy traits updated to be faithful to PTX without permutations. Permuted variant is renamed to LdMatrix16x8x8bOp. + - Grouped GEMM example takes the argument --host_problem_shape_available. If the argument is provided, grid is carved out based upon the host problem shapes, otherwise, we launch maximum possible SMs. + - hardware_info.get_max_active_cluster support pass in specific stream to query. Useful for green context based SM partition. - group_bulk_copy_modes in async bulk copy example is now deprecated, use group_modes directly instead. + - Deprecate nvvm wrapper from using nvvm enum, use str instead. - cute.arch.calc_packed_f32x2_op default enable ftz to default disable ftz + - In CuTe DSL with CTK 13.1, following APIs in cutlass.cute.arch now require string literal instead of enum as argument: + + fence_proxy + + fence_view_async_tmem_op + + calc_packed_f32x2_op + + warp_redux_sync + + atomic_add + + atomic_and + + atomic_or + + atomic_xor + + atomic_max + + atomic_min + + atomic_exch + + atomic_cas + + store + + load + +* Advanced compiler control +Use 'Advanced compiler control' for mixed input gemm examples for better performance. +Advanced compiler control is an experimental feature of CUDA compiler. The controls file contains internal compiler settings tuned for specific kernels with a specific version of CUDA toolkit to get better GPU kernel code. More details and documentation on how to create these controls files will be provided in future CUDA toolkit release. + +Note: The advanced compiler control file is not expected to work for kernels that it was not tuned for. There is no compatibility guarantee, and the controls file will not work for CUDA toolkit with a different version. ### CUTLASS C++ * Add Hopper e2m1 to fp32 optimized conversion and e2m1 * TF32 tensor core GEMM. @@ -54,12 +103,19 @@ - Fix a TMA descriptor bug where the CUDA driver is not properly setting the OOB address gen mode correctly. - Fix memory fence for clc scheduler in Blackwell SM120 pingpong kernel. - Fix missing SMEM alignment in Blackwell SM120 scale factors. + - Fix a PDL issue for grouped gemm. + - Fix divide-by-zero issue in canimplement for sm100 implicit gemm kernels. + - Fix cluster swizzle for Grouped GEMMs. + + Move host-side swizzling heuristics to device. + + Apply swizzle per group based on problem shape and max swizzle size. + + Improve examples and unit tests. * Fix some profiler issues: - Refactor L1 functional test generation logic to reduce the L1 test cases to avoid timeout. - Fix a core dump issue for nvfp4 grouped GEMM kernel. - Fix inconsistent GEMM verification logic. - Rework grouped gemm verification logic for different types. -* Fix some broken links under `media/docs`. + - Fix api break change in libheuristics. +* Fix some failed links under `media/docs`. * Various improvements and fixes from the community and CUTLASS team. Thanks to everyone who submitted PRs! * Optimal code generation with CUDA toolkit versions 13.1. diff --git a/README.md b/README.md index b7b32ee7..869c885d 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ # CUTLASS 4.4.0 -_CUTLASS 4.4.0 - Jan 2026_ +_CUTLASS 4.4.0 - Feb 2026_ CUTLASS is a collection of abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and related computations at all levels and scales within CUDA. It incorporates strategies for @@ -45,8 +45,21 @@ To get started quickly - please refer : # What's New in CUTLASS 4.4 -### CuTe DSL +## CuTe DSL * New features + - CuTe DSL now supports CUDA toolkit 13.1! + + Set up with cutlass/python/CuTeDSL/setup.sh --cu13 + + Refer to https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/quick_start.html for more details + - GB300 is now supported in CuTe DSL with CTK 13.1 + + Refer to [SM103 batched 3xFP4 blockscaled GEMM kernel](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py) for example kernel + - cute.experimental: introduce a higher-level, composable layer on top of existing CuTe DSL APIs (not a separate abstraction), which can be mixed with existing Cute DSL building blocks. + + Fragment-free programming model: copy/dot APIs take memrefs directly instead of descriptors/fragments. + + Automatic TMA descriptor generation and update insertion. + + Automatic vectorization and predication for SIMT copies. + + New pipeline abstraction with convenience wrappers + + New Partition ops to simplify partitioning logic. + + Device-side TMA descriptor allocation, initialization, and management + + These examples can be found here https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/experimental - Ahead of Time (AoT) compilation is now available! + Refer to files under https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/cute/export for example usage - JAX support - you can now use CuTeDSL along with JAX @@ -55,20 +68,56 @@ To get started quickly - please refer : + cutlass.__version__ for a string representation of DSL version + cutlass.CUDA_VERSION for a version class to tell the CUDA version used for DSL - Added CopyDsmemStoreOp to store data to distributed shared memory with explicit synchronization. + - Grouped GEMM example now supports device-only problem shapes. + - We allow grid carve-out without problem shapes being available on host. + - Tma+LdMatrix features for loading+unpacking narrow-width types (refer to mixed_input_fmha_decode.py for example usage). + - It is possible now to have customized epilogue fusion for persistent dense GEMM through a Python Epilogue Fusion Configuration (EFC) function, somewhat similar to CUTLASS C++ EVT. It also provides a PyTorch evaluator to compare the results. + +* More examples of authorizing peak-performance kernels + - [SM103 batched 3xFP4 blockscaled GEMM kernel](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py) + - Mixed input FMHA decode example with support for int4 KV (int8 KV supported in 4.3) + - New acc_scale grouped mixed input gemm kernel variant is introduced to deliver better performance for decoding cases. + - All mixed_input_gemm examples are moved into a separate folder `mixed_input_gemm`. Common utility functions are also extracted into mixed_input_host_utils.py under the same folder. * Bug fixing and improvements + - Fixed an issue that both branches of if are executed - Fixed `cute.printf` with f-string - - Fixed an issue that cutlass.cuda.initialize_cuda_context() silently kills python + - Fixed an indexing issue of scalar tensor + - Fixed small K reference check error for cta_tile_n = 256 case with overlapping accumulator optimization in [Blackwell SM100 persistent dense blockscaled GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py). * API changes - Deprecate get_num_tmem_alloc_cols from blackwell_helpers.py. Use the one from tmem_allocator.py instead. - Deprecate SM100_TMEM_CAPACITY_COLUMNS and SM100_TMEM_MIN_ALLOC_COLUMNS. - LdMatrix16x16x8bOp and StMatrix16x8x8bOp now require explicit transpose=True when calling __init__, to avoid ambiguity in data transposition. - LdMatrix16x16x8bOp copy traits updated to be faithful to PTX without permutations. Permuted variant is renamed to LdMatrix16x8x8bOp. + - Grouped GEMM example takes the argument --host_problem_shape_available. If the argument is provided, grid is carved out based upon the host problem shapes, otherwise, we launch maximum possible SMs. + - hardware_info.get_max_active_cluster support pass in specific stream to query. Useful for green context based SM partition. - group_bulk_copy_modes in async bulk copy example is now deprecated, use group_modes directly instead. + - Deprecate nvvm wrapper from using nvvm enum, use str instead. - cute.arch.calc_packed_f32x2_op default enable ftz to default disable ftz + - In CuTe DSL with CTK 13.1, following APIs in cutlass.cute.arch now require string literal instead of enum as argument: + + fence_proxy + + fence_view_async_tmem_op + + calc_packed_f32x2_op + + warp_redux_sync + + atomic_add + + atomic_and + + atomic_or + + atomic_xor + + atomic_max + + atomic_min + + atomic_exch + + atomic_cas + + store + + load -### CUTLASS C++ +* Advanced compiler control +Use 'Advanced compiler control' for mixed input gemm examples for better performance. +Advanced compiler control is an experimental feature of CUDA compiler. The controls file contains internal compiler settings tuned for specific kernels with a specific version of CUDA toolkit to get better GPU kernel code. More details and documentation on how to create these controls files will be provided in future CUDA toolkit release. + +Note: The advanced compiler control file is not expected to work for kernels that it was not tuned for. There is no compatibility guarantee, and the controls file will not work for CUDA toolkit with a different version. + +## CUTLASS C++ * Add Hopper e2m1 to fp32 optimized conversion and e2m1 * TF32 tensor core GEMM. - Set MmaType to tfloat32_t for FP32 mode. - TF32 provides FP32 inputs with reduced precision (19-bit vs 32-bit) @@ -95,12 +144,19 @@ To get started quickly - please refer : - Fix a TMA descriptor bug where the CUDA driver is not properly setting the OOB address gen mode correctly. - Fix memory fence for clc scheduler in Blackwell SM120 pingpong kernel. - Fix missing SMEM alignment in Blackwell SM120 scale factors. + - Fix a PDL issue for grouped gemm. + - Fix divide-by-zero issue in canimplement for sm100 implicit gemm kernels. + - Fix cluster swizzle for Grouped GEMMs. + + Move host-side swizzling heuristics to device. + + Apply swizzle per group based on problem shape and max swizzle size. + + Improve examples and unit tests. * Fix some profiler issues: - Refactor L1 functional test generation logic to reduce the L1 test cases to avoid timeout. - Fix a core dump issue for nvfp4 grouped GEMM kernel. - Fix inconsistent GEMM verification logic. - Rework grouped gemm verification logic for different types. -* Fix some broken links under `media/docs`. + - Fix api break change in libheuristics. +* Fix some failed links under `media/docs`. Note: CUTLASS 4.x builds are known to be down on Windows platforms for all CUDA toolkits. CUTLASS team is working on a fix. diff --git a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu index cde24323..52c06507 100644 --- a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu +++ b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu @@ -251,6 +251,8 @@ struct Options { dim3 cluster_shape = dim3(4,2,1); dim3 cluster_shape_fallback = dim3(2,1,1); RasterOrderOptions raster_order = RasterOrderOptions::AlongM; + char raster_char = 'M'; + int swizzle = 1; int max_sm_count = INT_MAX; std::string benchmark_path; std::vector problem_sizes_host; @@ -294,7 +296,6 @@ struct Options { randomize_problems(cmd); } - char raster_char; cmd.get_cmd_line_argument("raster", raster_char); if (raster_char == 'N' || raster_char == 'n') { @@ -303,6 +304,7 @@ struct Options { else if (raster_char == 'M' || raster_char == 'm') { raster_order = RasterOrderOptions::AlongM; } + cmd.get_cmd_line_argument("swizzle", swizzle, 1); } void randomize_problems(cutlass::CommandLine &cmd) { @@ -378,7 +380,8 @@ struct Options { << " --beta= Epilogue scalar beta\n\n" << " --cluster_m= and --cluster_n= Sets the X,Y dims of the preferred cluster shape\n" << " --cluster_fallback_m= and --cluster_fallback_n= Sets the X,Y dims of the fallback cluster shape\n\n" - << " --raster= CTA Rasterization direction (N for along N, M for along M)\n\n" + << " --raster= Cluster rasterization direction (N for along N, M for along M)\n" + << " --swizzle= Cluster swizzle (swizzle up to 8 and with the nearest multiple of 2)\n\n" << " --iterations= Number of profiling iterations to perform\n\n" << " --benchmark= Executes a benchmark problem size\n" << " --max_sm_count= Run kernels using only these number of SMs\n" @@ -615,6 +618,7 @@ typename Gemm::Arguments args_from_options(Options &options, bool host_problem_s typename Gemm::GemmKernel::TileSchedulerArguments scheduler; scheduler.raster_order = options.raster_order; + scheduler.max_swizzle_size = options.swizzle; if (host_problem_shapes_available) { arguments = typename Gemm::Arguments { @@ -685,7 +689,13 @@ int run(Options &options, bool host_problem_shapes_available = true) std::cout << " " << options.problem_sizes_host.at(i); std::cout << ", " << alpha_host.at(i) << ", " << beta_host.at(i) << std::endl; } - std::cout << " Groups : " << options.groups << std::endl; + std::cout << " Groups : " << options.groups << std::endl; + + std::cout << " Cluster Shape : " << options.cluster_shape.x << "x" << options.cluster_shape.y << std::endl; + std::cout << " Cluster Fallback Shape : " << options.cluster_shape_fallback.x << "x" << options.cluster_shape_fallback.y << std::endl; + + std::cout << " Raster Order : Along-" << options.raster_char << std::endl; + std::cout << " Max Swizzle Size : " << options.swizzle << std::endl; // Instantiate CUTLASS kernel depending on templates Gemm gemm; diff --git a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu index 13631595..1ff8f215 100644 --- a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu +++ b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu @@ -309,6 +309,8 @@ struct Options { dim3 cluster_shape = dim3(2,1,1); dim3 cluster_shape_fallback = dim3(2,1,1); RasterOrderOptions raster_order = RasterOrderOptions::AlongN; + char raster_char = 'N'; + int swizzle = 1; int max_sm_count = INT_MAX; std::string benchmark_path; std::vector problem_sizes_host; @@ -356,7 +358,6 @@ struct Options { randomize_problems(cmd); } - char raster_char; cmd.get_cmd_line_argument("raster", raster_char); if (raster_char == 'N' || raster_char == 'n') { @@ -365,6 +366,7 @@ struct Options { else if (raster_char == 'M' || raster_char == 'm') { raster_order = RasterOrderOptions::AlongM; } + cmd.get_cmd_line_argument("swizzle", swizzle, 1); } void randomize_problems(cutlass::CommandLine &cmd) { @@ -441,7 +443,8 @@ struct Options { << " --norm_constant= Epilogue scalar normalization constant for the output matrix\n\n" << " --cluster_m= and --cluster_n= Sets the X,Y dims of the preferred cluster shape\n" << " --cluster_fallback_m= and --cluster_fallback_n= Sets the X,Y dims of the fallback cluster shape\n\n" - << " --raster= CTA Rasterization direction (N for along N, M for along M)\n\n" + << " --raster= Cluster rasterization direction (N for along N, M for along M)\n" + << " --swizzle= Cluster swizzle (swizzle up to 8 and with the nearest multiple of 2)\n\n" << " --iterations= Number of profiling iterations to perform\n\n" << " --benchmark= Executes a benchmark problem size\n" << " --max_sm_count= Run kernels using only these number of SMs\n" @@ -722,6 +725,7 @@ typename Gemm::Arguments args_from_options(Options &options, bool host_problem_s typename Gemm::GemmKernel::TileSchedulerArguments scheduler; scheduler.raster_order = options.raster_order; + scheduler.max_swizzle_size = options.swizzle; if (host_problem_shapes_available) { arguments = typename Gemm::Arguments { @@ -813,7 +817,13 @@ int run(Options &options, bool host_problem_shapes_available = true) std::cout << " " << options.problem_sizes_host.at(i); std::cout << ", " << alpha_host.at(i) << ", " << beta_host.at(i) << std::endl; } - std::cout << " Groups : " << options.groups << std::endl; + std::cout << " Groups : " << options.groups << std::endl; + + std::cout << " Cluster Shape : " << options.cluster_shape.x << "x" << options.cluster_shape.y << std::endl; + std::cout << " Cluster Fallback Shape : " << options.cluster_shape_fallback.x << "x" << options.cluster_shape_fallback.y << std::endl; + + std::cout << " Raster Order : Along-" << options.raster_char << std::endl; + std::cout << " Max Swizzle Size : " << options.swizzle << std::endl; // Instantiate CUTLASS kernel depending on templates Gemm gemm; diff --git a/examples/75_blackwell_grouped_gemm/CMakeLists.txt b/examples/75_blackwell_grouped_gemm/CMakeLists.txt index e8a19aa4..b44659f7 100644 --- a/examples/75_blackwell_grouped_gemm/CMakeLists.txt +++ b/examples/75_blackwell_grouped_gemm/CMakeLists.txt @@ -31,16 +31,16 @@ -set(TEST_RANDOM --iterations=0) # Random problem sizes -set(TEST_RANDOM_LARGE_GROUP --groups=50 --iterations=0) # Random problem sizes +set(TEST_RANDOM --iterations=0 --raster=M --swizzle=4) # Random problem sizes +set(TEST_RANDOM_LARGE_GROUP --groups=50 --iterations=0 --raster=M --swizzle=2) # Random problem sizes set(TEST_EPILOGUE --alpha=0.5 --beta=0.5 --iterations=0) # Random problem sizes set(TEST_EPILOGUE_LARGE_GROUP --alpha=1.5 --beta=2.0 --groups=50 --iterations=0) # Random problem sizes -set(TEST_EPILOGUE_OP --beta=0.5 --iterations=1) # Random problem sizes +set(TEST_EPILOGUE_OP --beta=0.5 --iterations=1 --raster=N --swizzle=8) # Random problem sizes set(TEST_EPILOGUE_OP_LARGE_GROUP --alpha=1.5 --groups=50 --iterations=1) # Random problem sizes -set(TEST_FIXED --m=2048 --n=5120 --k=8192 --iterations=0) # Fixed problem sizes +set(TEST_FIXED --m=2048 --n=5120 --k=8192 --iterations=0 --raster=M --swizzle=8) # Fixed problem sizes set(TEST_FIXED_LARGE_GROUP --m=2048 --n=512 --k=512 --groups=51 --iterations=0) # Fixed problem sizes set(TEST_SMALL --m=256 --n=128 --iterations=0) # Small problem sizes diff --git a/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py b/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py index fc109f56..1a21d37a 100644 --- a/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py +++ b/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py @@ -29,7 +29,6 @@ import sys import os from typing import Tuple -import torch import cutlass import cutlass.cute as cute @@ -125,6 +124,8 @@ def tensor_op_gemm_wrapper( def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]): + import torch + print("\nRunning TensorOpGemm test with:") print(f"Tensor dimensions: {mnkl}") diff --git a/examples/python/CuTeDSL/ampere/call_from_jit.py b/examples/python/CuTeDSL/ampere/call_from_jit.py index 0e1ff2c0..f4fa4339 100644 --- a/examples/python/CuTeDSL/ampere/call_from_jit.py +++ b/examples/python/CuTeDSL/ampere/call_from_jit.py @@ -60,11 +60,8 @@ import os import sys from typing import Type, Tuple -import torch - import cutlass import cutlass.cute as cute -from cutlass.torch import dtype as torch_dtype from cutlass.cute.runtime import make_ptr if __name__ == "__main__": @@ -205,6 +202,9 @@ def tensor_op_gemm_wrapper( def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]): + import torch + from cutlass.torch import dtype as torch_dtype + print("\nRunning TensorOpGemm test with:") print(f"Tensor dimensions: {mnkl}") diff --git a/examples/python/CuTeDSL/ampere/cooperative_launch.py b/examples/python/CuTeDSL/ampere/cooperative_launch.py new file mode 100644 index 00000000..a447bf31 --- /dev/null +++ b/examples/python/CuTeDSL/ampere/cooperative_launch.py @@ -0,0 +1,627 @@ +# 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. + +""" +Cooperative Launch Example: + +This module demonstrates CUDA Cooperative Launch functionality. It implements a +global barrier that synchronizes ALL threads across the entire GPU grid. + +In traditional CUDA kernel launches, there is no guarantee that all thread blocks +will be resident on the GPU simultaneously. This means that thread blocks may +execute in waves (some finish before others start) and attempting to synchronize +across blocks can cause deadlock. + +**Cooperative Launch** solves this by guaranteeing that all thread blocks launch +atomically and simultaneously. + +For more details, see the CUDA Programming Guide official documentation: +https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/cooperative-groups.html#when-to-use-cudalaunchcooperativekernel + +Cooperative Launch Limitations: + +Cooperative launch has strict grid size constraints. +If you exceed this limit, cudaLaunchCooperativeKernel returns +cudaErrorCooperativeLaunchTooLarge. + +This example demonstrates both a successful cooperative launch with a small grid +and an expected failure when exceeding the grid size limit. + +Usage: + +Run directly: + $ python cooperative_launch.py + +This will: + 1. Demonstrate expected failure with too many thread blocks + 2. Successfully run a cooperative kernel with grid-wide barrier + 3. Print confirmation that all threads synchronized successfully + +""" + +from typing import List, Optional +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cutlass_dsl import ( + dsl_user_op, # Decorator for user-defined device operations + DSLCudaRuntimeError, # Exception type for CUDA runtime errors + extract_mlir_values, # Extract MLIR values from the object + new_from_mlir_values, # Create a new instance from MLIR values +) + +# Function to check cuda errors +from cutlass.base_dsl.runtime.cuda import checkCudaErrors + +# LLVM dialect for inline PTX assembly generation +from cutlass._mlir.dialects import llvm + +# CUDA Python bindings for runtime API (memory allocation, synchronization, etc.) +import cuda.bindings.runtime as cuda_runtime + + +class GlobalBarrier: + """ + A grid-wide barrier for synchronizing ALL thread blocks on the GPU. + + This class implements a cooperative barrier that enables grid-wide + synchronization. It requires cooperative launch to function correctly. + + Design Overview: + + The barrier uses a single 32-bit integer in global memory with the + following bit layout: + + ┌──────────────────────────────────────────────────────────────────┐ + │ Bit 31 │ Bits 30-0 │ + │ ────────── │ ───────────────────────────────────────────────────│ + │ Phase Bit │ Arrival Counter (supports up to 2^31 - 1 blocks) │ + └──────────────────────────────────────────────────────────────────┘ + + Capacity: + + - Maximum thread blocks: 2^31 - 1 = 2,147,483,647 blocks + + Memory Ordering: + + The barrier uses specific memory ordering semantics: + + - Release semantics on arrival (atom.add.release.gpu) + + - Acquire semantics on wait (ld.global.acquire.gpu) + + Usage Example: + + Host-side setup: + + >>> barrier_ptr = GlobalBarrier.allocate() # Allocate barrier memory + + Device-side usage (inside a kernel): + + >>> barrier = GlobalBarrier(barrier_ptr) + >>> + >>> # Do some work... + >>> + >>> barrier.arrive_and_wait() # Synchronize all blocks + >>> + >>> # All blocks proceed together after this point + + Warning: + + This barrier requires cooperative launch! Using it with a regular launch + can result in a deadlock because not all thread blocks may be resident + simultaneously. + """ + + @staticmethod + def allocate() -> cute.runtime.Pointer: + """ + Allocate and initialize barrier memory on the GPU. + + This function allocates device memory for the barrier. + It must be called before launching any kernel that uses the barrier. + """ + ptr = checkCudaErrors(cuda_runtime.cudaMalloc(4)) + + # This sets all 32 bits to 0: + # - Phase bit (bit 31) = 0 + # - Counter (bits 30-0) = 0 + checkCudaErrors(cuda_runtime.cudaMemset(ptr, 0, 4)) + + # Create a pointer with the following properties: + # - Type: Uint32 (32-bit unsigned integer) + # - Address: the allocated device pointer + # - Address Space: gmem (global memory) + barrier_ptr = cute.runtime.make_ptr( + cutlass.Uint32, # Element type + ptr, # Raw CUDA pointer + cute.AddressSpace.gmem, # Memory address space + ) + + return barrier_ptr + + @staticmethod + def free(barrier_ptr: cute.Pointer): + """ + Free the barrier memory on the GPU. + + This function frees the device memory for the barrier. + It must be called after the barrier is no longer needed. + """ + checkCudaErrors(cuda_runtime.cudaFree(barrier_ptr._pointer)) + + @dsl_user_op + def __init__( + self, + barrier_ptr: cute.Pointer, + *, + phase: Optional[cutlass.Uint32] = None, + is_leader: Optional[cutlass.Boolean] = None, + number_of_thread_blocks: Optional[cutlass.Uint32] = None, + loc=None, + ip=None, + ): + """ + Initialize a GlobalBarrier instance on the device. + + This constructor is called by each thread when the kernel + starts. It sets up the barrier state for this thread's participation + in grid-wide synchronization. + + Each thread stores the following: + + - A reference to the shared barrier memory + - Whether it's the leader thread of its block + - The current phase for barrier tracking + - The total number of thread blocks in the grid + """ + # The barrier is shared across ALL thread blocks, so it must be in + # global memory. Shared memory (smem) is block-local and wouldn't work. + if barrier_ptr.memspace != cute.AddressSpace.gmem: + raise ValueError( + "GlobalBarrier requires barrier_ptr to be in global memory (gmem)" + ) + + # Store barrier pointer reference + self.barrier_ptr = barrier_ptr + + # Initialize phase tracking + # Phase starts at 0 for the first barrier, then alternates: + # First barrier: wait for phase 1 + # Second barrier: wait for phase 0 + # Third barrier: wait for phase 1 + # ... and so on + if phase is not None: + self.phase = phase + else: + self.phase = cutlass.Uint32(0) + + if is_leader is not None: + self.is_leader = is_leader + else: + # Determine if this thread is the block leader + # Get this thread's position within its block + tidx, tidy, tidz = cute.arch.thread_idx() + + # Leader is the thread at position (0, 0, 0) in the block + # We use bitwise AND to combine the three conditions efficiently + self.is_leader = ( + cutlass.Boolean(tidx == 0) # First in X dimension + & cutlass.Boolean(tidy == 0) # First in Y dimension + & cutlass.Boolean(tidz == 0) # First in Z dimension + ) + + if number_of_thread_blocks is not None: + self.number_of_thread_blocks = number_of_thread_blocks + else: + # Calculate total number of thread blocks in the grid + # Get grid dimensions (how many blocks in each dimension) + gidx, gidy, gidz = cute.arch.grid_dim() + + # Total blocks = gridDim.x × gridDim.y × gridDim.z + # This is needed to know when ALL blocks have arrived + self.number_of_thread_blocks = cutlass.Uint32(gidx * gidy * gidz) + + @dsl_user_op + @cute.jit + def arrive(self, *, loc=None, ip=None): + """ + Arrive at the barrier without waiting. + + This signals that the calling thread block has reached the barrier + point, but does not wait for other blocks. Use this when you want + to overlap computation with barrier synchronization. + + This method must be called by ALL threads in the block, + not just the leader. The internal block-level sync ensures all + threads in the block agree before the leader signals arrival. + """ + # Ensure ALL threads in this block have reached this point before + # the leader signals arrival. This is critical for correctness! + cute.arch.sync_threads(loc=loc, ip=ip) + + # Only the leader thread performs atomic operations to minimize + # contention on the barrier memory location + if self.is_leader: + # Atomically increment the arrival counter by 1 + # The atomic add returns the value before the add, so we add 1 + # to get the current value after our arrival + barrier_value = ( + self._increment_barrier(cutlass.Uint32(1), loc=loc, ip=ip) + 1 + ) + + # Check if we're the last block to arrive + # Mask out the phase bit (bit 31) to get just the counter value + # Compare against total number of thread blocks + if (barrier_value & ~(1 << 31)) == self.number_of_thread_blocks: + # Flip phase and reset counter + # We add a value that simultaneously: + # 1. Flips bit 31 (adds 2^31) + # 2. Resets counter to 0 (subtracts N, where N was the count) + # + # Example with 8 blocks: + # Current: 0x00000008 (phase=0, counter=8) + # Add: 0x80000000 - 8 = 0x7FFFFFF8 + # Result: 0x80000000 (phase=1, counter=0) ✓ + # + # This works because we're doing modular arithmetic and the + # counter wraps correctly + self._increment_barrier( + cutlass.Uint32((1 << 31) - self.number_of_thread_blocks), + loc=loc, + ip=ip, + ) + + def _read_barrier(self, *, loc=None, ip=None) -> cutlass.Uint32: + """ + Read the barrier value with acquire memory semantics. + + This is an internal method that reads the 32-bit barrier value from + global memory using GPU-scope acquire semantics. + + Notes + ----- + PTX Instruction: + Uses ld.global.acquire.gpu.b32 which is a: + + - Global memory load (ld.global) + - With acquire semantics (.acquire) + - At GPU scope (.gpu) - visible across all thread blocks + - For 32-bit data (.b32) + + Inline Assembly: + We use LLVM inline assembly because CuTe DSL may not have a direct + high-level API for acquire loads. The assembly string format: + + - $0: Output operand (the loaded value) + - $1: Input operand (the address to load from) + + """ + # Use inline PTX assembly for the acquire-semantics load + return cutlass.Uint32( + llvm.inline_asm( + # Return type: 32-bit unsigned integer + cutlass.Uint32.mlir_type, + # Input arguments: barrier pointer address + # We convert the pointer to an integer (64-bit address) + [self.barrier_ptr.toint().ir_value(loc=loc, ip=ip)], + # PTX instruction + "ld.global.acquire.gpu.b32 $0, [$1];", + # Constraint string + # "=r" : Output is a 32-bit register (write-only) + # "l" : Input is a 64-bit register (pointer address) + "=r,l", + # Assembly attributes + # Mark as having side effects + has_side_effects=True, + # No special stack alignment needed + is_align_stack=False, + # Use AT&T syntax (required for LLVM inline asm) + asm_dialect=llvm.AsmDialect.AD_ATT, + # MLIR location and insertion point + loc=loc, + ip=ip, + ) + ) + + def _increment_barrier( + self, value: cutlass.Uint32, *, loc=None, ip=None + ) -> cutlass.Uint32: + """ + Atomically increment the barrier with release memory semantics. + + This is an internal method that performs an atomic add on the barrier + value using GPU-scope release semantics. + + Notes + ----- + PTX Instruction: + Uses atom.add.release.gpu.u32 which is a: + + - Atomic operation (atom) + - Addition (.add) + - With release semantics (.release) + - At GPU scope (.gpu) + - For unsigned 32-bit integers (.u32) + + Atomicity: + The atomic add is guaranteed to be indivisible - no other thread + can see a partial update or interleave with this operation. + + Return Value: + Atomic operations return the OLD value, not the new value. + This is why the caller adds 1 to get the current count. + """ + # Atomic add using inline PTX assembly with release semantics + return cutlass.Uint32( + llvm.inline_asm( + # Return type: 32-bit unsigned integer (the old value) + cutlass.Uint32.mlir_type, + # Input arguments: (barrier address, value to add) + [ + self.barrier_ptr.toint().ir_value( + loc=loc, ip=ip + ), # Barrier address + value.ir_value(loc=loc, ip=ip), # Value to add + ], + # PTX instruction + "atom.add.release.gpu.u32 $0, [$1], $2;", + # Constraint string + # "=r" : Output is a 32-bit register + # "l" : First input is 64-bit (pointer) + # "r" : Second input is 32-bit (value) + "=r,l,r", + # Assembly attributes + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + # MLIR metadata + loc=loc, + ip=ip, + ) + ) + + @dsl_user_op + @cute.jit + def wait(self, *, loc=None, ip=None): + """ + Wait for all thread blocks to arrive at the barrier. + + This method blocks (spins) until all thread blocks have called + arrive() on the barrier. It does NOT signal arrival itself - use + arrive_and_wait() if you need to both arrive and wait. + + IMPORTANT: This method MUST be called by ALL threads in the block. + The internal sync_threads ensures all threads proceed together. + + Algorithm: + - Leader thread spins, reading barrier with acquire semantics + - Waits until phase bit matches expected value + - Block-level sync ensures all threads proceed together + - Update local phase tracking for next barrier + + """ + # Leader thread: spin-wait for phase flip + if self.is_leader: + # Calculate expected phase (opposite of current phase) + # XOR with 1 flips: 0→1, 1→0 + expected = self.phase ^ 1 + + # Initial read of barrier value + barrier_value = self._read_barrier(loc=loc, ip=ip) + + # Spin loop: wait until phase matches expected + # Extract phase bit (bit 31) + # Compare against expected phase value + while (barrier_value >> 31) != expected: + # Keep reading barrier until phase flips + # The acquire semantics ensure memory ordering + barrier_value = self._read_barrier(loc=loc, ip=ip) + + # Block-level synchronization + # Ensure all threads in the block wait for the leader to see the + # phase flip before any thread proceeds + cute.arch.sync_threads(loc=loc, ip=ip) + + # Update phase for next barrier + # Flip local phase: 0→1 or 1→0 + # This prepares for the next barrier synchronization + self.phase = self.phase ^ 1 + + @dsl_user_op + def arrive_and_wait(self, *, loc=None, ip=None): + """ + Arrive at the barrier AND wait for all other thread blocks. + + This is the most common barrier operation - it combines arrive() + and wait() into a single call. All thread blocks will be synchronized + after this call returns. + + IMPORTANT: This method MUST be called by ALL threads in the block. + + Semantics + --------- + + Logically equivalent to: + + >>> barrier.arrive() # Signal we've reached this point + >>> barrier.wait() # Wait for everyone else + """ + # Execute both phases: arrive then wait + self.arrive(loc=loc, ip=ip) + self.wait(loc=loc, ip=ip) + + def __extract_mlir_values__(self) -> List[ir.Value]: + """ + Extract MLIR values from the GlobalBarrier instance. + """ + + assert len(extract_mlir_values(self.barrier_ptr)) == 1 + assert len(extract_mlir_values(self.is_leader)) == 1 + assert len(extract_mlir_values(self.phase)) == 1 + assert len(extract_mlir_values(self.number_of_thread_blocks)) == 1 + + return ( + extract_mlir_values(self.barrier_ptr) + + extract_mlir_values(self.is_leader) + + extract_mlir_values(self.phase) + + extract_mlir_values(self.number_of_thread_blocks) + ) + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GlobalBarrier": + """ + Create a new GlobalBarrier instance from MLIR values. + """ + assert len(values) == 4, f"Expected 4 IR values, but got {len(values)}" + return GlobalBarrier( + barrier_ptr=new_from_mlir_values(self.barrier_ptr, [values[0]]), + is_leader=new_from_mlir_values(self.is_leader, [values[1]]), + phase=new_from_mlir_values(self.phase, [values[2]]), + number_of_thread_blocks=new_from_mlir_values( + self.number_of_thread_blocks, [values[3]] + ), + ) + + +@cute.kernel +def cooperative_kernel(barrier_ptr: cute.Pointer): + """ + Example kernel demonstrating cooperative launch with grid-wide barrier. + + This kernel shows how to use the GlobalBarrier class to synchronize all + thread blocks in a grid. It performs 10 iterations, with a barrier + synchronization after each iteration. + + Launch Requirements: This kernel MUST be launched with cooperative=True. + + """ + # Initialize the barrier for this thread + # Each thread creates its own GlobalBarrier instance, all sharing the + # same underlying barrier_ptr in global memory + barrier = GlobalBarrier(barrier_ptr=barrier_ptr) + + for i in range(10): + # Synchronize all thread blocks across the entire grid + # After this call, ALL blocks have completed iterations 0..i + barrier.arrive_and_wait() + + # Get block and thread indices + bidx, bidy, bidz = cute.arch.block_idx() + tidx, tidy, tidz = cute.arch.thread_idx() + + # Check if this is the leader block (first block in the grid) + leader_cluster = bidx == 0 and bidy == 0 and bidz == 0 + + # Check if this is the leader thread (first thread in the block) + leader_thread = tidx == 0 and tidy == 0 and tidz == 0 + + # Only the single leader thread of the leader block prints + if leader_cluster and leader_thread: + cute.printf("All threads arrived at barrier for the %dth iteration\n", i) + + +# ============================================================================= +# KERNEL LAUNCH WRAPPERS +# ============================================================================= + + +@cute.jit +def run_cooperative_kernel(barrier_ptr: cute.runtime.Pointer): + """ + Launch the cooperative kernel with a reasonable grid size. + + This wrapper launches the cooperative_kernel with a grid of 8 thread blocks + (2×2×2), where each block contains 128 threads (32×2×2). + + Notes: The cooperative=True flag is ESSENTIAL. It tells CUDA to: + - Verify the grid fits within hardware limits + - Launch all blocks atomically + """ + cooperative_kernel(barrier_ptr).launch( + grid=(2, 2, 2), # 8 thread blocks + block=(32, 2, 2), # 128 threads per block + cooperative=True, # Enable cooperative launch semantics + ) + + +@cute.jit +def xfail_run_cooperative_kernel(barrier_ptr: cute.runtime.Pointer): + """ + Demonstrate cooperative launch failure with an oversized grid. + + This wrapper intentionally launches with a grid that exceeds + the limits, demonstrating how cooperative launch fails. + + This launch is expected to fail with cudaErrorCooperativeLaunchTooLarge. + + This demonstrates proper error handling for cooperative launch. + + See Also + -------- + The main() function shows how to properly catch and handle this error. + """ + # Attempt to launch with way too many blocks + cooperative_kernel(barrier_ptr).launch( + grid=(10000, 1, 1), # 10,000 blocks + block=(1024, 1, 1), # 1,024 threads per block + cooperative=True, # Cooperative launch will reject this + ) + + +if __name__ == "__main__": + # Initialize CUDA context + cutlass.cuda.initialize_cuda_context() + + # Allocate barrier memory + # Allocate 4 bytes in device global memory for the barrier state + barrier_ptr = GlobalBarrier.allocate() + + # Demonstrate expected failure (grid too large) + expectedly_failed = False + try: + # Attempt to launch with 10,000 blocks - this WILL fail + xfail_run_cooperative_kernel(barrier_ptr) + except DSLCudaRuntimeError as e: + # Verify we got the expected error code + assert ( + e.error_code == cuda_runtime.cudaError_t.cudaErrorCooperativeLaunchTooLarge + ) + expectedly_failed = True + finally: + # Ensure the failure actually happened (test validation) + assert expectedly_failed + + # Run successful cooperative kernel + # Launch with a reasonable grid size that fits hardware constraints + run_cooperative_kernel(barrier_ptr) + + # Synchronize and clean up + checkCudaErrors(cuda_runtime.cudaDeviceSynchronize()) + + # Free the barrier memory + GlobalBarrier.free(barrier_ptr) diff --git a/examples/python/CuTeDSL/ampere/elementwise_add.py b/examples/python/CuTeDSL/ampere/elementwise_add.py index 941bda44..92fd7e25 100644 --- a/examples/python/CuTeDSL/ampere/elementwise_add.py +++ b/examples/python/CuTeDSL/ampere/elementwise_add.py @@ -28,7 +28,6 @@ import argparse -import torch import time from typing import Type @@ -36,7 +35,6 @@ from typing import Type import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -import cutlass.torch as cutlass_torch from cutlass.cute.runtime import from_dlpack """ @@ -252,6 +250,8 @@ def elementwise_add(mA, mB, mC, copy_bits: cutlass.Constexpr = 128): cC = cute.zipped_divide(idC, tiler=tiler_mn) print(f"[DSL INFO] coord tensor = {cC.type}") + kernel_name = f"cutlass_dsl_elementwise_add_kernel" + elementwise_add_kernel.set_name_prefix(kernel_name) elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch( grid=[cute.size(gC, mode=[1]), 1, 1], block=[cute.size(tv_layout, mode=[0]), 1, 1], @@ -270,6 +270,12 @@ def run_elementwise_add( warmup_iterations=2, iterations=200, ): + import torch + import cutlass.torch as cutlass_torch + + if not torch.cuda.is_available(): + raise RuntimeError("Ampere GPU is required to run this example!") + print("\nRunning Elementwise Add test with:") print(f"Tensor dimensions: [{M}, {N}]") print(f"Input and Output Data type: {dtype}") @@ -304,6 +310,8 @@ def run_elementwise_add( else: c_tensor = c + elementwise_add.set_name_prefix("host_prefix") + print("Compiling kernel with cute.compile ...") start_time = time.time() compiled_func = cute.compile( @@ -386,9 +394,6 @@ if __name__ == "__main__": args = parser.parse_args() - if not torch.cuda.is_available(): - raise RuntimeError("Ampere GPU is required to run this example!") - run_elementwise_add( args.M, args.N, diff --git a/examples/python/CuTeDSL/ampere/elementwise_add_autotune.py b/examples/python/CuTeDSL/ampere/elementwise_add_autotune.py new file mode 100644 index 00000000..483c93ed --- /dev/null +++ b/examples/python/CuTeDSL/ampere/elementwise_add_autotune.py @@ -0,0 +1,372 @@ +# 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 argparse +from typing import Any, Callable, Type + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing + +""" +In this example we revisit the elementwise add example and use the autotune_jit decorator to +autotune the kernel. + +To run this example: + +.. code-block:: bash + + python examples/ampere/elementwise_add_autotune.py --M 3 --N 12 + python examples/ampere/elementwise_add_autotune.py --M 1024 --N 512 + python examples/ampere/elementwise_add_autotune.py --M 1024 --N 1024 --benchmark --warmup_iterations 2 --iterations 1000 + +""" + + +@cute.kernel +def elementwise_add_kernel( + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + cC: cute.Tensor, # coordinate tensor + shape: cute.Shape, + 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) + blkA = gA[blk_coord] # (TileM,TileN) + blkB = gB[blk_coord] # (TileM,TileN) + blkC = gC[blk_coord] # (TileM,TileN) + blkCrd = cC[blk_coord] # (TileM, TileN) + + # # declare the atoms which will be used later for memory copy + copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gA.element_type) + copy_atom_store = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gC.element_type) + + tiled_copy_A = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + tiled_copy_B = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + tiled_copy_C = cute.make_tiled_copy_tv(copy_atom_store, thr_layout, val_layout) + + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + thr_copy_C = tiled_copy_C.get_slice(tidx) + + thrA = thr_copy_A.partition_S(blkA) + thrB = thr_copy_B.partition_S(blkB) + thrC = thr_copy_C.partition_S(blkC) + + # allocate fragments for gmem->rmem + frgA = cute.make_rmem_tensor_like(thrA) + frgB = cute.make_rmem_tensor_like(thrB) + frgC = cute.make_rmem_tensor_like(thrC) + + thrCrd = thr_copy_C.partition_S(blkCrd) + frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean) + + for i in range(0, cute.size(frgPred), 1): + val = cute.elem_less(thrCrd[i], shape) + frgPred[i] = val + + # Print per thread predicate mask + # if tidx == 0 and bidx == 0: + # cute.printf("block_dim = {}", cute.arch.grid_dim()) + # cute.printf("shape = {}", shape) + # cute.print_tensor(thrA) + # cute.print_tensor(thrB) + # cute.print_tensor(frgPred) + + ########################################################## + # Move data to reg address space + ########################################################## + + cute.copy(copy_atom_load, thrA, frgA, pred=frgPred) + cute.copy(copy_atom_load, thrB, frgB, pred=frgPred) + + # if tidx == 0 and bidx == 0: + # cute.print_tensor(frgA) + # cute.print_tensor(frgB) + + # Load data before use. The compiler will optimize the copy and load + # operations to convert some memory ld/st into register uses. + result = frgA.load() + frgB.load() + + # Save the results back to registers. Here we reuse b's registers. + frgC.store(result) + + # Copy the results back to c + cute.copy(copy_atom_store, frgC, thrC, pred=frgPred) + + +@testing.autotune_jit( + params_dict={"copy_bits": [64, 128]}, + update_on_change=["M", "N"], + warmup_iterations=100, + iterations=100, +) +@cute.jit +def elementwise_add_autotune(mA, mB, mC, M, N, copy_bits: cutlass.Constexpr = 128): + dtype = mA.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) + + gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + idC = cute.make_identity_tensor(mC.shape) + cC = cute.zipped_divide(idC, tiler=tiler_mn) + + elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch( + grid=[cute.size(gC, mode=[1]), 1, 1], + block=[cute.size(tv_layout, mode=[0]), 1, 1], + ) + + +class ElementwiseAddWrapper: + """ + This class mimics more advanced kernel development, where a class encapsulates + pieces of the kernel implementation. + + The can_implement method can be used to check if the kernel can be implemented + for the given arguments. + + The __call__ method is the actual cute.jit function. + + """ + + def __init__(self, copy_bits: cutlass.Constexpr = 128): + self.copy_bits = copy_bits + + def can_implement(self, mA, mB, mC, M, N): + return self.copy_bits in [64, 128] + + @cute.jit + def __call__(self, mA, mB, mC, M, N): + dtype = mA.element_type + vector_size = self.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) + + gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + idC = cute.make_identity_tensor(mC.shape) + cC = cute.zipped_divide(idC, tiler=tiler_mn) + + elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch( + grid=[cute.size(gC, mode=[1]), 1, 1], + block=[cute.size(tv_layout, mode=[0]), 1, 1], + ) + + +def tune_class(mA, mB, mC, M, N): + """ + This function is used to autotune the elementwise add kernel which is wrapped in a class. + An internal function is defined to compile the class with the given arguments. + The internal function is then passed to the benchmarking.tune function to autotune. + The best parameters are then used to instantiate the class. + + :param mA: Input tensor A + :type mA: cute.Tensor + :param mB: Input tensor B + :type mB: cute.Tensor + :param mC: Output tensor C + :type mC: cute.Tensor + :param M: Number of rows in the input tensors + :type M: int + :param N: Number of columns in the input tensors + :type N: int + :return: An instance of the ElementwiseAddWrapper class with the best parameters + :rtype: ElementwiseAddWrapper + """ + + def compile_class(a, b, c, M, N, copy_bits=128) -> Callable[[], Any]: + kernel = ElementwiseAddWrapper(copy_bits) + if not kernel.can_implement(a, b, c, M, N): + raise ValueError(f"Cannot implement kernel for copy_bits={copy_bits}") + compiled_kernel = cute.compile(kernel, a, b, c, M, N) + return lambda: compiled_kernel(a, b, c, M, N) + + params = testing.tune( + compile_class, + params_dict={"copy_bits": [1, 64, 128]}, + kernel_arguments=testing.JitArguments(mA, mB, mC, M, N), + ) + return ElementwiseAddWrapper(**params) + + +def run_elementwise_add( + M_start, + M_range, + M_step, + N_start, + N_range, + N_step, + dtype: Type[cutlass.Numeric], + skip_ref_check=False, + warmup_iterations=2, + iterations=200, +): + import torch + import cutlass.torch as cutlass_torch + + if not torch.cuda.is_available(): + raise RuntimeError("Ampere GPU is required to run this example!") + + for M in range(M_start, M_start + M_range + 1, M_step): + for N in range(N_start, N_start + N_range + 1, N_step): + print("\nRunning Elementwise Add test with:") + print(f"Tensor dimensions: [{M}, {N}]") + print(f"Input and Output Data type: {dtype}") + + torch_dtype = cutlass_torch.dtype(dtype) + if dtype.is_integer: + a = torch.randint( + 0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype + ) + b = torch.randint( + 0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype + ) + else: + a = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype) + b = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype) + + c = torch.zeros_like(a) + + print("Input tensor shapes:") + print(f"a: {a.shape}, dtype: {a.dtype}") + print(f"b: {b.shape}, dtype: {b.dtype}") + print(f"c: {c.shape}, dtype: {c.dtype}\n") + + elementwise_class = tune_class(a, b, c, M, N) + + if not skip_ref_check: + print("Verifying results for class ...") + torch.testing.assert_close(a + b, c) + print("Results verified successfully!") + c = torch.zeros_like(a) + + elementwise_add_autotune(a, b, c, M, N) + + if not skip_ref_check: + print("Verifying results for autotuned function ...") + torch.testing.assert_close(a + b, c) + print("Results verified successfully!") + + def generate_kernel_arguments(): + if dtype.is_integer: + a = torch.randint( + 0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype + ) + b = torch.randint( + 0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype + ) + else: + a = torch.randn( + M, N, device=torch.device("cuda"), dtype=torch_dtype + ) + b = torch.randn( + M, N, device=torch.device("cuda"), dtype=torch_dtype + ) + + c = torch.zeros_like(a) + + return testing.JitArguments(a, b, c, M, N) + + avg_time_us = testing.benchmark( + elementwise_add_autotune, + workspace_generator=generate_kernel_arguments, + workspace_count=10, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + # Print execution results + print( + f"Kernel execution time for cute.jit kernel with M={M}, N={N}: {avg_time_us / 1e3:.4f} ms" + ) + print( + f"Achieved memory throughput for M={M}, N={N}: {(3 * a.numel() * dtype.width // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s" + ) + + compiled_class = cute.compile(elementwise_class, a, b, c, M, N) + + avg_time_us = testing.benchmark( + compiled_class, + workspace_generator=generate_kernel_arguments, + workspace_count=10, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print( + f"Kernel execution time for Class Wrapper with M={M}, N={N}: {avg_time_us / 1e3:.4f} ms" + ) + print( + f"Achieved memory throughput for M={M}, N={N}: {(3 * a.numel() * dtype.width // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s" + ) + + +if __name__ == "__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("--M_range", default=0, type=int) + parser.add_argument("--M_step", default=1024, type=int) + parser.add_argument("--N", default=1024, type=int) + parser.add_argument("--N_range", default=0, type=int) + parser.add_argument("--N_step", default=1024, type=int) + parser.add_argument("--warmup_iterations", default=2, type=int) + parser.add_argument("--iterations", default=100, type=int) + parser.add_argument("--skip_ref_check", action="store_true") + + args = parser.parse_args() + run_elementwise_add( + args.M, + args.M_range, + args.M_step, + args.N, + args.N_range, + args.N_step, + dtype=cutlass.Float32, + skip_ref_check=args.skip_ref_check, + warmup_iterations=args.warmup_iterations, + iterations=args.iterations, + ) + print("\nPASS") diff --git a/examples/python/CuTeDSL/ampere/elementwise_apply.py b/examples/python/CuTeDSL/ampere/elementwise_apply.py index 12f93df8..2d39cbce 100644 --- a/examples/python/CuTeDSL/ampere/elementwise_apply.py +++ b/examples/python/CuTeDSL/ampere/elementwise_apply.py @@ -36,8 +36,6 @@ from typing import List, Type import cuda.bindings.driver as cuda import cutlass.cute as cute import cutlass.cute.testing as testing -import cutlass.torch as cutlass_torch -import torch from cutlass.cute.runtime import from_dlpack import cutlass @@ -274,6 +272,8 @@ def leaky_relu(x, alpha, *, loc=None, ip=None): def leaky_relu_ref(x, alpha): + import torch + return torch.where(x > 0, x, alpha * x) @@ -287,6 +287,9 @@ def run_and_verify( warmup_iterations=2, iterations=100, ): + import torch + import cutlass.torch as cutlass_torch + if not torch.cuda.is_available(): raise RuntimeError("NVIDIA GPU is required to run this example!") diff --git a/examples/python/CuTeDSL/ampere/flash_attention_v2.py b/examples/python/CuTeDSL/ampere/flash_attention_v2.py index 5f936e72..95e3497b 100644 --- a/examples/python/CuTeDSL/ampere/flash_attention_v2.py +++ b/examples/python/CuTeDSL/ampere/flash_attention_v2.py @@ -30,13 +30,11 @@ import argparse from types import SimpleNamespace from typing import Type, Callable -import torch import cuda.bindings.driver as cuda import cutlass.cute.testing as testing import cutlass import cutlass.cute as cute from cutlass.cute.nvgpu import cpasync, warp -import cutlass.torch as cutlass_torch from cutlass.cute.runtime import from_dlpack import cutlass.pipeline as pipeline import cutlass.utils as utils @@ -1162,6 +1160,9 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + # Skip unsupported testcase if not FlashAttentionForwardAmpere.can_implement( dtype, @@ -1237,8 +1238,12 @@ def run( torch_stream = torch.cuda.current_stream() # Get the raw stream pointer as a CUstream current_stream = cuda.CUstream(torch_stream.cuda_stream) + # Pass compile options if needed + compile_options = "" # compile the fa2 forward pass - compiled_fa2_fwd = cute.compile(fa2_fwd, q, k, v, o, softmax_scale, current_stream) + compiled_fa2_fwd = cute.compile( + fa2_fwd, q, k, v, o, softmax_scale, current_stream, options=compile_options + ) if not skip_ref_check: compiled_fa2_fwd(q, k, v, o, softmax_scale, current_stream) @@ -1317,7 +1322,6 @@ if __name__ == "__main__": default=False, help="Use circular buffer tensor sets to ensure L2 cold cache", ) - args = parser.parse_args() run( args.dtype, diff --git a/examples/python/CuTeDSL/ampere/hstu_attention.py b/examples/python/CuTeDSL/ampere/hstu_attention.py index b32600e8..3b537792 100644 --- a/examples/python/CuTeDSL/ampere/hstu_attention.py +++ b/examples/python/CuTeDSL/ampere/hstu_attention.py @@ -29,10 +29,8 @@ from typing import Type import argparse -import torch import cuda.bindings.driver as cuda import cutlass -import cutlass.torch as cutlass_torch import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack from cutlass._mlir.dialects import llvm @@ -838,28 +836,25 @@ class HSTUAttentionForwardAmpere(object): def run_pytorch_hstu_test( - dtype: torch.dtype, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - rab: torch.Tensor, + dtype, + q, + k, + v, + rab, is_causal: bool, ): """Generate the reference output of the HSTU attention with Pytorch. :param dtype: data type of the input tensors - :type dtype: torch.dtype :param q: query tensor - :type q: torch.Tensor :param k: key tensor - :type k: torch.Tensor :param v: value tensor - :type v: torch.Tensor :param rab: RAB tensor - :type rab: torch.Tensor :param is_causal: whether to use causal masking :type is_causal: bool """ + import torch + q = q.to(dtype) k = k.to(dtype) v = v.to(dtype) @@ -922,6 +917,9 @@ def run( """ assert dtype == cutlass.Float16 or dtype == cutlass.BFloat16 + import torch + import cutlass.torch as cutlass_torch + torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) diff --git a/examples/python/CuTeDSL/ampere/inline_ptx.py b/examples/python/CuTeDSL/ampere/inline_ptx.py index 73d54dee..0eeedeb3 100644 --- a/examples/python/CuTeDSL/ampere/inline_ptx.py +++ b/examples/python/CuTeDSL/ampere/inline_ptx.py @@ -29,8 +29,6 @@ from functools import partial from typing import Union -import torch - import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack from cutlass._mlir.dialects import llvm @@ -49,7 +47,7 @@ Situations like: motivate developers to inline PTX themselves. In this example, we inline the vote.sync.ballot.b32, vote.sync.any.pred, vote.sync.all.pred, -vote.sync.uni.pred, and use the corresponding ops in nvvm_wrappers.py for the test. +vote.sync.uni.pred, and use the corresponding ops in nvvm dialect for the test. You can refer to the documentation of `inline_asm op in llvm dialect `_ and `vote.sync `_ @@ -61,8 +59,8 @@ To run this example: python examples/ampere/inline_ptx.py -The example will run the vote kernel with inline PTX and nvvm dialect separately. -The results from inline PTX and nvvm dialect will be verified correspondingly. +The example will run the vote kernel with inline ptx and nvvm dialect separately. +The results from inline ptx and nvvm dialect will be verified correspondingly. """ @@ -184,6 +182,8 @@ def vote( def run(): + import torch + ballot_ptx = torch.randint( 0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32 ) @@ -230,14 +230,11 @@ def run(): torch.testing.assert_close(ballot_ptx, ballot_nvvm) print("Verifying any results...") torch.testing.assert_close(any_ptx, any_nvvm) - print(torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE)))) - assert torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE))) print("Verifying all results...") torch.testing.assert_close(all_ptx, all_nvvm) - assert torch.all(all_ptx == all(i < 10 for i in range(WARP_SIZE))) print("Verifying uni results...") torch.testing.assert_close(uni_ptx, uni_nvvm) - assert torch.all(uni_ptx == (len(set(i < 10 for i in range(WARP_SIZE))) == 1)) + print("Results verified successfully!") diff --git a/examples/python/CuTeDSL/ampere/sgemm.py b/examples/python/CuTeDSL/ampere/sgemm.py index 17835dd0..3455e9b7 100644 --- a/examples/python/CuTeDSL/ampere/sgemm.py +++ b/examples/python/CuTeDSL/ampere/sgemm.py @@ -31,7 +31,6 @@ import time from typing import Tuple import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute @@ -643,6 +642,8 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + """Execute SIMT GEMM operation and benchmark performance. :param mnk: GEMM problem size (M, N, K, L) @@ -666,6 +667,7 @@ def run( :return: Execution time of the GEMM kernel in microseconds :rtype: float """ + torch.manual_seed(1024) print("Running Ampere SIMT GEMM example:") print(f"mnk: {mnk}") print(f"A major: {a_major}, B major: {b_major}, C major: {c_major}") @@ -851,8 +853,6 @@ if __name__ == "__main__": args = parser.parse_args() print("Running SIMT GEMM example:") - torch.manual_seed(1024) - run( args.mnk, args.a_major, diff --git a/examples/python/CuTeDSL/ampere/smem_allocator.py b/examples/python/CuTeDSL/ampere/smem_allocator.py index 3728db43..ea004ead 100644 --- a/examples/python/CuTeDSL/ampere/smem_allocator.py +++ b/examples/python/CuTeDSL/ampere/smem_allocator.py @@ -28,7 +28,6 @@ import cutlass.cute as cute import cutlass -import torch import numpy as np from cutlass.cute.runtime import from_dlpack @@ -175,6 +174,8 @@ def host( def run_and_verify(const_a, const_b, const_c): + import torch + dst_a = torch.zeros((8, 4), dtype=torch.float32, device="cuda") dst_b = torch.zeros((8, 2), dtype=torch.float32, device="cuda") dst_c = torch.zeros((16, 2), dtype=torch.float32, device="cuda") diff --git a/examples/python/CuTeDSL/ampere/tensorop_gemm.py b/examples/python/CuTeDSL/ampere/tensorop_gemm.py index ed9a455c..ad3fe0f0 100644 --- a/examples/python/CuTeDSL/ampere/tensorop_gemm.py +++ b/examples/python/CuTeDSL/ampere/tensorop_gemm.py @@ -30,12 +30,9 @@ import argparse import math from typing import Tuple, Type -import torch - import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -import cutlass.torch as cutlass_torch import cutlass.utils as utils from cutlass.cute.runtime import from_dlpack @@ -849,6 +846,9 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + print("Running Ampere tensor core GEMM example:") print(f"mnkl: {mnkl}") print( diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py index c6f7b5e0..a54e0fb0 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py @@ -30,13 +30,11 @@ import argparse from typing import Type, Tuple, Union import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @@ -64,7 +62,7 @@ This GEMM kernel supports the following features: 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. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations. +2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy 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. @@ -1008,7 +1006,10 @@ class BlockwiseGemmKernel: ) # fence view async shared - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1123,7 +1124,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1295,7 +1299,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1450,7 +1457,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1684,7 +1694,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1851,7 +1864,10 @@ class BlockwiseGemmKernel: 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("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1881,7 +1897,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -2553,6 +2572,9 @@ class BlockwiseGemmKernel: def create_tensors( l, m, n, k, a_major, b_major, cd_major, ab_dtype, c_dtype, scale_dtype ): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) @@ -2613,6 +2635,9 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + """ Prepare A/B/C tensors, launch GPU kernel, and reference checking. """ @@ -2688,6 +2713,7 @@ def run( # try to check CUDA version to decide the opt level try: from cutlass import CUDA_VERSION + opt_level = ( 3 if CUDA_VERSION.major < 13 diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py index 233dda4a..bc99a15c 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py @@ -30,13 +30,11 @@ import argparse from typing import Type, Tuple, Union import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @@ -80,7 +78,7 @@ This GEMM kernel supports the following features: 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. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations. +2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy 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. @@ -1034,7 +1032,10 @@ class BlockwiseContiguousGroupedGemmKernel: ) # fence view async shared - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1150,7 +1151,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1322,7 +1326,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1479,7 +1486,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1715,7 +1725,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1884,7 +1897,10 @@ class BlockwiseContiguousGroupedGemmKernel: 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("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1914,7 +1930,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -2595,6 +2614,8 @@ class BlockwiseContiguousGroupedGemmKernel: def create_mask(num_groups, expect_m, fixed_m=False, m_aligned=128): + import torch + valid_m = 0 group_m_list = [] gidx_mapping = [] @@ -2632,6 +2653,9 @@ def create_tensors( scale_dtype, fixed_m=False, ): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) valid_m, group_m_list, _gidx_mapping = create_mask(l, m, fixed_m) @@ -2702,6 +2726,9 @@ def run( fixed_m: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + """ Prepare A/B/C tensors, launch GPU kernel, and reference checking. """ diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py index a0b655ae..316a5e9a 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py @@ -30,13 +30,11 @@ import argparse from typing import Type, Tuple, Union import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @@ -79,7 +77,7 @@ Matrix A/C Memory Layout Diagrams: 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. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations. +2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy 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. @@ -1041,7 +1039,10 @@ class BlockwiseMaskedGroupedGemmKernel: ) # fence view async shared - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1156,7 +1157,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1328,7 +1332,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1483,7 +1490,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1717,7 +1727,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1884,7 +1897,10 @@ class BlockwiseMaskedGroupedGemmKernel: 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("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1914,7 +1930,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -2586,6 +2605,8 @@ class BlockwiseMaskedGroupedGemmKernel: def create_mask(num_groups: int, m: int, fixed_m=False, tile_m=128): + import torch + # align with block_m (or block_n if swapAB) masked_m_candidates = list( filter( @@ -2617,6 +2638,9 @@ def create_tensors( scale_dtype, fixed_m=False, ): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) _gidx_mapping, masked_m = create_mask(l, m, fixed_m) @@ -2684,6 +2708,9 @@ def run( fixed_m: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + """ Prepare A/B/C tensors, launch GPU kernel, and reference checking. """ diff --git a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py index f3573fa7..b74d35dd 100644 --- a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py @@ -27,7 +27,7 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import argparse -from typing import Type, Tuple, Union +from typing import Type, Tuple, Union, Literal import cuda.bindings.driver as cuda import torch @@ -41,7 +41,7 @@ import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils -from cutlass.cute.runtime import from_dlpack +from cutlass.cute.runtime import make_ptr """ This example provides an experimental implementation of the SM100 batched dense blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. @@ -117,10 +117,6 @@ Constraints: """ -def ceil_div(a, b): - return (a + b - 1) // b - - class Sm100BlockScaledPersistentDenseGemmKernel: """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. @@ -224,8 +220,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: num_threads=self.threads_per_warp * len((self.mma_warp_id, *self.epilog_warp_id)), ) self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") - SM100_TMEM_CAPACITY_COLUMNS = 512 - self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -375,24 +370,36 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Compute number of TMEM columns for SFA/SFB/Accumulator sf_atom_mn = 32 - self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k - self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k + self.num_sfa_tmem_cols = ( + self.cta_tile_shape_mnk[0] // sf_atom_mn + ) * mma_inst_tile_k + self.num_sfb_tmem_cols = ( + self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn + ) * mma_inst_tile_k self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols - self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage + if not self.overlapping_accum + else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols + ) # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue - # Use -1 since at that iteration the pipeline is updated after the tmem -> reg copy - num_subtiles_in_overlap_region = ceil_div(self.num_sf_tmem_cols, self.epi_tile_n) - self.iter_acc_early_release_in_epilogue = num_subtiles_in_overlap_region - 1 + self.iter_acc_early_release_in_epilogue = ( + self.num_sf_tmem_cols // self.epi_tile_n + ) @cute.jit def __call__( self, - a_tensor: cute.Tensor, - b_tensor: cute.Tensor, - sfa_tensor: cute.Tensor, - sfb_tensor: cute.Tensor, - c_tensor: cute.Tensor, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + layouts: cutlass.Constexpr[ + Tuple[tcgen05.OperandMajorMode, tcgen05.OperandMajorMode, utils.LayoutEnum] + ], + problem_mnkl: Tuple[int, int, int, int], max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x, @@ -423,13 +430,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel: :raises TypeError: If input data types are incompatible with the MMA instruction. """ # Setup static attributes before smem/grid/tma computation - self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type - self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type - self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type - self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type - self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() - self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() - self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) + self.a_dtype: Type[cutlass.Numeric] = a_ptr.value_type + self.b_dtype: Type[cutlass.Numeric] = b_ptr.value_type + self.sf_dtype: Type[cutlass.Numeric] = sfa_ptr.value_type + self.c_dtype: Type[cutlass.Numeric] = c_ptr.value_type + + m, n, k, l = problem_mnkl + self.a_major_mode, self.b_major_mode, self.c_layout = layouts # Check if input data types are compatible with MMA instruction if cutlass.const_expr(self.a_dtype != self.b_dtype): @@ -438,18 +445,37 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Setup attributes that dependent on gemm inputs self._setup_attributes() + a_layout = cute.make_ordered_layout((m, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.a_major_mode == tcgen05.OperandMajorMode.K): + a_layout = cute.make_ordered_layout( + (cute.assume(m, 32), k, l), order=(1, 0, 2) + ) + b_layout = cute.make_ordered_layout((n, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.b_major_mode == tcgen05.OperandMajorMode.K): + b_layout = cute.make_ordered_layout( + (cute.assume(n, 32), k, l), order=(1, 0, 2) + ) + c_layout = cute.make_ordered_layout((cute.assume(m, 32), n, l), order=(0, 1, 2)) + if cutlass.const_expr(self.c_layout == utils.LayoutEnum.ROW_MAJOR): + c_layout = cute.make_ordered_layout( + (m, cute.assume(n, 32), l), order=(1, 0, 2) + ) + a_tensor = cute.make_tensor(a_ptr, a_layout) + b_tensor = cute.make_tensor(b_ptr, b_layout) + c_tensor = cute.make_tensor(c_ptr, c_layout) + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( a_tensor.shape, self.sf_vec_size ) - sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( b_tensor.shape, self.sf_vec_size ) - sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( self.a_dtype, @@ -539,25 +565,21 @@ class Sm100BlockScaledPersistentDenseGemmKernel: y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) new_shape = ( - ( - tma_tensor_sfb.shape[0][0], - ((2, 2), y) - ), + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), tma_tensor_sfb.shape[1], - tma_tensor_sfb.shape[2] + tma_tensor_sfb.shape[2], ) # Use right multiplication for ScaledBasis (3 * x instead of x * 3) x_times_3 = 3 * x new_stride = ( - ( - tma_tensor_sfb.stride[0][0], - ((x, x), x_times_3) - ), + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), tma_tensor_sfb.stride[1], - tma_tensor_sfb.stride[2] + tma_tensor_sfb.stride[2], ) tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) - tma_tensor_sfb = cute.make_tensor(tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout) + tma_tensor_sfb = cute.make_tensor( + tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout + ) 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) @@ -896,7 +918,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: cute.group_modes(tCgB, 0, 3), ) - # TMA load SFA partition_S/D + # TMA load scaled factor A partition_S/D sfa_cta_layout = a_cta_layout # ((atom_v, rest_v), STAGE) # ((atom_v, rest_v), RestM, RestK, RestL) @@ -910,7 +932,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tAsSFA = cute.filter_zeros(tAsSFA) tAgSFA = cute.filter_zeros(tAgSFA) - # TMA load SFB partition_S/D + # TMA load scaled factor B partition_S/D sfb_cta_layout = cute.make_layout( cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape ) @@ -945,13 +967,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tCtAcc_fake.iterator, cute.make_layout( tCtAcc_fake.shape, - stride = ( + stride=( tCtAcc_fake.stride[0], tCtAcc_fake.stride[1], tCtAcc_fake.stride[2], - (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1] - ) - ) + (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1], + ), + ), ) else: # (MMA, MMA_M, MMA_N, STAGE) @@ -1010,9 +1032,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): slice_n = mma_tile_coord_mnl[1] // 2 # ((atom_v, rest_v), RestK) - tBgSFB_slice = tBgSFB[ - (None, slice_n, None, mma_tile_coord_mnl[2]) - ] + tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])] # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt ab_producer_state.reset_count() @@ -1188,11 +1208,15 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tCtSFB_mma = tCtSFB if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB) - offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0) + offset = ( + cutlass.Int32(2) + if mma_tile_coord_mnl[1] % 2 == 1 + else cutlass.Int32(0) + ) shifted_ptr = cute.recast_ptr( acc_tmem_ptr - + self.num_accumulator_tmem_cols - + self.num_sfa_tmem_cols + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + offset, dtype=self.sf_dtype, ) @@ -1201,7 +1225,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Move in increments of 64 columns of SFB offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) shifted_ptr = cute.recast_ptr( - acc_tmem_ptr + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset, @@ -1398,7 +1422,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Get accumulator stage index if cutlass.const_expr(self.overlapping_accum): acc_stage_index = acc_consumer_state.phase - reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False) + reverse_subtile = True if acc_stage_index == 0 else False else: acc_stage_index = acc_consumer_state.index @@ -1425,7 +1449,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: real_subtile_idx = subtile_idx if cutlass.const_expr(self.overlapping_accum): if reverse_subtile: - real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx + real_subtile_idx = ( + self.cta_tile_shape_mnk[1] // self.epi_tile_n + - 1 + - subtile_idx + ) # # Load accumulator from tensor memory buffer to register # @@ -1459,7 +1487,10 @@ class Sm100BlockScaledPersistentDenseGemmKernel: 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("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1900,9 +1931,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel: def is_valid_layouts( ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], ) -> bool: """ Check if layouts and dtypes are valid combinations @@ -1912,11 +1943,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: :param c_dtype: The data type of the output tensor :type c_dtype: Type[cutlass.Numeric] :param a_major: The major dimension of the A tensor - :type a_major: str + :type a_major: Literal["m", "k"] :param b_major: The major dimension of the B tensor - :type b_major: str + :type b_major: Literal["n", "k"] :param c_major: The major dimension of the C tensor - :type c_major: str + :type c_major: Literal["m", "n"] :return: True if the layouts are valid, False otherwise :rtype: bool @@ -1976,9 +2007,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel: l: int, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], ) -> bool: """ Check if the tensor alignment is valid @@ -1996,11 +2027,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: :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 + :type a_major: Literal["m", "k"] :param b_major: The major axis of the B tensor - :type b_major: str + :type b_major: Literal["n", "k"] :param c_major: The major axis of the C tensor - :type c_major: str + :type c_major: Literal["m", "n"] :return: True if the problem shape is valid, False otherwise :rtype: bool @@ -2023,27 +2054,32 @@ class Sm100BlockScaledPersistentDenseGemmKernel: @staticmethod def can_implement( + mnkl: Tuple[int, int, int, int], ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], - sf_vec_size: int, c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + sf_vec_size: int, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], - m: int, - n: int, - k: int, - l: int, - a_major: str, - b_major: str, - c_major: str, ) -> bool: """ Check if the gemm can be implemented + :param mnkl: The problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] :param ab_dtype: The data type of the A and B operands :type ab_dtype: Type[cutlass.Numeric] :param sf_dtype: The data type of the scale factor tensor :type sf_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: Literal["m", "k"] + :param b_major: The major axis of the B tensor + :type b_major: Literal["n", "k"] + :param c_major: The major axis of the C tensor + :type c_major: Literal["m", "n"] :param sf_vec_size: The vector size :type sf_vec_size: int :param c_dtype: The data type of the output tensor @@ -2052,24 +2088,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: :type mma_tiler_mn: Tuple[int, int] :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster :type cluster_shape_mn: Tuple[int, int] - :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 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 gemm can be implemented, False otherwise :rtype: bool """ + # Unpack parameters + m, n, k, l = mnkl can_implement = True # Skip unsupported types if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size( @@ -2094,30 +2117,448 @@ class Sm100BlockScaledPersistentDenseGemmKernel: return can_implement +# Helper function to convert scale factor tensor from MKL layout to (32, 4, restM, 4, restK, l) format @cute.jit def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( - sf_ref_tensor: cute.Tensor, - sf_mma_tensor: cute.Tensor, + sf_ref_ptr: cute.Pointer, + sf_mma_ptr: cute.Pointer, + mn: int, + sf_k: int, + l: int, + mma_shape: tuple, ): - """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout""" - # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) - # group to ((32, 4, rest_m), (4, rest_k), l) + mma_permute_order = (3, 4, 1, 5, 2, 0) + permuted_shape = tuple(mma_shape[i] for i in mma_permute_order) + cute_layout = cute.make_ordered_layout(permuted_shape, order=(2, 1, 4, 0, 3, 5)) + + sf_ref_tensor = cute.make_tensor( + sf_ref_ptr, cute.make_layout((mn, sf_k, l), stride=(sf_k, 1, mn * sf_k)) + ) + sf_mma_tensor = cute.make_tensor(sf_mma_ptr, cute_layout) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) for i in cutlass.range(cute.size(sf_ref_tensor)): mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + pass -def run( +# Helper function for ceil division +def ceil_div(a, b): + return (a + b - 1) // b + + +# Convert scale factor tensors from (m, k, l) to (32, 4, restM, 4, restK, l) format +def create_and_reorder_scale_factor_tensor( + l, mn, k, sf_vec_size, sf_dtype, torch_tensor +): + """ + Create the CUTE-format scale factor tensor on CUDA based on the reference tensor. + """ + sf_k = ceil_div(k, sf_vec_size) + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, # batch size + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + # Generate a random int8 tensor, then convert to float8_e4m3fn + cute_tensor = torch.ones(mma_shape, dtype=cutlass_torch.dtype(sf_dtype)).permute( + 3, 4, 1, 5, 2, 0 + ) + + # Call the helper function to do layout conversion + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + make_ptr( + sf_dtype, + torch_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr( + sf_dtype, + cute_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + mn, + sf_k, + l, + mma_shape, + ) + return cute_tensor.cuda() + + +# Compile the persistent dense blockscaled GEMM operation +def scaled_mm( + gemm_obj: Sm100BlockScaledPersistentDenseGemmKernel, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + options: str = "", +): + # Construct CuTe Pointers + a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + + a_major_mode = ( + tcgen05.OperandMajorMode.K if a_major == "k" else tcgen05.OperandMajorMode.MN + ) + b_major_mode = ( + tcgen05.OperandMajorMode.K if b_major == "k" else tcgen05.OperandMajorMode.MN + ) + c_layout = ( + utils.LayoutEnum.ROW_MAJOR if c_major == "n" else utils.LayoutEnum.COL_MAJOR + ) + return cute.compile( + gemm_obj, + a_ptr, + b_ptr, + sfa_ptr, + sfb_ptr, + c_ptr, + (a_major_mode, b_major_mode, c_layout), + (cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + max_active_clusters, + stream, + epilogue_op, + options=options, + ) + + +def is_emulated_dtype( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +) -> bool: + if c_dtype in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + }: + if ab_dtype == cutlass.Float4E2M1FN and sf_dtype == cutlass.Float8E4M3FN: + return False + if ab_dtype == cutlass.Float8E4M3FN and sf_dtype == cutlass.Float8E8M0FNU: + return False + + return True + + +# Convert scale factor tensor from MKL layout to blocked layout +def to_blocked(input_matrix): + rows, cols = input_matrix.shape + # Please ensure rows and cols are multiples of 128 and 4 respectively + n_row_blocks = ceil_div(rows, 128) + n_col_blocks = ceil_div(cols, 4) + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + # Pad the input matrix if necessary + if padded_rows != rows or padded_cols != cols: + # For FP8 types, convert to float32 for padding, then convert back + original_dtype = input_matrix.dtype + input_float32 = input_matrix.to(torch.float32) + padded = torch.nn.functional.pad( + input_float32, + (0, padded_cols - cols, 0, padded_rows - rows), + mode="constant", + value=0, + ) + # Convert back to original dtype if needed + if original_dtype != input_float32.dtype: + padded = padded.to(original_dtype) + else: + padded = input_matrix + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) + return rearranged.flatten() + + +# Reference implementation of the persistent dense blockscaled GEMM operation (emulated version) +def reference_scaled_mm_emulated( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + mnkl: Tuple[int, int, int, int], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], +): + m, n, k, l = mnkl + sfa_expanded = ( + torch.repeat_interleave(sfa, sf_vec_size, dim=1)[:, :k, :] + .to(dtype=torch.float32) + .cuda() + ) + sfb_expanded = ( + torch.repeat_interleave(sfb, sf_vec_size, dim=1)[:, :k, :] + .to(dtype=torch.float32) + .cuda() + ) + res_a = torch.einsum("mkl,mkl->mkl", a, sfa_expanded) + res_b = torch.einsum("nkl,nkl->nkl", b, sfb_expanded) + # Cast res_a and res_b to float32 for einsum to avoid NotImplementedError on 'Byte' + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + c_ref = ref.to(dtype=cutlass_torch.dtype(c_dtype)) + return c_ref + + +# Reference implementation of the persistent dense blockscaled GEMM operation (non-emulated version) +def reference_scaled_mm( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + mnkl: Tuple[int, int, int, int], + c_dtype: Type[cutlass.Numeric], +): + m, n, k, l = mnkl + c_ref = torch.clone(c) + for l_idx in range(l): + # Convert the scale factor tensor to blocked format + scale_a = to_blocked(sfa[:, :, l_idx]) + scale_b = to_blocked(sfb[:, :, l_idx]) + # Ensure a_slice is row-major (M, K) with stride (K, 1) + a_slice = a[:, :, l_idx].contiguous() + # Ensure b_slice is row-major (N, K) so that transpose gives column-major (K, N) + b_slice = b[:, :, l_idx].contiguous() + # (m, k) @ (n, k).T -> (m, n) + res = torch._scaled_mm( + a_slice, + b_slice.transpose(0, 1), + scale_a.cuda(), + scale_b.cuda(), + bias=None, + out_dtype=c_ref.dtype, + ) + c_ref[:, :, l_idx] = res + return c_ref + + +# Construct CuTe Pointers for the persistent dense blockscaled GEMM operation (emulated version) +def construct_cute_pointers_emulated( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +): + a_cute, _ = cutlass_torch.cute_tensor_like( + a.cpu(), + ab_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + a_cute = cutlass_torch.convert_cute_tensor( + a, + a_cute, + ab_dtype, + is_dynamic_layout=True, + ) + b_cute, _ = cutlass_torch.cute_tensor_like( + b.cpu(), + ab_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + b_cute = cutlass_torch.convert_cute_tensor( + b, + b_cute, + ab_dtype, + is_dynamic_layout=True, + ) + a_ptr = a_cute.iterator + b_ptr = b_cute.iterator + + sfa_ptr = make_ptr( + sf_dtype, sfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfb_ptr = make_ptr( + sf_dtype, sfb.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + return a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute + + +# Construct CuTe Pointers for the persistent dense blockscaled GEMM operation (non-emulated version) +def construct_cute_pointers( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +): + a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr( + sf_dtype, sfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfb_ptr = make_ptr( + sf_dtype, sfb.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + return a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr + + +# Use uint8 and uint32 to emulate unsupported +# dtype in torch +def prepare_tensors_emulated( mnkl: Tuple[int, int, int, int], ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], +): + m, n, k, l = mnkl + sf_k = ceil_div(k, sf_vec_size) + + # Create tensor SFA/SFB with values in [1, 3) + sfa = ( + torch.randint(0, 3, (l, m, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + sfb = ( + torch.randint(0, 3, (l, n, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + + # Create tensor A/B with values in [0, 2) + if a_major == "k": + a = torch.randint(-2, 2, (l, m, k), dtype=torch.float32, device="cuda").permute( + 1, 2, 0 + ) + else: + a = torch.randint(-2, 2, (l, k, m), dtype=torch.float32, device="cuda").permute( + 2, 1, 0 + ) + if b_major == "k": + b = torch.randint(-2, 2, (l, n, k), dtype=torch.float32, device="cuda").permute( + 1, 2, 0 + ) + else: + b = torch.randint(-2, 2, (l, k, n), dtype=torch.float32, device="cuda").permute( + 2, 1, 0 + ) + if c_major == "n": + c = torch.empty( + (l, m, n), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(1, 2, 0) + else: + c = torch.empty( + (l, n, m), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(2, 1, 0) + return a, b, c, sfa, sfb + + +def prepare_tensors( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], +): + m, n, k, l = mnkl + + if ab_dtype == cutlass.Float4E2M1FN: + # Using int8 for torch.float4_e2m1fn_x2 tensor allocation + # Thus the size of k needs to be halved in this case. + k_fct = 2 + else: + k_fct = 1 + + sf_k = ceil_div(k, sf_vec_size) + + # Create tensor SFA/SFB + sfa = ( + torch.randint(0, 3, (l, m, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + sfb = ( + torch.randint(0, 3, (l, n, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + + # Create tensor A/B/C + if a_major == "k": + a = torch.randint( + -2, 2, (l, m, k // k_fct), dtype=torch.int8, device="cuda" + ).permute(1, 2, 0) + else: + a = torch.randint(-2, 2, (l, k, m), dtype=torch.int8, device="cuda").permute( + 2, 1, 0 + ) + if b_major == "k": + b = torch.randint( + -2, 2, (l, n, k // k_fct), dtype=torch.int8, device="cuda" + ).permute(1, 2, 0) + else: + b = torch.randint(-2, 2, (l, k, n), dtype=torch.int8, device="cuda").permute( + 2, 1, 0 + ) + if c_major == "n": + c = torch.randint( + -2, 2, (l, m, n), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(1, 2, 0) + else: + c = torch.randint( + -2, 2, (l, n, m), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(2, 1, 0) + + if ab_dtype == cutlass.Float4E2M1FN: + a = a.view(dtype=torch.float4_e2m1fn_x2) + b = b.view(dtype=torch.float4_e2m1fn_x2) + else: + a = a.to(dtype=cutlass_torch.dtype(ab_dtype)) + b = b.to(dtype=cutlass_torch.dtype(ab_dtype)) + + c = c.to(dtype=cutlass_torch.dtype(c_dtype)) + return a, b, c, sfa, sfb + + +# This will show how to covert torch tensor +# and pass to CuTe kernel +def run_scaled_mm( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], tolerance: float = 1e-01, @@ -2127,7 +2568,7 @@ def run( use_cold_l2: bool = False, **kwargs, ): - """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking. + """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking (non-emulated dtypes). This function prepares input tensors, configures and launches the persistent GEMM kernel, optionally performs reference validation, and benchmarks the execution performance. @@ -2143,7 +2584,7 @@ def run( :param c_dtype: Data type for output tensor C :type c_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 + :type a_major/b_major/c_major: Literal["m", "k", "n"] :param mma_tiler_mn: MMA tiling size. :type mma_tiler_mn: Tuple[int, int] :param cluster_shape_mn: Cluster shape. @@ -2178,21 +2619,25 @@ def run( # Unpack parameters m, n, k, l = mnkl - # Skip unsupported testcase - if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement( - ab_dtype, - sf_dtype, + # Configure gemm kernel + gemm = Sm100BlockScaledPersistentDenseGemmKernel( sf_vec_size, - c_dtype, mma_tiler_mn, cluster_shape_mn, - m, - n, - k, - l, + ) + + # Skip unsupported testcase + if not gemm.can_implement( + mnkl, + ab_dtype, + sf_dtype, + c_dtype, a_major, b_major, c_major, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, ): raise TypeError( f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" @@ -2203,122 +2648,197 @@ def run( torch.manual_seed(1111) - # Create tensor A/B/C - a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) - b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) - c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) + # 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, a_torch = cutlass_torch.cute_tensor_like( - a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - b_tensor, b_torch = cutlass_torch.cute_tensor_like( - b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, c_torch = cutlass_torch.cute_tensor_like( - c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + # Check if configuration can be implemented + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] ) - # Mark tensor with element divisibility for 16B alignment - a_tensor.mark_compact_shape_dynamic( - mode=1 if a_major == "k" else 0, - stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), - divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, - ) - b_tensor.mark_compact_shape_dynamic( - mode=1 if b_major == "k" else 0, - stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), - divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, - ) - c_tensor.mark_compact_shape_dynamic( - mode=1 if c_major == "n" else 0, - stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), - divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + # Compile gemm kernel with fake tensors + compiled_gemm = scaled_mm( + gemm, + ab_dtype, + c_dtype, + sf_dtype, + a_major, + b_major, + c_major, + max_active_clusters, + current_stream, + options=f"--opt-level 2", ) - # Create scale factor tensor SFA/SFB - def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): - sf_k = ceil_div(k, sf_vec_size) - ref_shape = (l, mn, sf_k) + # Create Torch Tensors for A, scale factor A, B, scale factor B, C + a, b, c, sfa, sfb = prepare_tensors( + mnkl, ab_dtype, sf_dtype, sf_vec_size, c_dtype, a_major, b_major, c_major + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa + ) + sfb_reordered = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb + ) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr = construct_cute_pointers( + a, + b, + sfa_reordered, + sfb_reordered, + c, + ab_dtype, + sf_dtype, + c_dtype, + ) - atom_m = (32, 4) - atom_k = 4 - mma_shape = ( - l, - ceil_div(mn, atom_m[0] * atom_m[1]), - ceil_div(sf_k, atom_k), - atom_m[0], - atom_m[1], - atom_k, + # Compute reference result + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_gemm( + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream + ) + c_ref = reference_scaled_mm(a, b, sfa, sfb, c, (m, n, k, l), c_dtype) + if c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Rtol=0.001 and atol=0.1 are not supported for bitwise comparison of + # low dimensional floats. Please use rtol=0.0 and atol=0.0. + tolerance = 0.0 + torch.testing.assert_close(c, c_ref, atol=tolerance, rtol=tolerance) + + def generate_inputs(): + a, b, c, sfa, sfb = prepare_tensors( + mnkl, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa + ) + sfb_reordered = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb + ) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr = construct_cute_pointers( + a, + b, + sfa_reordered, + sfb_reordered, + c, + ab_dtype, + sf_dtype, + c_dtype, + ) + jit_args = cute.testing.JitArguments( + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream + ) + # Keep references to external variables (e.g., Torch tensors when taking a view) + jit_args.add_to_scope([a, b, sfa_reordered, sfb_reordered, c]) + return jit_args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a.numel() * a.element_size() + + b.numel() * b.element_size() + + sfa.numel() * sfa.element_size() + + sfb.numel() * sfb.element_size() + + c.numel() * c.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations ) - ref_permute_order = (1, 2, 0) - mma_permute_order = (3, 4, 1, 5, 2, 0) - - # Create f32 ref torch tensor (cpu) - ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( - ref_shape, - torch.float32, - permute_order=ref_permute_order, - init_type=cutlass_torch.TensorInitType.RANDOM, - init_config=cutlass_torch.RandomInitConfig( - min_val=1, - max_val=3, - ), - ) - - # Create f32 cute torch tensor (cpu) - cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( - mma_shape, - torch.float32, - permute_order=mma_permute_order, - init_type=cutlass_torch.TensorInitType.RANDOM, - init_config=cutlass_torch.RandomInitConfig( - min_val=0, - max_val=1, - ), - ) - - # convert ref f32 tensor to cute f32 tensor - cvt_sf_MKL_to_M32x4xrm_K4xrk_L( - from_dlpack(ref_f32_torch_tensor_cpu), - from_dlpack(cute_f32_torch_tensor_cpu), - ) - cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() - - # reshape makes memory contiguous - ref_f32_torch_tensor_cpu = ( - ref_f32_torch_tensor_cpu.permute(2, 0, 1) - .unsqueeze(-1) - .expand(l, mn, sf_k, sf_vec_size) - .reshape(l, mn, sf_k * sf_vec_size) - .permute(*ref_permute_order) - ) - # prune to mkl for reference check. - ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] - - # Create dtype cute torch tensor (cpu) - cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( - cute_f32_torch_tensor_cpu, - dtype, - is_dynamic_layout=True, - assumed_align=16, - ) - - # Convert f32 cute tensor to dtype cute tensor - cute_tensor = cutlass_torch.convert_cute_tensor( - cute_f32_torch_tensor, - cute_tensor, - dtype, - is_dynamic_layout=True, - ) - return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor - - sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( - l, m, k, sf_vec_size, sf_dtype - ) - sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( - l, n, k, sf_vec_size, sf_dtype + exec_time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_inputs, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, ) + return exec_time # Return execution time in microseconds + + +# This is to compatible with the other narrow +# precision combinations are not supported in either +# torch or dlpack. For example, Float4E2M1FN with Float8E8M0FNU. +def run_scaled_mm_with_emulated_dtype( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking (emulated dtypes). + + 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 sf_dtype: Data type for scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: Vector size for scale factor tensor + :type sf_vec_size: int + :param c_dtype: Data type for output tensor C + :type c_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: Literal["m", "n","k"] + :param mma_tiler_mn: MMA tiling size. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :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 Sm100 Persistent Dense BlockScaled GEMM test (Emulated) with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}") + print(f"C dtype: {c_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"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 # Configure gemm kernel gemm = Sm100BlockScaledPersistentDenseGemmKernel( @@ -2327,110 +2847,139 @@ def run( cluster_shape_mn, ) - # Compute max active clusters on current device - hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters( + # Skip unsupported testcase + if not gemm.can_implement( + mnkl, + ab_dtype, + sf_dtype, + c_dtype, + a_major, + b_major, + c_major, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + ): + raise TypeError( + f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" + ) + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # 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) + + # Check if configuration can be implemented + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) - # Initialize Stream - current_stream = cutlass_torch.default_stream() - - # Compile gemm kernel - compiled_gemm = cute.compile( + # Compile gemm kernel with fake tensors + compiled_gemm = scaled_mm( gemm, - a_tensor, - b_tensor, - sfa_tensor, - sfb_tensor, - c_tensor, + ab_dtype, + c_dtype, + sf_dtype, + a_major, + b_major, + c_major, max_active_clusters, current_stream, options=f"--opt-level 2", ) + # Create Torch Tensors for A, scale factor A, B, scale factor B, C + a, b, c, sfa, sfb = prepare_tensors_emulated( + mnkl, ab_dtype, sf_dtype, sf_vec_size, c_dtype, a_major, b_major, c_major + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa + ) + sfb_reordered = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb + ) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute = ( + construct_cute_pointers_emulated( + a, + b, + sfa_reordered, + sfb_reordered, + c, + ab_dtype, + sf_dtype, + c_dtype, + ) + ) + # Compute reference result if not skip_ref_check: # Execute kernel once for reference checking compiled_gemm( - a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, current_stream + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream ) - print("Verifying results...") - res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) - res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) - ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) - - # Convert c back to f32 for comparison. - c_ref_device = c_ref.cuda() - cute.testing.convert( - c_tensor, - from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( - leading_dim=(1 if c_major == "n" else 0) - ), + c_ref = reference_scaled_mm_emulated( + a, b, sfa, sfb, c, (m, n, k, l), sf_vec_size, c_dtype ) - c_ref = c_ref_device.cpu() + if c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Rtol=0.001 and atol=0.1 are not supported for bitwise comparison of + # low dimensional floats. Please use rtol=0.0 and atol=0.0. + tolerance = 0.0 + torch.testing.assert_close(c, c_ref, atol=tolerance, rtol=tolerance) - if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): - torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) - elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): - # Convert ref : f32 -> f8 -> f32 - ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute( - 1, 2, 0 + def generate_inputs(): + a, b, c, sfa, sfb = prepare_tensors_emulated( + mnkl, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa + ) + sfb_reordered = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb + ) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute = ( + construct_cute_pointers_emulated( + a, + b, + sfa_reordered, + sfb_reordered, + c, + ab_dtype, + sf_dtype, + c_dtype, ) - ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic( - leading_dim=1 - ) - ref_f8.element_type = c_dtype - ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() - ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic( - leading_dim=1 - ) - cute.testing.convert(ref_tensor, ref_f8) - cute.testing.convert(ref_f8, ref_tensor) - ref = ref_device.cpu() - torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) - def generate_tensors(): - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, _ = cutlass_torch.cute_tensor_like( - c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + jit_args = cute.testing.JitArguments( + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream ) + # Keep references to external variables (e.g., Torch tensors when taking a view) + jit_args.add_to_scope([a, b, sfa_reordered, sfb_reordered, c, a_cute, b_cute]) + return jit_args - # Mark tensor to be byte aligned - a_tensor.mark_compact_shape_dynamic( - mode=1 if a_major == "k" else 0, - stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), - divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1, - ) - b_tensor.mark_compact_shape_dynamic( - mode=1 if b_major == "k" else 0, - stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), - divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1, - ) - c_tensor.mark_compact_shape_dynamic( - mode=1 if c_major == "n" else 0, - stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), - divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, - ) - - _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype) - _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype) - return cute.testing.JitArguments( - a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, current_stream - ) workspace_count = 1 if use_cold_l2: one_workspace_bytes = ( - a_torch.numel() * a_torch.element_size() - + b_torch.numel() * b_torch.element_size() - + sfa_torch.numel() * sfa_torch.element_size() - + sfb_torch.numel() * sfb_torch.element_size() - + c_torch.numel() * c_torch.element_size() + a.numel() * a.element_size() + + b.numel() * b.element_size() + + sfa.numel() * sfa.element_size() + + sfb.numel() * sfb.element_size() + + c.numel() * c.element_size() ) workspace_count = cute.testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations @@ -2438,16 +2987,78 @@ def run( exec_time = cute.testing.benchmark( compiled_gemm, - workspace_generator=generate_tensors, + workspace_generator=generate_inputs, workspace_count=workspace_count, stream=current_stream, warmup_iterations=warmup_iterations, iterations=iterations, ) - return exec_time # Return execution time in microseconds +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """ + Execute the appropriate GEMM function based on dtype. + + Routes to either run_scaled_mm_with_emulated_dtype or run_scaled_mm + depending on whether the dtypes require emulation. + """ + if is_emulated_dtype(ab_dtype, sf_dtype, c_dtype): + exec_time = run_scaled_mm_with_emulated_dtype( + mnkl, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + tolerance, + warmup_iterations, + iterations, + skip_ref_check, + use_cold_l2, + ) + else: + exec_time = run_scaled_mm( + mnkl, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + tolerance, + warmup_iterations, + iterations, + skip_ref_check, + use_cold_l2, + ) + return exec_time + + if __name__ == "__main__": def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: @@ -2481,7 +3092,7 @@ if __name__ == "__main__": help="Cluster shape (comma-separated)", ) parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) - parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E4M3FN) parser.add_argument("--sf_vec_size", type=int, default=16) parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") @@ -2520,6 +3131,7 @@ if __name__ == "__main__": if len(args.cluster_shape_mn) != 2: parser.error("--cluster_shape_mn must contain exactly 2 values") + # Execute GEMM with appropriate function based on dtype run( args.mnkl, args.ab_dtype, diff --git a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py new file mode 100644 index 00000000..f423fcce --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py @@ -0,0 +1,2576 @@ +# 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 argparse +from typing import Type, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass._mlir.dialects import math +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import from_dlpack + +""" +This example provides an experimental implementation of the SM100 batched dense blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. + +A high-performance persistent batched dense blockscaled 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") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") +- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has M×ceil_div(K, sf_vec_size)×L elements respectively +- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has N×ceil_div(K, sf_vec_size)×L elements respectively + +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: + - Load scale factor A/B from shared memory (SMEM) to tensor memory (TMEM) using tcgen05.cp instruction. + - 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.kind.block_scale instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Read scalefactor A from TMEM +- Read scalefactor B from TMEM +- 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 shown below: + +.. code-block:: bash + + python examples/blackwell/dense_blockscaled_gemm_persistent.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,1024,1 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/dense_blockscaled_gemm_persistent.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,1024,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + + +Constraints: +* Supported input data types: mxf8, mxf4, nvf4 + see detailed valid dtype combinations in below Sm100BlockScaledPersistentDenseGemmKernel class documentation +* A/B tensor must have the same data type, mixed data type is not supported (e.g., mxf8 x mxf4) +* Mma tiler M must be 128 or 256(use_2cta_instrs) +* Mma tiler N must be 128 or 256 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs) +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 16 and 32 for Float8 and Float4, respectively. +""" + + +class Sm100BlockScaledPersistentDenseGemmKernel: + """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :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] + + :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 combinations of A/B data types, SF data typs and SF vector size: + - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16 + + :note: Supported accumulator data types: + - Float32 + + :note: Supported C data types: + - Float32 + - Float16/BFloat16 + - Float8E4M3FN/Float8E5M2 + :note: Constraints: + - MMA tiler M must be 128 or 256 (use_2cta_instrs) + - MMA tiler N must be 128/256 + - Cluster shape M must be multiple of 2 if Mma tiler M is 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors + + Example: + >>> gemm = Sm100BlockScaledPersistentDenseGemmKernel( + ... sf_vec_size=16, + ... mma_tiler_mn=(256, 128), + ... cluster_shape_mn=(2, 1) + ... ) + >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, amax_tensor, max_active_clusters, stream) + """ + + def __init__( + self, + sf_vec_size: int, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ): + """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, always set to Float32 + - sf_vec_size: Scalefactor A/B vector size. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + """ + + self.acc_dtype = cutlass.Float32 + self.sf_vec_size = sf_vec_size + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.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.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + + # Amax reduction configuration + self.num_epilog_warps = len(self.epilog_warp_id) + + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("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/SFA/SFB + - Computing epilogue subtile + - Setting up A/B/SFA/SFB/C stage counts in shared memory + - Computing A/B/SFA/SFB/C shared memory layout + """ + # Compute mma instruction shapes + # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) + self.mma_inst_shape_mn = ( + self.mma_tiler[0], + self.mma_tiler[1], + ) + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + + # 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_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[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,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.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.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute epilogue subtile + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + # 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.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + ) + + # Compute A/B/SFA/SFB/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.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + 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, + ) + + @cute.jit + def __call__( + self, + a_tensor: cute.Tensor, + b_tensor: cute.Tensor, + sfa_tensor: cute.Tensor, + sfb_tensor: cute.Tensor, + c_tensor: cute.Tensor, + amax_tensor: 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_tensor: Input tensor A + :type a_tensor: cute.Tensor + :param b_tensor: Input tensor B + :type b_tensor: cute.Tensor + :param sfa_tensor: Scale factor tensor A + :type sfa_tensor: cute.Tensor + :param sfb_tensor: Scale factor tensor B + :type sfb_tensor: cute.Tensor + :param c_tensor: Output tensor C + :type c_tensor: cute.Tensor + :param amax_tensor: Output tensor for absolute maximum value + :type amax_tensor: 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. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type + self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type + self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type + self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) + + # 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() + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, self.sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, self.sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + 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_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # 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_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Setup TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb_tensor, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + 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) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Setup TMA store for C + 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_tensor, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c_tensor, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + self.buffer_align_bytes = 1024 + + # 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, + cute.cosize(self.c_smem_layout_staged.outer), + ], + 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, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sSFA: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sSFB: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + # Amax reduction shared memory (one FP32 per epilogue warp) + # Use smaller alignment for amax since it's only 16 bytes + sAmax: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.num_epilog_warps], + 16, + ] + + self.shared_storage = SharedStorage + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + amax_tensor, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_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, + min_blocks_per_mp=1, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + mAmax: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout], + 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) + cpasync.prefetch_descriptor(tma_atom_sfa) + cpasync.prefetch_descriptor(tma_atom_sfb) + 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 + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_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) + + # 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, + defer_sync=True, + ) + + # 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, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.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=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/SFA/SFB/C + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + # (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 + ) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # Shared memory for amax reduction (one FP32 per epilogue warp) + # Simple 1D layout + amax_layout = cute.make_layout((self.num_epilog_warps,)) + sAmax = storage.sAmax.get_tensor(amax_layout) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_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 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_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, bK, RestM, RestK, RestL) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gSFB_nkl = cute.local_tile( + mSFB_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) + thr_mma_sfb = tiled_mma_sfb.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_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma_sfb.partition_B(gSFB_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), RestN, 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), + ) + + # TMA load scaled factor A partition_S/D + sfa_cta_layout = a_cta_layout + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load scaled factor B partition_S/D + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # 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=self.cluster_shape_mn) + + # + # 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]) + ] + + # ((atom_v, rest_v), RestK) + tAgSFA_slice = tAgSFA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[ + (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/SFA/SFB + 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, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfb_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.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # Make accumulator tmem tensor + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Make SFA tmem tensor + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + dtype=self.sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # Make SFB tmem tensor + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=self.sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + # + # Partition for S2T copy of SFA/SFB + # + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + # + # 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 + ) + + # Copy SFA/SFB from smem to tmem + s2t_stage_coord = ( + None, + None, + None, + None, + ab_consumer_state.index, + ) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + 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, + ) + + # Set SFA/SFB tensor to tiled_mma + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set( + tcgen05.Field.SFA, + tCtSFA[sf_kblock_coord].iterator, + ) + tiled_mma.set( + tcgen05.Field.SFB, + tCtSFB[sf_kblock_coord].iterator, + ) + + 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 + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_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 = 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 + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, 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 + ) + + # 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 + # + # ((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 + + # Initialize thread-local amax accumulator for this tile + # Use 0.0 as initial value since we're computing absolute maximum + thread_tile_amax = cutlass.Float32(0.0) + + 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) + + # Accumulate thread-level amax across all subtiles in this tile + # Note: We need absolute value maximum, so take abs first + acc_values = tTR_rAcc.load() + # Apply element-wise absolute value using math.absf (supports vectors) + abs_acc_values_ir = math.absf( + acc_values.ir_value() # operand (positional) + ) + abs_acc_values = type(acc_values)( + abs_acc_values_ir, acc_values.shape, acc_values.dtype + ) + subtile_amax = abs_acc_values.reduce( + cute.ReductionOp.MAX, + cutlass.Float32(0.0), + 0, # Use 0.0 as init for abs values + ) + thread_tile_amax = cute.arch.fmax(thread_tile_amax, subtile_amax) + + # + # 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( + "async.shared", + space="cta", + ) + self.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() + self.epilog_sync_barrier.arrive_and_wait() + + # Perform amax reduction after all subtiles are processed + # Warp-level reduction using wrapper function + warp_amax = cute.arch.warp_redux_sync( + value=thread_tile_amax, + kind="fmax", + mask_and_clamp=0xFFFFFFFF, + nan=True, + ) + # Each epilogue warp's lane 0 writes warp amax to shared memory + if cute.arch.lane_idx() == 0: + sAmax[warp_idx] = cutlass.Float32(warp_amax) + + # Ensure all epilogue warps complete their writes before block reduction + self.epilog_sync_barrier.arrive_and_wait() + + # Block-level reduction: only first epilogue warp's lane 0 handles this + if warp_idx == self.epilog_warp_id[0] and cute.arch.lane_idx() == 0: + block_amax = cutlass.Float32( + 0.0 + ) # Initial value for absolute maximum + for i in cutlass.range(self.num_epilog_warps): + warp_amax_val = sAmax[i] + block_amax = cute.arch.fmax(block_amax, warp_amax_val) + + # Global atomic max (accumulates across all tiles for final tensor amax) + # Since we compute absolute values, all values are non-negative + # Use wrapper function for atomic max operation + _ = cute.arch.atomic_max_float32( + ptr=mAmax.iterator.llvm_ptr, value=block_amax + ) + # + # 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() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(acc_tmem_ptr) + # + # Wait for C store complete + # + c_pipeline.producer_tail() + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). + + :param sSF: The scale factor tensor in smem + :type sSF: cute.Tensor + :param tSF: The scale factor tensor in tmem + :type tSF: cute.Tensor + + :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: + - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) + - tCsSF_compact_s2t: The partitioned scale factor tensor in smem + - tSF_compact_s2t: The partitioned scale factor tensor in tmem + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSF_compact = cute.filter_zeros(sSF) + # (MMA, MMA_MN, MMA_K) + tCtSF_compact = cute.filter_zeros(tSF) + + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + 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 + + 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 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 (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 + :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 + ) + + 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 + + @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, + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + smem_capacity: int, + occupancy: int, + ) -> 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 sf_dtype: Data type of Scale factor. + :type sf_dtype: type[cutlass.Numeric] + :param sf_vec_size: Scale factor vector size. + :type sf_vec_size: int + :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 + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # ACC stages + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + + # Default C stages + num_c_stage = 2 + + # Calculate smem layout and size for one stage of A, B, SFA, SFB 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 + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 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, + ) + + 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) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + amax_bytes = 16 + + # Calculate A/B/SFA/SFB 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/SFA/SFB stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes + amax_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes + # Add remaining unused smem to epilogue + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes + amax_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 is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + ) -> bool: + """ + Check if the dtypes and sf_vec_size are valid combinations + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes and sf_vec_size are valid, False otherwise + :rtype: bool + """ + is_valid = True + + # Check valid ab_dtype + if ab_dtype not in { + cutlass.Float4E2M1FN, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + # Check valid sf_vec_size + if sf_vec_size not in {16, 32}: + is_valid = False + + # Check valid sf_dtype + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + is_valid = False + + # Check valid sf_dtype and sf_vec_size combinations + if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32: + is_valid = False + if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16: + is_valid = False + + # Check valid c_dtype + if c_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + return is_valid + + @staticmethod + def is_valid_layouts( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if layouts and dtypes are valid combinations + + :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 dimension of the A tensor + :type a_major: str + :param b_major: The major dimension of the B tensor + :type b_major: str + :param c_major: The major dimension of the C tensor + :type c_major: str + + :return: True if the layouts are valid, False otherwise + :rtype: bool + """ + is_valid = True + + if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"): + is_valid = False + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ) -> bool: + """ + Check if the mma tiler and cluster shape are valid + + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if mma_tiler_mn[0] not in [128, 256]: + is_valid = False + if mma_tiler_mn[1] not in [128, 256]: + is_valid = False + # Skip illegal cluster shape + if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + # Special cluster shape check for scale factor multicasts. + # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + 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 + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the gemm can be implemented + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + :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 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 gemm can be implemented, False otherwise + :rtype: bool + """ + can_implement = True + # Skip unsupported types + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype, sf_dtype, sf_vec_size, c_dtype + ): + can_implement = False + # Skip unsupported layouts + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_layouts( + ab_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn, cluster_shape_mn + ): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_tensor_alignment( + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + return can_implement + + +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_tensor: cute.Tensor, + sf_mma_tensor: cute.Tensor, +): + """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout""" + # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) + # group to ((32, 4, rest_m), (4, rest_k), l) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + + +def compute_reference_amax(output_tensor) -> float: + import torch + + """ + Compute reference amax value on CPU. + + Args: + output_tensor: torch.Tensor, GEMM output result (CPU tensor) + + Returns: + float: reference amax value + """ + # Ensure FP32 for computation + if output_tensor.dtype != torch.float32: + output_fp32 = output_tensor.float() + else: + output_fp32 = output_tensor + + # Compute absolute maximum value + reference_amax = torch.amax(torch.abs(output_fp32)) + + return reference_amax.item() + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """Execute a persistent batched dense blockscaled 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 sf_dtype: Data type for scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: Vector size for scale factor tensor + :type sf_vec_size: int + :param c_dtype: Data type for output tensor C + :type c_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. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :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 Sm100 Persistent Dense BlockScaled GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}") + print(f"C dtype: {c_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"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'}") + import torch + import cutlass.torch as cutlass_torch + + # Unpack parameters + m, n, k, l = mnkl + + # Skip unsupported testcase + if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement( + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + c_major, + ): + raise TypeError( + f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" + ) + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Create tensor A/B/C + a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) + b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) + c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) + + a_tensor, a_torch = cutlass_torch.cute_tensor_like( + a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, b_torch = cutlass_torch.cute_tensor_like( + b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Create amax tensor (single FP32 value, initialized to -inf) + amax_ref = cutlass_torch.matrix( + 1, + 1, + 1, + False, + cutlass.Float32, + init_type=cutlass_torch.TensorInitType.SCALAR, + init_config=cutlass_torch.ScalarInitConfig(-float("inf")), + ) + amax_tensor, amax_torch = cutlass_torch.cute_tensor_like( + amax_ref, cutlass.Float32, is_dynamic_layout=True, assumed_align=16 + ) + + # Mark tensor with element divisibility for 16B alignment + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + + # Create scale factor tensor SFA/SFB + def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + ref_shape = (l, mn, sf_k) + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + ref_permute_order = (1, 2, 0) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + # Create f32 ref torch tensor (cpu) + ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + ref_shape, + torch.float32, + permute_order=ref_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=1, + max_val=3, + ), + ) + + # Create f32 cute torch tensor (cpu) + cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=0, + max_val=1, + ), + ) + + # convert ref f32 tensor to cute f32 tensor + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + from_dlpack(ref_f32_torch_tensor_cpu), + from_dlpack(cute_f32_torch_tensor_cpu), + ) + cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() + + # reshape makes memory contiguous + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(l, mn, sf_k, sf_vec_size) + .reshape(l, mn, sf_k * sf_vec_size) + .permute(*ref_permute_order) + ) + # prune to mkl for reference check. + ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] + + # Create dtype cute torch tensor (cpu) + cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( + cute_f32_torch_tensor_cpu, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + # Convert f32 cute tensor to dtype cute tensor + cute_tensor = cutlass_torch.convert_cute_tensor( + cute_f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=True, + ) + return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor + + sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype + ) + sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype + ) + + # Configure gemm kernel + gemm = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + ) + + # Compute max active clusters on current device + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + # Compile gemm kernel + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + amax_tensor, + max_active_clusters, + current_stream, + options=f"--opt-level 2", + ) + + # Compute reference result + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_gemm( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + amax_tensor, + current_stream, + ) + print("Verifying results...") + res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) + res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + + # Save original Float32 ref for amax computation (before quantization) + ref_for_amax = ref.clone() + + # Convert c back to f32 for comparison. + c_ref_device = c_ref.cuda() + cute.testing.convert( + c_tensor, + from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=(1 if c_major == "n" else 0) + ), + ) + c_ref = c_ref_device.cpu() + + if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Convert ref : f32 -> f8 -> f32 + ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute( + 1, 2, 0 + ) + ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + ref_f8.element_type = c_dtype + ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() + ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + cute.testing.convert(ref_tensor, ref_f8) + cute.testing.convert(ref_f8, ref_tensor) + ref = ref_device.cpu() + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + # Verify amax result + device_amax = amax_torch.cpu().squeeze() # Remove dimensions to make it scalar + # Use original Float32 ref (before quantization) for amax computation + reference_amax = torch.tensor(compute_reference_amax(ref_for_amax)) + + # AMAX validation using same approach as GEMM result validation + torch.testing.assert_close( + device_amax, reference_amax, atol=tolerance, rtol=1e-02 + ) + + def generate_tensors(): + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, _ = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Mark tensor to be byte aligned + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, + ) + + _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype) + _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype) + + # Create amax tensor for benchmarking (reset to -inf for each iteration) + amax_ref_bench = cutlass_torch.matrix( + 1, + 1, + 1, + False, + cutlass.Float32, + init_type=cutlass_torch.TensorInitType.SCALAR, + init_config=cutlass_torch.ScalarInitConfig(-float("inf")), + ) + amax_tensor_bench, _ = cutlass_torch.cute_tensor_like( + amax_ref_bench, cutlass.Float32, is_dynamic_layout=True, assumed_align=16 + ) + + return cute.testing.JitArguments( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + amax_tensor_bench, + current_stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch.numel() * a_torch.element_size() + + b_torch.numel() * b_torch.element_size() + + sfa_torch.numel() * sfa_torch.element_size() + + sfb_torch.numel() * sfb_torch.element_size() + + c_torch.numel() * c_torch.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +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 Sm100 Dense Persistent BlockScaled GEMM." + ) + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(512, 256, 256, 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.Float4E2M1FN) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) + parser.add_argument("--sf_vec_size", type=int, default=16) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) + 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( + "--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") + + run( + args.mnkl, + args.ab_dtype, + args.sf_dtype, + args.sf_vec_size, + args.c_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm.py b/examples/python/CuTeDSL/blackwell/dense_gemm.py index 1e347bd1..94495fb1 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm.py @@ -30,7 +30,6 @@ import argparse from typing import Optional, Type, Tuple, Union import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute @@ -38,7 +37,6 @@ import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.testing as testing @@ -215,7 +213,6 @@ class DenseGemmKernel: self.occupancy = 1 self.threads_per_cta = 128 - self.smem_capacity = utils.get_smem_capacity_in_bytes() def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -278,6 +275,8 @@ class DenseGemmKernel: else: self.epi_tile = self.cta_tile_shape_mnk[:2] + self.smem_capacity = utils.get_smem_capacity_in_bytes() + # Setup A/B/C stage count in shared memory self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( tiled_mma, @@ -1030,7 +1029,10 @@ class DenseGemmKernel: c_buffer = 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("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) pipeline.sync(barrier_id=1) # TMA store C to global memory @@ -1261,7 +1263,7 @@ class DenseGemmKernel: """ acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) - return sm100_utils.get_num_tmem_alloc_cols(tCtAcc_fake) + return utils.get_num_tmem_alloc_cols(tCtAcc_fake) def is_valid_dtypes( self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] @@ -1494,6 +1496,9 @@ class DenseGemmKernel: def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) @@ -1522,6 +1527,9 @@ def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): def compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance): + import torch + import cutlass.torch as cutlass_torch + # Copy gpu result back kernel_result = c_torch_gpu.cpu() @@ -1615,6 +1623,7 @@ def run( print(f"Iterations: {iterations}") print(f"Skip reference checking: {skip_ref_check}") print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + import torch # Unpack parameters m, n, k, l = mnkl @@ -1654,6 +1663,8 @@ def run( compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance) def generate_tensors(): + import cutlass.torch as cutlass_torch + a_tensor, _ = cutlass_torch.cute_tensor_like( a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py b/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py index 2b7d7d84..d8881325 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py @@ -29,13 +29,11 @@ import argparse from typing import Optional, Tuple, Type, Union -import torch import cuda.bindings.driver as cuda 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 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @@ -1176,7 +1174,10 @@ class SM100PersistentDenseGemmAlphaBetaKernel: tSR_sC[(None, None, None, c_pipeline_consumer_state.index)], tSR_rC, ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) c_pipeline.consumer_release(c_pipeline_consumer_state) # Advance pipeline states @@ -1203,7 +1204,10 @@ class SM100PersistentDenseGemmAlphaBetaKernel: tiled_copy_r2s, tRS_rD, tRS_sD[(None, None, None, d_buffer)] ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) epilog_sync_barrier.arrive_and_wait() # @@ -1875,6 +1879,9 @@ class SM100PersistentDenseGemmAlphaBetaKernel: def create_tensors(l, m, n, k, a_major, b_major, cd_major, ab_dtype, c_dtype, d_dtype): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) @@ -1991,6 +1998,9 @@ def compare( beta, tolerance, ): + import torch + import cutlass.torch as cutlass_torch + # Copy gpu result back kernel_result = d_torch_gpu.cpu() @@ -2057,6 +2067,8 @@ def run_dense_gemm( # Unpack parameters m, n, k, l = mnkl + import torch + if not torch.cuda.is_available(): raise RuntimeError("GPU is required to run this example!") diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py index e54b1412..d08ea62b 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py @@ -28,14 +28,14 @@ import argparse from typing import Optional, Tuple, Type, Union - +from functools import lru_cache import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -from cutlass.cute.runtime import from_dlpack import cutlass.utils as utils +from cutlass.utils import is_fp8_dtype, create_cute_tensor_for_fp8 import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 @@ -281,6 +281,7 @@ class PersistentDenseGemmKernel: self.mma_tiler_mn = mma_tiler_mn self.mma_tiler = (*mma_tiler_mn, 1) self.use_tma_store = use_tma_store + self.arch = "sm_100" self.cta_group = ( tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE @@ -288,17 +289,26 @@ class PersistentDenseGemmKernel: self.occupancy = 1 # Set specialized warp ids - self.epilog_warp_id = (0, 1, 2, 3) + self.epilogue_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) + (self.mma_warp_id, self.tma_warp_id, *self.epilogue_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") + + def _create_tiled_mma(self): + return utils.sm100.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], + ) def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -315,14 +325,7 @@ class PersistentDenseGemmKernel: - Computing tensor memory allocation columns """ # Configure tiled mma - tiled_mma = utils.sm100.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], - ) + tiled_mma = self._create_tiled_mma() # Compute mma/cluster/tile shapes mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) @@ -367,6 +370,8 @@ class PersistentDenseGemmKernel: self.c_dtype, self.c_layout, self.epi_tile, 1 ) + self.smem_capacity = utils.get_smem_capacity_in_bytes() + # 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, @@ -396,7 +401,7 @@ class PersistentDenseGemmKernel: # 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 + tiled_mma, self.mma_tiler, self.num_acc_stage, self.arch ) @cute.jit @@ -443,17 +448,11 @@ class PersistentDenseGemmKernel: if cutlass.const_expr(self.a_dtype != self.b_dtype): raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + tiled_mma = self._create_tiled_mma() + # Setup attributes that dependent on gemm inputs self._setup_attributes() - tiled_mma = utils.sm100.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 @@ -618,7 +617,7 @@ class PersistentDenseGemmKernel: # Initialize acc_pipeline (barrier) and states acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - num_acc_consumer_threads = len(self.epilog_warp_id) * ( + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( 2 if use_2cta_instrs else 1 ) acc_pipeline_consumer_group = pipeline.CooperativeGroup( @@ -635,19 +634,19 @@ class PersistentDenseGemmKernel: 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)), + 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.epilog_warp_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.epilog_warp_id[0], + 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, ) @@ -763,6 +762,16 @@ class PersistentDenseGemmKernel: # pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + # + # Construct the scheduler + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + ) + work_tile = tile_sched.initial_work_tile_info() + # # Specialized TMA load warp # @@ -771,10 +780,6 @@ class PersistentDenseGemmKernel: # # 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 @@ -855,10 +860,6 @@ class PersistentDenseGemmKernel: # # 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 @@ -973,37 +974,70 @@ class PersistentDenseGemmKernel: # # Persistent tile scheduling loop for epilogue # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage ) if cutlass.const_expr(self.use_tma_store): 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, + 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 + ) + 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], + ) + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete + c_pipeline.producer_tail() else: - self.epilogue( - tidx, - acc_pipeline, - tiled_mma, - tCtAcc_base, - tCgC, - epi_tile, - tile_sched, - epilogue_op, - tmem_dealloc_barrier, - ) + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() # # Dealloc the tensor memory buffer @@ -1011,354 +1045,6 @@ class PersistentDenseGemmKernel: 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("async.shared", space="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 = utils.sm100.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 = utils.sm100.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, @@ -1401,6 +1087,7 @@ class PersistentDenseGemmKernel: tiled_mma: cute.TiledMma, mma_tiler: Tuple[int, int, int], num_acc_stage: int, + arch: str, ) -> int: """ Compute the number of tensor memory allocation columns. @@ -1417,25 +1104,29 @@ class PersistentDenseGemmKernel: """ 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) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch) return num_tmem_alloc_cols - def is_valid_dtypes( - self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] - ) -> bool: + def check_supported_dtypes( + self, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + ): """ 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 a_dtype: The data type of the A operands + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operands + :type b_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 + :raises testing.CantImplementError: If the dtypes are invalid """ valid_ab_dtypes = { cutlass.Float16, @@ -1446,11 +1137,15 @@ class PersistentDenseGemmKernel: cutlass.Float8E4M3FN, cutlass.Float8E5M2, } - if ab_dtype not in valid_ab_dtypes: - return False + if a_dtype not in valid_ab_dtypes or b_dtype not in valid_ab_dtypes: + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype}" + ) if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: - return False + raise testing.CantImplementError( + f"Unsupported accumulator dtype: {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and AB type acc_ab_compatibility = { @@ -1469,8 +1164,13 @@ class PersistentDenseGemmKernel: 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 + if ( + a_dtype not in acc_ab_compatibility[self.acc_dtype] + or b_dtype not in acc_ab_compatibility[self.acc_dtype] + ): + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype} for accumulator dtype: {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and C type acc_c_compatibility = { @@ -1499,28 +1199,32 @@ class PersistentDenseGemmKernel: } # Check compatibility between accumulator type and C type if c_dtype not in acc_c_compatibility[self.acc_dtype]: - return False + raise testing.CantImplementError( + f"Unsupported C dtype: {c_dtype} for accumulator dtype: {self.acc_dtype}" + ) - return True - - def is_valid_mma_tiler_and_cluster_shape(self) -> bool: + def check_mma_tiler_and_cluster_shape(self): """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 + :raises testing.CantImplementError: If the mma tiler and cluster shape are invalid """ - 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 + raise testing.CantImplementError( + f"Invalid mma tiler & use_2cta_instrs: {self.mma_tiler_mn}, {self.use_2cta_instrs}" + ) if self.mma_tiler_mn[1] not in range(32, 257, 32): - is_valid = False + raise testing.CantImplementError( + f"Invalid mma tiler N: {self.mma_tiler_mn[1]}" + ) # Skip illegal cluster shape if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: - is_valid = False + raise testing.CantImplementError( + f"Invalid cluster shape M: {self.cluster_shape_mn[0]}" + ) # Skip invalid cluster shape is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 if ( @@ -1530,21 +1234,23 @@ class PersistentDenseGemmKernel: 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 + raise testing.CantImplementError( + f"Invalid cluster shape: {self.cluster_shape_mn}" + ) - def is_valid_tensor_alignment( + def check_tensor_alignment( self, m: int, n: int, k: int, l: int, - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_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 @@ -1556,8 +1262,10 @@ class PersistentDenseGemmKernel: :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 a_dtype: The data type of the A operands + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operands + :type b_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 @@ -1567,10 +1275,8 @@ class PersistentDenseGemmKernel: :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 + :raises testing.CantImplementError: If the tensor alignment is invalid """ - is_valid = True # TODO: move to utils def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): @@ -1580,14 +1286,15 @@ class PersistentDenseGemmKernel: 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)) + not check_contiguous_16B_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(b_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 + raise testing.CantImplementError( + f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {a_dtype}, {b_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}" + ) - def is_valid_epilog_store_option(self, m: int, n: int) -> bool: + def check_epilog_store_option(self, m: int, n: int): """ Check if the epilogue store option is valid @@ -1596,11 +1303,8 @@ class PersistentDenseGemmKernel: :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 + :raises testing.CantImplementError: If the epilogue store option is invalid """ - - 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), @@ -1608,13 +1312,15 @@ class PersistentDenseGemmKernel: ) 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 + raise testing.CantImplementError( + f"Invalid epilog store option: {m}, {n}" + ) def can_implement( self, mnkl: Tuple[int, int, int, int], - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, @@ -1625,8 +1331,10 @@ class PersistentDenseGemmKernel: :param mnkl: Problem size as a tuple (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 a_dtype: Data type for input tensors A. + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Data type for input tensors B. + :type b_dtype: Type[cutlass.Numeric] :param c_dtype: Data type for output tensor C. :type c_dtype: Type[cutlass.Numeric] :param a_major: Major dimension of the A tensor layout ("m" or "k"). @@ -1639,25 +1347,20 @@ class PersistentDenseGemmKernel: :rtype: bool """ - # Skip unsupported types - if not self.is_valid_dtypes(ab_dtype, c_dtype): - return False + try: + # Skip unsupported types + self.check_supported_dtypes(a_dtype, b_dtype, c_dtype) - # Skip invalid mma tile shape and cluster shape - if not self.is_valid_mma_tiler_and_cluster_shape(): - return False + # Skip invalid mma tile shape and cluster shape + self.check_mma_tiler_and_cluster_shape() - # Unpack mnkl for clarity in calling the epilog check - m, n, k, l = mnkl - # Skip illegal problem shape for load/store alignment - if not self.is_valid_tensor_alignment( - m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major - ): + m, n, k, l = mnkl + self.check_tensor_alignment( + m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ) + self.check_epilog_store_option(m, n) + except testing.CantImplementError: return False - # Skip invalid epilogue store option - if not self.is_valid_epilog_store_option(m, n): - return False - return True @@ -1702,45 +1405,117 @@ def bmm( gemm_op(a, b, c, max_active_clusters, stream, epilogue_op) +@lru_cache(maxsize=1) def prepare_tensors( mnkl: Tuple[int, int, int, int], - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, c_major: str, init_random: bool = True, + normal_mean: float = 0.0, + normal_std: float = 1.0, ): + """Prepare tensors for GEMM. + + Returns: + Tuple of (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage): + - *_f32: Float32 tensors with the logical data (for reference and fp8 conversion) + - *_storage: Storage tensors for DLPack (uint8 for fp8, otherwise the target dtype) + """ import torch from cutlass.torch import dtype as torch_dtype m, n, k, l = mnkl if a_major == "k": - a = torch.empty((l, m, k), dtype=torch.float32, device="cuda") + a_f32 = torch.empty((l, m, k), dtype=torch.float32, device="cuda") elif a_major == "m": - a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + a_f32 = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if b_major == "n": - b = torch.empty((l, k, n), dtype=torch.float32, device="cuda") + b_f32 = torch.empty((l, k, n), dtype=torch.float32, device="cuda") elif b_major == "k": - b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1) + b_f32 = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if c_major == "n": - c = torch.empty((l, m, n), dtype=torch.float32, device="cuda") + c_f32 = torch.empty((l, m, n), dtype=torch.float32, device="cuda") elif c_major == "m": - c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + c_f32 = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if init_random: - a.random_(-2, 3) - b.random_(-2, 3) - c.random_(-2, 3) + # Uniform random initialization in range [-2, 3) + a_f32.random_(-2, 3) + b_f32.random_(-2, 3) + c_f32.random_(-2, 3) - return ( - a.to(dtype=torch_dtype(ab_dtype)), - b.to(dtype=torch_dtype(ab_dtype)), - c.to(dtype=torch_dtype(c_dtype)), + else: + # Normal (Gaussian) initialization with user-specified mean and std + a_f32.normal_(mean=normal_mean, std=normal_std) + b_f32.normal_(mean=normal_mean, std=normal_std) + c_f32.normal_(mean=normal_mean, std=normal_std) + + # For float8 types, use uint8 as storage type to avoid dlpack limitation + # (dlpack doesn't support float8 types) + # For other types, convert to the target dtype + a_storage_dtype = torch.uint8 if is_fp8_dtype(a_dtype) else torch_dtype(a_dtype) + b_storage_dtype = torch.uint8 if is_fp8_dtype(b_dtype) else torch_dtype(b_dtype) + c_storage_dtype = torch.uint8 if is_fp8_dtype(c_dtype) else torch_dtype(c_dtype) + + a_storage = a_f32.to(dtype=a_storage_dtype) + b_storage = b_f32.to(dtype=b_storage_dtype) + c_storage = c_f32.to(dtype=c_storage_dtype) + + return (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage) + + +@lru_cache(maxsize=1) +def compile_bmm( + mnkl: Tuple[int, int, int, int], + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + 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), + max_active_clusters: cutlass.Constexpr = None, + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + epilogue_op: cutlass.Constexpr = lambda x: x, +): + from cutlass.cute.runtime import make_fake_stream + + gemm = PersistentDenseGemmKernel( + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, ) + # Check if configuration can be implemented + can_implement = gemm.can_implement( + mnkl, a.element_type, b.element_type, c.element_type, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + 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}" + ) + + stream = make_fake_stream() + return cute.compile(bmm, gemm, a, b, c, max_active_clusters, stream, epilogue_op) def run( @@ -1808,36 +1583,9 @@ def run( :return: Execution time of the GEMM kernel. :rtype: float """ - print("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'}") - import torch from cutlass.torch import dtype as torch_dtype - # Build GEMM object - gemm_op = PersistentDenseGemmKernel( - acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store - ) - can_implement = gemm_op.can_implement( - mnkl, ab_dtype, c_dtype, a_major, b_major, c_major - ) - if not can_implement: - raise testing.CantImplementError( - 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}" - ) - if not torch.cuda.is_available(): raise RuntimeError("GPU is required to run this example!") @@ -1852,67 +1600,96 @@ def run( ) # Run and verify BMM with torch - a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major) + a_f32, b_f32, c_f32, a_storage, b_storage, c_storage = prepare_tensors( + mnkl, ab_dtype, ab_dtype, c_dtype, a_major, b_major, c_major + ) - # Leading dim is 2 leading_dim_a = 2 if a_major == "k" else 1 leading_dim_b = 1 if b_major == "k" else 2 leading_dim_c = 2 if c_major == "n" else 1 - a_ = from_dlpack(a).mark_layout_dynamic(leading_dim=leading_dim_a) - b_ = from_dlpack(b).mark_layout_dynamic(leading_dim=leading_dim_b) - c_ = from_dlpack(c).mark_layout_dynamic(leading_dim=leading_dim_c) + # Create CuTe tensors, passing float32 source for fp8 conversion + a_ = create_cute_tensor_for_fp8( + a_storage, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_ = create_cute_tensor_for_fp8( + b_storage, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_ = create_cute_tensor_for_fp8( + c_storage, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) - compiled_fn = cute.compile( - bmm, - gemm_op, + compiled_fn = compile_bmm( + mnkl, a_, b_, c_, + acc_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, max_active_clusters, - current_stream, + use_2cta_instrs, + use_tma_store, epilogue_op=lambda x: x, ) + print("Running Blackwell Persistent Dense GEMM test with:") + print(f"mnkl: {mnkl}") + 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'}") + if not skip_ref_check: # Use small random number for deterministic result for reference check compiled_fn(a_, b_, c_, current_stream) # Manually quantize to be comparable + # Use float32 source data for reference calculation ref = ( - torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32)) + torch.bmm(a_f32, b_f32) .to(dtype=torch_dtype(c_dtype)) .to(dtype=torch.float32) ) + # Read back the result from CuTe tensor (c_storage was updated in-place) torch.testing.assert_close( - c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 + c_storage.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 ) if not benchmark: return 0 def generate_tensors(): - init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8] - a, b, c = prepare_tensors( + a_f32, b_f32, c_f32, a_st, b_st, c_st = prepare_tensors( mnkl, ab_dtype, + ab_dtype, c_dtype, a_major, b_major, c_major, - init_random=not init_normal, ) - a_ = from_dlpack(a).mark_layout_dynamic(leading_dim=leading_dim_a) - b_ = from_dlpack(b).mark_layout_dynamic(leading_dim=leading_dim_b) - c_ = from_dlpack(c).mark_layout_dynamic(leading_dim=leading_dim_c) + a_ = create_cute_tensor_for_fp8( + a_st, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_ = create_cute_tensor_for_fp8( + b_st, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_ = create_cute_tensor_for_fp8( + c_st, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) return testing.JitArguments(a_, b_, c_, current_stream) workspace_count = 1 if use_cold_l2: one_workspace_bytes = ( - a.numel() * a.element_size() - + b.numel() * b.element_size() - + c.numel() * c.element_size() + a_storage.numel() * a_storage.element_size() + + b_storage.numel() * b_storage.element_size() + + c_storage.numel() * c_storage.element_size() ) workspace_count = testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations @@ -1929,6 +1706,11 @@ def run( ) +def compute_tflops(time_ns, m, n, k): + return 2.0 * m * n * k / time_ns / 1000.0 + + + def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: try: return tuple(int(x.strip()) for x in s.split(",")) @@ -1939,7 +1721,6 @@ def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: def prepare_parser(): - parser = argparse.ArgumentParser( description="Example of Dense Persistent GEMM on Blackwell." ) @@ -1974,7 +1755,14 @@ def prepare_parser(): "--tolerance", type=float, default=1e-01, help="Tolerance for validation" ) parser.add_argument( - "--benchmark", action="store_true", help="Only benchmark the kernel" + "--benchmark", + type=str, + default="default", + choices=[ + "default", + "none", + ], + help="Benchmark the kernel with nsight or default (cute.testing.benchmark) or none", ) parser.add_argument( "--warmup_iterations", type=int, default=0, help="Warmup iterations" @@ -2018,6 +1806,20 @@ if __name__ == "__main__": if len(args.cluster_shape_mn) != 2: parser.error("--cluster_shape_mn must contain exactly 2 values") + print(f"[DSL INFO] Compiling Blackwell Persistent Dense GEMM with:") + print( + f"[DSL INFO] A dtype: {args.ab_dtype}, B dtype: {args.c_dtype}, C dtype: {args.acc_dtype}, Acc dtype: {args.acc_dtype}" + ) + print( + f"[DSL INFO] Matrix majors - A: {args.a_major}, B: {args.b_major}, C: {args.c_major}" + ) + print(f"[DSL INFO] Mma Tiler (M, N): {args.mma_tiler_mn}") + print(f"[DSL INFO] Cluster Shape (M, N): {args.cluster_shape_mn}") + print( + f"[DSL INFO] 2CTA MMA instructions: {'True' if args.use_2cta_instrs else 'False'}" + ) + print(f"[DSL INFO] Use TMA Store: {'True' if args.use_tma_store else 'False'}") + run( args.mnkl, args.ab_dtype, @@ -2035,6 +1837,6 @@ if __name__ == "__main__": args.iterations, args.skip_ref_check, args.use_cold_l2, - args.benchmark, + args.benchmark == "default", ) print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py index d68d9fc9..7a259295 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py @@ -34,8 +34,8 @@ import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -from cutlass.cute.runtime import from_dlpack import cutlass.utils as utils +from cutlass.utils import is_fp8_dtype, create_cute_tensor_for_fp8 import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 @@ -401,6 +401,7 @@ class PersistentDenseGemmKernel: # Setup clc stage by default self.num_clc_stage = 1 + assert self.num_clc_stage == 1, "Only single-stage CLC pipeline is supported" # Compute A/B/C shared memory layout self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( @@ -615,8 +616,8 @@ class PersistentDenseGemmKernel: ] tmem_dealloc_mbar_ptr: cutlass.Int64 tmem_holding_buf: cutlass.Int32 - clc_ptr: cute.struct.MemRange[cutlass.Int64, self.num_clc_stage * 2] - clc_response_ptr: cute.struct.MemRange[cutlass.Int32, 1] + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + clc_response: cute.struct.MemRange[cutlass.Int32, 4] smem = utils.SmemAllocator() storage = smem.allocate(SharedStorage) @@ -664,7 +665,7 @@ class PersistentDenseGemmKernel: pipeline.Agent.Thread, num_clc_consumer_threads ) clc_pipeline = pipeline.PipelineClcFetchAsync.create( - barrier_storage=storage.clc_ptr.data_ptr(), + barrier_storage=storage.clc_mbar_ptr.data_ptr(), num_stages=self.num_clc_stage, producer_group=clc_pipeline_producer_group, consumer_group=clc_pipeline_consumer_group, @@ -696,7 +697,7 @@ class PersistentDenseGemmKernel: pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) # Initial clc response pointer - clc_response_ptr = storage.clc_response_ptr.data_ptr() + clc_response_ptr = storage.clc_response.data_ptr() clc_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_clc_stage @@ -1046,45 +1047,70 @@ class PersistentDenseGemmKernel: # (MMA, MMA_M, MMA_N, STAGE) tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - # - # Persistent tile scheduling loop for epilogue - # + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) if cutlass.const_expr(self.use_tma_store): assert tma_atom_c is not None and sC is not None - utils.gemm.sm100.epilogue_tma_store( - self, - tidx, - warp_idx, - acc_pipeline, - tiled_mma, - tma_atom_c, - tCtAcc_base, - sC, - tCgC, - epi_tile, - tile_sched, - epilogue_op, - clc_pipeline, - clc_consumer_state, + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), ) - else: - utils.gemm.sm100.epilogue( - self, - tidx, - acc_pipeline, - tiled_mma, - tCtAcc_base, - tCgC, - epi_tile, - tile_sched, - epilogue_op, - tmem_dealloc_barrier, - None, - None, - clc_pipeline, - clc_consumer_state, + 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], + ) + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + # + # Advance to next tile + # + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete + c_pipeline.producer_tail() + else: + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() # # Dealloc the tensor memory buffer # @@ -1150,8 +1176,11 @@ class PersistentDenseGemmKernel: return num_tmem_alloc_cols def check_supported_dtypes( - self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] - ) -> bool: + self, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + ): """ Check if the dtypes are valid @@ -1173,8 +1202,10 @@ class PersistentDenseGemmKernel: cutlass.Float8E4M3FN, cutlass.Float8E5M2, } - if ab_dtype not in valid_ab_dtypes: - raise testing.CantImplementError(f"Unsupported AB dtype: {ab_dtype}") + if a_dtype not in valid_ab_dtypes or b_dtype not in valid_ab_dtypes: + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype}" + ) if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: raise testing.CantImplementError( @@ -1198,8 +1229,13 @@ class PersistentDenseGemmKernel: 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 + if ( + a_dtype not in acc_ab_compatibility[self.acc_dtype] + or b_dtype not in acc_ab_compatibility[self.acc_dtype] + ): + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype} for accumulator dtype: {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and C type acc_c_compatibility = { @@ -1228,11 +1264,11 @@ class PersistentDenseGemmKernel: } # Check compatibility between accumulator type and C type if c_dtype not in acc_c_compatibility[self.acc_dtype]: - return False + raise testing.CantImplementError( + f"Unsupported C dtype: {c_dtype} for accumulator dtype: {self.acc_dtype}" + ) - return True - - def check_mma_tiler_and_cluster_shape(self) -> bool: + def check_mma_tiler_and_cluster_shape(self): """Check if the mma tiler and cluster shape are valid. :raises testing.CantImplementError: If the mma tiler and cluster shape are invalid @@ -1273,12 +1309,13 @@ class PersistentDenseGemmKernel: n: int, k: int, l: int, - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_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 @@ -1290,8 +1327,10 @@ class PersistentDenseGemmKernel: :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 a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_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 @@ -1301,8 +1340,7 @@ class PersistentDenseGemmKernel: :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 + :raises testing.CantImplementError: If the tensor alignment is invalid """ # TODO: move to utils @@ -1313,15 +1351,15 @@ class PersistentDenseGemmKernel: 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)) + not check_contiguous_16B_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(b_dtype, b_major == "n", (n, k, l)) or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) ): raise testing.CantImplementError( - f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {ab_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}" + f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {a_dtype}, {b_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}" ) - def check_epilog_store_option(self, m: int, n: int) -> bool: + def check_epilog_store_option(self, m: int, n: int): """ Check if the epilogue store option is valid @@ -1346,7 +1384,8 @@ class PersistentDenseGemmKernel: def can_implement( self, mnkl: Tuple[int, int, int, int], - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, @@ -1357,8 +1396,10 @@ class PersistentDenseGemmKernel: :param mnkl: Problem size as a tuple (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 a_dtype: Data type for input tensors A. + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Data type for input tensors B. + :type b_dtype: Type[cutlass.Numeric] :param c_dtype: Data type for output tensor C. :type c_dtype: Type[cutlass.Numeric] :param a_major: Major dimension of the A tensor layout ("m" or "k"). @@ -1373,14 +1414,14 @@ class PersistentDenseGemmKernel: try: # Skip unsupported types - self.check_supported_dtypes(ab_dtype, c_dtype) + self.check_supported_dtypes(a_dtype, b_dtype, c_dtype) # Skip invalid mma tile shape and cluster shape self.check_mma_tiler_and_cluster_shape() m, n, k, l = mnkl self.check_tensor_alignment( - m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major ) self.check_epilog_store_option(m, n) except testing.CantImplementError: @@ -1432,43 +1473,72 @@ def bmm( @lru_cache(maxsize=1) def prepare_tensors( mnkl: Tuple[int, int, int, int], - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, c_major: str, init_random: bool = True, + normal_mean: float = 0.0, + normal_std: float = 1.0, ): + """Prepare tensors for GEMM. + + Returns: + Tuple of (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage): + - *_f32: Float32 tensors with the logical data (for reference and fp8 conversion) + - *_storage: Storage tensors for DLPack (uint8 for fp8, otherwise the target dtype) + """ import torch from cutlass.torch import dtype as torch_dtype m, n, k, l = mnkl if a_major == "k": - a = torch.empty((l, m, k), dtype=torch.float32, device="cuda") + a_f32 = torch.empty((l, m, k), dtype=torch.float32, device="cuda") elif a_major == "m": - a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + a_f32 = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if b_major == "n": - b = torch.empty((l, k, n), dtype=torch.float32, device="cuda") + b_f32 = torch.empty((l, k, n), dtype=torch.float32, device="cuda") elif b_major == "k": - b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1) + b_f32 = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if c_major == "n": - c = torch.empty((l, m, n), dtype=torch.float32, device="cuda") + c_f32 = torch.empty((l, m, n), dtype=torch.float32, device="cuda") elif c_major == "m": - c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + c_f32 = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if init_random: - a.random_(-2, 3) - b.random_(-2, 3) - c.random_(-2, 3) + # Uniform random initialization in range [-2, 3) + a_f32.random_(-2, 3) + b_f32.random_(-2, 3) + c_f32.random_(-2, 3) + else: + # Normal (Gaussian) initialization with user-specified mean and std + a_f32.normal_(mean=normal_mean, std=normal_std) + b_f32.normal_(mean=normal_mean, std=normal_std) + c_f32.normal_(mean=normal_mean, std=normal_std) - return ( - a.to(dtype=torch_dtype(ab_dtype)), - b.to(dtype=torch_dtype(ab_dtype)), - c.to(dtype=torch_dtype(c_dtype)), - ) + # For float8 types, use uint8 as storage type to avoid dlpack limitation + # (dlpack doesn't support float8 types) + # For other types, convert to the target dtype + a_storage_dtype = torch.uint8 if is_fp8_dtype(a_dtype) else torch_dtype(a_dtype) + b_storage_dtype = torch.uint8 if is_fp8_dtype(b_dtype) else torch_dtype(b_dtype) + c_storage_dtype = torch.uint8 if is_fp8_dtype(c_dtype) else torch_dtype(c_dtype) + + a_storage = a_f32.to(dtype=a_storage_dtype) + b_storage = b_f32.to(dtype=b_storage_dtype) + c_storage = c_f32.to(dtype=c_storage_dtype) + + return (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage) @lru_cache(maxsize=1) @@ -1499,7 +1569,7 @@ def compile_bmm( ) # Check if configuration can be implemented can_implement = gemm.can_implement( - mnkl, a.element_type, c.element_type, a_major, b_major, c_major + mnkl, a.element_type, b.element_type, c.element_type, a_major, b_major, c_major ) if not can_implement: raise testing.CantImplementError( @@ -1594,21 +1664,30 @@ def run( ) # Run and verify BMM with torch - a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major) + a_f32, b_f32, c_f32, a_storage, b_storage, c_storage = prepare_tensors( + mnkl, ab_dtype, ab_dtype, c_dtype, a_major, b_major, c_major + ) leading_dim_a = 2 if a_major == "k" else 1 leading_dim_b = 1 if b_major == "k" else 2 leading_dim_c = 2 if c_major == "n" else 1 - a_ = from_dlpack( - a, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32 - ).mark_layout_dynamic(leading_dim=leading_dim_a) - b_ = from_dlpack( - b, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32 - ).mark_layout_dynamic(leading_dim=leading_dim_b) - c_ = from_dlpack( - c, assumed_align=16, force_tf32=c_dtype == cutlass.TFloat32 - ).mark_layout_dynamic(leading_dim=leading_dim_c) + # Create CuTe tensors, passing float32 source for fp8 conversion + a_ = create_cute_tensor_for_fp8( + a_storage, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_ = create_cute_tensor_for_fp8( + b_storage, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_ = create_cute_tensor_for_fp8( + c_storage, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) + + print("Compile Blackwell Persistent Dense GEMM with:") + print(f"ab_dtype: {ab_dtype}, c_dtype: {c_dtype}, acc_dtype: {acc_dtype}") + print(f"a_major: {a_major}, b_major: {b_major}, c_major: {c_major}") + print(f"mma_tiler_mn: {mma_tiler_mn}, cluster_shape_mn: {cluster_shape_mn}") + print(f"use_2cta_instrs: {use_2cta_instrs}, use_tma_store: {use_tma_store}") compiled_fn = compile_bmm( mnkl, @@ -1640,46 +1719,47 @@ def run( compiled_fn(a_, b_, c_, current_stream) # Manually quantize to be comparable + # Use float32 source data for reference calculation ref = ( - torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32)) + torch.bmm(a_f32, b_f32) .to(dtype=torch_dtype(c_dtype)) .to(dtype=torch.float32) ) + # Read back the result from CuTe tensor (c_storage was updated in-place) torch.testing.assert_close( - c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 + c_storage.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 ) if not benchmark: return 0 def generate_tensors(): - init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8] - a, b, c = prepare_tensors( + a_f32, b_f32, c_f32, a_st, b_st, c_st = prepare_tensors( mnkl, ab_dtype, + ab_dtype, c_dtype, a_major, b_major, c_major, - init_random=not init_normal, ) - a_ = from_dlpack( - a, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32 - ).mark_layout_dynamic(leading_dim=leading_dim_a) - b_ = from_dlpack( - b, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32 - ).mark_layout_dynamic(leading_dim=leading_dim_b) - c_ = from_dlpack( - c, assumed_align=16, force_tf32=c_dtype == cutlass.TFloat32 - ).mark_layout_dynamic(leading_dim=leading_dim_c) + a_ = create_cute_tensor_for_fp8( + a_st, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_ = create_cute_tensor_for_fp8( + b_st, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_ = create_cute_tensor_for_fp8( + c_st, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) return testing.JitArguments(a_, b_, c_, current_stream) workspace_count = 1 if use_cold_l2: one_workspace_bytes = ( - a.numel() * a.element_size() - + b.numel() * b.element_size() - + c.numel() * c.element_size() + a_storage.numel() * a_storage.element_size() + + b_storage.numel() * b_storage.element_size() + + c_storage.numel() * c_storage.element_size() ) workspace_count = testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations @@ -1735,6 +1815,7 @@ def prepare_parser(): 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") diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py b/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py index 3952cd46..3622b211 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py @@ -30,15 +30,12 @@ import argparse from typing import Optional, Type, Tuple, Union import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.testing as testing @@ -767,6 +764,14 @@ class DenseGemmKernel: acc_pipeline.consumer_wait(acc_consumer_state) if cutlass.const_expr(self.use_tma_store): + # TODO: segment fault if I put here + # sC = smem.allocate_tensor( + # element_type=self.c_dtype, + # layout=c_smem_layout_staged.outer, + # byte_alignment=128, + # swizzle=c_smem_layout_staged.inner, + # ) + assert tma_atom_c is not None and sC is not None self.epilogue_tma_store( tidx, @@ -983,7 +988,10 @@ class DenseGemmKernel: c_buffer = 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("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) pipeline.sync(barrier_id=1) # TMA store C to global memory @@ -1214,7 +1222,7 @@ class DenseGemmKernel: """ acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) - return sm100_utils.get_num_tmem_alloc_cols(tCtAcc_fake) + return utils.get_num_tmem_alloc_cols(tCtAcc_fake) def is_valid_dtypes( self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] @@ -1365,6 +1373,7 @@ class DenseGemmKernel: """ 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] @@ -1446,6 +1455,9 @@ class DenseGemmKernel: def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) @@ -1474,6 +1486,9 @@ def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): def compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance): + import torch + import cutlass.torch as cutlass_torch + # Copy gpu result back kernel_result = c_torch_gpu.cpu() @@ -1563,6 +1578,7 @@ def run( print(f"Iterations: {iterations}") print(f"Skip reference checking: {skip_ref_check}") print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + import torch # Unpack parameters m, n, k, l = mnkl @@ -1600,6 +1616,8 @@ def run( compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance) def generate_tensors(): + import cutlass.torch as cutlass_torch + a_tensor, _ = cutlass_torch.cute_tensor_like( a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) diff --git a/examples/python/CuTeDSL/blackwell/epilogue/activation_custom_epilogue_dense_gemm.py b/examples/python/CuTeDSL/blackwell/epilogue/activation_custom_epilogue_dense_gemm.py new file mode 100644 index 00000000..5981ec7a --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/activation_custom_epilogue_dense_gemm.py @@ -0,0 +1,754 @@ +# 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 +from common_efc import ACTIVATION_FUNCTIONS + +""" +A high-performance persistent batched dense GEMM with activation functions in custom epilogue fusion +for the NVIDIA Blackwell SM100 architecture using CUTE DSL and Epilogue Fusion Configuration (EFC). + +This example demonstrates GEMMs with custom fused epilogues inspired by Ada FP8 GEMM epilogue +from https://github.com/NVIDIA/cutlass/blob/main/examples/58_ada_fp8_gemm/ada_fp8_gemm.cu : + Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias + D = activation(Aux) + +The scale factors (scale_a, scale_b, scale_c) default to 1.0 but can be customized via CLI: + +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, "source"), C can be row-major("N") or column-major("M") +- Matrix Aux is MxNxL (auxiliary output, pre-activation), same layout as C/D +- Matrix D is MxNxL (final output, post-activation), same layout as C +- alpha, beta are scalar scale factors +- scale_a, scale_b, scale_c are scalar scale factors for A, B, and C matrices +- bias is a scalar bias term + +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 with activation functions + +Supported activation functions: + - identity: f(x) = x + - relu: f(x) = max(0, x) + - leaky_relu: f(x) = max(0, x) + negative_slope * min(0, x) + - tanh: f(x) = tanh(x) + - sigmoid: f(x) = 1 / (1 + exp(-x)) + - silu: f(x) = x * sigmoid(x) + - hardswish: f(x) = x * relu6(x + 3) / 6 + - gelu: f(x) = 0.5 * x * (1 + erf(x / sqrt(2))) + +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 (source) matrix from global memory (GMEM) to shared memory (SMEM) using TMA, then to registers (RMEM). + - Compute Aux = (alpha * scale_a * scale_b) * accumulator + (beta * scale_c) * C + bias + - Compute D = activation(Aux) + - Type convert Aux and D matrices to output types. + - Store Aux and D 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 activation_custom_epilogue_dense_gemm.py \ + --activation relu \ + --ab_dtype Float16 --c_dtype Float16 --aux_dtype Float16 --d_dtype Float16 \ + --acc_dtype Float32 --epi_dtype Float32 \ + --mma_tiler_mn 128,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 1.0 --beta 1.0 --bias 0.0 \ + --scale_a 1.0 --scale_b 1.0 --scale_c 1.0 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python activation_custom_epilogue_dense_gemm.py \ + --activation gelu \ + --ab_dtype Float16 --c_dtype Float16 --aux_dtype Float16 --d_dtype Float16 \ + --acc_dtype Float32 --epi_dtype Float32 \ + --mma_tiler_mn 128,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 1.0 --beta 1.0 --bias 0.0 \ + --scale_a 1.0 --scale_b 1.0 --scale_c 1.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/Aux 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 DenseGemmActivation(DenseGemmEFC): + """Implements batched GEMM with activation function in epilogue using EFC. + + This class extends DenseGemmEFC to provide a fused epilogue inspired by + Ada FP8 GEMM that: + - Reads from input tensor C (source) + - Writes to output tensors Aux (auxiliary, pre-activation) and D (final, post-activation) + - Performs: Aux = alpha * accumulator + beta * C + bias + D = activation(Aux) + + The class provides CLI argument parsing and tensor creation for the + specific epilogue configuration with C, Aux, D tensors and alpha, + beta, bias scalar parameters. + """ + + def __init__( + self, + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue_fn, + activation_name, + ): + """Initialize the GEMM with activation epilogue. + + :param acc_dtype: Accumulator data type + :param epi_dtype: Epilogue computation data type + :param use_2cta_instrs: Whether to use 2-CTA MMA instructions + :param mma_tiler_mn: MMA tile shape (M, N) + :param cluster_shape_mn: Cluster shape (M, N) + :param epilogue_fn: Epilogue function to use + :param activation_name: Name of the activation function + """ + super().__init__( + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue_fn, + ) + self.activation_name = activation_name + + class CLIParser(DenseGemmEFC.CLIParser): + @override + def more_parsing(self): + self.parser.add_argument( + "--activation", + type=str, + default="relu", + choices=ACTIVATION_FUNCTIONS, + help="Activation function to use in epilogue", + ) + self.parser.add_argument( + "--alpha", + type=float, + default=1.0, + help="alpha scale factor for accumulator", + ) + self.parser.add_argument( + "--beta", type=float, default=1.0, help="beta scale factor for source" + ) + self.parser.add_argument( + "--bias", type=float, default=0.0, help="bias term to add" + ) + self.parser.add_argument( + "--scale_a", type=float, default=1.0, help="scale factor for matrix A" + ) + self.parser.add_argument( + "--scale_b", type=float, default=1.0, help="scale factor for matrix B" + ) + self.parser.add_argument( + "--scale_c", type=float, default=1.0, help="scale factor for source C" + ) + self.parser.add_argument( + "--c_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="C tensor dtype", + ) + self.parser.add_argument( + "--aux_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="Aux tensor dtype", + ) + self.parser.add_argument( + "--d_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="D tensor dtype", + ) + self.parser.add_argument( + "--leaky_relu_alpha", + type=float, + default=0.01, + help="negative slope for leaky_relu", + ) + + @override + def create_arguments( + self, + l, + m, + n, + k, + a_major, + b_major, + cd_major, + ab_dtype, + # For the supplemental tensors. + c_dtype, + aux_dtype, + d_dtype, + ): + """Create arguments for GEMM operations with epilogue tensors. + + Creates tensors for A, B (from parent class) and epilogue-specific + tensors C, Aux, D 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, + aux_tensor, aux_torch_cpu, aux_torch_gpu, + d_tensor, d_torch_cpu, d_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 + ) + + # Create C tensor (source for epilogue) + 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 + ) + + # Create Aux tensor (auxiliary/pre-activation output) + aux_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", aux_dtype) + aux_tensor, aux_torch_gpu = cutlass_torch.cute_tensor_like( + aux_torch_cpu, aux_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Create D tensor (final/post-activation output) + 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 + ) + + return ( + *std_args, + c_tensor, + c_torch_cpu, + c_torch_gpu, + aux_tensor, + aux_torch_cpu, + aux_torch_gpu, + d_tensor, + d_torch_cpu, + d_torch_gpu, + ) + + def compare( + self, + a_torch_cpu, + b_torch_cpu, + epi_dtype, + tolerance, + # For the tensor check. + c_torch_gpu, + aux_torch_gpu, + d_torch_gpu, + # The EFC epilogue arguments. + c_torch_cpu, + aux_torch_cpu, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_torch_cpu, + leaky_relu_alpha=0.01, + ): + """Compare GPU results against CPU reference implementation. + + :param a_torch_cpu: Input tensor A on CPU + :param b_torch_cpu: Input tensor B on CPU + :param epi_dtype: Epilogue data type + :param tolerance: Comparison tolerance + :param c_torch_gpu: GPU result for C + :param aux_torch_gpu: GPU result for Aux + :param d_torch_gpu: GPU result for D + :param c_torch_cpu: CPU reference for C + :param aux_torch_cpu: CPU reference for Aux + :param alpha: Alpha scale factor + :param beta: Beta scale factor + :param bias: Bias term + :param scale_a: Scale factor for matrix A + :param scale_b: Scale factor for matrix B + :param scale_c: Scale factor for source C + :param d_torch_cpu: CPU reference for D + :param leaky_relu_alpha: Negative slope for leaky_relu + """ + # Compute reference result + self.evaluate_on_cpu( + a_torch_cpu, + b_torch_cpu, + epi_dtype, + c_torch_cpu, + aux_torch_cpu, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_torch_cpu, + leaky_relu_alpha, + ) + # Assert close results for output tensors + torch.testing.assert_close( + aux_torch_gpu.cpu(), aux_torch_cpu, atol=tolerance, rtol=1e-03 + ) + torch.testing.assert_close( + d_torch_gpu.cpu(), d_torch_cpu, atol=tolerance, rtol=1e-03 + ) + # Assert that the read tensor has not been changed + torch.testing.assert_close( + c_torch_gpu.cpu(), c_torch_cpu, atol=tolerance, rtol=1e-03 + ) + + @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], + aux_dtype: typing.Type[cutlass.Numeric], + d_dtype: typing.Type[cutlass.Numeric], + alpha: float, + beta: float, + bias: float, + scale_a: float, + scale_b: float, + scale_c: float, + activation: str, + leaky_relu_alpha: float, + tolerance: float, + ) -> str: + """Format test parameters as CLI arguments for activation_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, aux_dtype, d_dtype, + alpha, beta, bias, activation). + + :return: Formatted CLI command string + """ + # Get base command from parent class + base_cmd = DenseGemmEFC.format_as_cli_args( + "activation_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 arguments + epilogue_args = ( + f" --activation {activation}" + f" --c_dtype {c_dtype.__name__}" + f" --aux_dtype {aux_dtype.__name__}" + f" --d_dtype {d_dtype.__name__}" + f" --alpha {alpha}" + f" --beta {beta}" + f" --bias {bias}" + f" --scale_a {scale_a}" + f" --scale_b {scale_b}" + f" --scale_c {scale_c}" + f" --leaky_relu_alpha {leaky_relu_alpha}" + ) + + return base_cmd + epilogue_args + + +def create_epilogue_function(activation_name: str): + """Create an epilogue function with the specified activation. + + :param activation_name: Name of the activation function to use + :return: Epilogue function + """ + # Validate activation name + if activation_name not in ACTIVATION_FUNCTIONS: + raise ValueError(f"Unsupported activation: {activation_name}") + + def epilogue( + efc_config, + C, + Aux, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + D, + leaky_relu_alpha, + ): + # Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias + # Following Ada FP8 GEMM epilogue pattern + aux_val = ( + (alpha * scale_a * scale_b) * efc_config.accum() + + (beta * scale_c) * C.load() + + bias + ) + Aux.store(aux_val) + # D = activation(Aux) + activation_fn = getattr(efc_config, activation_name) + # leaky_relu needs an extra parameter, others don't + if activation_name == "leaky_relu": + D.store(activation_fn(aux_val, leaky_relu_alpha)) + else: + D.store(activation_fn(aux_val)) + + return epilogue + + +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, + c_dtype: typing.Type[cutlass.Numeric], + aux_dtype: typing.Type[cutlass.Numeric], + d_dtype: typing.Type[cutlass.Numeric], + alpha: float, + beta: float, + bias: float, + scale_a: float, + scale_b: float, + scale_c: float, + activation: str, + leaky_relu_alpha: float, + mma_tiler_mn: typing.Tuple[int, int], + cluster_shape_mn: typing.Tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + warmup_iterations: int = 3, + iterations: int = 100, + skip_ref_check: bool = False, +): + """Run GEMM with activation function in epilogue. + + :param mnkl: Tuple of (M, N, K, L) dimensions + :param ab_dtype: Data type for A and B tensors + :param acc_dtype: Accumulator data type + :param epi_dtype: Epilogue computation data type + :param a_major: Major dimension for A ("m" or "k") + :param b_major: Major dimension for B ("n" or "k") + :param cd_major: Major dimension for C/D/Aux ("m" or "n") + :param c_dtype: Data type for C tensor + :param aux_dtype: Data type for Aux tensor + :param d_dtype: Data type for D tensor + :param alpha: Alpha scale factor + :param beta: Beta scale factor + :param bias: Bias term + :param scale_a: Scale factor for matrix A + :param scale_b: Scale factor for matrix B + :param scale_c: Scale factor for source C + :param activation: Activation function name + :param leaky_relu_alpha: Negative slope for leaky_relu + :param mma_tiler_mn: MMA tile shape (M, N) + :param cluster_shape_mn: Cluster shape (M, N) + :param use_2cta_instrs: Whether to use 2-CTA MMA instructions + :param tolerance: Comparison tolerance + :param warmup_iterations: Number of warmup iterations + :param iterations: Number of benchmark iterations + :param skip_ref_check: Whether to skip reference check + """ + 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}, {aux_dtype = !s}, {d_dtype = !s}") + print(f"\t{alpha = }, {beta = }, {bias = }") + print(f"\t{scale_a = }, {scale_b = }, {scale_c = }") + print(f"\t{activation = !s}") + if activation == "leaky_relu": + print(f"\t{leaky_relu_alpha = }") + + # 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) + + # Create the epilogue function with the specified activation + epilogue_fn = create_epilogue_function(activation) + + # Build GEMM object with EFC configuration + gemm = DenseGemmActivation( + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue_fn, + activation, + ) + + ( + a_tensor, + b_tensor, + a_torch_cpu, + b_torch_cpu, + # The supplemental tensors. + c_tensor, + c_torch_cpu, + c_torch_gpu, + aux_tensor, + aux_torch_cpu, + aux_torch_gpu, + d_tensor, + d_torch_cpu, + d_torch_gpu, + ) = gemm.create_arguments( + l, + m, + n, + k, + a_major, + b_major, + cd_major, + ab_dtype, + # For the supplemental tensors. + c_dtype, + aux_dtype, + d_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, + aux_tensor, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_tensor, + leaky_relu_alpha, + ) + + 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, + aux_tensor, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_tensor, + leaky_relu_alpha, + ) + + # 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, + aux_tensor, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_tensor, + leaky_relu_alpha, + ), + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print(f"Execution time: {exec_time} us") + + # Compute reference result + if not skip_ref_check: + print("Checking results against CPU reference...") + gemm.compare( + # The usual arguments. + a_torch_cpu, + b_torch_cpu, + epi_dtype, + tolerance, + # For the tensor check. + c_torch_gpu, + aux_torch_gpu, + d_torch_gpu, + # The EFC epilogue arguments. + c_torch_cpu, + aux_torch_cpu, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_torch_cpu, + leaky_relu_alpha, + ) + print("Results match CPU reference!") + + +if __name__ == "__main__": + args = DenseGemmActivation.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.aux_dtype, + args.d_dtype, + args.alpha, + args.beta, + args.bias, + args.scale_a, + args.scale_b, + args.scale_c, + args.activation, + args.leaky_relu_alpha, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + ) + print("\n" + "=" * 80) + print( + f"PASS - {args.activation.upper()} activation test completed successfully!" + ) + print("=" * 80 + "\n") + except Exception as exc: + traceback.print_exception(exc) + raise diff --git a/examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py b/examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py new file mode 100644 index 00000000..0a4fa2d7 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py @@ -0,0 +1,2011 @@ +# 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 argparse +import logging +import types +import typing +from typing import Tuple, Type, Union + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.nvgpu import cpasync, tcgen05 + +import common_efc +from common_efc import log + +""" +Common base infrastructure for 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 module provides the DenseGemmEFC base class that implements the core GEMM functionality with +support for custom epilogue operations. Subclasses define specific epilogue configurations by +providing an epilogue function that operates on the accumulator and supplemental tensors. + +Key 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 + +GEMM Execution Flow: +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 (customizable via EFC): + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Load supplemental input tensors from GMEM to SMEM using TMA, then to RMEM. + - Execute custom epilogue function (defined by subclass) that: + * Accesses the accumulator via efc_config.accum() + * Reads from supplemental input tensors via tensor.load() + * Writes to supplemental output tensors via tensor.store() + * Can apply arbitrary element-wise operations, scaling, and fusion + - Store result tensors from RMEM to SMEM to 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. + +Base 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") +- Supplemental tensors are MxNxL with layout matching the epilogue configuration + +Common Constraints: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2) +* A/B tensors 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 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 + +Subclass Examples: +- custom_epilogue_dense_gemm.py: Custom fused epilogue with multiple read/write tensors +- synthetic_custom_epilogue_dense_gemm.py: Synthetic epilogue for testing with configurable tensor counts +""" + + +class DenseGemmEFC: + """Base class for batched GEMM with custom epilogue fusion using EFC. + + This class provides the core infrastructure for persistent batched GEMM operations + with customizable epilogue fusion. Subclasses define specific epilogue behaviors + by providing an epilogue configuration function that describes operations on the + accumulator and supplemental tensors. + + The class handles: + - GEMM mainloop (A * B computation) + - TMA-based memory operations + - Warp specialization + - Persistent tile scheduling + - EFC (Epilogue Fusion Configuration) integration + - CLI argument parsing (extensible via CLIParser.more_parsing()) + - Tensor creation and validation + + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param epi_dtype: Data type for epilogue operation + :type epi_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 epilogue_function_configuration: Function defining the epilogue behavior via EFC + :type epilogue_function_configuration: Callable + + :note: Supported A/B data types: + - TFloat32 + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + (A and B must have the same data type) + + :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 supplemental tensor data types (epilogue-dependent): + - 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: + >>> def my_epilogue(efc_config, alpha, beta, output_tensor, input_tensor): + ... result = efc_config.accum() * alpha + input_tensor.load() * beta + ... output_tensor.store(result) + ... + >>> gemm = DenseGemmEFC( + ... acc_dtype=cutlass.Float32, + ... epi_dtype=cutlass.Float32, + ... use_2cta_instrs=True, + ... mma_tiler_mn=(128, 128), + ... cluster_shape_mn=(2, 2), + ... epilogue_function_configuration=my_epilogue + ... ) + """ + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + epi_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + epilogue_function_configuration: typing.Callable, + ): + """Initializes the configuration for a Blackwell dense GEMM kernel with EFC. + + 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. Epilogue Configuration: + - epilogue_function_configuration: Defines custom epilogue behavior + that operates on accumulator and supplemental tensors. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param epi_dtype: Data type of the epilogue. + :type epi_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + :param epilogue_function_configuration: Function defining epilogue behavior via EFC. + :type epilogue_function_configuration: Callable + """ + + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.epi_dtype: Type[cutlass.Numeric] = epi_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.arch = "sm_100" + + self.c_dtype = self.epi_dtype + + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids: + + # The warps responsible for computing the epilogue function and storing + # the results. + self.epilogue_warp_id = (0, 1, 2, 3) + # The warp responsible for computing the matrix multiplication. + self.mma_warp_id = 4 + # The warp responsible for loading the tensors A & B to feed the MMA. + self.tma_warp_id = 5 + # The warp responsible for loading the auxiliary tensors used in the epilogue. + self.epilogue_load_warp_id = 6 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + *self.epilogue_warp_id, + self.epilogue_load_warp_id, + ) + ) + # Barrier ids for cta sync, epilogue sync and tmem ptr sync. + self.cta_sync_bar_id = 1 + self.epilogue_sync_bar_id = 2 + self.tmem_alloc_sync_bar_id = 3 + # Amount of available shared memory. + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + + # Setup the EFC from the given function representing the epilogue + # configuration. + self.efc = common_efc.EFC(self, epilogue_function_configuration) + + def _create_tiled_mma(self): + """Make a tiled MMA atom with given data type, leading dimension, CTA + group and MMA tile shape. Use SMEM operand source for A.""" + return utils.sm100.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], + ) + + 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/D stage counts in shared memory + - Computing A/B/C/D shared memory layout + - Computing tensor memory allocation columns + """ + # Get the right tiled MMA. + self._tiled_mma = self._create_tiled_mma() + log(f"{self._tiled_mma = !s}") + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(self._tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + # Extend mma_tiler with k-dimension (MMA_M, MMA_N, MMA_K) + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + log(f"{self.mma_tiler = !s}") + # CTA tiler with the 2CTA instruction correction. + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(self._tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + log(f"{self.cta_tile_shape_mnk = !s}") + # Compute cluster layout, V for the 2CTA instructions. + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (self._tiled_mma.thr_id.shape,), + ) + log(f"{cute.make_layout((*self.cluster_shape_mn, 1)) = !s}") + log(f"{self.cluster_layout_vmnk = !s}") + # 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 + log(f"{self.num_mcast_ctas_a = }, {self.num_mcast_ctas_b = }") + log(f"{self.is_a_mcast = }, {self.is_b_mcast = }") + + # Compute epilogue (EPI_TILE_M, EPI_TILE_N) subtile of cta_tile_shape_mnk + # according to some heuristics. + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + layout_d=self.d_layout, + elem_ty_d=self.d_dtype, + layout_c=self.c_layout, + elem_ty_c=self.c_dtype, + ) + log(f"{self.epi_tile = !s}") + + # Setup A/B/C/D pipeline stage count in shared memory and ACC stage + # count in tensor memory. + self.compute_stages() + log(f"{self.num_acc_stage = }, {self.num_ab_stage = }, {self.num_c_stage = }") + # Compute A/B shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + self._tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + log(f"{self.a_smem_layout_staged = !s}") + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + self._tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + log(f"{self.b_smem_layout_staged = !s}") + # Get the smem_layout for the tensors used in the EFC. + self.efc.jit.smem_layout() + + # Compute the number of tensor memory allocation columns + self.compute_num_tmem_alloc_cols() + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + supplemental_parameters: Tuple, + ): + """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 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 supplemental_parameters: Tuple or None used to pass variadic values + :type supplemental_parameters: Tuple + :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. + """ + # Process the variadic parameters. + self.efc.jit.unpack_parameters(supplemental_parameters) + + # 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.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + + # Gather all the auxiliary tensor element data types. + self.efc.jit.record_tensor_dtypes() + + # There is no D tensor to be used as a returned tensor. In the + # following, D is used more like of a "store" concept. So use the + # written tensor with the biggest element_type to set up all the tiling + # heuristics and epilogue store pipeline. + self.d_name_bigger = self.efc.jit.written_tensor_name_with_bigger_element_type() + d = self.efc.jit.parameter[self.d_name_bigger] + self.d_dtype: Type[cutlass.Numeric] = d.element_type + self.d_layout = utils.LayoutEnum.from_tensor(d) + log(f"d{self.d_name_bigger} = {d!s}") + + # C is the read tensor with the biggest element_type, if any, used by + # some heuristics for tiling. + self.c_dtype = None + self.c_layout = None + self.c_name_bigger = self.efc.jit.read_tensor_name_with_bigger_element_type() + if cutlass.const_expr(self.c_name_bigger): + c = self.efc.jit.parameter[self.c_name_bigger] + log(f"{self.c_name_bigger = } -> {c = !s}") + self.c_dtype = c.element_type + 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 depend on gemm inputs + self._setup_attributes() + + atom_thr_size = cute.size(self._tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, self._tiled_mma.thr_id + ) + log(f"{a_op = !s}") + # Get read of the pipeline dimension. + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + log(f"{a_smem_layout = !s}") + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + self._tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + log(f"{tma_atom_a = !s}") + log(f"{tma_tensor_a = !s}") + # Setup TMA load for B + b_op = utils.sm100.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, self._tiled_mma.thr_id + ) + log(f"{b_op = !s}") + # Get written of the pipeline dimension. + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + log(f"{b_smem_layout = !s}") + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + self._tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + log(f"{tma_atom_b = !s}") + log(f"{tma_tensor_b = !s}") + 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 + log(f"{self.num_tma_load_bytes = }") + + # Set the TMA related arguments for the tensors used in the EFC. + self.efc.jit.create_tma_arguments() + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + d, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + + self.efc.jit.create_supplemental_arguments_for_kernel() + + # Launch the kernel synchronously + self.kernel( + self._tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + self.efc.kernel.pack_arguments(), + ).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, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + supplemental_parameters: Tuple, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + # Process the variadic parameters. + self.efc.kernel.unpack_parameters(supplemental_parameters) + + 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) + # Prefetch the TMA descriptors for all the supplemental tensors. + self.efc.kernel.prefetch_tma_descriptors() + + # + # 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 + # + self.buffer_align_bytes = 1024 + + # 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] + # Barriers used by the supplemental load tensor pipeline. + c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage] + c_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.a_dtype, cute.cosize(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(b_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + self.smem = utils.SmemAllocator() + storage = self.smem.allocate(self.shared_storage) + + # Allocate the shared memory for all the supplemental tensors. + self.efc.kernel.allocate_smem() + + # 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.epilogue_warp_id) * ( + 2 if self.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, + ) + + # Load pipeline, used to load all the supplemental tensors of the + # epilogue. + c_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + c_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.c_full_mbar_ptr.data_ptr(), + num_stages=self.num_c_stage, + producer_group=c_producer_group, + consumer_group=c_consumer_group, + # Unlock the barrier when all the tensor bytes have been loaded. + tx_count=self.efc.jit.total_tma_load_bytes, + ) + + 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)), + ) + # 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=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + # + # Setup smem tensor A/B + # + # (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 self.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, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/D + # + self.thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = self.thr_mma.partition_A(gA_mkl) + log(f"{tCgA = !s}") + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = self.thr_mma.partition_B(gB_nkl) + log(f"{tCgB = !s}") + # Create the local_tile gX_mnl for all the EFC supplemental tensors. + self.efc.kernel.partition_global_tensors_for_tiled_mma() + + # + # 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), tiles_m, tiles_k, tiles_l) + 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), tiles_n, tiles_k, tiles_l) + 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/D + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + log(f"{tCrA = !s}") + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + log(f"{tCrB = !s}") + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + log(f"{acc_shape = !s}") + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + log(f"{tCtAcc_fake = !s}") + + # Named barriers + # + cta_sync_barrier = pipeline.NamedBarrier( + self.cta_sync_bar_id, self.threads_per_cta + ) + epilogue_sync_barrier = pipeline.NamedBarrier( + self.epilogue_sync_bar_id, 32 * len(self.epilogue_warp_id) + ) + + # + # Cluster wait before tensor memory alloc + # + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + cta_sync_barrier.arrive_and_wait() + + # + # 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), loopK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), loopK) + 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.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + 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() + + 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_k_blocks = cute.size(tCrA, mode=[2]) + for k_block_idx in cutlass.range( + num_k_blocks, unroll_full=True + ): + k_block_coord = ( + None, + None, + k_block_idx, + ab_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[k_block_coord], + tCrB[k_block_coord], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first k_block + 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 + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + log(f"tmem_ptr = {tmem_ptr!s}") + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + log(f"tCtAcc_base = {tCtAcc_base!s}") + # + # Partition for epilogue + # + epi_tidx = tidx + tCgD = self.efc.kernel.tCgD_written[self.d_name_bigger] + log(f"tCgD (aka tCgD_written[{self.d_name_bigger}])= {tCgD!s}") + + ( + tiled_copy_t2r, # (EPI_TILE_M, EPI_TILE_N) + tTR_tAcc_base, # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_rAcc, # (T2R, T2R_M, T2R_N) + ) = self.epilogue_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgD, epi_tile + ) + log(f"{tiled_copy_t2r = !s}") + log(f"{tTR_tAcc_base = !s}") + log(f"{tTR_rAcc = !s}") + # Copy and partition for the supplemental EFC tensors. + self.efc.kernel.copy_and_partition_supplemental_rmem_tensors( + tiled_copy_t2r, tTR_rAcc, epi_tidx, epi_tile + ) + + # + # 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 + ) + + # Store D pipeline used for all the written tensors in the epilogue. + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=d_producer_group, + ) + + c_pipeline_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_c_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 the supplemental written tensors per MMA tile index. + self.efc.kernel.slice_written_tensors_per_mma_tile_index( + 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) + ] + log(f"tTR_tAcc = {tTR_tAcc!s}") + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + # Group together the EPI_M, EPI_M which are starting at group 3. + # (T2R, T2R_M, T2R_N, (EPI_M, EPI_M)) + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + log(f"group_modes tTR_tAcc = {tTR_tAcc!s}") + # + # Store accumulator to global memory in subtiles + # + # Use EPI_M*EPI_M to iterate using the 1-D coordinate. + 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) + log(f"cute.copy tiled_copy_t2r = {tiled_copy_t2r!s}") + log(f"cute.copy tTR_tAcc_mn = {tTR_tAcc_mn!s}") + log(f"cute.copy tTR_rAcc = {tTR_rAcc!s}") + + # Wait for the EFC tensor loads to complete. + if cutlass.const_expr(self.efc.read_tensor_names): + # The wait is blocking even if the tx_count is 0 when + # there is no tensor to load. So need to skip it when + # there is no tensor to read. + c_pipeline.consumer_wait(c_pipeline_consumer_state) + + # Load supplemental tensors from shared memory to register. + self.efc.kernel.load_tensors_from_smem_to_register( + c_pipeline_consumer_state.index + ) + + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + c_pipeline.consumer_release(c_pipeline_consumer_state) + + # Advance pipeline states + c_pipeline_consumer_state.advance() + + # + # Perform epilogue op on accumulator. + # + tiled_copy_r2s = self.efc.kernel.tiled_copy_r2s[self.d_name_bigger] + log(f"tiled_copy_r2s = {tiled_copy_r2s!s}") + # Use a SimpleNamespace to pass easily some local content as + # an extensible class compatible with CuTe DSL + # implementation. + epilogue_context = types.SimpleNamespace() + # Load the accumulator cast to the epi_dtype used to do all + # the computations in the epilogue. + # Retile the accumulator subtile to fit the destination + # subtile vector TV layout. + epilogue_context.acc_vec = ( + tiled_copy_r2s.retile(tTR_rAcc).load().to(self.epi_dtype) + ) + log(f"before .retile tTR_rAcc = {tTR_rAcc!s}") + log( + f"tiled_copy_r2s.retile(tTR_rAcc) = {tiled_copy_r2s.retile(tTR_rAcc)!s}" + ) + log( + f"tiled_copy_r2s.retile(tTR_rAcc).load() = {tiled_copy_r2s.retile(tTR_rAcc).load()!s}" + ) + log(f"epilogue_context.acc_vec = {epilogue_context.acc_vec!s}") + + # Execute the EFC epilogue. + self.efc.kernel.epilogue_computation(epilogue_context) + d_buffer = (num_prev_subtiles + subtile_idx) % self.num_d_stage + + # Store the EFC written tensors to shared memory. + self.efc.kernel.store_written_tensors_to_smem(d_buffer) + + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + epilogue_sync_barrier.arrive_and_wait() + + # + # TMA store D to global memory + # + if warp_idx == self.epilogue_warp_id[0]: + # Store with TMA the written EFC tensors to global memory. + self.efc.kernel.tma_store_written_tensors_to_gmem( + d_buffer, subtile_idx + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + d_pipeline.producer_commit() + d_pipeline.producer_acquire() + + 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() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + epilogue_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for D store complete + # + d_pipeline.producer_tail() + + # + # Specialized epilog load warp + # + if warp_idx == self.epilogue_load_warp_id: + # Create the tiled tensors to be loaded in the epilogue. + self.efc.kernel.create_epilogue_subtile_tensors(tidx, epi_tile) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + c_pipeline_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_c_stage + ) + + # Setup the pipelines reading the EFC supplemental tensors. + + 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], + ) + # Prepare the EFC tensors to be loaded by the subtiles. + subtile_cnt = self.efc.kernel.prepare_tensor_load_for_subtiles( + mma_tile_coord_mnl, + ) + + # Assume the pipeline can work even in the case there is no + # tensor to load and so subtile_cnt is 0. + for subtile_idx in cutlass.range(subtile_cnt): + # Load C from global memory to shared memory. + c_pipeline.producer_acquire(c_pipeline_producer_state) + + # Load the subtiles of EFC tensors. + self.efc.kernel.load_tensor_subtiles( + subtile_idx, c_pipeline, c_pipeline_producer_state + ) + + c_pipeline_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for the load buffer to be empty. + # + c_pipeline.producer_tail(c_pipeline_producer_state) + + def epilogue_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + tCgC: cute.Tensor, + epi_tile: cute.Tile, + ) -> 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). + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, # Take this as the reference layout for the epilogue tile. + self.epi_dtype, # But we get the accumulator as epi_dtype in the epilogue. + self.acc_dtype, + epi_tile, + self.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) + tCgC_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) + tTR_gC = thr_copy_t2r.partition_D(tCgC_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_load( + 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 load, then use it to partition register array (destination) and shared memory (source). + + :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 + + :return: A tuple containing (tiled_copy_s2r, tSR_rC, tSR_sC) where: + - tiled_copy_s2r: The tiled copy operation for smem to register copy(s2r) + - tSR_rC: The partitioned tensor C (register destination) + - tSR_sC: The partitioned tensor C (smem source) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_s2r = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + tiled_copy_s2r = cute.make_tiled_copy_D(copy_atom_s2r, tiled_copy_t2r) + # (S2R, S2R_M, S2R_N, PIPE_C) + thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) + tSR_sC = thr_copy_s2r.partition_D(sC) + # (S2R, S2R_M, S2R_N) + tSR_rC = tiled_copy_s2r.retile(tTR_rC) + return tiled_copy_s2r, tSR_rC, tSR_sC + + def epilogue_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, + dtype: Type[cutlass.Numeric], + ) -> 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, tiles_m, tiles_n, tiles_l) + gC_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + + 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, tiles_m, tiles_n, tiles_l) + 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 + + def compute_stages(self) -> None: + """Compute and set the number of stages for A/B/C/D operands. + + Uses instance attributes to compute and assign: + `self.num_acc_stage`, `self.num_ab_stage`, `self.num_c_stage`, + and `self.num_d_stage`. + """ + # Defaults + self.num_acc_stage = 2 + # To read the tensors needed for the epilogue: + self.num_c_stage = 2 + # To write the tensors produced by the epilogue: + self.num_d_stage = 2 + + # Calculate smem layout and size for one stage of A, B, C, and D + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + self._tiled_mma, self.mma_tiler, self.a_dtype, 1 + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + self._tiled_mma, self.mma_tiler, self.b_dtype, 1 + ) + + # Get the contribution from the tensors used in the EFC. + self.efc.jit.compute_stage() + + ab_bytes_per_stage = cute.size_in_bytes( + self.a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(self.b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + # Contribution from the tensors loaded in the EFC. + c_bytes_per_stage = self.efc.jit.smem_size_in_bytes_of_read_tensors() + c_bytes = c_bytes_per_stage * self.num_c_stage + # Contribution from the tensors stored in the EFC. There is at least 1 + # written tensor, so the following is strictly positive. + d_bytes_per_stage = self.efc.jit.smem_size_in_bytes_of_written_tensors() + d_bytes = d_bytes_per_stage * self.num_d_stage + + # Calculate A/B stages + self.num_ab_stage = ( + self.smem_capacity // self.occupancy + - (mbar_helpers_bytes + c_bytes + d_bytes) + ) // ab_bytes_per_stage + log(f"\t{self.num_ab_stage = }") + + if self.num_ab_stage <= 0: + raise MemoryError("Not enough smem capacity to allocate all the tensors.") + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes. + # Add remaining unused smem to epilogue. + self.num_d_stage += ( + self.smem_capacity + - self.occupancy * ab_bytes_per_stage * self.num_ab_stage + - self.occupancy * (mbar_helpers_bytes + c_bytes + d_bytes) + ) // (self.occupancy * d_bytes_per_stage) + log(f"\tnew {self.num_d_stage = }") + + @staticmethod + def _compute_grid( + d: 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 D. + + :param d: The output tensor D + :type d: 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]] + """ + log(f"compute_grid: {max_active_clusters = }") + d_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gd = cute.zipped_divide(d, tiler=d_shape) + num_ctas_mnl = gd[(0, (None, None, None))].shape + common_efc.if_debug( + lambda: cute.printf("compute_grid: num_ctas_mnl = {}", num_ctas_mnl) + ) + 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 + ) + common_efc.if_debug(lambda: cute.printf("compute_grid: grid = {}", grid)) + return tile_sched_params, grid + + def compute_num_tmem_alloc_cols(self) -> None: + """Compute and set the number of tensor memory allocation columns. + + This method uses the instance attributes computed during setup to + determine the number of tensor memory allocation columns and stores + the result in `self.num_tmem_alloc_cols`. + """ + acc_shape = self._tiled_mma.partition_shape_C(self.mma_tiler[:2]) + log(f"compute_num_tmem_alloc_cols: {acc_shape = !s}") + tCtAcc_fake = self._tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + log(f"compute_num_tmem_alloc_cols: {tCtAcc_fake = !s}") + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + log(f"compute_num_tmem_alloc_cols: {self.num_tmem_alloc_cols = }") + + def check_valid_dtypes( + self, + ab_dtype: Type[cutlass.Numeric], + ): + """ + 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] + + :raises ValueError: If the dtypes are invalid or incompatible + """ + 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: + raise ValueError( + f"Invalid A/B dtype: {ab_dtype}. " + f"Supported dtypes: {', '.join(str(dt) for dt in valid_ab_dtypes)}" + ) + + valid_acc_dtypes = {cutlass.Float32, cutlass.Float16, cutlass.Int32} + if self.acc_dtype not in valid_acc_dtypes: + raise ValueError( + f"Invalid accumulator dtype: {self.acc_dtype}. " + f"Supported dtypes: {', '.join(str(dt) for dt in valid_acc_dtypes)}" + ) + + # 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]: + compatible_types = acc_ab_compatibility[self.acc_dtype] + raise ValueError( + f"Incompatible dtype combination: A/B dtype {ab_dtype} is not compatible " + f"with accumulator dtype {self.acc_dtype}. " + f"Compatible A/B dtypes for {self.acc_dtype}: {', '.join(str(dt) for dt in compatible_types)}" + ) + + def check_valid_mma_tiler_and_cluster_shape(self): + """Check if the mma tiler and cluster shape are valid. + + :raises ValueError: If the mma tiler or cluster shape is invalid + """ + # Check invalid mma tile shape M dimension + 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]) + ): + expected = [128, 256] if self.use_2cta_instrs else [64, 128] + raise ValueError( + f"Invalid MMA tile M dimension: {self.mma_tiler_mn[0]}. " + f"Expected one of {expected} (use_2cta_instrs={self.use_2cta_instrs})" + ) + + # Check invalid mma tile shape N dimension + if self.mma_tiler_mn[1] not in range(32, 257, 32): + raise ValueError( + f"Invalid MMA tile N dimension: {self.mma_tiler_mn[1]}. " + f"Expected a multiple of 32 in range [32, 256]" + ) + # Check illegal cluster shape M dimension + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + divisor = 2 if self.use_2cta_instrs else 1 + raise ValueError( + f"Invalid cluster shape M dimension: {self.cluster_shape_mn[0]}. " + f"Must be divisible by {divisor} (use_2cta_instrs={self.use_2cta_instrs})" + ) + + def is_power_of_2(x): + return x > 0 and (x & (x - 1)) == 0 + + # Check invalid cluster shape constraints + if self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16: + raise ValueError( + f"Invalid cluster shape: {self.cluster_shape_mn}. " + f"Product {self.cluster_shape_mn[0]} * {self.cluster_shape_mn[1]} = " + f"{self.cluster_shape_mn[0] * self.cluster_shape_mn[1]} exceeds maximum of 16" + ) + if self.cluster_shape_mn[0] <= 0 or self.cluster_shape_mn[1] <= 0: + raise ValueError( + f"Invalid cluster shape: {self.cluster_shape_mn}. " + f"Both dimensions must be positive" + ) + if not is_power_of_2(self.cluster_shape_mn[0]) or not is_power_of_2( + self.cluster_shape_mn[1] + ): + raise ValueError( + f"Invalid cluster shape: {self.cluster_shape_mn}. " + f"Both dimensions must be powers of 2" + ) + + def check_valid_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + ): + """ + 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 d_dtype: The data type of the D tensor + :type d_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 cd_major: The major axis of the C/D tensor + :type cd_major: str + + :raises ValueError: If the tensor alignment is invalid + """ + + def check_contigous_16B_alignment( + dtype, is_mode0_major, tensor_shape, tensor_name + ): + 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 + if num_major_elements % num_contiguous_elements != 0: + raise ValueError( + f"Invalid alignment for tensor {tensor_name}. " + f"Major dimension has {num_major_elements} elements, " + f"but requires alignment to {num_contiguous_elements} elements (16 bytes). " + f"Dtype: {dtype}, width: {dtype.width} bits" + ) + + check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l), "A") + check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l), "B") + check_contigous_16B_alignment(d_dtype, cd_major == "m", (m, n, l), "D") + + def check_implementable(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor): + """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 d: One of the tensor used as some output + :type d: cute.Tensor + + :raises CantImplementError: If the configuration is not implementable + """ + m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] + + # infer a_major, b_major, cd_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_d = utils.LayoutEnum.from_tensor(d).is_m_major_c() + a_major = "m" if is_m_major_a else "k" + b_major = "n" if is_n_major_b else "k" + cd_major = "m" if is_m_major_d else "n" + + try: + # Check dtypes (raises ValueError if invalid) + self.check_valid_dtypes(a.element_type) + + # Check mma tile shape and cluster shape (raises ValueError if invalid) + self.check_valid_mma_tiler_and_cluster_shape() + + # Check problem shape for load/store alignment (raises ValueError if invalid) + self.check_valid_tensor_alignment( + m, + n, + k, + l, + a.element_type, + d.element_type, + a_major, + b_major, + cd_major, + ) + except ValueError as e: + raise cute.testing.CantImplementError(f"Configuration error: {e}") + + class CLIParser: + """Parse command-line arguments for the Blackwell Dense GEMM example.""" + + def __init__(self): + self.parser = argparse.ArgumentParser( + description="Example of Dense Persistent GEMM on Blackwell." + ) + self.parser.add_argument( + "--mnkl", + type=self.parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", + ) + self.parser.add_argument( + "--mma_tiler_mn", + type=self.parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + self.parser.add_argument( + "--cluster_shape_mn", + type=self.parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + self.parser.add_argument( + "--ab_dtype", type=cutlass.dtype, default=cutlass.TFloat32 + ) + self.parser.add_argument( + "--acc_dtype", type=cutlass.dtype, default=cutlass.Float32 + ) + self.parser.add_argument( + "--epi_dtype", type=cutlass.dtype, default=cutlass.Float32 + ) + self.parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + self.parser.add_argument( + "--a_major", choices=["k", "m"], type=str, default="k" + ) + self.parser.add_argument( + "--b_major", choices=["k", "n"], type=str, default="k" + ) + self.parser.add_argument( + "--cd_major", choices=["n", "m"], type=str, default="n" + ) + self.parser.add_argument( + "--tolerance", + type=float, + default=1e-01, + help="Tolerance for validation", + ) + self.parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + self.parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + self.parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + + # A children class may add more things to parse. + self.more_parsing() + + def parse(self): + """Parse the command-line arguments.""" + args = self.parser.parse_args() + + if len(args.mnkl) != 4: + self.parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + self.parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + self.parser.error("--cluster_shape_mn must contain exactly 2 values") + + return args + + @staticmethod + 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." + ) + + def more_parsing(self): + """To override to add more stuff to self.parser""" + + @staticmethod + def dtype_name(dtype: Type[cutlass.Numeric]) -> str: + """Convert a CUTLASS dtype object to its clean string name. + + This is needed to format dtype objects into CLI arguments without + full module paths. CUTLASS dtype objects have different representations: + some have a __name__ attribute while others need string parsing. + + We want "Float16" not "cutlass.Float16" or "". + + :param dtype: CUTLASS numeric data type + :return: Clean type name string (e.g., "Float16", "BFloat16") + + Example: + >>> DenseGemmEFC.dtype_name(cutlass.Float16) + 'Float16' + """ + return ( + dtype.__name__ if hasattr(dtype, "__name__") else str(dtype).split(".")[-1] + ) + + @staticmethod + def format_as_cli_args( + script_name: str, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + epi_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + ) -> str: + """Format common test parameters as CLI arguments base. + + This method formats the common parameters shared across different GEMM examples. + Subclass-specific parameters (like alpha, beta, etc.) should be added by overriding methods. + + :param script_name: Name of the Python script + :param mnkl: Matrix dimensions (M, N, K, L) + :param ab_dtype: Data type for A and B matrices + :param acc_dtype: Data type for accumulation + :param epi_dtype: Data type for epilogue + :param a_major: Major order for matrix A + :param b_major: Major order for matrix B + :param cd_major: Major order for matrices C and D + :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 + :return: Formatted CLI command string + """ + # Format tuples as comma-separated values + mnkl_str = ",".join(map(str, mnkl)) + mma_tiler_str = ",".join(map(str, mma_tiler_mn)) + cluster_shape_str = ",".join(map(str, cluster_shape_mn)) + + cmd = ( + f"python {script_name} " + f"--mnkl {mnkl_str} " + f"--ab_dtype {DenseGemmEFC.dtype_name(ab_dtype)} " + f"--acc_dtype {DenseGemmEFC.dtype_name(acc_dtype)} " + f"--epi_dtype {DenseGemmEFC.dtype_name(epi_dtype)} " + f"--a_major {a_major} " + f"--b_major {b_major} " + f"--cd_major {cd_major} " + f"--mma_tiler_mn {mma_tiler_str} " + f"--cluster_shape_mn {cluster_shape_str} " + f"{'--use_2cta_instrs ' if use_2cta_instrs else ''}" + f"--tolerance {tolerance}" + ) + return cmd + + def create_arguments(self, l, m, n, k, a_major, b_major, cd_major, ab_dtype): + """Create base GEMM input tensors A and B. + + This method creates the input matrices for GEMM computation. Subclasses + typically override this method to create additional supplemental tensors + for the epilogue. + + :param l: Batch dimension + :param m: M dimension (rows of A and output) + :param n: N dimension (columns of B and output) + :param k: K dimension (inner dimension) + :param a_major: Major order for A matrix ('m' or 'k') + :param b_major: Major order for B matrix ('n' or 'k') + :param cd_major: Major order for supplemental tensors ('m' or 'n') + :param ab_dtype: Data type for A and B matrices + :return: Tuple of (a_tensor, b_tensor, a_torch_cpu, b_torch_cpu) + """ + torch.manual_seed(1111) + + a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_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 + ) + + return ( + a_tensor, + b_tensor, + a_torch_cpu, + b_torch_cpu, + ) + + def evaluate_on_cpu( + self, + a_torch_cpu, + b_torch_cpu, + epi_dtype, + *epilogue_args, + ): + """Evaluate the GEMM and epilogue computation on CPU for validation. + + Computes the reference result by performing A*B using einsum, then + evaluates the epilogue function with the accumulator and supplemental + arguments. This updates any output tensors in epilogue_args. + + :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 epilogue_args: Supplemental arguments for the epilogue (tensors and scalars) + """ + # Compute reference result + ref = torch.einsum( + "mkl,nkl->mnl", + a_torch_cpu.to(dtype=torch.float32), + b_torch_cpu.to(dtype=torch.float32), + ) + + self.efc.evaluate_on_cpu(ref, *epilogue_args) + + def compile( + self, + a_tensor, + b_tensor, + max_active_clusters, + current_stream, + *supplemental_arguments, + **compiler_options, + ): + """Compile the GEMM kernel with epilogue fusion. + + Compiles the kernel using CUTE DSL compilation, incorporating the EFC + (Epilogue Fusion Configuration) with all supplemental arguments. Returns + a callable function that accepts the same arguments for execution. + + :param a_tensor: Input tensor A + :param b_tensor: Input tensor B + :param max_active_clusters: Maximum number of active clusters + :param current_stream: CUDA stream for execution + :param supplemental_arguments: Additional arguments for the epilogue (tensors and scalars) + :param compiler_options: Keywords arguments passed to CuTe DSL compiler. "options = " for now + :return: Compiled callable function that executes the GEMM with the same signature + """ + self.efc.compile(supplemental_arguments) + compiled = cutlass.cute.compile( + self, + a_tensor, + b_tensor, + max_active_clusters, + current_stream, + self.efc.jit.pack_arguments(*supplemental_arguments), + **compiler_options, + ) + + def inject_supplemental_arguments( + a_tensor, + b_tensor, + current_stream, + *supplemental_arguments, + ): + # Run the compiled code, do not pass the constexpr parameters nor + # the compiler options. + return compiled( + a_tensor, + b_tensor, + current_stream, + self.efc.jit.pack_arguments(*supplemental_arguments), + ) + + return inject_supplemental_arguments diff --git a/examples/python/CuTeDSL/blackwell/epilogue/common_efc.py b/examples/python/CuTeDSL/blackwell/epilogue/common_efc.py new file mode 100644 index 00000000..684364be --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/common_efc.py @@ -0,0 +1,1512 @@ +# 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. + +# This not to use module annotations from future version but to change the type system to postpone the evaluation of annotations, +# about forward declaration and lazy type checking. +# See https://docs.python.org/3/library/__future__.html#future__.annotations and https://peps.python.org/pep-0563/. +from __future__ import annotations + +import dataclasses +import enum +import functools +import inspect +import logging +import os +import types +import typing + +import cutlass +import torch + +# To have some verbosity, set the CUTE_DSL_EFC_LOG_LEVEL environment variable to +# INFO or even DEBUG before launching this program. +if log_level := os.environ.get("CUTE_DSL_EFC_LOG_LEVEL", None): + logging.basicConfig(level=getattr(logging, log_level)) +logger = logging.getLogger(__name__) + + +def log(message: str): + """Helper function to log messages. Change logger.info to another level here + if needed.""" + logger.info(message) + + +""" +CUTLASS EFC Framework +""" + +# Available activation functions in the EFC Configuration class +ACTIVATION_FUNCTIONS = [ + "identity", + "relu", + "leaky_relu", + "tanh", + "sigmoid", + "silu", + "hardswish", + "gelu", +] + + +def if_debug(function): + """Execute a function if in debug mode.""" + if logger.isEnabledFor(logging.DEBUG): + function() + + +def mark_mlir(message: str): + """Put some message in MLIR output to make MLIR assembly clearer or trace execution.""" + if_debug(lambda: cutlass.cute.printf(f"mark_mlir: {message}")) + + +def trace_in_mlir(func): + """Decorator to trace function entry and exit in MLIR.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + function_name = func.__name__ + mark_mlir(f"entering {function_name}") + result = func(*args, **kwargs) + mark_mlir(f"leaving {function_name}") + return result + + return wrapper + + +def create_named_epilogue(param_names, func): + """Create a wrapper function with specific parameter names that delegates to an implementation function. + + This function solves a common problem in the EFC (Epilogue Fusion Configuration) framework: + epilogue functions must have parameters with specific names (e.g., "alpha", "beta", "C", "D") + to match the EFC calling convention, but you may want to generate these functions + programmatically from a generic implementation. + + Instead of using string manipulation with exec() or eval() (which is insecure and breaks + tooling), this function uses Python's inspect module to create a proper function signature + that tools like debuggers, type checkers, and IDEs can understand. + + Args: + param_names: List of parameter names for the generated function + (e.g., ["alpha", "beta", "C", "x_factor"]) + func: Implementation function that accepts the same number of arguments as param_names. + The arguments will be passed in the order specified by param_names. + + Returns: + A new function with the specified parameter names that calls func with those + parameters in order. The wrapper preserves func's name and docstring, and + has a proper signature for introspection. + + Example: + # Generic implementation that doesn't care about parameter names + def compute(a, b, c): + return a + b * c + + # Create EFC-compliant function with required parameter names + epilogue = create_named_epilogue(["alpha", "X", "Y"], compute) + # Now epilogue(alpha=1, X=2, Y=3) calls compute(1, 2, 3) + # and inspect.signature(epilogue) shows the correct parameter names + + Use Case: + When programmatically generating epilogue functions with different tensor + configurations, you need each function to have the right parameter names + for the EFC framework to call them correctly with keyword arguments. + + """ + # Create Parameter objects for each parameter name, using standard Python argument binding. + parameters = [ + inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD) + for name in param_names + ] + + # Create a new signature with the custom parameter names + new_signature = inspect.Signature(parameters) + + # Create a wrapper function that accepts arguments according to the new signature + def wrapper(*args, **kwargs): + # Bind the provided arguments to our custom signature + bound = new_signature.bind(*args, **kwargs) + bound.apply_defaults() + + # Extract argument values in the order specified by param_names + ordered_args = [bound.arguments[name] for name in param_names] + + # Call the original function with the properly ordered arguments + return func(*ordered_args) + + # Assign the custom signature to the wrapper so introspection works correctly + wrapper.__signature__ = new_signature + wrapper.__name__ = getattr(func, "__name__", "generated_function") + wrapper.__doc__ = func.__doc__ + + return wrapper + + +class VariadicParameters: + """Minimal mixin wrapper for variadic parameters for @cute.jit/@cute.kernel + functions taking advantage that the DSL to can ingest a recursive + combination of tuples and lists.""" + + def __init__(self, efc, parameter_names): + # Add local shortcuts to the efc and gemm objects + self.efc = efc + self.gemm = efc.gemm + # Create a dataclass to have an aggregate initializer. + # Use __slots__ to avoid assigning wrong argument by error. + fields = [(name, typing.Any) for name in parameter_names] + self._parameter_class = dataclasses.make_dataclass( + "Parameter", fields, slots=True + ) + + # Add some methods to the dataclass so we can access for example arg.a + # and parameter.b also with arg["a"] or parameter["b"]. + def getitem(self, name): + """Access the dataclass attribute by name.""" + return getattr(self, name) + + self._parameter_class.__getitem__ = getitem + + def setitem(self, name, value): + """Set the dataclass attribute by name.""" + setattr(self, name, value) + + self._parameter_class.__setitem__ = setitem + + self.instantiate_args() + logger.info(f"Initial {self.arg = }") + + def pack_arguments(self, *args, **kwargs): + """Pack the arguments to pass them through a @cute.jit/@cute.kernel + call. + + If some arguments are provided, pack them, otherwise just use the + self.arg object by default. + + Return a tuple as an interface object since a @cute.jit/@cute.kernel + can ingest a recursive combination of tuples and lists.""" + if args or kwargs: + # Override the current argument object from the provided arguments, if any. + self.arg = self._parameter_class(*args, **kwargs) + # dataclasses.astuple(self.arg) breaks because it is recursive and + # applies a deepcopy incompatible with the DSL magic. Just generate 1 + # level of tuple of object references. + r = tuple(self.arg[f.name] for f in dataclasses.fields(self.arg)) + logger.info(f"pack_arguments {args = } {kwargs = } {self.arg = } {r = }") + # The DSL does not accept an empty tuple but can handle None. So + # remap to None in that case. + if not r: + return None + return r + + def unpack_parameters(self, p: typing.Tuple): + """Unpack the parameters inside a @cute.jit/@cute.kernel function. + + Assign all the self.parameter attributes.""" + # Do the opposite mapping of None to an empty tuple to have the + # parameter constructor happy. + if p is None: + p = () + # Instantiate the dataclass holding the parameters from the + # individual parameter values. + self.parameter = self._parameter_class(*p) + logger.info(f"unpack_parameters {p = } {self.parameter = }") + + def instantiate_args(self): + """Create an arg attribute from the Parameter class to be used + as an alternative way to pass the arguments instead of using an + explicit pack_arguments(). + + All the arg attributes are initialized to a noticeable name so that + any forgotten field will trigger an error.""" + + class _UnassignedArgument: + """Sentinel class to detect uninitialized arguments""" + + def __repr__(self): + return "" + + self.arg = self._parameter_class( + *([_UnassignedArgument] * len(dataclasses.fields(self._parameter_class))) + ) + + +class EFC: + """Epilogue Fusion Configuration.""" + + # Helper functions for CuTe operations + @staticmethod + def maximum(x, y): + """Element-wise maximum of 2 CuTe tensors""" + x_type = x.element_type + y_type = y.element_type + assert x_type is y_type, f"Type mismatch: x is {x_type}, y is {y_type}" + return cutlass.cute.where(x > y, x, y) + + @staticmethod + def minimum(x, y): + """Element-wise minimum of 2 CuTe tensors""" + x_type = x.element_type + y_type = y.element_type + assert x_type is y_type, f"Type mismatch: x is {x_type}, y is {y_type}" + return cutlass.cute.where(x < y, x, y) + + class JIT(VariadicParameters): + """Handle Python/@cute.jit and its boundaries with Host.""" + + # All the following customization functions should go somewhere else in + # the long term, as part of a refactoring similar to CUTLASS + # collective/main loop/epilogue... + + @trace_in_mlir + def record_tensor_dtypes(self): + """It does not seem that the tma_tensor and tma_atom carry over the + element type, so, store it here for later use.""" + self.tensor_dtype = {} + + def f(tensor_name, attributes): + tensor = self.parameter[tensor_name] + self.tensor_dtype[tensor_name] = tensor.element_type + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def written_tensor_name_with_bigger_element_type(self): + """The type of the written tensor is used to compute a lot of + implementation details about tiling and so on in the kernel. + + The compilation phase has checked already there is at least 1 + written tensor name. + + Return the name of the written tensor with the biggest + element_type. + + """ + return max( + (tensor_name for tensor_name in self.efc.written_tensor_names), + key=lambda tensor_name: self.tensor_dtype[tensor_name].width, + ) + + @trace_in_mlir + def read_tensor_name_with_bigger_element_type(self): + """The type of the read tensor is used to compute a lot of + implementation details about tiling and so on in the kernel. Return the name of the read tensor with the biggest element_type, or None if there is no read tensor.""" + if self.efc.read_tensor_names: + return max( + (tensor_name for tensor_name in self.efc.read_tensor_names), + key=lambda tensor_name: self.tensor_dtype[tensor_name].width, + ) + return None + + @trace_in_mlir + def compute_stage(self): + """Get the contribution from the tensors used in the EFC to the + pipeline stage numbers.""" + self.smem_size_of_read_tensors = 0 + self.smem_size_of_written_tensors = 0 + self.tensor_dtype = {} + + def f(tensor_name, attributes): + tensor = self.parameter[tensor_name] + tensor_layout = cutlass.utils.LayoutEnum.from_tensor(tensor) + if cutlass.const_expr(self.gemm.d_layout != tensor_layout): + error_msg = ( + f"The tensor {tensor_name} has layout {tensor_layout} which is " + f"different from C/D specified layout {self.gemm.d_layout}." + ) + raise ValueError(error_msg) + + # It does not seem that the tma_tensor and tma_atom carry over + # the element type, so, store it here for later use. + self.tensor_dtype[tensor_name] = tensor.element_type + + smem_size_in_bytes_of_a_pipeline_stage = cutlass.cute.size_in_bytes( + tensor.element_type, + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, self.gemm.d_layout, self.gemm.epi_tile, 1 + ), + ) + # Prepare the information to be asked soon, to recycle this + # loop. + if attributes.is_read: + self.smem_size_of_read_tensors += ( + smem_size_in_bytes_of_a_pipeline_stage + ) + if attributes.is_written: + self.smem_size_of_written_tensors += ( + smem_size_in_bytes_of_a_pipeline_stage + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def smem_size_in_bytes_of_read_tensors(self): + """Get the contribution in a smem pipeline stage from the tensors + loaded in the EFC.""" + logger.info(f"\t{self.smem_size_of_read_tensors = }") + return self.smem_size_of_read_tensors + + @trace_in_mlir + def smem_size_in_bytes_of_written_tensors(self): + """Get the contribution in a smem pipeline stage from the tensors + stored in the EFC.""" + logger.info(f"\t{self.smem_size_of_written_tensors = }") + return self.smem_size_of_written_tensors + + @trace_in_mlir + def smem_layout(self): + """Get the smem_layout for the tensors used in the EFC.""" + self.smem_layout_staged_read = {} + self.smem_layout_staged_written = {} + + def f(tensor_name, attributes): + tensor = self.parameter[tensor_name] + tensor_layout = cutlass.utils.LayoutEnum.from_tensor(tensor) + log(f"JIT.smem_layout {tensor_name} = {tensor!s}") + log(f"JIT.smem_layout tensor_layout[{tensor_name}] = {tensor_layout!s}") + + if attributes.is_read: + self.smem_layout_staged_read[tensor_name] = ( + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, + tensor_layout, + self.gemm.epi_tile, + self.gemm.num_c_stage, + ) + ) + log(f"JIT.smem_layout read {self.gemm.num_c_stage = }") + log( + f"JIT.smem_layout read self.smem_layout_staged_read[{tensor_name}] = {self.smem_layout_staged_read[tensor_name]!s}" + ) + if attributes.is_written: + self.smem_layout_staged_written[tensor_name] = ( + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, + tensor_layout, + self.gemm.epi_tile, + self.gemm.num_d_stage, + ) + ) + log(f"JIT.smem_layout written {self.gemm.num_d_stage = }") + log( + f"JIT.smem_layout written self.smem_layout_staged_written[{tensor_name}] = {self.smem_layout_staged_written[tensor_name]!s}" + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def create_tma_arguments(self): + """Set the TMA related arguments for the tensors used in the EFC.""" + # Make the difference for read/written to handle the case a tensor + # is both read and written. + self.total_tma_load_bytes = 0 # Used by the PipelineTmaAsync + self.tma_atom_read = {} + self.tma_tensor_read = {} + self.tma_atom_written = {} + self.tma_tensor_written = {} + + def f(tensor_name, attributes): + tensor = self.parameter[tensor_name] + + if attributes.is_read: + smem_layout = cutlass.cute.slice_( + self.smem_layout_staged_read[tensor_name], (None, None, 0) + ) + self.total_tma_load_bytes += cutlass.cute.size_in_bytes( + tensor.element_type, smem_layout + ) + ( + self.tma_atom_read[tensor_name], + self.tma_tensor_read[tensor_name], + ) = cutlass.cute.nvgpu.cpasync.make_tiled_tma_atom( + cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(), + tensor, + smem_layout, + self.gemm.epi_tile, + ) + log( + f"JIT.tma_atom_read[{tensor_name}] = {self.tma_atom_read[tensor_name]!s}" + ) + log( + f"JIT.tma_tensor_read[{tensor_name}] = {self.tma_tensor_read[tensor_name]!s}" + ) + + if attributes.is_written: + smem_layout = cutlass.cute.slice_( + self.smem_layout_staged_written[tensor_name], (None, None, 0) + ) + ( + self.tma_atom_written[tensor_name], + self.tma_tensor_written[tensor_name], + ) = cutlass.cute.nvgpu.cpasync.make_tiled_tma_atom( + cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(), + tensor, + smem_layout, + self.gemm.epi_tile, + ) + log( + f"JIT.tma_atom_written[{tensor_name}] = {self.tma_atom_written[tensor_name]!s}" + ) + log( + f"JIT.tma_tensor_written[{tensor_name}] = {self.tma_tensor_written[tensor_name]!s}" + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def create_supplemental_arguments_for_kernel(self): + """Executed before launching the @cute.kernel function to set up the + supplemental arguments to pass to the @cute.kernel function. + + In the @cute.kernel example, the parameters like `X_tma_tensor_read` + or `Y_tma_tensor_written` correspond to `mX_mnl` and `mY_mnl`.""" + argument_names = [] + + def compute_argument_names(name, attributes): + if not attributes.is_tensor: + # Just propagate the dynamic scalar with the same name. + argument_names.append(name) + else: + if attributes.is_read: + argument_names.append(f"{name}_tma_atom_read") + argument_names.append(f"{name}_tma_tensor_read") + argument_names.append(f"{name}_smem_layout_staged_read") + if attributes.is_written: + argument_names.append(f"{name}_tma_atom_written") + argument_names.append(f"{name}_tma_tensor_written") + argument_names.append(f"{name}_smem_layout_staged_written") + + self.efc.foreach_argument(compute_argument_names) + # Create the @cute.kernel-side meta-programming infrastructure + # handling also the supplemental argument handling. + self.efc.kernel = EFC.Kernel(self.efc, argument_names) + + arg = self.efc.kernel.arg + + def populate_the_kernel_arguments(name, attributes): + if not attributes.is_tensor: + # Just propagate the dynamic scalar with the same name. + arg[name] = self.parameter[name] + else: + if attributes.is_read: + arg[f"{name}_tma_atom_read"] = self.tma_atom_read[name] + arg[f"{name}_tma_tensor_read"] = self.tma_tensor_read[name] + arg[f"{name}_smem_layout_staged_read"] = ( + self.smem_layout_staged_read[name] + ) + if attributes.is_written: + arg[f"{name}_tma_atom_written"] = self.tma_atom_written[name] + arg[f"{name}_tma_tensor_written"] = self.tma_tensor_written[ + name + ] + arg[f"{name}_smem_layout_staged_written"] = ( + self.smem_layout_staged_written[name] + ) + + self.efc.foreach_argument(populate_the_kernel_arguments) + + class Kernel(VariadicParameters): + """Handle kernel part and @cute.jit/@cute.kernel boundaries.""" + + @trace_in_mlir + def prefetch_tma_descriptors(self): + """Prefetch the TMA descriptors for the tensors used in the EFC.""" + + def f(tensor_name, attributes): + if attributes.is_read: + cutlass.cute.nvgpu.cpasync.prefetch_descriptor( + self.parameter[f"{tensor_name}_tma_atom_read"] + ) + + if attributes.is_written: + cutlass.cute.nvgpu.cpasync.prefetch_descriptor( + self.parameter[f"{tensor_name}_tma_atom_written"] + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def allocate_smem(self): + """Allocate the shared memory for all the supplemental tensors.""" + self.smem_read = {} + self.smem_written = {} + + def f(tensor_name, attributes): + element_type = self.efc.jit.tensor_dtype[tensor_name] + if attributes.is_read: + smem_layout_staged = self.parameter[ + f"{tensor_name}_smem_layout_staged_read" + ] + self.smem_read[tensor_name] = self.gemm.smem.allocate_tensor( + element_type=element_type, + layout=smem_layout_staged.outer, + byte_alignment=self.gemm.buffer_align_bytes, + swizzle=smem_layout_staged.inner, + ) + if attributes.is_written: + smem_layout_staged = self.parameter[ + f"{tensor_name}_smem_layout_staged_written" + ] + self.smem_written[tensor_name] = self.gemm.smem.allocate_tensor( + element_type=element_type, + layout=smem_layout_staged.outer, + byte_alignment=self.gemm.buffer_align_bytes, + swizzle=smem_layout_staged.inner, + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def partition_global_tensors_for_tiled_mma(self): + """Partition the global supplemental tensors for TiledMMA_C/D.""" + self.tCgC_read = {} + self.tCgD_written = {} + + def f(tensor_name, attributes): + if attributes.is_read: + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cutlass.cute.local_tile( + self.parameter[f"{tensor_name}_tma_tensor_read"], + cutlass.cute.slice_(self.gemm.mma_tiler, (None, None, 0)), + (None, None, None), + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: gC_mnl[{tensor_name}] = {gC_mnl!s}" + ) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + self.tCgC_read[tensor_name] = self.gemm.thr_mma.partition_C(gC_mnl) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: self.tCgC_read[{tensor_name}] = {self.tCgC_read[tensor_name]!s}" + ) + + if attributes.is_written: + # (bM, bN, loopM, loopN, loopL) + gD_mnl = cutlass.cute.local_tile( + self.parameter[f"{tensor_name}_tma_tensor_written"], + cutlass.cute.slice_(self.gemm.mma_tiler, (None, None, 0)), + (None, None, None), + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: gD_mnl[{tensor_name}] = {gD_mnl!s}" + ) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + self.tCgD_written[tensor_name] = self.gemm.thr_mma.partition_C( + gD_mnl + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: self.tCgD_written[{tensor_name}] = {self.tCgD_written[tensor_name]!s}" + ) + + self.efc.foreach_tensor(f) + + # The following functions are executed by the specialized warps for + # epilogue computation. + + @trace_in_mlir + def copy_and_partition_supplemental_rmem_tensors( + self, tiled_copy_t2r, tTR_rAcc, epi_tidx, epi_tile + ): + # Load tensor. + self.tiled_copy_s2r = {} + self.tSR_rC = {} + self.tSR_sC = {} + + # Store tensor. + self.tiled_copy_r2s = {} + self.tRS_rD = {} + self.tRS_sD = {} + self.bSG_sD = {} # ((ATOM_V, REST_V), EPI_M, EPI_N) + self.bSG_gD_partitioned = {} # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: tiled_copy_t2r = {tiled_copy_t2r!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: tTR_rAcc = {tTR_rAcc!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: epi_tile = {epi_tile!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: epi_tidx = {epi_tidx!s}" + ) + + def f(tensor_name, attributes): + element_type = self.efc.jit.tensor_dtype[tensor_name] + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: element_type[{tensor_name}] = {element_type!s}" + ) + + if attributes.is_read: + tTR_rC = cutlass.cute.make_rmem_tensor(tTR_rAcc.shape, element_type) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: tTR_rC[{tensor_name}] = {tTR_rC!s}" + ) + + ( + self.tiled_copy_s2r[tensor_name], + self.tSR_rC[tensor_name], + self.tSR_sC[tensor_name], + ) = self.gemm.epilogue_smem_copy_and_partition_load( + tiled_copy_t2r, + tTR_rC, + epi_tidx, + self.smem_read[tensor_name], + ) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tiled_copy_s2r[{tensor_name}] = {self.tiled_copy_s2r[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tSR_sC[{tensor_name}] = {self.tSR_sC[tensor_name]!s}" + ) + + if attributes.is_written: + # (T2R, T2R_M, T2R_N) + tTR_rD = cutlass.cute.make_rmem_tensor(tTR_rAcc.shape, element_type) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: tTR_rD[{tensor_name}] = {tTR_rD!s}" + ) + + # utils.gemm.sm100.epilogue_smem_copy_and_partition uses + # explicitly "C" as the output matrix and introspects the + # gemm object while in this kernel "C" is used for read but + # "D" is for output according to the BLAS convention. + # So construct a minimal mock-up with the required + # information. + faux_gemm = types.SimpleNamespace() + faux_gemm.c_layout = self.gemm.d_layout + faux_gemm.c_dtype = self.gemm.d_dtype + faux_gemm.acc_dtype = self.gemm.acc_dtype + ( + self.tiled_copy_r2s[tensor_name], + self.tRS_rD[tensor_name], # (R2S, R2S_M, R2S_N) + self.tRS_sD[tensor_name], # (R2S, R2S_M, R2S_N) + ) = cutlass.utils.gemm.sm100.epilogue_smem_copy_and_partition( + faux_gemm, + tiled_copy_t2r, # (EPI_TILE_M, EPI_TILE_N) + tTR_rD, + epi_tidx, + self.smem_written[tensor_name], + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.smem_written[{tensor_name}] = {self.smem_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tiled_copy_r2s[{tensor_name}] = {self.tiled_copy_r2s[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tRS_rD[{tensor_name}] = {self.tRS_rD[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tRS_sD[{tensor_name}] = {self.tRS_sD[tensor_name]!s}" + ) + ( + _, + self.bSG_sD[tensor_name], # ((ATOM_V, REST_V), EPI_M, EPI_N) + self.bSG_gD_partitioned[ + tensor_name + ], # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + ) = self.gemm.epilogue_gmem_copy_and_partition( + epi_tidx, + self.parameter[f"{tensor_name}_tma_atom_written"], + self.tCgD_written[tensor_name], + epi_tile, + self.smem_written[tensor_name], + element_type, + ) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.parameter[{tensor_name}_tma_atom_written] = {self.parameter[f'{tensor_name}_tma_atom_written']!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tCgD_written[{tensor_name}] = {self.tCgD_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.smem_written[{tensor_name}] = {self.smem_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: element_type = {element_type!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.bSG_sD[{tensor_name}] = {self.bSG_sD[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.bSG_gD_partitioned[{tensor_name}] = {self.bSG_gD_partitioned[tensor_name]!s}" + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def slice_written_tensors_per_mma_tile_index(self, mma_tile_coord_mnl): + """Slice the supplemental written tensors per MMA tile index.""" + self.bSG_gD = {} # ((ATOM_V, REST_V), (EPI_M, EPI_N)) + + def f(tensor_name, attributes): + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gD = self.bSG_gD_partitioned[tensor_name][ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + log( + f"Kernel.slice_written_tensors_per_mma_tile_index: bSG_gD[{tensor_name}] = {bSG_gD!s}" + ) + # Group the 2 last modes so the subtile_idx loop can iterate + # through it using 1-D indexing. + # ((ATOM_V, REST_V), (EPI_M, EPI_N)) + self.bSG_gD[tensor_name] = cutlass.cute.group_modes( + bSG_gD, 1, cutlass.cute.rank(bSG_gD) + ) + log( + f"Kernel.slice_written_tensors_per_mma_tile_index: self.bSG_gD[{tensor_name}] = {self.bSG_gD[tensor_name]!s}" + ) + + self.efc.foreach_written_tensor(f) + + @trace_in_mlir + def load_tensors_from_smem_to_register(self, index): + """Load supplemental tensors from shared memory to register.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.tiled_copy_s2r[tensor_name], + self.tSR_sC[tensor_name][ + ( + None, + None, + None, + index, + ) + ], + self.tSR_rC[tensor_name], + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tiled_copy_s2r[{tensor_name}] = {self.tiled_copy_s2r[tensor_name]!s}" + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tSR_sC[{tensor_name}] = {self.tSR_sC[tensor_name]!s}" + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + + self.efc.foreach_read_tensor(f) + + @trace_in_mlir + def epilogue_computation(self, epilogue_context): + """Execute the EFC epilogue.""" + + epilogue_context.load = {} + epilogue_context.store = {} + + def load_setup(tensor_name, attributes): + # Retile the read subtile to fit the accumulator subtile vector + # TV layout. + epilogue_context.load[tensor_name] = ( + self.tiled_copy_r2s[self.gemm.d_name_bigger] + .retile(self.tSR_rC[tensor_name]) + .load() + ) + log( + f"Kernel.epilogue_computation load_setup: {self.tiled_copy_r2s[self.gemm.d_name_bigger] = !s}" + ) + log( + f"Kernel.epilogue_computation load_setup: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + log( + f"Kernel.epilogue_computation load_setup: self.tiled_copy_r2s[self.gemm.d_name_bigger].retile(self.tSR_rC[{tensor_name}]) = {self.tiled_copy_r2s[self.gemm.d_name_bigger].retile(self.tSR_rC[tensor_name])!s}" + ) + log( + f"Kernel.epilogue_computation load_setup: epilogue_context.load[{tensor_name}] = {epilogue_context.load[tensor_name]!s}" + ) + + self.efc.foreach_read_tensor(load_setup) + + def store_setup(tensor_name, attributes): + epilogue_context.store[tensor_name] = self.tRS_rD[tensor_name] + log( + f"Kernel.epilogue_computation store_setup: epilogue_context.store[{tensor_name}] = {epilogue_context.store[tensor_name]!s}" + ) + + self.efc.foreach_written_tensor(store_setup) + + self.efc.specialized_epilogue(EFC.Phase.ThreadOperation, epilogue_context)() + + @trace_in_mlir + def store_written_tensors_to_smem(self, d_buffer): + """Store the EFC written tensors to shared memory.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.tiled_copy_r2s[tensor_name], + self.tRS_rD[tensor_name], + self.tRS_sD[tensor_name][(None, None, None, d_buffer)], + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tiled_copy_r2s[{tensor_name}] = {self.tiled_copy_r2s[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_rD[{tensor_name}] = {self.tRS_rD[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_sD[{tensor_name}] = {self.tRS_sD[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_sD[{tensor_name}][(None, None, None, d_buffer)] = {self.tRS_sD[tensor_name][(None, None, None, d_buffer)]!s}" + ) + + self.efc.foreach_written_tensor(f) + + @trace_in_mlir + def tma_store_written_tensors_to_gmem(self, d_buffer, subtile_idx): + """Store with TMA the written EFC tensors to global memory.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.parameter[f"{tensor_name}_tma_atom_written"], + self.bSG_sD[tensor_name][(None, d_buffer)], + self.bSG_gD[tensor_name][(None, subtile_idx)], + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.parameter[{tensor_name}_tma_atom_written] = {self.parameter[f'{tensor_name}_tma_atom_written']!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_sD[{tensor_name}] = {self.bSG_sD[tensor_name]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_sD[{tensor_name}][(None, d_buffer)] = {self.bSG_sD[tensor_name][(None, d_buffer)]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_gD[{tensor_name}] = {self.bSG_gD[tensor_name]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_gD[{tensor_name}][(None, subtile_idx)] = {self.bSG_gD[tensor_name][(None, subtile_idx)]!s}" + ) + + self.efc.foreach_written_tensor(f) + + # The following functions are executed by the specialized warp for the + # epilogue load. + + @trace_in_mlir + def create_epilogue_subtile_tensors(self, tidx, epi_tile): + """Setup the pipelines reading the EFC supplemental tensors.""" + self.bGS_sC = {} + self.bGS_gC_partitioned = {} + + def f(tensor_name, attributes): + ( + _, + self.bGS_sC[tensor_name], + self.bGS_gC_partitioned[tensor_name], + ) = self.gemm.epilogue_gmem_copy_and_partition( + tidx, + self.parameter[f"{tensor_name}_tma_atom_read"], + self.tCgC_read[tensor_name], + epi_tile, + self.smem_read[tensor_name], + self.efc.jit.tensor_dtype[tensor_name], + ) + + self.efc.foreach_read_tensor(f) + + @trace_in_mlir + def prepare_tensor_load_for_subtiles( + self, + mma_tile_coord_mnl, + ): + """Prepare the EFC tensors to be loaded by the subtiles and return the number of subtiles to compute.""" + self.bGS_gC = {} + # In the case there is no supplemental tensor to load in the + # epilogue: + self._subtile_cnt = 0 + + def f(tensor_name, attributes): + self.bGS_gC[tensor_name] = self.bGS_gC_partitioned[tensor_name][ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + self.bGS_gC[tensor_name] = cutlass.cute.group_modes( + self.bGS_gC[tensor_name], + 1, + cutlass.cute.rank(self.bGS_gC[tensor_name]), + ) + st_cnt = cutlass.cute.size(self.bGS_gC[tensor_name].shape, mode=[1]) + if self._subtile_cnt == 0: + # Keep the first loaded tensor as a reference. + self._subtile_cnt = st_cnt + if st_cnt != self._subtile_cnt: + raise NotImplementedError( + f"Subtile count mismatch: tensor '{self.efc.read_tensor_names[0]}' has {self._subtile_cnt} subtiles, " + f"but tensor '{tensor_name}' has {st_cnt} subtiles. All tensors must have the same subtile count." + ) + + self.efc.foreach_read_tensor(f) + + return self._subtile_cnt + + @trace_in_mlir + def load_tensor_subtiles( + self, subtile_idx, c_pipeline, c_pipeline_producer_state + ): + """Load the subtiles of the EFC tensors.""" + + def f(tensor_name, attributes): + # Load supplemental tensor from global memory to shared memory. + cutlass.cute.copy( + self.parameter[f"{tensor_name}_tma_atom_read"], + self.bGS_gC[tensor_name][(None, subtile_idx)], + self.bGS_sC[tensor_name][(None, c_pipeline_producer_state.index)], + tma_bar_ptr=c_pipeline.producer_get_barrier( + c_pipeline_producer_state + ), + ) + + self.efc.foreach_read_tensor(f) + + class Phase(enum.Enum): + ParameterAnalysis = enum.auto() + """Epilogue function during analysis of its parameters.""" + + ThreadOperation = enum.auto() + """Epilogue function used for computation.""" + + PyTorchEvaluation = enum.auto() + """Epilogue function used for verification on CPU with PyTorch.""" + + class Tensor: + """A proxy object to be used as an argument to introspect or execute the + epilogue configuration function in a given phase.""" + + @dataclasses.dataclass + class ParameterAttributes: + """Store some characteristics of the epilogue parameters""" + + is_tensor: bool # Tensor or scalar + is_read: bool = False + is_written: bool = False + + def __init__( + self, + phase: typing.ForwardRef("EFC.Phase"), + name: str, + efc: EFC, + configuration, + ): + self.phase = phase + self.name = name + self.efc = efc + self.configuration = configuration + self.attributes: EFC.Tensor.ParameterAttributes = efc.parameter_attributes[ + name + ] + logger.info(f"Tensor {self.name = }") + + def load(self): + """""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Record that the tensor is read: + self.attributes.is_read = True + # Some value to have expression evaluation happy + return 1 + + case EFC.Phase.ThreadOperation: + # arg[0] is the epilogue_context from epilogue_computation(). + return ( + self.configuration.args[0] + .load[self.name] + .to(self.efc.gemm.epi_dtype) + ) + + case EFC.Phase.PyTorchEvaluation: + # args[1] is VariadicParameters constructed in + # evaluate_on_cpu(). Use .arg and not .parameter since it is + # not used actually to handle variadic parameter passing + # here. Just return the PyTorch tensor. + # TODO: Need to map to matching cutlass type. + return self.configuration.args[1].arg[self.name] + + case _: + raise NotImplementedError( + f"load({self.name}) not implemented for phase {self.phase}" + ) + + def store(self, value): + """""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Record that the tensor is written: + self.attributes.is_written = True + + case EFC.Phase.ThreadOperation: + # arg[0] is the epilogue_context from epilogue_computation(). + tRS_rD = self.configuration.args[0].store[self.name] + tRS_rD.store(value.to(self.efc.jit.tensor_dtype[self.name])) + + case EFC.Phase.PyTorchEvaluation: + # args[1] is VariadicParameters constructed in + # evaluate_on_cpu(). Use .arg and not .parameter since it is + # not used actually to handle variadic parameter passing + # here. Assign the PyTorch tensor target with the given + # value. + self.configuration.args[1].arg[self.name].copy_(value) + + case _: + raise NotImplementedError( + f"store({self.name}) not implemented for phase {self.phase}" + ) + + class Configuration: + """Specialize the epilogue provided by the user to be called in the + compute kernel customization point at a given phase.""" + + def __init__(self, efc: EFC, phase: EFC.Phase, *args): + """""" + self.efc = efc + self.phase = phase + # args[0] is the epilogue_context from the kernel for EFC.Phase.ThreadOperation. + self.args = args + self.arguments = [ + self._argument(name) for name in efc.epilogue_parameter_names + ] + + def _argument(self, name): + """Generate the argument used by the specialized epilogue with the + given name""" + if self.efc.parameter_attributes[name].is_tensor: + # Delegate the phase-related behavior to the Tensor object + # itself. + return EFC.Tensor(self.phase, name, self.efc, self) + # Otherwise, we have a dynamic scalar parameter. + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Use some dummy value during introspection phase. + return cutlass.Float32(42).to(self.efc.gemm.epi_dtype) + + case EFC.Phase.ThreadOperation: + # TODO: Need to map to matching cutlass type. + # Return directly the real kernel parameter with the same name. + return cutlass.Float32(self.efc.kernel.parameter[name]).to( + self.efc.gemm.epi_dtype + ) + + case EFC.Phase.PyTorchEvaluation: + # args[1] is VariadicParameters constructed in + # evaluate_on_cpu(). Use .arg and not .parameter since it is + # not used actually to handle variadic parameter passing + # here. + # TODO: Need to map to matching cutlass type. + return self.args[1].arg[name] + + case _: + raise NotImplementedError( + f"argument({name}) not implemented for phase {self.phase}" + ) + + def __call__(self): + """Execute the epilogue provided by the end-user with some specific + arguments crafted for the current phase. + + Pass self as an argument, to be seen as `efc_config`, a way to + access the EFC instance and its properties.""" + return self.efc.epilogue_function_configuration(self, *self.arguments) + + def accum(self): + """Provide the accumulator value to the user.""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + # The answer to anything. + return cutlass.Float32(42) + + case EFC.Phase.ThreadOperation: + # args[0] is epilogue_context passed to + # efc.kernel.epilogue_computation(). + return self.args[0].acc_vec + + case EFC.Phase.PyTorchEvaluation: + # Return matrix_multiplication_ref from evaluate_on_cpu(). + return self.args[0] + + case _: + raise NotImplementedError( + f"accum() not implemented for phase {self.phase}" + ) + + # Some helper functions for common operations. + + def maximum(self, x, y): + """Element-wise maximum of 2 tensors""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.maximum(x, y) + case EFC.Phase.PyTorchEvaluation: + return torch.maximum(x, y) + case _: + raise NotImplementedError( + f"maximum() not implemented for phase {self.phase}" + ) + + def minimum(self, x, y): + """Element-wise minimum of 2 tensors""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.minimum(x, y) + case EFC.Phase.PyTorchEvaluation: + return torch.minimum(x, y) + case _: + raise NotImplementedError( + f"minimum() not implemented for phase {self.phase}" + ) + + # Define some activation functions inspired by: + # - cutlass/python/cutlass_cppgen/epilogue/epilogue.py + # - cutlass/python/cutlass_cppgen/backend/epilogue.py + + def identity(self, x): + """Identity activation function: f(x) = x""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return x + case EFC.Phase.PyTorchEvaluation: + return x + case _: + raise NotImplementedError( + f"identity() not implemented for phase {self.phase}" + ) + + def relu(self, x): + """ReLU activation function: f(x) = maximum(0, x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.maximum(x, self.full_like(x, 0)) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.relu(x) + case _: + raise NotImplementedError( + f"relu() not implemented for phase {self.phase}" + ) + + def leaky_relu(self, x, negative_slope=0.01): + """Leaky ReLU activation function: f(x) = maximum(0, x) + negative_slope * minimum(0, x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # same type as x element type. + zero = self.full_like(x, 0) + return EFC.maximum(x, zero) + EFC.minimum(x, zero) * self.full_like( + x, negative_slope + ) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.leaky_relu(x, negative_slope) + case _: + raise NotImplementedError( + f"leaky_relu() not implemented for phase {self.phase}" + ) + + def tanh(self, x): + """Hyperbolic tangent activation function""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return cutlass.cute.tanh(x) + case EFC.Phase.PyTorchEvaluation: + return torch.tanh(x) + case _: + raise NotImplementedError( + f"tanh() not implemented for phase {self.phase}" + ) + + def sigmoid(self, x): + """Sigmoid activation function: f(x) = 1 / (1 + exp(-x))""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # same type as x element type. + # sigmoid(x) = 1 / (1 + exp(-x)) + return self.full_like(x, 1) / ( + self.full_like(x, 1) + cutlass.cute.exp(-x) + ) + case EFC.Phase.PyTorchEvaluation: + return torch.sigmoid(x) + case _: + raise NotImplementedError( + f"sigmoid() not implemented for phase {self.phase}" + ) + + def silu(self, x): + """SiLU (Swish) activation function: f(x) = x * sigmoid(x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # silu(x) = x * sigmoid(x) + return x * self.sigmoid(x) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.silu(x) + case _: + raise NotImplementedError( + f"silu() not implemented for phase {self.phase}" + ) + + def hardswish(self, x): + """Hard Swish activation function: f(x) = x * relu6(x + 3) / 6""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # same type as x element type. + # hardswish(x) = x * minimum(maximum(x + 3, 0), 6) / 6 + relu6 = EFC.minimum( + EFC.maximum(x + self.full_like(x, 3), self.full_like(x, 0)), + self.full_like(x, 6), + ) + return x * relu6 / self.full_like(x, 6) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.hardswish(x) + case _: + raise NotImplementedError( + f"hardswish() not implemented for phase {self.phase}" + ) + + def gelu(self, x): + """GELU (Gaussian Error Linear Unit) activation function.""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # same type as x element type. + # GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) + # Using a simpler approximation for CuTe + sqrt_2_over_pi = self.full_like(x, 0.7978845608028654) + return ( + self.full_like(x, 0.5) + * x + * ( + self.full_like(x, 1) + + cutlass.cute.tanh( + sqrt_2_over_pi + * (x + self.full_like(x, 0.044715) * x * x * x) + ) + ) + ) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.gelu(x) + case _: + raise NotImplementedError( + f"gelu() not implemented for phase {self.phase}" + ) + + def __getattr__(self, name): + """Called when an attribute or method is not found. + + Hijack this mechanism to dispatch/emulate functions like + cute.full_like() or torch.full_like() provided inside the epilogue + function as self.full_like(). + + This is required since the epilogue is used not only in a @cute.jit + or @cute.kernel but also executed in a normal context for analyzing + the epilogue content and even run in emulation with frameworks like + PyTorch.""" + + def chameleon(self, *args, **kwargs): + """The great impostor method. + + TODO: add some level of configuration to tweak the CuTe/Python + name mapping, handle some specific default values for some + parameters...""" + + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Just return a value to go on with the fake evaluation, in + # the case the function is expected to return a result. It + # will be ignored anyway in the opposite case. + return 1 + + case EFC.Phase.ThreadOperation: + # In the @cute.kernel context, just use the normal CuTe + # implementation. + return getattr(cutlass.cute, name)(*args, **kwargs) + + case EFC.Phase.PyTorchEvaluation: + # In the PyTorch context, call the equivalent function + # with the same name. + return getattr(torch, name)(*args, **kwargs) + + case _: + raise NotImplementedError( + f"self.{name} not implemented for phase {self.phase}" + ) + + # Update the function name to match the requested attribute name. + chameleon.__name__ = name + # Return chameleon blessed as a bound method of self. + return types.MethodType(chameleon, self) + + def __init__( + self, + gemm, + epilogue_function_configuration, + ): + """Construct an EFC instance.""" + self.gemm = gemm + self.epilogue_function_configuration = epilogue_function_configuration + self.analyze_epilogue(epilogue_function_configuration) + + def analyze_epilogue(self, epilogue_function_configuration): + """Analyze the epilogue configuration function to extract its parameter + names.""" + sig = inspect.signature(epilogue_function_configuration) + names = [name for name in sig.parameters.keys()] + # Impose to have the first parameter named "efc_config". This is very + # intrusive but at the same time some people got confused when they + # forgot this parameter. + if names[0] != "efc_config": + raise RuntimeError( + "The epilogue configuration function must take efc_config as an argument" + ) + + # Keep all the argument names but the first "efc_config" one. + self.epilogue_parameter_names = names[1:] + logger.info(f"{self.epilogue_parameter_names = }") + + def compile(self, supplemental_arguments): + """Compile with all the arguments to know the types during compilation + while hiding the epilogue detail1s.""" + assert len(supplemental_arguments) == len(self.epilogue_parameter_names) + # Update the active epilogue instance to use the new Parameter class + self.analyze_epilogue_with_arguments(supplemental_arguments) + # Create the metaprogramming objects for the @cute.jit and @cute.kernel + # parts. For now just forward all the parameters as is. + self.jit = EFC.JIT(self, self.epilogue_parameter_names) + if not self.written_tensor_names: + raise NotImplementedError( + "The epilogue requires at least one written tensor to do something useful." + ) + + def analyze_epilogue_with_arguments(self, supplemental_arguments): + self.parameter_attributes = {} + logger.info(f"{self.analyze_epilogue_with_arguments}:") + for name, a in zip(self.epilogue_parameter_names, supplemental_arguments): + logger.info(f"{name = } {a = }, {type(a) = }") + self.parameter_attributes[name] = EFC.Tensor.ParameterAttributes( + is_tensor=isinstance(a, cutlass.cute.Tensor) + ) + + # Evaluate the epilogue function for parameter analysis + self.specialized_epilogue(EFC.Phase.ParameterAnalysis)() + logger.info(f"\t{self.parameter_attributes = }") + # Keep track of all the epilogue tensor use cases per name: + self.used_tensor_names = [] + self.read_tensor_names = [] + self.written_tensor_names = [] + for name in self.epilogue_parameter_names: + q = self.parameter_attributes[name] + if not q.is_tensor: + continue + if q.is_read or q.is_written: + self.used_tensor_names.append(name) + if q.is_read: + self.read_tensor_names.append(name) + if q.is_written: + self.written_tensor_names.append(name) + logger.info( + f"\t{self.used_tensor_names = }\n\t{self.read_tensor_names = }\n\t{self.written_tensor_names = }" + ) + + def specialized_epilogue(self, phase: typing.ForwardRef("EFC.Phase"), *args): + """Construct a configuration of the epilogue specialized for a given + phase. The arguments are opaque and depend on the actual phase use.""" + return EFC.Configuration(self, phase, *args) + + def foreach_argument(self, function): + """Execute the given function for each supplemental argument of the epilogue.""" + for name in self.epilogue_parameter_names: + attributes = self.parameter_attributes[name] + function(name, attributes) + + def foreach_tensor(self, function): + """Execute the given function for each supplemental tensor.""" + for tensor_name in self.used_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def foreach_read_tensor(self, function): + """Execute the given function for each supplemental read tensor.""" + + for tensor_name in self.read_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def foreach_written_tensor(self, function): + """Execute the given function for each supplemental written tensor.""" + + for tensor_name in self.written_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def evaluate_on_cpu(self, matrix_multiplication_ref, *args): + """Evaluate the epilogue fusion configuration function on CPU for + validation using the precomputed matrix multiplication result. + + Use PyTorch for now but could be whatever.""" + # Recycle the VariadicParameters class to map the arguments according to + # their names: + epilogue_args = VariadicParameters(self, self.epilogue_parameter_names) + epilogue_args.pack_arguments(*args) + # Evaluate the epilogue with PyTorch. The tensor arguments which are + # stored are also evaluated and this is how some results are returned. + self.specialized_epilogue( + EFC.Phase.PyTorchEvaluation, + matrix_multiplication_ref, + epilogue_args, + )() diff --git a/examples/python/CuTeDSL/blackwell/epilogue/custom_epilogue_dense_gemm.py b/examples/python/CuTeDSL/blackwell/epilogue/custom_epilogue_dense_gemm.py new file mode 100644 index 00000000..954cf480 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/custom_epilogue_dense_gemm.py @@ -0,0 +1,625 @@ +# 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 diff --git a/examples/python/CuTeDSL/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py b/examples/python/CuTeDSL/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py new file mode 100644 index 00000000..0516cebb --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py @@ -0,0 +1,403 @@ +# 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 +import torch + +import cutlass +import cutlass.cute.testing as testing +import cutlass.torch as cutlass_torch + +from common_dense_gemm_efc import DenseGemmEFC +import common_efc + +""" +A high-performance persistent batched dense GEMM (D = alpha * A * B + beta * C) 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 D is MxNxL, L is batch dimension, D can be row-major("N") or column-major("M") +- alpha and beta are float scalars + +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. + - Load C matrix from global memory (GMEM) to shared memory (SMEM) using TMA operations and then copied to registers (RMEM). + - Compute D = alpha * accumulator + beta * C. + - Type convert D matrix to output type. + - Store D matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with 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.py. + +.. code-block:: bash + + python examples/internal/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py \ + --ab_dtype Float16 --c_dtype Float16 --d_dtype Float16 --acc_dtype Float32 --epi_dtype Float32 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 2.0 --beta 1.0 --t_dtype Float32 --read_tensors 2 --written_tensors 3 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/internal/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py \ + --ab_dtype Float16 --c_dtype Float16 --d_dtype Float16 --acc_dtype Float32 --epi_dtype Float32 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 2.0 --beta 1.0 --t_dtype Float32 --read_tensors 2 --written_tensors 3 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + + +Constraints are same as dense_gemm.py: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), + see detailed valid dtype combinations in below SM100PersistentDenseGemmAlphaBetaKernel class documentation +* A/B tensor must have the same data type +* C/D tensor 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 A/B/C/D 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 +""" + + +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, + t_dtype: typing.Type[cutlass.Numeric], + alpha: float, + beta: float, + read_tensors: int, + written_tensors: int, + tolerance: float, +) -> str: + """Format test parameters as CLI arguments for synthetic_custom_epilogue_dense_gemm.py""" + + # Get base command from DenseGemmEFC class + base_cmd = DenseGemmEFC.format_as_cli_args( + "synthetic_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 synthetic epilogue-specific parameters + specific_args = ( + f" --alpha {alpha} " + f"--beta {beta} " + f"--t_dtype {DenseGemmEFC.dtype_name(t_dtype)} " + f"--read_tensors {read_tensors} " + f"--written_tensors {written_tensors}" + ) + + 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, + alpha: float, + beta: float, + t_dtype: typing.Type[cutlass.Numeric], + 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, + read_tensors: int = 1, + written_tensors: int = 1, + verbose: bool = False, +): + """ + Prepare A/B/C/D tensors, launch GPU kernel, and 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{alpha = }, {beta = }") + print(f"\t{t_dtype = !s}") + print(f"\t{read_tensors = }, {written_tensors = }") + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("A 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) + + def meta_epilogue(read_tensors, written_tensors): + """Build a synthetic epilogue function with parameters + (self, alpha, beta, read_t0, read_t1,..., read_t{read__tensors-1}, written_t0, written_t1,..., written_t{written_tensors-1}""" + + param_names = ( + ["efc_config", "alpha", "beta"] + + [f"read_t{i}" for i in range(read_tensors)] + + [f"written_t{i}" for i in range(written_tensors)] + ) + + assert written_tensors > 0, ( + "At least one tensor must be written in the epilogue." + ) + + def computation_impl(efc_config, alpha, beta, *tensors): + """Implementation of the epilogue computation.""" + read = beta + for tensor in tensors[:read_tensors]: + read += tensor.load() * alpha + if read_tensors > 0: + # Can use some CuTe/PyTorch-like functions exposed under + # efc_config namespace for portability: + read = efc_config.where( + read < 1, read, read * efc_config.full_like(read, 2) + ) + + t = efc_config.accum() + for tensor in tensors[read_tensors:]: + t = t * alpha + read + 5000 + tensor.store(t) + + # Wrap the implementation with a function with the correct parameter + # names. + return common_efc.create_named_epilogue(param_names, computation_impl) + + epilogue = meta_epilogue(read_tensors, written_tensors) + + # Build GEMM object with EFC configuration: + # TODO: generalize acc_dtype and epi_dtype + gemm = DenseGemmEFC( + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue, + ) + ( + a_tensor, + b_tensor, + a_torch_cpu, + b_torch_cpu, + ) = gemm.create_arguments(l, m, n, k, a_major, b_major, cd_major, ab_dtype) + + # Create all the supplemental tensors. + t_torch_cpu, t_torch_gpu, t_tensor = ([], [], []) + for i in range(read_tensors + written_tensors): + t_torch_cpu.append(cutlass_torch.matrix(l, m, n, cd_major == "m", t_dtype)) + tensor, torch_gpu = cutlass_torch.cute_tensor_like( + t_torch_cpu[i], t_dtype, is_dynamic_layout=True, assumed_align=16 + ) + t_tensor.append(tensor) + t_torch_gpu.append(torch_gpu) + + # Check if configuration can be implemented + gemm.check_implementable(a_tensor, b_tensor, t_tensor[0]) + + 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. + alpha, + beta, + *t_tensor, + ) + torch.cuda.synchronize() + + # 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. + alpha, + beta, + *t_tensor, + ), + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print(f"Execution time: {exec_time} us") + + # Evaluate the epilogue on the host: + gemm.evaluate_on_cpu( + a_torch_cpu, + b_torch_cpu, + epi_dtype, + # The EFC arguments: + alpha, + beta, + *t_torch_cpu, + ) + + # Print tensors if verbose mode is enabled + if verbose: + print("\n=== Read Tensors ===") + for i in range(read_tensors): + print(f"\nRead Tensor {i} (GPU):") + print(t_torch_gpu[i].cpu()) + + print("\n=== Written Tensors ===") + for i in range(written_tensors): + idx = read_tensors + i + print(f"\nWritten Tensor {i} (GPU):") + print(t_torch_gpu[idx].cpu()) + print() + + # Assert close results between the values computed on GPU and CPU. + for torch_gpu, torch_cpu in zip(t_torch_gpu, t_torch_cpu): + torch.testing.assert_close( + torch_gpu.cpu(), torch_cpu, atol=tolerance, rtol=1e-03 + ) + + +if __name__ == "__main__": + cli = DenseGemmEFC.CLIParser() + cli.parser.add_argument( + "--alpha", type=float, default=1.0, help="alpha scale factor" + ) + cli.parser.add_argument("--beta", type=float, default=0.0, help="beta scale factor") + cli.parser.add_argument("--t_dtype", type=cutlass.dtype, default=cutlass.Float32) + cli.parser.add_argument( + "--read_tensors", + type=int, + default=1, + help="number of tensors to read inside the epilogue", + ) + cli.parser.add_argument( + "--written_tensors", + type=int, + default=1, + help="number of tensors to write inside the epilogue", + ) + cli.parser.add_argument( + "--verbose", + action="store_true", + help="print read and written tensors", + ) + args = cli.parse() + + try: + run( + args.mnkl, + args.ab_dtype, + args.acc_dtype, + args.epi_dtype, + args.a_major, + args.b_major, + args.cd_major, + args.alpha, + args.beta, + args.t_dtype, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.read_tensors, + args.written_tensors, + args.verbose, + ) + print("PASS") + except Exception as exc: + traceback.print_exception(exc) + raise diff --git a/examples/python/CuTeDSL/blackwell/fmha.py b/examples/python/CuTeDSL/blackwell/fmha.py index f52d9538..7b91051a 100644 --- a/examples/python/CuTeDSL/blackwell/fmha.py +++ b/examples/python/CuTeDSL/blackwell/fmha.py @@ -33,16 +33,14 @@ import sys import time from typing import Type, Tuple, Union, Optional -import torch -import torch.nn.functional as F import cuda.bindings.driver as cuda +import torch import cutlass import cutlass.cute as cute import cutlass.cute.nvgpu.tcgen05 as tcgen05 import cutlass.utils as utils import cutlass.pipeline as pipeline -import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.testing as testing from cutlass.cute.runtime import from_dlpack @@ -174,8 +172,7 @@ class BlackwellFusedMultiHeadAttentionForward: self.load_warp_id = 13 self.epilogue_warp_id = 14 self.empty_warp_id = 15 - SM100_TMEM_CAPACITY_COLUMNS = 512 - self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") self.threads_per_warp = 32 self.threads_per_cta = self.threads_per_warp * len( @@ -1637,6 +1634,8 @@ class BlackwellFusedMultiHeadAttentionForward: :type atom_args: tuple :param tensor_args: Tuple containing softmax related tensors :type tensor_args: tuple + :param fused_mask: Compute trip counts and apply masking for attention blocks + :type fused_mask: fmha_utils.FusedMask :return: Updated state values (row_max, row_sum, and pipeline related arguments) :rtype: tuple """ @@ -1729,20 +1728,16 @@ class BlackwellFusedMultiHeadAttentionForward: tTMEM_STORErS_x4_e, cute.make_layout(frg_tile) ) for j in range(frg_cnt): - for k in range(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2): - tTMEM_LOADrS_frg[k, j], tTMEM_LOADrS_frg[k + 1, j] = ( - cute.arch.fma_packed_f32x2( - (tTMEM_LOADrS_frg[k, j], tTMEM_LOADrS_frg[k + 1, j]), - (scale, scale), - (minus_row_max_scale, minus_row_max_scale), - ) + for k in cutlass.range( + cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True + ): + tTMEM_LOADrS_frg[k, j] = ( + tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale ) tTMEM_LOADrS_frg[k, j] = cute.math.exp2( tTMEM_LOADrS_frg[k, j], fastmath=True ) - tTMEM_LOADrS_frg[k + 1, j] = cute.math.exp2( - tTMEM_LOADrS_frg[k + 1, j], fastmath=True - ) + s_vec = tTMEM_LOADrS_frg[None, j].load() tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) # Sequence barrier arrive @@ -1859,6 +1854,8 @@ class BlackwellFusedMultiHeadAttentionForward: :type s0_s1_sequence_pipeline: pipeline.PipelineAsync :param tile_sched_params: Parameters for tile scheduling :type tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams + :param fused_mask: Compute trip counts and apply masking for attention blocks + :type fused_mask: fmha_utils.FusedMask """ tidx, _, _ = cute.arch.thread_idx() thread_idx = tidx % ( @@ -2204,11 +2201,8 @@ class BlackwellFusedMultiHeadAttentionForward: ) cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in range(0, cute.size(tTMrO_i), 2): - tTMrO_i[j], tTMrO_i[j + 1] = cute.arch.mul_packed_f32x2( - (tTMrO_i[j], tTMrO_i[j + 1]), - (scale, scale), - ) + for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): + tTMrO_i[j] = tTMrO_i[j] * scale cute.copy(tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) @cute.jit @@ -2310,11 +2304,8 @@ class BlackwellFusedMultiHeadAttentionForward: tTMEM_LOADoO[None, 0, 0, i].shape, self.pv_acc_dtype ) cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) - for j in range(0, cute.size(tTMrO), 2): - tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( - (tTMrO[j], tTMrO[j + 1]), - (scale, scale), - ) + for j in range(cute.size(tTMrO), vectorize=True): + tTMrO[j] = tTMrO[j] * scale tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) o_vec = tTMrO.load() tSMrO.store(o_vec.to(self.o_dtype)) @@ -2327,7 +2318,10 @@ class BlackwellFusedMultiHeadAttentionForward: mLSE[row_idx + cuseqlen_q, blk_coord[2]] = lse # fence view async shared - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) def run( @@ -2442,6 +2436,7 @@ def run( print(f" iterations: {iterations}") print(f" skip_ref_check: {skip_ref_check}") print(f" use_cold_l2: {use_cold_l2}") + import cutlass.torch as cutlass_torch # Unpack parameters b, s_q, h_q, d = q_shape @@ -2950,7 +2945,7 @@ def run( else: lse_tensor = None - return testing.JitArguments( + args = testing.JitArguments( q_tensor_workspace.iterator, k_tensor_workspace.iterator, v_tensor_workspace.iterator, @@ -2970,6 +2965,15 @@ def run( ), current_stream, ) + args.add_to_scope( + [ + q_tensor_workspace, + k_tensor_workspace, + v_tensor_workspace, + o_tensor_workspace, + ] + ) + return args workspace_count = 1 if use_cold_l2: diff --git a/examples/python/CuTeDSL/blackwell/fmha_bwd.py b/examples/python/CuTeDSL/blackwell/fmha_bwd.py index f40af955..5725d5e8 100644 --- a/examples/python/CuTeDSL/blackwell/fmha_bwd.py +++ b/examples/python/CuTeDSL/blackwell/fmha_bwd.py @@ -166,8 +166,7 @@ class BlackwellFusedMultiHeadAttentionBackward: self.num_reduce_warps = 4 self.num_compute_warps = 8 - SM100_TMEM_CAPACITY_COLUMNS = 512 - self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") self.threads_per_warp = 32 self.threads_per_cta = self.threads_per_warp * ( @@ -2093,7 +2092,6 @@ class BlackwellFusedMultiHeadAttentionBackward: cute.arch.fence_view_async_tmem_load() self.compute_sync_barrier.arrive_and_wait() - cute.arch.fence_view_async_tmem_load() cute.copy(tiled_r2t, tRT_rST_reshaped, tRT_tP) @@ -2161,7 +2159,10 @@ class BlackwellFusedMultiHeadAttentionBackward: cute.autovec_copy(tTR_rdST, sdS_slice) # Notify for dS - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) compute_mma_dS_pipeline.producer_commit(compute_mma_dS_producer_state) compute_mma_dS_producer_state.advance() @@ -2279,7 +2280,10 @@ class BlackwellFusedMultiHeadAttentionBackward: ) # Wait for the stores to all be visible to the TMA - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.reduce_sync_barrier.arrive_and_wait() if warp_idx == 0: @@ -2489,11 +2493,13 @@ class BlackwellFusedMultiHeadAttentionBackward: mma_compute_dKdV_pipeline.consumer_wait(mma_compute_dKdV_consumer_state) + # Load tdKtdK cute.copy(tiled_t2r_dK, tTR_tdK, tTR_rdK) for i in cutlass.range(cute.size(tTR_rdK), unroll_full=True): tTR_rdK[i] = scale_softmax * tTR_rdK[i] + # Store tdKgdK self.store(tTR_gdK, tTR_rdK, tTR_cdK, (K, D)) cute.arch.fence_view_async_tmem_load() diff --git a/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py index 7cc59105..0f2dd5fa 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py @@ -193,8 +193,7 @@ class Sm100GroupedBlockScaledGemmKernel: num_threads=64, ) self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") - SM100_TMEM_CAPACITY_COLUMNS = 512 - self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") # Set up configurations that dependent on gemm inputs. def _setup_attributes(self): @@ -423,6 +422,7 @@ class Sm100GroupedBlockScaledGemmKernel: self.b_dtype = initial_b.element_type self.sf_dtype = initial_sfa.element_type self.c_dtype = initial_c.element_type + self.is_nvfp4_output = self.c_dtype is cutlass.Float4E2M1FN self.a_major_mode = utils.LayoutEnum.from_tensor(initial_a).mma_major_mode() self.b_major_mode = utils.LayoutEnum.from_tensor(initial_b).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(initial_c) @@ -893,7 +893,7 @@ class Sm100GroupedBlockScaledGemmKernel: cute.group_modes(tCgB, 0, 3), ) - # TMA Load SFA partition_S/D + # TMA load scaled factor A partition_S/D sfa_cta_layout = a_cta_layout # ((atom_v, rest_v), STAGE) # ((atom_v, rest_v), RestM, RestK, RestL) @@ -907,7 +907,7 @@ class Sm100GroupedBlockScaledGemmKernel: tAsSFA = cute.filter_zeros(tAsSFA) tAgSFA = cute.filter_zeros(tAgSFA) - # TMA Load SFB partition_S/D + # TMA load scaled factor B partition_S/D sfb_cta_layout = cute.make_layout( cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape ) @@ -970,222 +970,239 @@ class Sm100GroupedBlockScaledGemmKernel: tensormaps[(tensormap_workspace_idx, 4, None)].iterator ) + # + # Persistent tile scheduling loop + # + # When the problem shapes are on device, we launch one CTA per SM. + # The if condition later prevents the warps from extra CTAs from doing any work. + tile_sched = utils.StaticPersistentGroupTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + grid_dim, + self.cluster_tile_shape_mnk, + utils.create_initial_search_state(), + group_count, + problem_sizes_mnkl, + ) + initial_work_tile_info = tile_sched.initial_work_tile_info() + # # Specialized TMA load warp # - if warp_idx == self.tma_warp_id: + if warp_idx == self.tma_warp_id and initial_work_tile_info.is_valid_tile: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) + work_tile = initial_work_tile_info + tensormap_init_done = cutlass.Boolean(False) # group index of last tile last_group_idx = cutlass.Int32(-1) - 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: - cur_tile_coord = work_tile.tile_idx - grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z( - cur_tile_coord, - problem_sizes_mnkl, - ) + grouped_gemm_cta_tile_info = work_tile.group_search_result cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k cur_group_idx = grouped_gemm_cta_tile_info.group_idx - is_group_changed = cur_group_idx != last_group_idx - # skip tensormap update if we're working on the same group - if is_group_changed: - real_tensor_a = self.make_tensor_abc_for_tensormap_update( - cur_group_idx, - self.a_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - strides_abc, - ptrs_abc, - 0, # 0 for tensor A - ) - real_tensor_b = self.make_tensor_abc_for_tensormap_update( - cur_group_idx, - self.b_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - strides_abc, - ptrs_abc, - 1, # 1 for tensor B - ) - real_tensor_sfa = self.make_tensor_sfasfb_for_tensormap_update( - cur_group_idx, - self.sf_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - ptrs_sfasfb, - 0, # 0 for tensor SFA - ) - real_tensor_sfb = self.make_tensor_sfasfb_for_tensormap_update( - cur_group_idx, - self.sf_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - ptrs_sfasfb, - 1, # 1 for tensor SFB - ) - if tensormap_init_done == False: - # wait tensormap initialization complete - self.tensormap_ab_init_barrier.arrive_and_wait() - tensormap_init_done = True + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 + # Do not load any data if cur_k_tile_cnt is 0 + if not is_k_tile_cnt_zero: + is_group_changed = cur_group_idx != last_group_idx + # skip tensormap update if we're working on the same group + if is_group_changed: + real_tensor_a = self.make_tensor_abc_for_tensormap_update( + cur_group_idx, + self.a_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 0, # 0 for tensor A + ) + real_tensor_b = self.make_tensor_abc_for_tensormap_update( + cur_group_idx, + self.b_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 1, # 1 for tensor B + ) + real_tensor_sfa = self.make_tensor_sfasfb_for_tensormap_update( + cur_group_idx, + self.sf_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + ptrs_sfasfb, + 0, # 0 for tensor SFA + ) + real_tensor_sfb = self.make_tensor_sfasfb_for_tensormap_update( + cur_group_idx, + self.sf_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + ptrs_sfasfb, + 1, # 1 for tensor SFB + ) + if not tensormap_init_done: + # wait tensormap initialization complete + self.tensormap_ab_init_barrier.arrive_and_wait() + tensormap_init_done = True - tensormap_manager.update_tensormap( - ( - real_tensor_a, - real_tensor_b, - real_tensor_sfa, - real_tensor_sfb, - ), - (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb), - ( - tensormap_a_gmem_ptr, - tensormap_b_gmem_ptr, - tensormap_sfa_gmem_ptr, - tensormap_sfb_gmem_ptr, - ), - self.tma_warp_id, - ( - tensormap_a_smem_ptr, - tensormap_b_smem_ptr, - tensormap_sfa_smem_ptr, - tensormap_sfb_smem_ptr, - ), + tensormap_manager.update_tensormap( + ( + real_tensor_a, + real_tensor_b, + real_tensor_sfa, + real_tensor_sfb, + ), + (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb), + ( + tensormap_a_gmem_ptr, + tensormap_b_gmem_ptr, + tensormap_sfa_gmem_ptr, + tensormap_sfb_gmem_ptr, + ), + self.tma_warp_id, + ( + tensormap_a_smem_ptr, + tensormap_b_smem_ptr, + tensormap_sfa_smem_ptr, + tensormap_sfb_smem_ptr, + ), + ) + + mma_tile_coord_mnl = ( + grouped_gemm_cta_tile_info.cta_tile_idx_m + // cute.size(tiled_mma.thr_id.shape), + grouped_gemm_cta_tile_info.cta_tile_idx_n, + 0, ) - mma_tile_coord_mnl = ( - grouped_gemm_cta_tile_info.cta_tile_idx_m - // cute.size(tiled_mma.thr_id.shape), - grouped_gemm_cta_tile_info.cta_tile_idx_n, - 0, - ) + # + # 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]) + ] - # - # 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]) - ] + # ((atom_v, rest_v), RestK) + tAgSFA_slice = tAgSFA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] - # ((atom_v, rest_v), RestK) - tAgSFA_slice = tAgSFA[ - (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) - ] - # ((atom_v, rest_v), RestK) - tBgSFB_slice = tBgSFB[ - (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 < cur_k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire( - ab_producer_state - ) - - if is_group_changed: - tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr) - tensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr) - tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr) - tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr) - # - # Tma load loop - # - for k_tile in cutlass.range(0, cur_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/SFA/SFB - 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, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_a_gmem_ptr, - cute.AddressSpace.generic, - ), - ) - 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, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_b_gmem_ptr, - cute.AddressSpace.generic, - ), - ) - cute.copy( - tma_atom_sfa, - tAgSFA_slice[(None, ab_producer_state.count)], - tAsSFA[(None, ab_producer_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), - mcast_mask=sfa_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_sfa_gmem_ptr, - cute.AddressSpace.generic, - ), - ) - cute.copy( - tma_atom_sfb, - tBgSFB_slice[(None, ab_producer_state.count)], - tBsSFB[(None, ab_producer_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), - mcast_mask=sfb_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_sfb_gmem_ptr, - cute.AddressSpace.generic, - ), - ) - - # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 - ab_producer_state.advance() + # 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 < cur_k_tile_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire( ab_producer_state ) + if is_group_changed: + tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr) + # + # Tma load loop + # + for k_tile in cutlass.range(0, cur_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/SFA/SFB + 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, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_a_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + 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, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_b_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier( + ab_producer_state + ), + mcast_mask=sfa_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_sfa_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier( + ab_producer_state + ), + mcast_mask=sfb_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_sfb_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + + # 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 < cur_k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + else: + if not tensormap_init_done: + # wait tensormap initialization complete + self.tensormap_ab_init_barrier.arrive_and_wait() + tensormap_init_done = True # # Advance to next tile # @@ -1201,7 +1218,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Specialized MMA warp # - if warp_idx == self.mma_warp_id: + if warp_idx == self.mma_warp_id and initial_work_tile_info.is_valid_tile: # # Initialize tensormaps for A, B, SFA and SFB # @@ -1279,18 +1296,8 @@ class Sm100GroupedBlockScaledGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) + work_tile = initial_work_tile_info - work_tile = tile_sched.initial_work_tile_info() ab_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_ab_stage ) @@ -1298,15 +1305,14 @@ class Sm100GroupedBlockScaledGemmKernel: pipeline.PipelineUserType.Producer, self.num_acc_stage ) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx + cur_group_idx = work_tile.group_search_result.group_idx + problem_shape_k = work_tile.group_search_result.problem_shape_k + # MMA warp is only interested in number of tiles along K dimension - ( - cur_k_tile_cnt, - cur_group_idx, - ) = group_gemm_ts_helper.search_cluster_tile_count_k( - cur_tile_coord, - problem_sizes_mnkl, - ) + cur_k_tile_cnt = ( + problem_shape_k + self.cluster_tile_shape_mnk[2] - 1 + ) // self.cluster_tile_shape_mnk[2] + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 # (MMA, MMA_M, MMA_N) tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] @@ -1322,7 +1328,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Wait for accumulator buffer empty # - if is_leader_cta: + if is_leader_cta and not is_k_tile_cnt_zero: acc_pipeline.producer_acquire(acc_producer_state) # @@ -1408,9 +1414,10 @@ class Sm100GroupedBlockScaledGemmKernel: # # Async arrive accumulator buffer full # - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() + if not is_k_tile_cnt_zero: + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() # # Advance to next tile @@ -1426,7 +1433,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Specialized epilogue warps # - if warp_idx < self.mma_warp_id: + if warp_idx < self.mma_warp_id and initial_work_tile_info.is_valid_tile: # initialize tensorap for C tensormap_manager.init_tensormap_from_atom( tma_atom_c, @@ -1483,18 +1490,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) - - work_tile = tile_sched.initial_work_tile_info() + work_tile = initial_work_tile_info acc_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_acc_stage @@ -1513,14 +1509,13 @@ class Sm100GroupedBlockScaledGemmKernel: last_group_idx = cutlass.Int32(-1) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z( - cur_tile_coord, - problem_sizes_mnkl, - ) + grouped_gemm_cta_tile_info = work_tile.group_search_result cur_group_idx = grouped_gemm_cta_tile_info.group_idx + cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 is_group_changed = cur_group_idx != last_group_idx + # We still need to store 0s when k_tile_cnt is 0 if is_group_changed: # construct tensor c based on real shape, stride information real_tensor_c = self.make_tensor_abc_for_tensormap_update( @@ -1549,7 +1544,6 @@ class Sm100GroupedBlockScaledGemmKernel: grouped_gemm_cta_tile_info.cta_tile_idx_n, 0, ) - cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k # # Slice to per mma tile index @@ -1573,7 +1567,8 @@ class Sm100GroupedBlockScaledGemmKernel: # # Wait for accumulator buffer full # - acc_pipeline.consumer_wait(acc_consumer_state) + if not is_k_tile_cnt_zero: + 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)) @@ -1588,17 +1583,34 @@ class Sm100GroupedBlockScaledGemmKernel: subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt for subtile_idx in 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 not is_k_tile_cnt_zero: + # + # 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() - tRS_rC.store(acc_vec.to(self.c_dtype)) + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + tRS_rC.store(acc_vec.to(self.c_dtype)) + else: + if cutlass.const_expr(self.is_nvfp4_output): + zeros_i8 = cute.make_rmem_tensor( + cute.recast_layout( + cutlass.Int8.width, + self.c_dtype.width, + tRS_rC.layout, + ), + cutlass.Int8, + ) + zeros_i8.fill(0) + tRS_rC.store( + cute.recast_tensor(zeros_i8, self.c_dtype).load() + ) + else: + tRS_rC.fill(0) # # Store C to shared memory @@ -1610,7 +1622,10 @@ class Sm100GroupedBlockScaledGemmKernel: 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("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1633,9 +1648,10 @@ class Sm100GroupedBlockScaledGemmKernel: # # Async arrive accumulator buffer empty # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() + if not is_k_tile_cnt_zero: + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() # # Advance to next tile @@ -2442,13 +2458,6 @@ def create_tensor_and_stride( torch_tensor_cpu, dtype, is_dynamic_layout, assumed_align=16 ) - # Mark tensor with element divisibility for 16B alignment - cute_tensor.mark_compact_shape_dynamic( - mode=0 if is_mode0_major else 1, - stride_order=(2, 1, 0) if is_mode0_major else (2, 0, 1), - divisibility=32 if dtype == cutlass.Float4E2M1FN else 16, - ) - # omit stride for L mode as it is always 1 stride = (1, mode0) if is_mode0_major else (mode1, 1) @@ -2552,7 +2561,7 @@ def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): def ceil_div(a, b): return (a + b - 1) // b - sf_k = ceil_div(k, sf_vec_size) + sf_k = max(1, ceil_div(k, sf_vec_size)) ref_shape = (l, mn, sf_k) atom_m = (32, 4) @@ -2675,6 +2684,7 @@ def create_tensors_sfasfb_for_all_groups( def run( num_groups: int, problem_sizes_mnkl: List[Tuple[int, int, int, int]], + host_problem_shape_available: bool, ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, @@ -2761,22 +2771,40 @@ def run( sf_vec_size, ) - # Choose A, B, C, SFA, SFB with the smallest size to create initial tensormaps - key_size_a = lambda item: item[1][0] * item[1][2] - key_size_b = lambda item: item[1][1] * item[1][2] - key_size_c = lambda item: item[1][0] * item[1][1] - # Find the indices of the groups with the smallest tensor sizes - min_a_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_a) - min_b_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_b) - min_c_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_c) + # Setup inital tensors for TMA of A,B and C + alignment = 16 # 16 bytes aligned + divisibility_ab = 32 if ab_dtype == cutlass.Float4E2M1FN else 16 + divisibility_c = 32 if c_dtype == cutlass.Float4E2M1FN else 16 + divisibility_sf = 32 if sf_dtype == cutlass.Float4E2M1FN else 16 + + min_ab_size = alignment * 8 // ab_dtype.width # alignment bytes of width + div_mul_ab = (divisibility_ab + min_ab_size - 1) // min_ab_size + min_ab_size = min_ab_size * div_mul_ab + + min_c_size = alignment * 8 // c_dtype.width + div_mul_c = (divisibility_c + min_c_size - 1) // min_c_size + min_c_size = min_c_size * div_mul_c + + min_sf_size = alignment * 8 // sf_dtype.width + div_mul_sf = (divisibility_sf + min_sf_size - 1) // min_sf_size + min_sf_size = min_sf_size * div_mul_sf + initial_cute_tensors_abc = [ - cute_tensors_abc[min_a_idx][0], # A with smallest (m, k) - cute_tensors_abc[min_b_idx][1], # B with smallest (n, k) - cute_tensors_abc[min_c_idx][2], # C with smallest (m, n) + create_tensor_and_stride(1, min_ab_size, min_ab_size, a_major == "m", ab_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_ab_size, min_ab_size, b_major == "n", ab_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_c_size, min_c_size, c_major == "m", c_dtype)[2], ] initial_cute_tensors_sfasfb = [ - cute_tensors_sfasfb[min_a_idx][0], # SFA with smallest (m, k)'s group - cute_tensors_sfasfb[min_b_idx][1], # SFB with smallest (n, k)'s group + create_tensor_and_stride(1, min_sf_size, min_sf_size, a_major == "m", sf_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_sf_size, min_sf_size, b_major == "n", sf_dtype)[ + 2 + ], ] hardware_info = cutlass.utils.HardwareInfo() @@ -2867,6 +2895,19 @@ def run( # Initialize Stream current_stream = cutlass_torch.default_stream() + # If the host problem shape is available, we will launch the grid with only + # the necessary clusters. The function compute_total_num_clusters() does that. + # If the problem shape only exists on device, we will need to launch all active + # clusters possible on a device. + if host_problem_shape_available: + print("Problem shapes available on host and device") + total_num_clusters = compute_total_num_clusters( + problem_sizes_mnkl, cluster_tile_shape_mn + ) + else: + print("Problem shapes available only on device") + total_num_clusters = max_active_clusters + # Compile grouped GEMM kernel compiled_grouped_gemm = cute.compile( grouped_blockscaled_gemm, @@ -2980,18 +3021,23 @@ def run( ) initial_cute_tensors_abc_workspace = [ - cute_tensors_abc_workspace[min_a_idx][0], # A with smallest (m, k) - cute_tensors_abc_workspace[min_b_idx][1], # B with smallest (n, k) - cute_tensors_abc_workspace[min_c_idx][2], # C with smallest (m, n) + create_tensor_and_stride( + 1, min_ab_size, min_ab_size, a_major == "m", ab_dtype + )[2], + create_tensor_and_stride( + 1, min_ab_size, min_ab_size, b_major == "n", ab_dtype + )[2], + create_tensor_and_stride( + 1, min_c_size, min_c_size, c_major == "m", c_dtype + )[2], ] - initial_cute_tensors_sfasfb_workspace = [ - cute_tensors_sfasfb_workspace[min_a_idx][ - 0 - ], # SFA with smallest (m, k)'s group - cute_tensors_sfasfb_workspace[min_b_idx][ - 1 - ], # SFB with smallest (n, k)'s group + create_tensor_and_stride( + 1, min_sf_size, min_sf_size, a_major == "m", sf_dtype + )[2], + create_tensor_and_stride( + 1, min_sf_size, min_sf_size, b_major == "n", sf_dtype + )[2], ] # Create new tensors for this workspace @@ -3022,7 +3068,7 @@ def run( is_dynamic_layout=False, ) - return cute.testing.JitArguments( + args = cute.testing.JitArguments( initial_cute_tensors_abc_workspace[0], initial_cute_tensors_abc_workspace[1], initial_cute_tensors_abc_workspace[2], @@ -3035,6 +3081,8 @@ def run( tensormap_workspace, current_stream, ) + args.add_to_scope([torch_tensors_abc_workspace, torch_tensors_sfasfb_workspace]) + return args workspace_count = 1 if use_cold_l2: @@ -3078,6 +3126,18 @@ def run( iterations=iterations, ) + runtime_s = exec_time / 1.0e6 + fmas = 0 + for group in range(num_groups): + [M, N, K, _] = problem_sizes_mnkl[group] + fmas += M * N * K + flop = 2 * fmas + gflop = flop / 1.0e9 + gflops = gflop / runtime_s + + print("Average Runtime : ", exec_time / 1000, "ms") + print("GFLOPS : ", gflops) + return exec_time # Return execution time in microseconds @@ -3138,6 +3198,11 @@ if __name__ == "__main__": default=(128, 128), help="Mma tile shape (comma-separated)", ) + parser.add_argument( + "--host_problem_shape_available", + action="store_true", + help="Enable the compute of grid based upon host problem shape", + ) parser.add_argument( "--cluster_shape_mn", type=parse_comma_separated_ints, @@ -3195,6 +3260,7 @@ if __name__ == "__main__": run( args.num_groups, args.problem_sizes_mnkl, + args.host_problem_shape_available, args.ab_dtype, args.sf_dtype, args.sf_vec_size, diff --git a/examples/python/CuTeDSL/blackwell/grouped_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_gemm.py index 1c63ac8b..50c208f2 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_gemm.py @@ -335,9 +335,11 @@ class GroupedGemmKernel: :type stream: cuda.CUstream :raises TypeError: If A and B data types do not match. """ + self.a_dtype = initial_a.element_type self.b_dtype = initial_b.element_type self.c_dtype = initial_c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(initial_a).mma_major_mode() self.b_major_mode = utils.LayoutEnum.from_tensor(initial_b).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(initial_c) @@ -475,6 +477,7 @@ class GroupedGemmKernel: block=[self.threads_per_cta, 1, 1], cluster=(*self.cluster_shape_mn, 1), stream=stream, + min_blocks_per_mp=1, ) return @@ -553,28 +556,38 @@ class GroupedGemmKernel: tensormap_c_smem_ptr = ( tensormap_b_smem_ptr + GroupedGemmKernel.bytes_per_tensormap // 8 ) - ab_full_mbar_ptr = storage.ab_full_mbar_ptr.data_ptr() - ab_empty_mbar_ptr = storage.ab_empty_mbar_ptr.data_ptr() - acc_full_mbar_ptr = storage.acc_full_mbar_ptr.data_ptr() - acc_empty_mbar_ptr = storage.acc_empty_mbar_ptr.data_ptr() # init barrier for loading A, B with TMA - if warp_idx == self.epilog_warp_id[0]: - for k_stage in range(self.num_ab_stage): - num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 - with cute.arch.elect_one(): - cute.arch.mbarrier_init(ab_full_mbar_ptr + k_stage, 1) - cute.arch.mbarrier_init( - ab_empty_mbar_ptr + k_stage, num_tma_producer - ) + 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, + defer_sync=True, + ) # Accumulator barrier init - if warp_idx == self.mma_warp_id: - for acc_stage in range(self.num_acc_stage): - with cute.arch.elect_one(): - cute.arch.mbarrier_init(acc_full_mbar_ptr + acc_stage, 1) - cute.arch.mbarrier_init( - acc_empty_mbar_ptr + acc_stage, 8 if use_2cta_instrs else 4 - ) + 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, + ) # Tensor memory dealloc barrier init tmem = utils.TmemAllocator( storage.tmem_holding_buf, @@ -747,10 +760,26 @@ class GroupedGemmKernel: tensormap_b_init_ptr = tensormap_b_ptr tensormap_c_init_ptr = tensormap_c_ptr + # + # Persistent tile scheduling loop + # + # When the problem shapes are on device, we launch one CTA per SM. + # The if condition later prevents the warps from extra CTAs from doing any work. + tile_sched = utils.StaticPersistentGroupTileScheduler.create( + tile_sched_params, + bid, + grid_dim, + self.cluster_tile_shape_mnk, + utils.create_initial_search_state(), + group_count, + problem_sizes_mnkl, + ) + initial_work_tile_info = tile_sched.initial_work_tile_info() + # # Specialized TMA load warp # - if warp_idx == self.tma_warp_id: + if warp_idx == self.tma_warp_id and initial_work_tile_info.is_valid_tile: # Initialize tensormaps for A, B if cutlass.const_expr(self.delegate_tensormap_ab_init == False): tensormap_manager.init_tensormap_from_atom( @@ -759,185 +788,161 @@ class GroupedGemmKernel: tensormap_manager.init_tensormap_from_atom( tma_atom_b, tensormap_b_init_ptr, self.tma_warp_id ) - # - # Persistent tile scheduling loop - # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, bid, grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) + tensormap_init_done = cutlass.Boolean(False) - # tile count we have searched - total_k_tile_cnt = cutlass.Int32(0) # group index of last tile last_group_idx = cutlass.Int32(-1) - work_tile = tile_sched.initial_work_tile_info() + + work_tile = initial_work_tile_info + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z( - cur_tile_coord, - problem_sizes_mnkl, - ) + grouped_gemm_cta_tile_info = work_tile.group_search_result + cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 cur_group_idx = grouped_gemm_cta_tile_info.group_idx - is_group_changed = cur_group_idx != last_group_idx - # skip tensormap update if we're working on the same group - if is_group_changed: - real_tensor_a = self.make_tensor_for_tensormap_update( - cur_group_idx, - self.a_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - strides_abc, - ptrs_abc, - 0, # 0 for tensor A + # Do not load any data if cur_k_tile_cnt is 0 + if not is_k_tile_cnt_zero: + is_group_changed = cur_group_idx != last_group_idx + # skip tensormap update if we're working on the same group + if is_group_changed: + real_tensor_a = self.make_tensor_for_tensormap_update( + cur_group_idx, + self.a_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 0, # 0 for tensor A + ) + real_tensor_b = self.make_tensor_for_tensormap_update( + cur_group_idx, + self.b_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 1, # 1 for tensor B + ) + # wait tensormap initialization complete before update + if not tensormap_init_done: + if cutlass.const_expr(self.delegate_tensormap_ab_init): + self.tensormap_ab_init_barrier.arrive_and_wait() + tensormap_manager.fence_tensormap_initialization() + tensormap_init_done = True + + tensormap_manager.update_tensormap( + (real_tensor_a, real_tensor_b), + (tma_atom_a, tma_atom_b), + (tensormap_a_ptr, tensormap_b_ptr), + self.tma_warp_id, + (tensormap_a_smem_ptr, tensormap_b_smem_ptr), + ) + + mma_tile_coord_mnl = ( + grouped_gemm_cta_tile_info.cta_tile_idx_m + // cute.size(tiled_mma.thr_id.shape), + grouped_gemm_cta_tile_info.cta_tile_idx_n, + 0, ) - real_tensor_b = self.make_tensor_for_tensormap_update( - cur_group_idx, - self.b_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - strides_abc, - ptrs_abc, - 1, # 1 for tensor B - ) - # wait tensormap initialization complete before update - if tensormap_init_done == False: + + # + # 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 < cur_k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + # ensure the update to tensormap has completed before using it + if is_group_changed: + tensormap_manager.fence_tensormap_update(tensormap_a_ptr) + tensormap_manager.fence_tensormap_update(tensormap_b_ptr) + # + # Tma load loop + # + for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): + # Wait for AB buffer empty + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) + + # Load A/B with TMA + 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, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_a_ptr, + cute.AddressSpace.generic, + ), + ) + 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, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_b_ptr, + cute.AddressSpace.generic, + ), + ) + + # 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 < cur_k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + else: + # If tensormap initialization is not done, wait for it to complete + if not tensormap_init_done: if cutlass.const_expr(self.delegate_tensormap_ab_init): self.tensormap_ab_init_barrier.arrive_and_wait() tensormap_manager.fence_tensormap_initialization() tensormap_init_done = True - - tensormap_manager.update_tensormap( - (real_tensor_a, real_tensor_b), - (tma_atom_a, tma_atom_b), - (tensormap_a_ptr, tensormap_b_ptr), - self.tma_warp_id, - (tensormap_a_smem_ptr, tensormap_b_smem_ptr), - ) - - mma_tile_coord_mnl = ( - grouped_gemm_cta_tile_info.cta_tile_idx_m - // cute.size(tiled_mma.thr_id.shape), - grouped_gemm_cta_tile_info.cta_tile_idx_n, - 0, - ) - - # - # 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]) - ] - - num_prev_k_blk = total_k_tile_cnt - total_k_tile_cnt += cur_k_tile_cnt - - # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt - tma_wr_k_tile = cutlass.Int32(0) - smem_wr_buffer = (num_prev_k_blk + tma_wr_k_tile) % self.num_ab_stage - tma_wr_ab_empty_phase = ( - num_prev_k_blk + tma_wr_k_tile - ) // self.num_ab_stage % 2 ^ 1 - peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait( - tma_wr_k_tile < cur_k_tile_cnt, - ab_empty_mbar_ptr + smem_wr_buffer, - tma_wr_ab_empty_phase, - ) - # ensure the update to tensormap has completed before using it - if is_group_changed: - tensormap_manager.fence_tensormap_update(tensormap_a_ptr) - tensormap_manager.fence_tensormap_update(tensormap_b_ptr) - # - # Tma load loop - # - for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): - tma_wr_k_tile_next = tma_wr_k_tile + 1 - smem_wr_buffer_next = ( - num_prev_k_blk + tma_wr_k_tile_next - ) % self.num_ab_stage - tma_wr_ab_empty_phase_next = ( - tma_wr_ab_empty_phase ^ 1 - if smem_wr_buffer_next == 0 - else tma_wr_ab_empty_phase - ) - - smem_full_mbar_ptr = ab_full_mbar_ptr + smem_wr_buffer - - # Wait for AB buffer empty - if peek_ab_empty_status == 0: - cute.arch.mbarrier_wait( - ab_empty_mbar_ptr + smem_wr_buffer, tma_wr_ab_empty_phase - ) - - # Arrive AB buffer and expect full transaction bytes - if is_leader_cta: - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - smem_full_mbar_ptr, self.num_tma_load_bytes - ) - - # Load A/B with TMA - cute.copy( - tma_atom_a, - tAgA_slice[(None, tma_wr_k_tile)], - tAsA[(None, smem_wr_buffer)], - tma_bar_ptr=smem_full_mbar_ptr, - mcast_mask=a_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_a_ptr, - cute.AddressSpace.generic, - ), - ) - cute.copy( - tma_atom_b, - tBgB_slice[(None, tma_wr_k_tile)], - tBsB[(None, smem_wr_buffer)], - tma_bar_ptr=smem_full_mbar_ptr, - mcast_mask=b_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_b_ptr, - cute.AddressSpace.generic, - ), - ) - - # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 - peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait( - tma_wr_k_tile_next < cur_k_tile_cnt, - ab_empty_mbar_ptr + smem_wr_buffer_next, - tma_wr_ab_empty_phase_next, - ) - - tma_wr_k_tile = tma_wr_k_tile_next - smem_wr_buffer = smem_wr_buffer_next - tma_wr_ab_empty_phase = tma_wr_ab_empty_phase_next - # Advance to next tile tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() last_group_idx = cur_group_idx + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + # # Specialized MMA warp # - if warp_idx == self.mma_warp_id: + if warp_idx == self.mma_warp_id and initial_work_tile_info.is_valid_tile: # Bar sync for retrieve tmem ptr from shared mem tmem.wait_for_alloc() @@ -951,63 +956,42 @@ class GroupedGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, bid, grid_dim + work_tile = initial_work_tile_info + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage ) - work_tile = tile_sched.initial_work_tile_info() # tile count we have searched - total_k_tile_cnt = cutlass.Int32(0) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - # MMA warp is only interested in number of tiles along K dimension - ( - cur_k_tile_cnt, - cur_group_idx, - ) = group_gemm_ts_helper.search_cluster_tile_count_k( - cur_tile_coord, - problem_sizes_mnkl, - ) - # Set tensor memory buffer for current tile - acc_buf_idx = tile_sched.num_tiles_executed % self.num_acc_stage - # (MMA, MMA_M, MMA_N) - tCtAcc = tCtAcc_base[(None, None, None, acc_buf_idx)] + cur_group_idx = work_tile.group_search_result.group_idx + problem_shape_k = work_tile.group_search_result.problem_shape_k - num_prev_k_blk = total_k_tile_cnt - total_k_tile_cnt += cur_k_tile_cnt + # MMA warp is only interested in number of tiles along K dimension + cur_k_tile_cnt = ( + problem_shape_k + self.cluster_tile_shape_mnk[2] - 1 + ) // self.cluster_tile_shape_mnk[2] + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 + + # (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 - mma_rd_k_tile = cutlass.Int32(0) - smem_rd_buffer = (num_prev_k_blk + mma_rd_k_tile) % self.num_ab_stage + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) if is_leader_cta: - need_check_rd_buffer_full = ( - mma_rd_k_tile < cur_k_tile_cnt and is_leader_cta - ) - mma_rd_ab_full_phase = ( - (num_prev_k_blk + mma_rd_k_tile) // self.num_ab_stage % 2 - ) - peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait( - need_check_rd_buffer_full, - ab_full_mbar_ptr + smem_rd_buffer, - mma_rd_ab_full_phase, - ) + if ab_consumer_state.count < cur_k_tile_cnt: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) # # Wait for accumulator buffer empty # - acc_empty_phase = ( - tile_sched.num_tiles_executed // self.num_acc_stage % 2 ^ 1 - ) - cute.arch.mbarrier_wait( - acc_empty_mbar_ptr + acc_buf_idx, acc_empty_phase - ) + if not is_k_tile_cnt_zero: + acc_pipeline.producer_acquire(acc_producer_state) # # Reset the ACCUMULATE field for each tile @@ -1017,26 +1001,20 @@ class GroupedGemmKernel: # # Mma mainloop # - for k_tile in range(cur_k_tile_cnt): - mma_rd_k_tile_next = cutlass.Int32(k_tile + 1) - smem_rd_buffer_next = ( - num_prev_k_blk + mma_rd_k_tile_next - ) % self.num_ab_stage - mma_rd_ab_full_phase_next = ( - mma_rd_ab_full_phase ^ 1 - if smem_rd_buffer_next == 0 - else mma_rd_ab_full_phase - ) + for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): # Wait for AB buffer full - if peek_ab_full_status == 0: - cute.arch.mbarrier_wait( - ab_full_mbar_ptr + smem_rd_buffer, mma_rd_ab_full_phase - ) - + 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, smem_rd_buffer) + kblock_coord = ( + None, + None, + kblock_idx, + ab_consumer_state.index, + ) cute.gemm( tiled_mma, @@ -1049,48 +1027,37 @@ class GroupedGemmKernel: tiled_mma.set(tcgen05.Field.ACCUMULATE, True) # Async arrive AB buffer empty - with cute.arch.elect_one(): - tcgen05.commit( - ab_empty_mbar_ptr + smem_rd_buffer, - ab_empty_mcast_mask, - self.cta_group, - ) + ab_pipeline.consumer_release(ab_consumer_state) # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 - need_check_rd_buffer_full = ( - mma_rd_k_tile_next < cur_k_tile_cnt and is_leader_cta - ) - - peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait( - need_check_rd_buffer_full, - ab_full_mbar_ptr + smem_rd_buffer_next, - mma_rd_ab_full_phase_next, - ) - - mma_rd_k_tile = mma_rd_k_tile_next - smem_rd_buffer = smem_rd_buffer_next - mma_rd_ab_full_phase = mma_rd_ab_full_phase_next + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < cur_k_tile_cnt: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) # # Async arrive accumulator buffer full # - with cute.arch.elect_one(): - tcgen05.commit( - acc_full_mbar_ptr + acc_buf_idx, - acc_full_mcast_mask, - self.cta_group, - ) + if not is_k_tile_cnt_zero: + 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: + if warp_idx < self.mma_warp_id and initial_work_tile_info.is_valid_tile: # initialize tensormap A, B for TMA warp if cutlass.const_expr(self.delegate_tensormap_ab_init): tensormap_manager.init_tensormap_from_atom( @@ -1147,32 +1114,32 @@ class GroupedGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, bid, grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) - work_tile = tile_sched.initial_work_tile_info() + work_tile = initial_work_tile_info + # wait tensormap initialization complete before update tensormap_manager.fence_tensormap_initialization() - # tile count we have searched - total_k_tile_cnt = cutlass.Int32(0) + 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_epi_stage, + producer_group=c_producer_group, + ) # group index of last tile last_group_idx = cutlass.Int32(-1) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z( - cur_tile_coord, - problem_sizes_mnkl, - ) + grouped_gemm_cta_tile_info = work_tile.group_search_result cur_group_idx = grouped_gemm_cta_tile_info.group_idx + cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 is_group_changed = cur_group_idx != last_group_idx + # We still need to store 0s when k_tile_cnt is 0 if is_group_changed: # construct tensor C based on real address, shape and stride information real_tensor_c = self.make_tensor_for_tensormap_update( @@ -1201,8 +1168,6 @@ class GroupedGemmKernel: grouped_gemm_cta_tile_info.cta_tile_idx_n, 0, ) - cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k - total_k_tile_cnt += cur_k_tile_cnt # # Slice to per mma tile index @@ -1216,17 +1181,16 @@ class GroupedGemmKernel: *mma_tile_coord_mnl, ) ] - - # Set tensor memory buffer for current tile - acc_buf_idx = tile_sched.num_tiles_executed % self.num_acc_stage # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) - tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_buf_idx)] - + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] # # Wait for accumulator buffer full # - acc_full_phase = tile_sched.num_tiles_executed // self.num_acc_stage % 2 - cute.arch.mbarrier_wait(acc_full_mbar_ptr + acc_buf_idx, acc_full_phase) + # Not waiting for accumulator buffer full when k_tile_cnt is 0 + if not is_k_tile_cnt_zero: + 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)) @@ -1240,28 +1204,34 @@ class GroupedGemmKernel: subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt for subtile_idx in 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 output type - # - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - tRS_rC.store(acc_vec.to(self.c_dtype)) # # Store C to shared memory # epi_buffer = (num_prev_subtiles + subtile_idx) % self.num_epi_stage + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + if not is_k_tile_cnt_zero: + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to output type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + tRS_rC.store(acc_vec.to(self.c_dtype)) + else: + tRS_rC.fill(0) cute.copy( tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, epi_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # # store C to global memory with TMA @@ -1276,19 +1246,17 @@ class GroupedGemmKernel: cute.AddressSpace.generic, ), ) - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group( - self.num_epi_stage - 1, read=True - ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() self.epilog_sync_barrier.arrive_and_wait() # # Async arrive accumulator buffer empty # - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive( - acc_empty_mbar_ptr + acc_buf_idx, - cta_rank_in_cluster // 2 * 2 if use_2cta_instrs else None, - ) + if not is_k_tile_cnt_zero: + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() # # Advance to next tile @@ -1305,13 +1273,9 @@ class GroupedGemmKernel: tmem.free(tmem_ptr) # - # Wait a/b buffer empty + # Wait for C store complete # - if warp_idx == self.epilog_warp_id[0]: - cute.arch.mbarrier_wait( - (ab_empty_mbar_ptr + ((total_k_tile_cnt - 1) % self.num_ab_stage)), - (((total_k_tile_cnt - 1) // self.num_ab_stage) % 2), - ) + c_pipeline.producer_tail() @cute.jit def make_tensor_for_tensormap_update( @@ -1649,7 +1613,7 @@ class GroupedGemmKernel: problem_shape_ntile_mnl, (*cluster_shape_mn, 1) ) - grid = utils.StaticPersistentTileScheduler.get_grid_shape( + grid = utils.StaticPersistentGroupTileScheduler.get_grid_shape( tile_sched_params, max_active_clusters ) @@ -1866,6 +1830,7 @@ def create_tensors_for_all_groups( def run( num_groups: int, problem_sizes_mnkl: tuple[int, int, int, int], + host_problem_shape_available: bool, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], acc_dtype: Type[cutlass.Numeric], @@ -1975,18 +1940,18 @@ def run( c_major, ) - # Choose A, B, C with the smallest size to create initial tensormaps - key_size_a = lambda item: item[1][0] * item[1][2] - key_size_b = lambda item: item[1][1] * item[1][2] - key_size_c = lambda item: item[1][0] * item[1][1] - # Find the indices of the groups with the smallest tensor sizes - min_a_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_a) - min_b_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_b) - min_c_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_c) + # Setup inital tensors for TMA of A,B and C + alignment = 16 # 16 bytes aligned + min_ab_size = alignment * 8 // ab_dtype.width + min_c_size = alignment * 8 // c_dtype.width initial_cute_tensors_abc = [ - cute_tensors_abc[min_a_idx][0], # A with smallest (m, k) - cute_tensors_abc[min_b_idx][1], # B with smallest (n, k) - cute_tensors_abc[min_c_idx][2], # C with smallest (m, n) + create_tensor_and_stride(1, min_ab_size, min_ab_size, a_major == "m", ab_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_ab_size, min_ab_size, b_major == "n", ab_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_c_size, min_c_size, c_major == "m", c_dtype)[2], ] hardware_info = utils.HardwareInfo() @@ -1994,6 +1959,7 @@ def run( max_active_clusters = hardware_info.get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) + # Prepare tensormap buffer for each SM num_tensormap_buffers = sm_count tensormap_shape = ( @@ -2069,9 +2035,19 @@ def run( cluster_tile_shape_mn = compute_cluster_tile_shape( mma_tiler_mn, cluster_shape_mn, use_2cta_instrs ) - total_num_clusters = compute_total_num_clusters( - problem_sizes_mnkl, cluster_tile_shape_mn - ) + + # If the host problem shape is available, we will launch the grid with only + # the necessary clusters. The function compute_total_num_clusters() does that. + # If the problem shape only exists on device, we will need to launch all active + # clusters possible on a device. + if host_problem_shape_available: + print("Problem shapes available on host and device") + total_num_clusters = compute_total_num_clusters( + problem_sizes_mnkl, cluster_tile_shape_mn + ) + else: + print("Problem shapes available only on device") + total_num_clusters = max_active_clusters # Initialize Stream current_stream = cutlass_torch.default_stream() @@ -2079,6 +2055,7 @@ def run( # try to check CUDA version to decide the opt level try: from cutlass import CUDA_VERSION + opt_level = ( 3 if CUDA_VERSION.major < 13 @@ -2131,6 +2108,9 @@ def run( rtol=1e-05, ) + if iterations <= 0: + return 0 + def generate_tensors(): # Reuse existing CPU tensors and create new GPU tensors from them ( @@ -2150,9 +2130,15 @@ def run( ) initial_cute_tensors_abc_workspace = [ - cute_tensors_abc_workspace[min_a_idx][0], # A with smallest (m, k) - cute_tensors_abc_workspace[min_b_idx][1], # B with smallest (n, k) - cute_tensors_abc_workspace[min_c_idx][2], # C with smallest (m, n) + create_tensor_and_stride( + 1, min_ab_size, min_ab_size, a_major == "m", ab_dtype + )[2], + create_tensor_and_stride( + 1, min_ab_size, min_ab_size, b_major == "n", ab_dtype + )[2], + create_tensor_and_stride( + 1, min_c_size, min_c_size, c_major == "m", c_dtype + )[2], ] # Create new tensors for this workspace @@ -2176,7 +2162,7 @@ def run( is_dynamic_layout=False, ) - return testing.JitArguments( + args = testing.JitArguments( initial_cute_tensors_abc_workspace[0], initial_cute_tensors_abc_workspace[1], initial_cute_tensors_abc_workspace[2], @@ -2186,6 +2172,8 @@ def run( tensormap_workspace, current_stream, ) + args.add_to_scope([torch_tensors_abc_workspace]) + return args workspace_count = 1 if use_cold_l2: @@ -2225,6 +2213,18 @@ def run( iterations=iterations, ) + runtime_s = exec_time / 1.0e6 + fmas = 0 + for group in range(num_groups): + [M, N, K, _] = problem_sizes_mnkl[group] + fmas += M * N * K + flop = 2 * fmas + gflop = flop / 1.0e9 + gflops = gflop / runtime_s + + print("Average Runtime : ", exec_time / 1000, "ms") + print("GFLOPS : ", gflops) + return exec_time # Return execution time in microseconds @@ -2270,15 +2270,20 @@ if __name__ == "__main__": parser.add_argument( "--num_groups", type=int, - default=2, + default=3, help="Number of groups", ) parser.add_argument( "--problem_sizes_mnkl", type=parse_comma_separated_tuples, - default=((128, 128, 128, 1), (128, 128, 128, 1)), + default=((128, 128, 128, 1), (512, 128, 128, 1), (128, 256, 128, 1)), help="a tuple of problem sizes for each group (comma-separated tuples)", ) + parser.add_argument( + "--host_problem_shape_available", + action="store_true", + help="Enable the compute of grid based upon host problem shape", + ) parser.add_argument( "--mma_tiler_mn", type=parse_comma_separated_ints, @@ -2362,6 +2367,7 @@ if __name__ == "__main__": run( args.num_groups, args.problem_sizes_mnkl, + args.host_problem_shape_available, args.ab_dtype, args.c_dtype, args.acc_dtype, diff --git a/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py b/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py index c72c24c8..630de111 100644 --- a/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py +++ b/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py @@ -33,9 +33,6 @@ import argparse from typing import List, Type, Tuple, Optional import cuda.bindings.driver as cuda -import torch -import torch.nn.functional as F - import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing @@ -43,7 +40,6 @@ import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack @@ -702,7 +698,7 @@ class SSDKernel: G = cute.size(tma_tensor_b, mode=[3]) NGROUP_RATIO = EH // G - # Make TiledMma + # Make tiledMma ( tiled_mma_intra1, tiled_mma_intra2, @@ -1670,7 +1666,10 @@ class SSDKernel: cute.copy(tiled_r2s_p, tRS_rP, tRS_sP[inter2_p_coord]) # Fence for shared memory - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Async arrive INTER2_P buffer full inter2_p_pipeline.producer_commit(inter2_p_producer_state) # Advance INTER2_P producer state @@ -1700,7 +1699,10 @@ class SSDKernel: ] # Fence for shared memory - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Combine B/Delta/DeltaA/last_column tScaledB = self.pre_inter_scale_bt_with_delta( @@ -1716,7 +1718,10 @@ class SSDKernel: cute.copy(tiled_r2s_b, tBrB_r2s, tBsB_r2s[inter1_b_coord]) # Fence for shared memory - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Async arrive B/Delta/B_TMEM buffer empty/empty/full b_pipeline.consumer_release( @@ -1743,14 +1748,9 @@ class SSDKernel: # Combine INTER1_ACC/last_column/State exp_last_column = cute.math.exp(last_column, fastmath=True) - for reg_idx in range(0, cute.size(tTR_rP), 2): - ( - tTR_rP[reg_idx], - tTR_rP[reg_idx + 1], - ) = cute.arch.fma_packed_f32x2( - (exp_last_column, exp_last_column), - (tState[reg_idx], tState[reg_idx + 1]), - (tTR_rP[reg_idx], tTR_rP[reg_idx + 1]), + for reg_idx in cutlass.range(cute.size(tTR_rP), vectorize=True): + tTR_rP[reg_idx] = ( + exp_last_column * tState[reg_idx] + tTR_rP[reg_idx] ) # Store scaled P to tRS_rP @@ -1765,7 +1765,10 @@ class SSDKernel: cute.copy(tiled_r2s_p, tRS_rP, tRS_sP[inter2_p_coord]) # Fence for shared memory - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Async arrive INTER1_ACC/INTER2_P buffer empty/full inter1_acc_pipeline.consumer_release(inter1_acc_consumer_state) @@ -1798,8 +1801,11 @@ class SSDKernel: # END of for chunk_idx in cutlass.range(C, unroll=1) # Store last INTER2_P (State) from smem to gmem - # Wait for all previous stores to smem to be done - cute.arch.fence_proxy("async.shared", space="cta") + # Wait for all previous store to smem done + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.pre_inter_sync_barrier.arrive_and_wait() if local_warp_idx == 0: @@ -2245,58 +2251,26 @@ class SSDKernel: cute.copy(s2r_atom_d, tRS_sD[d_coord], tRS_rD) # Combine INTRA2_ACC/INTER2_ACC/Delta/X/D - for reg_idx in range(0, cute.size(tRS_rCompute), 2): - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ) = cute.arch.fma_packed_f32x2( - (tTR_rInter[reg_idx], tTR_rInter[reg_idx + 1]), - ( - cute.math.exp( - tTR_rDeltaA[reg_idx], fastmath=True - ), - cute.math.exp( - tTR_rDeltaA[reg_idx + 1], fastmath=True - ), - ), - (tTR_rIntra[reg_idx], tTR_rIntra[reg_idx + 1]), + for reg_idx in cutlass.range( + cute.size(tRS_rCompute), vectorize=True + ): + tRS_rCompute[reg_idx] = ( + tTR_rInter[reg_idx] + * cute.math.exp(tTR_rDeltaA[reg_idx], fastmath=True) + + tTR_rIntra[reg_idx] ) # Fuse Y += X * D if cutlass.const_expr(self.d_has_hdim): - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ) = cute.arch.fma_packed_f32x2( - ( - tRS_rD[reg_idx].to(self.acc_dtype), - tRS_rD[reg_idx + 1].to(self.acc_dtype), - ), - ( - tSR_rX[reg_idx].to(self.acc_dtype), - tSR_rX[reg_idx + 1].to(self.acc_dtype), - ), - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ), + tRS_rCompute[reg_idx] = ( + tRS_rD[reg_idx].to(self.acc_dtype) + * tSR_rX[reg_idx].to(self.acc_dtype) + + tRS_rCompute[reg_idx] ) elif cutlass.const_expr(self.has_d): - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ) = cute.arch.fma_packed_f32x2( - ( - tRS_rD.to(self.acc_dtype), - tRS_rD.to(self.acc_dtype), - ), - ( - tSR_rX[reg_idx].to(self.acc_dtype), - tSR_rX[reg_idx + 1].to(self.acc_dtype), - ), - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ), + tRS_rCompute[reg_idx] = ( + tRS_rD.to(self.acc_dtype) + * tSR_rX[reg_idx].to(self.acc_dtype) + + tRS_rCompute[reg_idx] ) tRS_rY.store(tRS_rCompute.load().to(self.io_dtype)) @@ -2309,7 +2283,10 @@ class SSDKernel: ) # Fence for R2S store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Sync before TMA store self.epilog_sync_barrier.arrive_and_wait() @@ -2426,7 +2403,6 @@ class SSDKernel: internal_stages, intra1_acc_stages, ): - SM100_TMEM_CAPACITY_COLUMNS = 512 BITS_PER_TMEM_COL = 32 # (MMA, MMA_M, MMA_N) acc_shape_intra1 = tiled_mma_intra1.partition_shape_C(tile_shape_mnk_intra1[:2]) @@ -2483,7 +2459,7 @@ class SSDKernel: num_tmem_cols_total = 1 while num_tmem_cols_total < num_tmem_cols_total_tmp: num_tmem_cols_total *= 2 - assert num_tmem_cols_total <= SM100_TMEM_CAPACITY_COLUMNS + assert num_tmem_cols_total <= cute.arch.get_max_tmem_alloc_cols("sm_100") return ( tmem_intra1_acc_offset, @@ -3036,41 +3012,26 @@ class SSDKernel: # SegSum # fadd2 + fsel + fmul2/mufu + fmul2 - for subtile_idx in cutlass.range(0, cute.size(tTR_rQ), 2, unroll_full=True): - ( - tCompute[subtile_idx], - tCompute[subtile_idx + 1], - ) = cute.arch.add_packed_f32x2( - (tCrDeltaA_Col[subtile_idx], tCrDeltaA_Col[subtile_idx + 1]), - (-tCrDeltaA_Row[subtile_idx], -tCrDeltaA_Row[subtile_idx + 1]), + for subtile_idx in cutlass.range( + cute.size(tTR_rQ), unroll_full=True, vectorize=True + ): + tCompute[subtile_idx] = tCrDeltaA_Col[subtile_idx] + ( + -tCrDeltaA_Row[subtile_idx] ) for subtile_idx in cutlass.range(cute.size(tTR_rQ), unroll_full=True): m, n = tCoord[subtile_idx] if m < n: tCompute[subtile_idx] = cutlass.Float32(-float("inf")) LOG2_E = cutlass.Float32(1.4426950408889634) - for subtile_idx in cutlass.range(0, cute.size(tTR_rQ), 2, unroll_full=True): + for subtile_idx in cutlass.range( + cute.size(tTR_rQ), unroll_full=True, vectorize=True + ): # TODO: use math.exp directly - tCompute_log2e = cute.arch.mul_packed_f32x2( - (tCompute[subtile_idx], tCompute[subtile_idx + 1]), (LOG2_E, LOG2_E) - ) - ( - tCompute[subtile_idx], - tCompute[subtile_idx + 1], - ) = cute.arch.mul_packed_f32x2( - ( - cute.math.exp2(tCompute_log2e[0], fastmath=True), - cute.math.exp2(tCompute_log2e[1], fastmath=True), - ), - (tCrDelta[subtile_idx], tCrDelta[subtile_idx + 1]), - ) - ( - tCompute[subtile_idx], - tCompute[subtile_idx + 1], - ) = cute.arch.mul_packed_f32x2( - (tCompute[subtile_idx], tCompute[subtile_idx + 1]), - (tTR_rQ[subtile_idx], tTR_rQ[subtile_idx + 1]), + tCompute_log2e = tCompute[subtile_idx] * LOG2_E + tCompute[subtile_idx] = ( + cute.math.exp2(tCompute_log2e, fastmath=True) * tCrDelta[subtile_idx] ) + tCompute[subtile_idx] = tCompute[subtile_idx] * tTR_rQ[subtile_idx] tRT_rQ.store(tCompute.load().to(self.io_dtype)) return tRT_rQ @@ -3211,6 +3172,7 @@ class SSDKernel: ) return sDeltaA + @cute.jit def pre_inter_scale_bt_with_delta( self, tBrB_s2r, tBrDelta_s2r, tBrDeltaA_s2r, last_column ): @@ -3223,22 +3185,15 @@ class SSDKernel: tBrDelta_Compute.store(tBrDelta_s2r.load().to(self.acc_dtype)) tBrDeltaA_Compute.store(tBrDeltaA_s2r.load().to(self.acc_dtype)) - for reg_idx in range(0, cute.size(tBrB_Compute), 2): - tCompute[reg_idx], tCompute[reg_idx + 1] = cute.arch.mul_packed_f32x2( - ( - cute.math.exp( - (last_column - tBrDeltaA_Compute[reg_idx]), fastmath=True - ), - cute.math.exp( - (last_column - tBrDeltaA_Compute[reg_idx + 1]), fastmath=True - ), - ), - (tBrDelta_Compute[reg_idx], tBrDelta_Compute[reg_idx + 1]), - ) - tCompute[reg_idx], tCompute[reg_idx + 1] = cute.arch.mul_packed_f32x2( - (tCompute[reg_idx], tCompute[reg_idx + 1]), - (tBrB_Compute[reg_idx], tBrB_Compute[reg_idx + 1]), + for reg_idx in cutlass.range( + cute.size(tBrB_Compute), vectorize=True, unroll_full=True + ): + tCompute[reg_idx] = ( + cute.math.exp((last_column - tBrDeltaA_Compute[reg_idx]), fastmath=True) + * tBrDelta_Compute[reg_idx] ) + + tCompute[reg_idx] = tCompute[reg_idx] * tBrB_Compute[reg_idx] return tCompute def epilog_make_delta(self, smem_cumsum_delta): @@ -3349,6 +3304,10 @@ def run( print(f"Skip reference checking: {skip_ref_check}") print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + import torch + import torch.nn.functional as F + import cutlass.torch as cutlass_torch + # Unpack parameters G, B, E, H, C, D, L, N = gbehcdln EH = E * H diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py index a1b14194..0fd5d365 100644 --- a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py @@ -49,53 +49,52 @@ import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.testing as testing from cutlass.cute.runtime import from_dlpack -from cutlass.cute.typing import Int32, Int64, Float32, Pointer, AddressSpace +from cutlass.cute.typing import * -from cutlass._mlir.dialects import nvvm, llvm +from cutlass._mlir.dialects import llvm +from cutlass.cute.arch.nvvm_wrappers import mapa # Kernel invariants mma_modes = (0, 1, 2) -mma_dice = (None, None, None) # (MMA, #MMA_M, #MMA_K) -cpy_dice = (None, *mma_dice) # (CPY, #CPY_MMA, #CPY_M, #CPY_K) +mma_dice = (None, None, None) # (MMA, #MMA_M, #MMA_K) +cpy_dice = (None,) + mma_dice # (CPY, #CPY_MMA, #CPY_M, #CPY_K) warp_threads = 32 warpgroup_warps = 4 warpgroup_threads = 128 # Math helpers -log2_e = math.log2(math.e) # change exponential base -use_tensor_ssa_math = False # experimental -fadd2 = partial(cute.arch.add_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) -fmul2 = partial(cute.arch.mul_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) -ffma2 = partial(cute.arch.fma_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) +log2_e = math.log2(math.e) # change exponential base +use_tensor_ssa_math = False # experimental +fadd2 = partial(cute.arch.add_packed_f32x2, ftz=False, rnd="rn") +fmul2 = partial(cute.arch.mul_packed_f32x2, ftz=False, rnd="rn") +ffma2 = partial(cute.arch.fma_packed_f32x2, ftz=False, rnd="rn") exp2 = partial(cute.math.exp2, fastmath=True) + class MixedInputFusedMultiHeadAttentionDecode: def __init__( self, headdim, block_scaledim, # headdim per scale factor; scale factor shape is (batches, heads_k, seqlen, headdim / block_scaledim) grouped_head_tile, # GQA packing tile size, can be less than group size - dual_convert = False, # Dual warpgroups pingponging on convert stages - deterministic = False, # If True, cluster reduction is disabled + convert_warpgroups = 1, # Multiple warpgroups striding on convert stages ): self.headdim = headdim self.grouped_head_tile = grouped_head_tile self.block_scaledim = block_scaledim self.scaledim = headdim // block_scaledim - self.dual_convert = dual_convert - self.deterministic = deterministic + self.convert_warpgroups = convert_warpgroups assert headdim % block_scaledim == 0 assert grouped_head_tile % 8 == 0 and 0 < grouped_head_tile <= 32 - + warpgroup_id = 0 self.softmax_warpgroup_id = warpgroup_id warpgroup_id += 1 - self.cvt_warpgroup_ids = (warpgroup_id, (warpgroup_id+1) if dual_convert else None) - cvt_warpgroups = 2 if dual_convert else 1 - warpgroup_id += cvt_warpgroups + self.cvt_warpgroup_ids = tuple(range(warpgroup_id, warpgroup_id+convert_warpgroups)) + warpgroup_id += convert_warpgroups # Why 2 MMA+TMA warps when not MMA bound? # Less register pressure per warp promotes concise SASS @@ -112,53 +111,65 @@ class MixedInputFusedMultiHeadAttentionDecode: self.threads_per_cta = warpgroup_id * warpgroup_threads self.use_reg_reconfig = grouped_head_tile > 16 - max_regs_per_wg_thread = (64*1024 // warpgroup_threads) # 64K regs per SM + max_regs_per_wg_thread = 64 * 1024 // warpgroup_threads # 64K regs per SM self.mma_tma_regs = 72 self.cvt_regs = 112 - self.softmax_regs = min(256, max_regs_per_wg_thread - self.mma_tma_regs - - self.cvt_regs * cvt_warpgroups) + self.softmax_regs = (max_regs_per_wg_thread + - self.mma_tma_regs + - self.cvt_regs * convert_warpgroups) + self.softmax_regs = max(128, min(256, self.softmax_regs)) assert (self.mma_tma_regs + self.softmax_regs + - self.cvt_regs * cvt_warpgroups) <= max_regs_per_wg_thread + self.cvt_regs * convert_warpgroups) <= max_regs_per_wg_thread or not self.use_reg_reconfig self.bs_stages = 2 self.sp_stages = 2 self.o_stages = 1 - def can_implement(self, problem_shape, kv_splits, kv_cluster_dim, q_dtype, kv_dtype, o_dtype, acc_dtype): + def can_implement( + self, + problem_shape, + kv_splits, + q_dtype, + kv_dtype, + o_dtype, + acc_dtype, + ): b, h_q, h_k, s_k, d = problem_shape + if kv_dtype is cutlass.Float8E4M3: + raise ValueError("use Float8E4M3FN instead of Float8E4M3") + + if d % 64 != 0: + raise ValueError(f"headdim({d}) must be multiple of 64") + if h_q % h_k != 0: raise ValueError(f"heads_q({h_q}) must be a multiple of heads_k({h_k})") - - if kv_splits % kv_cluster_dim != 0: - raise ValueError(f"kv_splits({kv_splits}) must be a multiple of kv_cluster_dim({kv_cluster_dim})") - if self.deterministic and kv_cluster_dim != 1: - raise ValueError(f"kv_cluster_dim({kv_cluster_dim}) must be 1 for determinism") - - align_scale_bits = 128 # TMA requirement + align_scale_bits = 128 # TMA requirement if self.scaledim * q_dtype.width < align_scale_bits: align_seq = align_scale_bits // (self.scaledim * q_dtype.width) if s_k % align_seq != 0: raise ValueError(f"seqlen({s_k}) must be a multiple of {align_seq}") + if kv_dtype.width < 8 and d % 128 != 0: # TMA requirement + raise ValueError(f"headdim({d}) must be multiple of 128 for {kv_dtype} KV") + @cute.jit def __call__( self, - problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32], # b, h_q, h_k, s_k, d - kv_splits: Int32, # threadblocks per sequence - kv_cluster_dim: Int32, # threadblocks per partial buffer (atomic reduction) + problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32], # b, h_q, h_k, s_k, d + kv_splits: Int32, # threadblocks per sequence q_iter: cute.Pointer, k_iter: cute.Pointer, v_iter: cute.Pointer, k_scale_iter: cute.Pointer, v_scale_iter: cute.Pointer, o_iter: cute.Pointer, - m_iter: cute.Pointer, # colmax_s, must be -inf initialized - l_iter: cute.Pointer, # logsumexp - o_partial_iter: cute.Pointer, # partial O per kv cluster, must be zero initialized if nondeterminism is enabled - m_partial_iter: cute.Pointer, # partial colmax_s per kv cluster - l_partial_iter: cute.Pointer, # partial colsum_p per kv cluster, must be zero initialized if nondeterminism is enabled + m_iter: cute.Pointer, # colmax_s, must be -inf initialized + l_iter: cute.Pointer, # logsumexp + o_partial_iter: cute.Pointer, # partial O per kv split + m_partial_iter: cute.Pointer, # partial colmax_s per kv split + l_partial_iter: cute.Pointer, # partial colsum_p per kv split scale_qs: Float32, scale_o: Float32, stream: cuda.CUstream, @@ -168,7 +179,7 @@ class MixedInputFusedMultiHeadAttentionDecode: ############################## mma_dtype = q_iter.dtype acc_dtype = o_partial_iter.dtype - assert acc_dtype is Float32 # don't support other acc types for now + assert acc_dtype is Float32 # don't support other acc types for now # Block tile sets the granularity at which threadblocks consume work blk_tile_s = 128 @@ -186,29 +197,30 @@ class MixedInputFusedMultiHeadAttentionDecode: # GEMM1: (S_K, H_R, D, (H_K, B)) tiled_mma_kq = sm100_utils.make_trivial_tiled_mma( mma_dtype, - tcgen05.OperandMajorMode.K, # K - tcgen05.OperandMajorMode.K, # Q + tcgen05.OperandMajorMode.K, # K + tcgen05.OperandMajorMode.K, # Q acc_dtype, tcgen05.CtaGroup.ONE, mma_tile_mnk[:2], - tcgen05.OperandSource.TMEM, # converted K in tmem + tcgen05.OperandSource.TMEM, # converted K in tmem ) # GEMM2: (D, H_R, S_K, (H_K, B)) - tiled_mma_vp = sm100_utils.make_trivial_tiled_mma(# + tiled_mma_vp = sm100_utils.make_trivial_tiled_mma( # mma_dtype, - tcgen05.OperandMajorMode.K, # V - tcgen05.OperandMajorMode.MN, # P + tcgen05.OperandMajorMode.K, # V + tcgen05.OperandMajorMode.MN, # P acc_dtype, tcgen05.CtaGroup.ONE, mma_tile_mnk[:2], - tcgen05.OperandSource.TMEM, # converted V in tmem + tcgen05.OperandSource.TMEM, # converted V in tmem ) # Calculate Q stages self.q_stages = blk_tile_d // mma_tile_k - # Heuristics to avoid power throttling + # Perf heuristics + cap_kv_stages = k_iter.dtype.width >= 8 max_cvt_stages = 4 if self.grouped_head_tile == 32 and mma_tile_k == 128 else 8 max_kv_stages = 8 if mma_tile_k == 128 else 14 @@ -218,25 +230,40 @@ class MixedInputFusedMultiHeadAttentionDecode: tmem_capacity = 512 cvt_stage_cols = mma_tile_k * mma_dtype.width // 32 self.cvt_stages = (tmem_capacity - tmem_alloc_cols) // cvt_stage_cols - self.cvt_stages = min(self.cvt_stages, max_cvt_stages) + self.cvt_stages = ( + min(self.cvt_stages, max_cvt_stages) if cap_kv_stages else self.cvt_stages + ) tmem_alloc_cols += self.cvt_stages * cvt_stage_cols - self.tmem_alloc_cols = 2 ** math.ceil(math.log2(tmem_alloc_cols)) # Tmem alloc must be PO2 + self.tmem_alloc_cols = 2 ** math.ceil( + math.log2(tmem_alloc_cols) + ) # Tmem alloc must be PO2 print(f"\tcvt stages: {self.cvt_stages}") # Calculate KV smem stages self.mbarrier_reserved_bytes = 768 - smem_alloc_bits = self.mbarrier_reserved_bytes*8 - smem_alloc_bits += mma_tile_n * 2 * acc_dtype.width # colmax + cluster colmax - smem_alloc_bits += self.scaledim * blk_tile_s * self.bs_stages * mma_dtype.width # block scale - smem_alloc_bits += mma_tile_n * warpgroup_warps * acc_dtype.width # colsum - smem_alloc_bits += mma_tile_n * mma_tile_k * self.q_stages * mma_dtype.width # Q - smem_alloc_bits += mma_tile_m * mma_tile_n * self.sp_stages * mma_dtype.width # P + smem_alloc_bits = self.mbarrier_reserved_bytes * 8 + smem_alloc_bits += mma_tile_n * acc_dtype.width # colmax + smem_alloc_bits += ( + self.scaledim * blk_tile_s * self.bs_stages * mma_dtype.width + ) # block scale + smem_alloc_bits += mma_tile_n * warpgroup_warps * acc_dtype.width # colsum + smem_alloc_bits += ( + mma_tile_n * mma_tile_k * self.q_stages * mma_dtype.width + ) # Q + smem_alloc_bits += ( + mma_tile_m * mma_tile_n * self.sp_stages * mma_dtype.width + ) # P smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") - self.kv_stages = (smem_capacity*8 - smem_alloc_bits) // (mma_tile_m * mma_tile_k * k_iter.dtype.width) - self.kv_stages = min(self.kv_stages, max_kv_stages) + kv_smem_dtype = cutlass.Int8 if k_iter.dtype.width < 8 else k_iter.dtype + self.kv_stages = (smem_capacity * 8 - smem_alloc_bits) // ( + mma_tile_m * mma_tile_k * kv_smem_dtype.width + ) + self.kv_stages = ( + min(self.kv_stages, max_kv_stages) if cap_kv_stages else self.kv_stages + ) print(f"\tkv stages: {self.kv_stages}") @@ -245,74 +272,68 @@ class MixedInputFusedMultiHeadAttentionDecode: ############################## b, h_q, h_k, s_k, d = problem_shape h_r = h_q // h_k - kv_clusters = kv_splits // kv_cluster_dim - q = cute.make_tensor(q_iter, - cute.make_ordered_layout( - shape=(h_r, d, (h_k, b)), - order=( 1, 0, ( 2, 3)) - ) + q = cute.make_tensor( + q_iter, + cute.make_ordered_layout(shape=(h_r, d, (h_k, b)), order=(1, 0, (2, 3))), ) - k = cute.make_tensor(k_iter, - cute.make_ordered_layout( - shape=(s_k, d, (h_k, b)), - order=( 1, 0, ( 2, 3)) - ) + k = cute.make_tensor( + k_iter, + cute.make_ordered_layout(shape=(s_k, d, (h_k, b)), order=(1, 0, (2, 3))), ) assert k_iter.dtype is not q_iter.dtype - v = cute.make_tensor(v_iter, - cute.make_ordered_layout( - shape=(d, s_k, (h_k, b)), - order=(0, 1, ( 2, 3)) - ) + v = cute.make_tensor( + v_iter, + cute.make_ordered_layout(shape=(d, s_k, (h_k, b)), order=(0, 1, (2, 3))), ) assert v_iter.dtype is k_iter.dtype - o_partial = cute.make_tensor(o_partial_iter, + o_partial = cute.make_tensor( + o_partial_iter, cute.make_ordered_layout( - shape=(d, h_r, (h_k, b), kv_clusters), - order=(0, 1, ( 2, 3), 4) - ) + shape=(d, h_r, (h_k, b), kv_splits), order=(0, 1, (2, 3), 4) + ), ) - m = cute.make_tensor(m_iter, + m = cute.make_tensor( + m_iter, cute.make_ordered_layout( shape=(h_r, (h_k, b)), - order=( 0, ( 1, 2)), - ) + order=(0, (1, 2)), + ), ) assert m_iter.dtype is acc_dtype - m_partial = cute.make_tensor(m_partial_iter, + m_partial = cute.make_tensor( + m_partial_iter, cute.make_ordered_layout( - shape=(h_r, (h_k, b), kv_clusters), - order=( 0, ( 1, 2), 3), - ) + shape=(h_r, (h_k, b), kv_splits), + order=(0, (1, 2), 3), + ), ) assert m_partial_iter.dtype is acc_dtype - l_partial = cute.make_tensor(l_partial_iter, + l_partial = cute.make_tensor( + l_partial_iter, cute.make_ordered_layout( - shape=(h_r, (h_k, b), kv_clusters), - order=( 0, ( 1, 2), 3), - ) + shape=(h_r, (h_k, b), kv_splits), + order=(0, (1, 2), 3), + ), ) assert l_partial_iter.dtype is acc_dtype - align_scale_bits = 128 # TMA requirement + align_scale_bits = 128 # TMA requirement if cutlass.const_expr(self.scaledim * mma_dtype.width >= align_scale_bits): scale_layout = cute.make_ordered_layout( - shape=(self.scaledim, s_k, (h_k, b)), - order=(0, 1, (2, 3)) + shape=(self.scaledim, s_k, (h_k, b)), order=(0, 1, (2, 3)) ) else: align_seq = align_scale_bits // (self.scaledim * mma_dtype.width) s_ks = (align_seq, s_k // align_seq) scale_layout = cute.make_ordered_layout( - shape=(self.scaledim, s_ks, (h_k, b)), - order=(0, (1, 2), (3, 4)) + shape=(self.scaledim, s_ks, (h_k, b)), order=(0, (1, 2), (3, 4)) ) k_scale = cute.make_tensor(k_scale_iter, scale_layout) @@ -321,40 +342,70 @@ class MixedInputFusedMultiHeadAttentionDecode: v_scale = cute.make_tensor(v_scale_iter, scale_layout) assert v_scale_iter.dtype is mma_dtype - # (MMA, MMA_M/N, MMA_K, Stages) - smem_layout_q = sm100_utils.make_smem_layout_b(tiled_mma_kq, mma_tile_mnk, q_iter.dtype, self.q_stages) - smem_layout_k = sm100_utils.make_smem_layout_a(tiled_mma_kq, mma_tile_mnk, k_iter.dtype, self.kv_stages) - smem_layout_v = sm100_utils.make_smem_layout_a(tiled_mma_vp, mma_tile_mnk, v_iter.dtype, self.kv_stages, - is_k_major=False) # V is always headdim-major (GEMM2 M-major) in gmem+smem + smem_layout_q = sm100_utils.make_smem_layout_b( + tiled_mma_kq, mma_tile_mnk, q_iter.dtype, self.q_stages + ) + smem_layout_k = sm100_utils.make_smem_layout_a( + tiled_mma_kq, mma_tile_mnk, kv_smem_dtype, self.kv_stages + ) + smem_layout_v = sm100_utils.make_smem_layout_a( + tiled_mma_vp, mma_tile_mnk, kv_smem_dtype, self.kv_stages, is_k_major=False + ) # V is always headdim-major (GEMM2 M-major) in gmem+smem smem_layout_bs = cute.make_layout((self.scaledim, blk_tile_s, self.bs_stages)) - smem_layout_atom_o = tcgen05.make_smem_layout_atom(tcgen05.mma.SmemLayoutAtomKind.MN_SW128, o_partial_iter.dtype) - smem_layout_o = cute.tile_to_shape(smem_layout_atom_o, (blk_tile_d, blk_tile_h), order=(1,0)) + smem_layout_atom_o = tcgen05.make_smem_layout_atom( + tcgen05.mma.SmemLayoutAtomKind.MN_SW128, o_partial_iter.dtype + ) + smem_layout_o = cute.tile_to_shape( + smem_layout_atom_o, (blk_tile_d, blk_tile_h), order=(1, 0) + ) smem_layout_o = cute.flat_divide(smem_layout_o, (mma_tile_m, mma_tile_n)) tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() - tma_store_op = (cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() if self.deterministic - else cute.nvgpu.cpasync.CopyReduceBulkTensorTileS2GOp()) + tma_store_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_B( - tma_load_op, q, cute.select(smem_layout_q, mma_modes), mma_tile_mnk, tiled_mma_kq + tma_load_op, + q, + cute.select(smem_layout_q, mma_modes), + mma_tile_mnk, + tiled_mma_kq, ) tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_A( - tma_load_op, k, cute.select(smem_layout_k, mma_modes), mma_tile_mnk, tiled_mma_kq + tma_load_op, + k, + cute.select(smem_layout_k, mma_modes), + mma_tile_mnk, + tiled_mma_kq, + internal_type=kv_smem_dtype, ) tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_A( - tma_load_op, v, cute.select(smem_layout_v, mma_modes), mma_tile_mnk, tiled_mma_vp + tma_load_op, + v, + cute.select(smem_layout_v, mma_modes), + mma_tile_mnk, + tiled_mma_vp, + internal_type=kv_smem_dtype, ) tma_atom_ks, tma_tensor_ks = cute.nvgpu.cpasync.make_tiled_tma_atom( - tma_load_op, k_scale, cute.select(smem_layout_bs, mode=[0,1]), smem_layout_bs.shape[:2] + tma_load_op, + k_scale, + cute.select(smem_layout_bs, mode=[0, 1]), + smem_layout_bs.shape[:2], ) tma_atom_vs, tma_tensor_vs = cute.nvgpu.cpasync.make_tiled_tma_atom( - tma_load_op, v_scale, cute.select(smem_layout_bs, mode=[0,1]), smem_layout_bs.shape[:2] + tma_load_op, + v_scale, + cute.select(smem_layout_bs, mode=[0, 1]), + smem_layout_bs.shape[:2], ) tma_atom_o, tma_tensor_o = cute.nvgpu.cpasync.make_tiled_tma_atom( - tma_store_op, o_partial, cute.select(smem_layout_o, mode=[0,1]), mma_tile_mnk[:2] + tma_store_op, + o_partial, + cute.select(smem_layout_o, mode=[0, 1]), + mma_tile_mnk[:2], ) # K scale and V scale will have the same TMA tensor (coord tensor) @@ -371,18 +422,39 @@ class MixedInputFusedMultiHeadAttentionDecode: grid = (kv_splits, n_tiles, l_tiles) self.decode( - blk_tile_shd, mma_tile_mnk, tiled_mma_kq, tiled_mma_vp, - q_iter.dtype, smem_layout_q, tma_atom_q, tma_tensor_q, - k_iter.dtype, smem_layout_k, tma_atom_k, tma_tensor_k, - v_iter.dtype, smem_layout_v, tma_atom_v, tma_tensor_v, - smem_layout_bs, tma_atom_ks, tma_atom_vs, tma_tensor_bs, - o_partial_iter.dtype, smem_layout_o, tma_atom_o, tma_tensor_o, - m, m_partial, l_partial, - scale_qs, scale_qs_log2_e, + blk_tile_shd, + mma_tile_mnk, + tiled_mma_kq, + tiled_mma_vp, + q_iter.dtype, + smem_layout_q, + tma_atom_q, + tma_tensor_q, + k_iter.dtype, + smem_layout_k, + tma_atom_k, + tma_tensor_k, + v_iter.dtype, + smem_layout_v, + tma_atom_v, + tma_tensor_v, + smem_layout_bs, + tma_atom_ks, + tma_atom_vs, + tma_tensor_bs, + o_partial_iter.dtype, + smem_layout_o, + tma_atom_o, + tma_tensor_o, + m, + m_partial, + l_partial, + scale_qs, + scale_qs_log2_e, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], - cluster=[kv_cluster_dim, 1, 1], + cluster=[1, 1, 1], stream=stream, min_blocks_per_mp=1, ) @@ -394,9 +466,15 @@ class MixedInputFusedMultiHeadAttentionDecode: m = cute.make_tensor(m_iter, cute.make_layout((h_q, b))) l = cute.make_tensor(l_iter, cute.make_layout((h_q, b))) - o_partial = cute.make_tensor(o_partial_iter, cute.make_layout((d, h_q, b, kv_clusters))) - m_partial = cute.make_tensor(m_partial_iter, cute.make_layout((h_q, b, kv_clusters))) - l_partial = cute.make_tensor(l_partial_iter, cute.make_layout((h_q, b, kv_clusters))) + o_partial = cute.make_tensor( + o_partial_iter, cute.make_layout((d, h_q, b, kv_splits)) + ) + m_partial = cute.make_tensor( + m_partial_iter, cute.make_layout((h_q, b, kv_splits)) + ) + l_partial = cute.make_tensor( + l_partial_iter, cute.make_layout((h_q, b, kv_splits)) + ) d_per_blk = 128 d_blks = cute.ceil_div(d, d_per_blk) @@ -451,9 +529,6 @@ class MixedInputFusedMultiHeadAttentionDecode: # Read special registers kv_splits, tiles_hr, tiles_hb = cute.arch.grid_dim() kv_split_idx, coord_hr, coord_hb = cute.arch.block_idx() - kv_split_in_cluster, _, _ = cute.arch.block_in_cluster_idx() - kv_cluster_dim, _, _ = cute.arch.block_in_cluster_dim() - kv_cluster_idx, _, _ = cute.arch.cluster_idx() tidx, _, _ = cute.arch.thread_idx() lane_idx = cute.arch.lane_idx() warp_idx = cute.arch.make_warp_uniform(tidx // warp_threads) @@ -464,17 +539,22 @@ class MixedInputFusedMultiHeadAttentionDecode: # No multicast mcast_coord = 0 - mcast_layout = cute.make_layout((1,1,1,1)) # vmnk + mcast_layout = cute.make_layout((1, 1, 1, 1)) # vmnk # Alias types mma_dtype = q_dtype acc_dtype = o_dtype + kv_smem_dtype = cutlass.Int8 if k_dtype.width < 8 else k_dtype # Shapes for MMA tile indexing (Read TMA partition for example) blk_tile_s, blk_tile_h, blk_tile_d = blk_tile_shd mma_tile_m, mma_tile_n, mma_tile_k = mma_tile_mnk - tiles_dm, tiles_sk = cute.ceil_div((blk_tile_d, blk_tile_s), (mma_tile_m, mma_tile_k)) - tiles_dk, tiles_sm = cute.ceil_div((blk_tile_d, blk_tile_s), (mma_tile_k, mma_tile_m)) + tiles_dm, tiles_sk = cute.ceil_div( + (blk_tile_d, blk_tile_s), (mma_tile_m, mma_tile_k) + ) + tiles_dk, tiles_sm = cute.ceil_div( + (blk_tile_d, blk_tile_s), (mma_tile_k, mma_tile_m) + ) tiles_s = cute.ceil_div(mK.shape[0], blk_tile_s) iters_s = cute.ceil_div(tiles_s - kv_split_idx, kv_splits) prefetch_iters = self.sp_stages - 1 @@ -484,7 +564,6 @@ class MixedInputFusedMultiHeadAttentionDecode: # Runtime checks exit_early = kv_split_idx >= tiles_s - do_cluster_reduction = not self.deterministic and kv_cluster_dim > 1 lane_store_max = mma_tile_n == warp_threads or lane_idx < mma_tile_n # Smem alloc helper @@ -523,7 +602,6 @@ class MixedInputFusedMultiHeadAttentionDecode: s_pipeline_ptr = smem.allocate_array(Int64, self.sp_stages * 2) p_pipeline_ptr = smem.allocate_array(Int64, self.sp_stages * 2) o_pipeline_ptr = smem.allocate_array(Int64, self.o_stages * 2) - m_cluster_full_ptr = smem.allocate_array(Int64) # signal cluster colmax is in split 0 smem assert smem._allocated_bytes <= self.mbarrier_reserved_bytes @@ -532,92 +610,75 @@ class MixedInputFusedMultiHeadAttentionDecode: mma_kq_nbar_id = 2 mma_vp_nbar_id = 3 - # Initialize cluster colmax + mbar (even if this split exits early) - sM_layout = cute.make_layout(shape=(mma_tile_m, mma_tile_n), stride=(0, 1)) - sM_cluster = smem.allocate_tensor(acc_dtype, sM_layout, svector_align) - if do_cluster_reduction: - if warp_idx == init_warp and kv_split_in_cluster == 0: - if lane_store_max: - sM_cluster[(0,lane_idx)] = -Float32.inf - cute.arch.fence_acq_rel_cluster() - init_warp += 1 - - if warp_idx == init_warp: - # split 0 waits for one arrive per colmax elt per split in cluster - # other splits in cluster wait for one arrive from split 0 - arrive_count = (kv_cluster_dim * mma_tile_n) if kv_split_in_cluster == 0 else 1 - cute.arch.mbarrier_init(m_cluster_full_ptr, arrive_count) - cute.arch.mbarrier_init_fence() - init_warp += 1 - - cute.arch.cluster_arrive_relaxed() - - # Setup up thread cooperatives + # Alias thread cooperatives elect_one_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread) warpgroup_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads) - dual_warpgroup_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads * 2) mma_group = elect_one_cooperative tma_group = elect_one_cooperative - tma_qo_group = elect_one_cooperative cvt_group = warpgroup_cooperative + cvt_groups = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads * self.convert_warpgroups) softmax_group = warpgroup_cooperative # Initialize pipelines q_producer, q_consumer = pipeline.PipelineTmaAsync.create( num_stages=self.q_stages, producer_group=tma_group, - consumer_group=softmax_group, # Reuse Q consumer mbarriers to sync O store + consumer_group=softmax_group, # Reuse Q consumer mbarriers to sync O store tx_count=cute.size_in_bytes(q_dtype, cute.select(smem_layout_q, mma_modes)), barrier_storage=q_pipeline_ptr, - tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + tidx=mcast_coord, + cta_layout_vmnk=mcast_layout, defer_sync=True, ).make_participants() kv_producer, kv_consumer = pipeline.PipelineTmaAsync.create( num_stages=self.kv_stages, - producer_group=tma_group, consumer_group=cvt_group, + producer_group=tma_group, + consumer_group=cvt_group, tx_count=cute.size_in_bytes(k_dtype, cute.select(smem_layout_k, mma_modes)), barrier_storage=kv_pipeline_ptr, - tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + tidx=mcast_coord, + cta_layout_vmnk=mcast_layout, defer_sync=True, ).make_participants() bs_producer, bs_consumer = pipeline.PipelineTmaAsync.create( num_stages=self.bs_stages, - producer_group=tma_qo_group, - consumer_group=(dual_warpgroup_cooperative if self.dual_convert else cvt_group), + producer_group=tma_group, + consumer_group=cvt_groups, tx_count=cute.size_in_bytes(mma_dtype, cute.select(smem_layout_bs, mode=[0,1])), barrier_storage=bs_pipeline_ptr, - tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + tidx=mcast_coord, + cta_layout_vmnk=mcast_layout, defer_sync=True, ).make_participants() cvt_producer, cvt_consumer = pipeline.PipelineAsyncUmma.create( num_stages=self.cvt_stages, - producer_group=cvt_group, consumer_group=mma_group, + producer_group=cvt_group, + consumer_group=mma_group, barrier_storage=cvt_pipeline_ptr, defer_sync=True, ).make_participants() s_producer, s_consumer = pipeline.PipelineUmmaAsync.create( num_stages=self.sp_stages, - producer_group=mma_group, consumer_group=softmax_group, + producer_group=mma_group, + consumer_group=softmax_group, barrier_storage=s_pipeline_ptr, defer_sync=True, ).make_participants() p_producer, p_consumer = pipeline.PipelineAsyncUmma.create( num_stages=self.sp_stages, - producer_group=softmax_group, consumer_group=mma_group, + producer_group=softmax_group, + consumer_group=mma_group, barrier_storage=p_pipeline_ptr, defer_sync=True, ).make_participants() o_producer, o_consumer = pipeline.PipelineUmmaAsync.create( num_stages=self.o_stages, - producer_group=mma_group, consumer_group=softmax_group, + producer_group=mma_group, + consumer_group=softmax_group, barrier_storage=o_pipeline_ptr, defer_sync=True, ).make_participants() - # Ensure visibility of cluster mbarrier + colmax init - if do_cluster_reduction: - cute.arch.cluster_wait() - # Ensure visibility of local mbarrier inits and tmem alloc cute.arch.sync_threads() @@ -629,41 +690,65 @@ class MixedInputFusedMultiHeadAttentionDecode: thrblk_mma_vp = tiled_mma_vp.get_slice(0) # M - colmax + sM_layout = cute.make_layout(shape=(mma_tile_m, mma_tile_n), stride=(0, 1)) sM = smem.allocate_tensor(acc_dtype, sM_layout, svector_align) tCsM = thrblk_mma_kq.partition_C(sM) - tCsM_cluster = thrblk_mma_kq.partition_C(sM_cluster) # L - colsum - sL_layout = cute.make_layout(shape=(mma_tile_m, mma_tile_n, warpgroup_warps), stride=(0, 1, mma_tile_n)) + sL_layout = cute.make_layout( + shape=(mma_tile_m, mma_tile_n, warpgroup_warps), stride=(0, 1, mma_tile_n) + ) sL = smem.allocate_tensor(acc_dtype, sL_layout, svector_align) tCsL = thrblk_mma_kq.partition_C(sL) # BS - block scale - sBS = smem.allocate_tensor(mma_dtype, smem_layout_bs, stensor_align) # (SCALE, TILE_S, bs_stages) + sBS = smem.allocate_tensor( + mma_dtype, smem_layout_bs, stensor_align + ) # (SCALE, TILE_S, bs_stages) # Q - tBsQ = smem.allocate_tensor(q_dtype, smem_layout_q.outer, stensor_align, smem_layout_q.inner) # (MMA, #MMA_N, #MMA_K, q_stages) + tBsQ = smem.allocate_tensor( + q_dtype, smem_layout_q.outer, stensor_align, smem_layout_q.inner + ) # (MMA, #MMA_N, #MMA_K, q_stages) # K - tAsK = smem.allocate_tensor(k_dtype, smem_layout_k.outer, stensor_align, smem_layout_k.inner) # (MMA, #MMA_M, #MMA_K, kv_stages) - tAtK_cvt_shape = tiled_mma_kq.partition_shape_A((mma_tile_m, mma_tile_k, self.cvt_stages)) # (MMA, #MMA_M, #MMA_K, cvt_stages) + tAsK = smem.allocate_tensor( + kv_smem_dtype, smem_layout_k.outer, stensor_align, smem_layout_k.inner + ) # (MMA, #MMA_M, #MMA_K, kv_stages) + tAtK_cvt_shape = tiled_mma_kq.partition_shape_A( + (mma_tile_m, mma_tile_k, self.cvt_stages) + ) # (MMA, #MMA_M, #MMA_K, cvt_stages) tAtK_cvt = thrblk_mma_kq.make_fragment_A(tAtK_cvt_shape) # V - tAsV_iterator = cute.recast_ptr(tAsK.iterator, smem_layout_v.inner, dtype=v_dtype) # KV share input buffers - tAsV = cute.make_tensor(tAsV_iterator, smem_layout_v.outer) # (MMA, #MMA_M, #MMA_K, kv_stages) - tAtV_cvt_shape = tiled_mma_vp.partition_shape_A((mma_tile_m, mma_tile_k, self.cvt_stages)) # (MMA, #MMA_M, #MMA_K, cvt_stages) + tAsV_iterator = cute.recast_ptr( + tAsK.iterator, smem_layout_v.inner, dtype=kv_smem_dtype + ) # KV share input buffers + tAsV = cute.make_tensor( + tAsV_iterator, smem_layout_v.outer + ) # (MMA, #MMA_M, #MMA_K, kv_stages) + tAtV_cvt_shape = tiled_mma_vp.partition_shape_A( + (mma_tile_m, mma_tile_k, self.cvt_stages) + ) # (MMA, #MMA_M, #MMA_K, cvt_stages) tAtV_cvt = thrblk_mma_vp.make_fragment_A(tAtV_cvt_shape) # S - tCtS_shape = tiled_mma_kq.partition_shape_C((mma_tile_m, mma_tile_n, self.sp_stages)) - tCtS = thrblk_mma_kq.make_fragment_C(tCtS_shape) # (MMA_MN, #MMA_M=1, #MMA_N=1, sp_stages) + tCtS_shape = tiled_mma_kq.partition_shape_C( + (mma_tile_m, mma_tile_n, self.sp_stages) + ) + tCtS = thrblk_mma_kq.make_fragment_C( + tCtS_shape + ) # (MMA_MN, #MMA_M=1, #MMA_N=1, sp_stages) # P - Treat MN C tile of BMM0 as NM B tile of BMM1 # (MMA_NK, #MMA_N, #MMA_K=MMA_TILE_M/MMA_K, sp_stages) mma_tile_nm = (None, mma_tile_n, mma_tile_m) - tBsP_nm_layout = sm100_utils.make_smem_layout_b(tiled_mma_vp, mma_tile_nm, mma_dtype, self.sp_stages) - tBsP_nm = smem.allocate_tensor(mma_dtype, tBsP_nm_layout.outer, stensor_align, tBsP_nm_layout.inner) + tBsP_nm_layout = sm100_utils.make_smem_layout_b( + tiled_mma_vp, mma_tile_nm, mma_dtype, self.sp_stages + ) + tBsP_nm = smem.allocate_tensor( + mma_dtype, tBsP_nm_layout.outer, stensor_align, tBsP_nm_layout.inner + ) # Tile for NK B tile iteration # (MMA_NK, #MMA_N, #MMA_K=MMA_TILE_K/MMA_K, #TILES_SK=MMA_TILE_M/MMA_TILE_K, sp_stages) @@ -677,39 +762,49 @@ class MixedInputFusedMultiHeadAttentionDecode: tCsP = cute.composition(tBsP_nm, tCsP_tile) # O - sO_iterator = cute.recast_ptr(tBsQ.iterator, smem_layout_o.inner, dtype=o_dtype) # Reuse QKV smem for O TMA store - sO_mma = cute.make_tensor(sO_iterator, smem_layout_o.outer) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) - tCsO = thrblk_mma_vp.partition_C(sO_mma) # (MMA, #MMA_M, #MMA_N, #TILE_DM, #TILE_HN) + sO_iterator = cute.recast_ptr( + tBsQ.iterator, smem_layout_o.inner, dtype=o_dtype + ) # Reuse QKV smem for O TMA store + sO_mma = cute.make_tensor( + sO_iterator, smem_layout_o.outer + ) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) + tCsO = thrblk_mma_vp.partition_C( + sO_mma + ) # (MMA, #MMA_M, #MMA_N, #TILE_DM, #TILE_HN) tCtO = thrblk_mma_vp.make_fragment_C(tCsO.shape) # Tmem tensor allocation tmem_ptr = cute.arch.retrieve_tmem_ptr(Int32, 16, tmem_ptr_smem_ptr) tmem_offset = 0 - tAtK_cvt = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtK_cvt.layout) - tAtV_cvt = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtV_cvt.layout) + tAtK_cvt = cute.make_tensor( + cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtK_cvt.layout + ) + tAtV_cvt = cute.make_tensor( + cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtV_cvt.layout + ) tmem_offset += tcgen05.find_tmem_tensor_col_offset(tAtK_cvt) - tCtS = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtS.layout) + tCtS = cute.make_tensor( + cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtS.layout + ) tmem_offset += tcgen05.find_tmem_tensor_col_offset(tCtS) - tCtO = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtO.layout) + tCtO = cute.make_tensor( + cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtO.layout + ) tmem_offset += tcgen05.find_tmem_tensor_col_offset(tCtO) - print(f"\t{tmem_offset} tmem cols used, {self.tmem_alloc_cols} tmem cols allocated") + print( + f"\t{tmem_offset} tmem cols used, {self.tmem_alloc_cols} tmem cols allocated" + ) assert tmem_offset <= self.tmem_alloc_cols ############################## # Exit early ############################## if exit_early: - if do_cluster_reduction and tidx == 0: - waiting_split_in_cluster = 0 - cute.arch.mbarrier_arrive( - m_cluster_full_ptr, - waiting_split_in_cluster, - arrive_count=mma_tile_n - ) + noop = None # early return not supported ############################## # TMA KV Dispatch @@ -720,14 +815,26 @@ class MixedInputFusedMultiHeadAttentionDecode: cute.arch.setmaxregister_decrease(self.mma_tma_regs) # Apply block tiler and slice - gK = cute.local_tile(mK, tiler=(blk_tile_s, blk_tile_d), coord=(None, 0, coord_hb)) # (TILE_S, TILE_D, #TILE_S) - gV = cute.local_tile(mV, tiler=(blk_tile_d, blk_tile_s), coord=(0, None, coord_hb)) # (TILE_D, TILE_S, #TILE_S) + gK = cute.local_tile( + mK, tiler=(blk_tile_s, blk_tile_d), coord=(None, 0, coord_hb) + ) # (TILE_S, TILE_D, #TILE_S) + gV = cute.local_tile( + mV, tiler=(blk_tile_d, blk_tile_s), coord=(0, None, coord_hb) + ) # (TILE_D, TILE_S, #TILE_S) # Apply MMA tiler and MMA partition - gK_mma = cute.flat_divide(gK, (mma_tile_m, mma_tile_k)) # (MMA_TILE_M, MMA_TILE_K, #TILE_SM, #TILE_DK, #TILE_S) - gV_mma = cute.flat_divide(gV, (mma_tile_m, mma_tile_k)) # (MMA_TILE_M, MMA_TILE_K, #TILE_DM, #TILE_SK, #TILE_S) - tAgK = thrblk_mma_kq.partition_A(gK_mma) # (MMA, #MMA_M, #MMA_K, #TILE_SM, #TILE_DK, #TILE_S) - tAgV = thrblk_mma_vp.partition_A(gV_mma) # (MMA, #MMA_M, #MMA_K, #TILE_DM, #TILE_SK, #TILE_S) + gK_mma = cute.flat_divide( + gK, (mma_tile_m, mma_tile_k) + ) # (MMA_TILE_M, MMA_TILE_K, #TILE_SM, #TILE_DK, #TILE_S) + gV_mma = cute.flat_divide( + gV, (mma_tile_m, mma_tile_k) + ) # (MMA_TILE_M, MMA_TILE_K, #TILE_DM, #TILE_SK, #TILE_S) + tAgK = thrblk_mma_kq.partition_A( + gK_mma + ) # (MMA, #MMA_M, #MMA_K, #TILE_SM, #TILE_DK, #TILE_S) + tAgV = thrblk_mma_vp.partition_A( + gV_mma + ) # (MMA, #MMA_M, #MMA_K, #TILE_DM, #TILE_SK, #TILE_S) # #TILE_SM=TILE_S/MMA_TILE_M, #TILE_HN=TILE_H/MMA_TILE_N, #TILE_DK=TILE_D/MMA_TILE_K # #TILE_DM=TILE_D/MMA_TILE_M, #TILE_HN=TILE_H/MMA_TILE_N, #TILE_SK=TILE_S/MMA_TILE_K @@ -739,14 +846,20 @@ class MixedInputFusedMultiHeadAttentionDecode: # TMA partition # (MMA, #MMA_M, #MMA_K, Rest...) -> (TMA, Rest...) tGSsK, tGSgK = cute.nvgpu.cpasync.tma_partition( - tma_atom_k, mcast_coord, mcast_layout, + tma_atom_k, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(tAsK, 0, 3), - gmem_tensor=cute.group_modes(tAgK, 0, 3)) + gmem_tensor=cute.group_modes(tAgK, 0, 3), + ) tGSsV, tGSgV = cute.nvgpu.cpasync.tma_partition( - tma_atom_v, mcast_coord, mcast_layout, + tma_atom_v, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(tAsV, 0, 3), - gmem_tensor=cute.group_modes(tAgV, 0, 3)) + gmem_tensor=cute.group_modes(tAgV, 0, 3), + ) # # Sequence loop @@ -762,11 +875,12 @@ class MixedInputFusedMultiHeadAttentionDecode: tma_atom_k, tGSgK_s[None, 0, dk], tGSsK[None, k_handle.index], - tma_bar_ptr=k_handle.barrier) + tma_bar_ptr=k_handle.barrier, + ) # Load V if s >= prefetch_tiles: - tGSgV_s = tGSgV[None, None, None, s-prefetch_tiles] + tGSgV_s = tGSgV[None, None, None, s - prefetch_tiles] for sk in cutlass.range_constexpr(tiles_sk): for dm in cutlass.range_constexpr(tiles_dm): v_handle = kv_producer.acquire_and_advance() @@ -774,7 +888,8 @@ class MixedInputFusedMultiHeadAttentionDecode: tma_atom_v, tGSgV_s[None, dm, sk], tGSsV[None, v_handle.index], - tma_bar_ptr=v_handle.barrier) + tma_bar_ptr=v_handle.barrier, + ) ############################## # TMA QO Dispatch @@ -785,39 +900,63 @@ class MixedInputFusedMultiHeadAttentionDecode: cute.arch.setmaxregister_decrease(self.mma_tma_regs) # Apply block tiler and slice - gQ = cute.local_tile(mQ, tiler=(blk_tile_h, blk_tile_d), coord=(coord_hr, 0, coord_hb)) # (TILE_H, TILE_D) - gO = cute.local_tile(mO, tiler=(blk_tile_d, blk_tile_h), coord=(0, coord_hr, coord_hb, kv_cluster_idx)) # (TILE_D, TILE_H) - gBS = cute.local_tile(mBS, tiler=(self.scaledim, blk_tile_s), coord=(0, None, coord_hb)) # (SCALE, TILE_S, #TILE_S) + gQ = cute.local_tile( + mQ, tiler=(blk_tile_h, blk_tile_d), coord=(coord_hr, 0, coord_hb) + ) # (TILE_H, TILE_D) + gO = cute.local_tile( + mO, + tiler=(blk_tile_d, blk_tile_h), + coord=(0, coord_hr, coord_hb, kv_split_idx), + ) # (TILE_D, TILE_H) + gBS = cute.local_tile( + mBS, tiler=(self.scaledim, blk_tile_s), coord=(0, None, coord_hb) + ) # (SCALE, TILE_S, #TILE_S) # Apply MMA tiler and MMA partition - gQ_mma = cute.flat_divide(gQ, (mma_tile_n, mma_tile_k)) # (MMA_TILE_N, MMA_TILE_K, #TILE_HN, #TILE_DK) - gO_mma = cute.flat_divide(gO, (mma_tile_m, mma_tile_n)) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) - tBgQ = thrblk_mma_kq.partition_B(gQ_mma) # (MMA, #MMA_N, #MMA_K, #TILE_HN, #TILE_DK) + gQ_mma = cute.flat_divide( + gQ, (mma_tile_n, mma_tile_k) + ) # (MMA_TILE_N, MMA_TILE_K, #TILE_HN, #TILE_DK) + gO_mma = cute.flat_divide( + gO, (mma_tile_m, mma_tile_n) + ) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) + tBgQ = thrblk_mma_kq.partition_B( + gQ_mma + ) # (MMA, #MMA_N, #MMA_K, #TILE_HN, #TILE_DK) # TMA partition tGSsQ, tGSgQ = cute.nvgpu.cpasync.tma_partition( - tma_atom_q, mcast_coord, mcast_layout, + tma_atom_q, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(tBsQ, 0, 3), - gmem_tensor=cute.group_modes(tBgQ, 0, 3)) + gmem_tensor=cute.group_modes(tBgQ, 0, 3), + ) tSGsO, tSGgO = cute.nvgpu.cpasync.tma_partition( - tma_atom_o, mcast_coord, mcast_layout, + tma_atom_o, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(sO_mma, 0, 2), - gmem_tensor=cute.group_modes(gO_mma, 0, 2)) + gmem_tensor=cute.group_modes(gO_mma, 0, 2), + ) tGSsBS, tGSgBS = cute.nvgpu.cpasync.tma_partition( - tma_atom_ks, mcast_coord, mcast_layout, + tma_atom_ks, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(sBS, 0, 2), - gmem_tensor=cute.group_modes(gBS, 0, 2)) + gmem_tensor=cute.group_modes(gBS, 0, 2), + ) # Load Q for dk in cutlass.range_constexpr(tiles_dk): q_handle = q_producer.acquire_and_advance() cute.copy( tma_atom_q, - tGSgQ[None, 0, dk], # stages_q == tiles_dk by construction + tGSgQ[None, 0, dk], # stages_q == tiles_dk by construction tGSsQ[None, dk], - tma_bar_ptr=q_handle.barrier) + tma_bar_ptr=q_handle.barrier, + ) # Sequence Loop prefetch_tiles = prefetch_iters * kv_splits @@ -828,23 +967,22 @@ class MixedInputFusedMultiHeadAttentionDecode: tma_atom_ks, tGSgBS[None, s], tGSsBS[None, bs_handle.index], - tma_bar_ptr=bs_handle.barrier) + tma_bar_ptr=bs_handle.barrier, + ) if s >= prefetch_tiles: bs_handle = bs_producer.acquire_and_advance() cute.copy( tma_atom_vs, - tGSgBS[None, s-prefetch_tiles], + tGSgBS[None, s - prefetch_tiles], tGSsBS[None, bs_handle.index], - tma_bar_ptr=bs_handle.barrier) - + tma_bar_ptr=bs_handle.barrier, + ) + # Store O for dm in cutlass.range_constexpr(tiles_dm): - q_producer.acquire_and_advance() # Reuse Q load barriers to sync O store - cute.copy( - tma_atom_o, - tSGsO[None, dm, 0], - tSGgO[None, dm, 0]) + q_producer.acquire_and_advance() # Reuse Q load barriers to sync O store + cute.copy(tma_atom_o, tSGsO[None, dm, 0], tSGgO[None, dm, 0]) ############################## # Convert Dispatch @@ -854,65 +992,113 @@ class MixedInputFusedMultiHeadAttentionDecode: if cutlass.const_expr(self.use_reg_reconfig): cute.arch.setmaxregister_decrease(self.cvt_regs) - # Initialize for dual convert if necessary - convert_warpgroups = 1 + # Intermediate convert type + cvt_type = Float32 + if cutlass.const_expr(mma_dtype is cutlass.BFloat16 and + k_dtype in (cutlass.Int4, cutlass.Int8)): + cvt_type = mma_dtype + + # Initialize for multiple warpgroups if necessary convert_phase = 0 - if cutlass.const_expr(self.dual_convert): - assert tiles_dk % 2 == 0 - assert (tiles_dm * tiles_sk) % 2 == 0 - convert_warpgroups = 2 - convert_phase = warpgroup_idx % convert_warpgroups - if convert_phase == 1: + if cutlass.const_expr(self.convert_warpgroups > 1): + assert tiles_dk % self.convert_warpgroups == 0 + assert (tiles_dm * tiles_sk) % self.convert_warpgroups == 0 + convert_phase = warpgroup_idx % self.convert_warpgroups + for _ in cutlass.range(convert_phase): kv_consumer.advance() cvt_producer.advance() # Construct tiled copy and partition K - tmem_op_width = 32 - tmem_op_repeat = tcgen05.Repetition(mma_tile_k * mma_dtype.width // tmem_op_width) - tmem_store_atom_k = cute.make_copy_atom(tcgen05.St32x32bOp(tmem_op_repeat), mma_dtype) - tmem_store_k = tcgen05.make_tmem_copy(tmem_store_atom_k, tAtK_cvt[*mma_dice, 0]) - thr_store_k = tmem_store_k.get_slice(warpgroup_tidx) # tmem copy is always 128 threads + mma_k_bits = mma_tile_k * mma_dtype.width + tmem_store_atom_k = cute.make_copy_atom( + tcgen05.St16x256bOp(tcgen05.Repetition(mma_k_bits // 256)), + mma_dtype, + ) + smem_load_atom_k = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x16x8bOp( + num_matrices=4, + unpack_bits=(k_dtype.width if k_dtype.width < 8 else None)), + kv_smem_dtype, + ) - tKsK = thr_store_k.partition_S(tAsK) + tmem_store_k = tcgen05.make_tmem_copy( + tmem_store_atom_k, tAtK_cvt[mma_dice+(0,)] + ) + thr_store_k = tmem_store_k.get_slice(warpgroup_tidx) + tKrK_cvt_shape = thr_store_k.partition_S(tAtK_cvt).shape[:-1] tKtK_cvt = thr_store_k.partition_D(tAtK_cvt) + smem_load_k = cute.make_tiled_copy_S(smem_load_atom_k, tmem_store_k) + thr_load_k = smem_load_k.get_slice(warpgroup_tidx) + tKsK = thr_load_k.partition_S(tAsK) + tKrK_shape = thr_load_k.partition_D(tAsK).shape[:-1] + # Construct tiled copy and partition V - tmem_op_width = 128 - tmem_op_repeat = tcgen05.Repetition(mma_tile_k * mma_dtype.width // tmem_op_width) - tmem_store_atom_v = cute.make_copy_atom(tcgen05.St16x128bOp(tmem_op_repeat), mma_dtype) - tmem_store_v = tcgen05.make_tmem_copy(tmem_store_atom_v, tAtV_cvt[*mma_dice, 0]) + tmem_store_atom_v = cute.make_copy_atom( + tcgen05.St16x256bOp(tcgen05.Repetition(mma_k_bits // 256)), mma_dtype + ) + smem_load_op_v = cute.nvgpu.warp.LdMatrix16x16x8bOp( + transpose=True, + num_matrices=2, + unpack_bits=(v_dtype.width if v_dtype.width < 8 else None), + ) + smem_load_atom_v = cute.make_copy_atom(smem_load_op_v, kv_smem_dtype) + + tmem_store_v = tcgen05.make_tmem_copy( + tmem_store_atom_v, tAtV_cvt[mma_dice+(0,)] + ) thr_store_v = tmem_store_v.get_slice(warpgroup_tidx) + tVrV_cvt_shape = thr_store_v.partition_S(tAtV_cvt).shape[:-1] + tVtV_cvt = thr_store_v.partition_D(tAtV_cvt) - smem_load_atom_v = cute.make_copy_atom(cute.nvgpu.warp.LdMatrix16x16x8bOp(num_matrices=2), v_dtype) smem_load_v = cute.make_tiled_copy_S(smem_load_atom_v, tmem_store_v) thr_load_v = smem_load_v.get_slice(warpgroup_tidx) - tVsV = thr_load_v.partition_S(tAsV) - tVrV_shape = smem_load_v.get_slice(0).partition_D(tAsV).shape[:-1] - tVrV_cvt_shape = thr_store_v.get_slice(0).partition_S(tAtV_cvt).shape[:-1] - tVtV_cvt = thr_store_v.partition_D(tAtV_cvt) + tVrV_shape = thr_load_v.partition_D(tAsV).shape[:-1] # Partition KS - K block scale sKS_layout = cute.make_layout( - shape=(blk_tile_s, (self.block_scaledim, self.scaledim), self.bs_stages), - stride=(self.scaledim, (0, 1), blk_tile_s * self.scaledim)) - sKS = cute.make_tensor(sBS.iterator, sKS_layout) # (TILE_S, TILE_D, bs_stages) - # (MMA_TILE_M, MMA_TILE_K, (#TILE_SM, #TILE_DK), bs_stages) - sKS_mma = cute.group_modes(cute.flat_divide(sKS, (mma_tile_m, mma_tile_k)), 2, 4) - tAsKS = thrblk_mma_kq.partition_A(sKS_mma) # (MMA, #MMA_M, #MMA_K, (#TILE_SM, #TILE_DK), bs_stages) - tKsKS = thr_store_k.partition_S(tAsKS) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) - tKrKS = cute.make_rmem_tensor_like(tKsKS[*cpy_dice, None, 0]) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE) + shape=( + blk_tile_s, + (self.block_scaledim, self.scaledim), + self.bs_stages, + ), + stride=(self.scaledim, (0, 1), blk_tile_s * self.scaledim), + ) + sKS = cute.make_tensor( + sBS.iterator, sKS_layout + ) # (TILE_S, TILE_D, bs_stages) + sKS_mma = cute.group_modes( + cute.flat_divide(sKS, (mma_tile_m, mma_tile_k)), 2, 4 + ) # (MMA_TILE_M, MMA_TILE_K, (#TILE_SM, #TILE_DK), bs_stages) + tAsKS = thrblk_mma_kq.partition_A( + sKS_mma + ) # (MMA, #MMA_M, #MMA_K, (#TILE_SM, #TILE_DK), bs_stages) + tKsKS = thr_load_k.partition_D( + tAsKS + ) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) # Partition VS - V block scale sVS_layout = cute.make_layout( - shape=((self.block_scaledim, self.scaledim), blk_tile_s, self.bs_stages), - stride=((0, 1), self.scaledim, blk_tile_s * self.scaledim)) - sVS = cute.make_tensor(sBS.iterator, sVS_layout) # (TILE_D, TILE_S, bs_stages) - # (MMA_TILE_M, MMA_TILE_K, (#TILE_DM, #TILE_SK), bs_stages) - sVS_mma = cute.group_modes(cute.flat_divide(sVS, (mma_tile_m, mma_tile_k)), 2, 4) - tAsVS = thrblk_mma_vp.partition_A(sVS_mma) # (MMA, #MMA_M, #MMA_K, (#TILE_DM, #TILE_SK), bs_stages) - tVsVS = thr_load_v.partition_D(tAsVS) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) - tVrVS = cute.make_rmem_tensor_like(tVsVS[*cpy_dice, None, 0]) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE) + shape=( + (self.block_scaledim, self.scaledim), + blk_tile_s, + self.bs_stages, + ), + stride=((0, 1), self.scaledim, blk_tile_s * self.scaledim), + ) + sVS = cute.make_tensor( + sBS.iterator, sVS_layout + ) # (TILE_D, TILE_S, bs_stages) + sVS_mma = cute.group_modes( + cute.flat_divide(sVS, (mma_tile_m, mma_tile_k)), 2, 4 + ) # (MMA_TILE_M, MMA_TILE_K, (#TILE_DM, #TILE_SK), bs_stages) + tAsVS = thrblk_mma_vp.partition_A( + sVS_mma + ) # (MMA, #MMA_M, #MMA_K, (#TILE_DM, #TILE_SK), bs_stages) + tVsVS = thr_load_v.partition_D( + tAsVS + ) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) # # Sequence loop @@ -921,62 +1107,94 @@ class MixedInputFusedMultiHeadAttentionDecode: if s < iters_s: # Load K scale bs_handle = bs_consumer.wait_and_advance() - cute.autovec_copy(tKsKS[*cpy_dice, None, bs_handle.index], tKrKS) + tKrKS = cute.make_rmem_tensor_like( + tKsKS[cpy_dice + (None, 0)] + ) # 'like' preserves 0 strides + cute.autovec_copy(tKsKS[cpy_dice + (None, bs_handle.index)], tKrKS) cute.arch.fence_view_async_shared() bs_handle.release() # Convert and scale K - for dk in cutlass.range(tiles_dk // convert_warpgroups, unroll=2): - tKrK = cute.make_rmem_tensor(tKsK.shape[:-1], k_dtype) - tKrK_cvt = cute.make_rmem_tensor(tKsK.shape[:-1], mma_dtype) + for dk in cutlass.range(tiles_dk // self.convert_warpgroups, unroll=2): + tKrK = cute.make_rmem_tensor(tKrK_shape, kv_smem_dtype) + tKrK_cvt = cute.make_rmem_tensor(tKrK_cvt_shape, mma_dtype) kv_handle = kv_consumer.wait_and_advance() - cute.autovec_copy(tKsK[*cpy_dice, kv_handle.index], tKrK) + cute.copy(thr_load_k, tKsK[cpy_dice + (kv_handle.index,)], tKrK) cute.arch.fence_view_async_shared() kv_handle.release() - coord_dk = dk * convert_warpgroups + convert_phase - scale_k = tKrKS[*cpy_dice, coord_dk].load() - tKrK_cvt.store(tKrK.load().to(mma_dtype) * scale_k) + # Sign extend unpacked int4 to int8 + if cutlass.const_expr(k_dtype is cutlass.Int4): + tKrK_unpacked_i4_vec = tKrK.load().maybe_downcast() + tKrK_i8_vec = cute.arch.sext_unpacked_i4_i8_intrinsic( + tKrK_unpacked_i4_vec, cute.size(tKrK_shape) + ) + tKrK.store( + cute.TensorSSA(tKrK_i8_vec, tKrK_shape, cutlass.Int8) + ) + + coord_dk = dk * self.convert_warpgroups + convert_phase + scale_k = tKrKS[cpy_dice + (coord_dk,)].load() + tKrK_ssa = tKrK.load().to(cvt_type).to(mma_dtype) * scale_k + tKrK_cvt.store(tKrK_ssa.reshape(tKrK_cvt_shape)) cvt_handle = cvt_producer.acquire_and_advance() - cute.copy(thr_store_k, tKrK_cvt, tKtK_cvt[*cpy_dice, cvt_handle.index]) + cute.copy( + thr_store_k, tKrK_cvt, tKtK_cvt[cpy_dice + (cvt_handle.index,)] + ) cute.arch.fence_view_async_tmem_store() cvt_handle.commit() - # Advance again for dual warpgroups - if cutlass.const_expr(self.dual_convert): + # Advance again for multiple warpgroups + for _ in cutlass.range_constexpr(self.convert_warpgroups - 1): kv_consumer.advance() cvt_producer.advance() if s >= prefetch_iters: # Load V scale bs_handle = bs_consumer.wait_and_advance() - cute.autovec_copy(tVsVS[*cpy_dice, None, bs_handle.index], tVrVS) + tVrVS = cute.make_rmem_tensor_like( + tVsVS[cpy_dice + (None, 0)] + ) # 'like' preserves 0 strides + cute.autovec_copy(tVsVS[cpy_dice + (None, bs_handle.index)], tVrVS) cute.arch.fence_view_async_shared() bs_handle.release() # Convert and scale V - for dmsk in cutlass.range(tiles_dm * tiles_sk // convert_warpgroups, unroll=2): - tVrV = cute.make_rmem_tensor(tVrV_shape, v_dtype) + for dmsk in cutlass.range(tiles_dm * tiles_sk // self.convert_warpgroups, unroll=2): + tVrV = cute.make_rmem_tensor(tVrV_shape, kv_smem_dtype) tVrV_cvt = cute.make_rmem_tensor(tVrV_cvt_shape, mma_dtype) kv_handle = kv_consumer.wait_and_advance() - cute.copy(thr_load_v, tVsV[*cpy_dice, kv_handle.index], tVrV) + cute.copy(thr_load_v, tVsV[cpy_dice + (kv_handle.index,)], tVrV) cute.arch.fence_view_async_shared() kv_handle.release() - coord_dmsk = dmsk * convert_warpgroups + convert_phase - scale_v = tVrVS[*cpy_dice, coord_dmsk].load() - tVrV_cvt.store(tVrV.load().to(mma_dtype) * scale_v) + # Sign extend unpacked int4 to int8 + if cutlass.const_expr(v_dtype is cutlass.Int4): + tVrV_unpacked_i4_vec = tVrV.load().maybe_downcast() + tVrV_i8_vec = cute.arch.sext_unpacked_i4_i8_intrinsic( + tVrV_unpacked_i4_vec, cute.size(tVrV_shape) + ) + tVrV.store( + cute.TensorSSA(tVrV_i8_vec, tVrV_shape, cutlass.Int8) + ) + + coord_dmsk = dmsk * self.convert_warpgroups + convert_phase + scale_v = tVrVS[cpy_dice + (coord_dmsk,)].load() + tVrV_ssa = tVrV.load().to(cvt_type).to(mma_dtype) * scale_v + tVrV_cvt.store(tVrV_ssa.reshape(tVrV_cvt_shape)) cvt_handle = cvt_producer.acquire_and_advance() - cute.copy(thr_store_v, tVrV_cvt, tVtV_cvt[*cpy_dice, cvt_handle.index]) + cute.copy( + thr_store_v, tVrV_cvt, tVtV_cvt[cpy_dice + (cvt_handle.index,)] + ) cute.arch.fence_view_async_tmem_store() cvt_handle.commit() - # Advance again for dual warpgroups - if cutlass.const_expr(self.dual_convert): + # Advance again for multiple warpgroups + for _ in cutlass.range_constexpr(self.convert_warpgroups - 1): kv_consumer.advance() cvt_producer.advance() @@ -996,7 +1214,7 @@ class MixedInputFusedMultiHeadAttentionDecode: q_consumer.wait_and_advance() # Sequence loop - s_token = True # Producer always acquires first + s_token = True # Producer always acquires first for s in cutlass.range(iters_s): # BMM1 k_token = cvt_consumer.try_wait() @@ -1007,14 +1225,16 @@ class MixedInputFusedMultiHeadAttentionDecode: k_handle = cvt_consumer.wait_and_advance(k_token) # Signal BMM2 to start if is_last_iter: - cute.arch.barrier_arrive(barrier_id=mma_kq_nbar_id, number_of_threads=64) + cute.arch.barrier_arrive( + barrier_id=mma_kq_nbar_id, number_of_threads=64 + ) for mma_k in cutlass.range_constexpr(tAtK_cvt.shape[2]): cute.gemm( tiled_mma_kq, - tCtS[*mma_dice, s_handle.index], + tCtS[mma_dice + (s_handle.index,)], tAtK_cvt[None, None, mma_k, k_handle.index], tBsQ_desc[None, None, mma_k, dk], - tCtS[*mma_dice, s_handle.index], + tCtS[mma_dice + (s_handle.index,)], ) if dk == 0 and mma_k == 0: tiled_mma_kq.set(tcgen05.Field.ACCUMULATE, True) @@ -1030,7 +1250,6 @@ class MixedInputFusedMultiHeadAttentionDecode: cute.arch.barrier(barrier_id=mma_vp_nbar_id, number_of_threads=64) s_token = s_producer.try_acquire() - ############################## # MMA VP Dispatch ############################## @@ -1043,13 +1262,14 @@ class MixedInputFusedMultiHeadAttentionDecode: tiled_mma_vp.set(tcgen05.Field.ACCUMULATE, True) tBsP_desc = thrblk_mma_vp.make_fragment_B(tBsP_nk) + # Advance and wait for BMM1 for _ in cutlass.range_constexpr(tiles_dk): cvt_consumer.advance() cute.arch.barrier(barrier_id=mma_kq_nbar_id, number_of_threads=64) # Sequence loop p_token = False - o_token = True # Producer always acquires first + o_token = True # Producer always acquires first for s in cutlass.range(iters_s): # Advance and wait for BMM1 if s < iters_s - 1: @@ -1068,14 +1288,16 @@ class MixedInputFusedMultiHeadAttentionDecode: v_handle = cvt_consumer.wait_and_advance(v_token) # Signal BMM1 to start if is_last_iter: - cute.arch.barrier_arrive(barrier_id=mma_vp_nbar_id, number_of_threads=64) + cute.arch.barrier_arrive( + barrier_id=mma_vp_nbar_id, number_of_threads=64 + ) for mma_k in cutlass.range_constexpr(tAtV_cvt.shape[2]): cute.gemm( tiled_mma_vp, - tCtO[*mma_dice, dm, 0], + tCtO[mma_dice + (dm, 0)], tAtV_cvt[None, None, mma_k, v_handle.index], tBsP_desc[None, None, mma_k, sk, p_handle.index], - tCtO[*mma_dice, dm, 0], + tCtO[mma_dice + (dm, 0)], ) v_handle.release() if not is_last_iter: @@ -1099,43 +1321,70 @@ class MixedInputFusedMultiHeadAttentionDecode: # Construct tiled copies tmem_op_width = 32 - tmem_op_repeat = tcgen05.Repetition(mma_tile_n * acc_dtype.width // tmem_op_width) - tmem_load_atom_s = cute.make_copy_atom(tcgen05.Ld32x32bOp(tmem_op_repeat), acc_dtype) - tmem_load_s = tcgen05.make_tmem_copy(tmem_load_atom_s, tCtS[*mma_dice, 0]) + tmem_op_repeat = tcgen05.Repetition( + mma_tile_n * acc_dtype.width // tmem_op_width + ) + tmem_load_atom_s = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tmem_op_repeat), acc_dtype + ) + tmem_load_s = tcgen05.make_tmem_copy(tmem_load_atom_s, tCtS[mma_dice + (0,)]) thr_load_s = tmem_load_s.get_slice(warpgroup_tidx) - tmem_store_atom_o = cute.make_copy_atom(tcgen05.St32x32bOp(tmem_op_repeat), o_dtype) - tmem_store_o = tcgen05.make_tmem_copy(tmem_store_atom_o, tCtO[*mma_dice, 0, 0]) + tmem_store_atom_o = cute.make_copy_atom( + tcgen05.St32x32bOp(tmem_op_repeat), o_dtype + ) + tmem_store_o = tcgen05.make_tmem_copy( + tmem_store_atom_o, tCtO[mma_dice + (0, 0)] + ) thr_store_o = tmem_store_o.get_slice(warpgroup_tidx) # Partition S and P - tStS = thr_load_s.partition_S(tCtS) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) - tSsP = thr_load_s.partition_D(tCsP) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) + tStS = thr_load_s.partition_S( + tCtS + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) + tSsP = thr_load_s.partition_D( + tCsP + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) # Partition O - tStO = thr_load_s.partition_S(tCtO) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) - tSsO = thr_load_s.partition_D(tCsO) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) + tStO = thr_load_s.partition_S( + tCtO + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) + tSsO = thr_load_s.partition_D( + tCsO + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) tSrO = cute.make_rmem_tensor_like(tSsO) # Partition colmax and initialize in RF - tSsM = thr_load_s.partition_D(tCsM) # (CPY, #CPY_MMA, #CPY_M, #CPY_N) - tSsM_cluster = thr_load_s.partition_D(tCsM_cluster) + tSsM = thr_load_s.partition_D(tCsM) # (CPY, #CPY_MMA, #CPY_M, #CPY_N) tSrM_prev = cute.make_rmem_tensor_like(tSsM) tSrM_prev.fill(-Float32.inf) # Partition colsum and initialize in RF # Each thread maintains a local colsum in RF, smem reduction happens after loop - tSsL = thr_load_s.partition_D(tCsL) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, WARPS) - tSrL = cute.make_rmem_tensor_like(tSsL[*cpy_dice, 0]) + tSsL = thr_load_s.partition_D( + tCsL + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, WARPS) + tSrL = cute.make_rmem_tensor_like(tSsL[cpy_dice + (0,)]) tSrL.fill(Float32(0)) assert warp_threads >= cute.size(tSsM) # get gmem colmax + colsum to store to inbound_hr = coord_hr * blk_tile_h + lane_idx < mM.shape[0] - gM = cute.local_tile(mM, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb)) # (TILE_H) = (MMA_TILE_N) - gM_partial = cute.local_tile(mM_partial, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb, kv_cluster_idx)) - gL_partial = cute.local_tile(mL_partial, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb, kv_cluster_idx)) + gM = cute.local_tile( + mM, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb) + ) # (TILE_H) = (MMA_TILE_N) + gM_partial = cute.local_tile( + mM_partial, + tiler=(mma_tile_n, 1), + coord=(coord_hr, coord_hb, kv_split_idx), + ) + gL_partial = cute.local_tile( + mL_partial, + tiler=(mma_tile_n, 1), + coord=(coord_hr, coord_hb, kv_split_idx), + ) # Initialize O tSrO.fill(Float32(0)) @@ -1146,7 +1395,9 @@ class MixedInputFusedMultiHeadAttentionDecode: tSsM[lane_idx] = -Float32.inf if warpgroup_widx == 1 and lane_store_max: tSsL[lane_idx] = Float32(0) - cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + cute.arch.barrier( + barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads + ) # # Sequence loop @@ -1155,15 +1406,15 @@ class MixedInputFusedMultiHeadAttentionDecode: # Load S from tmem s_handle = s_consumer.wait_and_advance() tSrS = cute.make_rmem_tensor(tSsP.shape[:-1], acc_dtype) - cute.copy(tmem_load_s, tStS[*cpy_dice, s_handle.index], tSrS) + cute.copy(tmem_load_s, tStS[cpy_dice + (s_handle.index,)], tSrS) cute.arch.fence_view_async_tmem_load() s_handle.release() # Reduce colmax in warp RF tSrM = cute.make_rmem_tensor_like(tSsM) - tSrM_lane = Float32(0) # Avoid dynamic register indexing + tSrM_lane = Float32(0) # Avoid dynamic register indexing for i in cutlass.range_constexpr(cute.size(tSrS)): - tSrM[i] = warp_fmax(tSrS[i]) + tSrM[i] = cute.arch.warp_redux_sync(tSrS[i], kind="fmax", nan=True) if i == lane_idx: tSrM_lane = tSrM[i] @@ -1172,44 +1423,56 @@ class MixedInputFusedMultiHeadAttentionDecode: self.smem_fmax(tSsM.iterator + tSsM.layout(lane_idx), tSrM_lane) # Wait for colmax then load - cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + cute.arch.barrier( + barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads + ) cute.autovec_copy(tSsM, tSrM) # Compute online softmax tSrP = cute.make_rmem_tensor(tSsP.shape[:-1], mma_dtype) if cutlass.const_expr(use_tensor_ssa_math): tSrP_f32 = exp2(scale_qs_log2_e * (tSrS.load() - tSrM.load())) - tSrP.store(tSrP_f32.to(mma_dtype)) # convert + tSrP.store(tSrP_f32.to(mma_dtype)) # convert else: tSrP_f32 = cute.make_rmem_tensor(tSrS.shape, acc_dtype) for i in cutlass.range_constexpr(0, cute.size(tSrS), 2): - p_f32x2 = fadd2((tSrS[i], tSrS[i+1]), (-tSrM[i], -tSrM[i+1])) + p_f32x2 = fadd2( + (tSrS[i], tSrS[i + 1]), (-tSrM[i], -tSrM[i + 1]) + ) p_f32x2 = fmul2(p_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) tSrP_f32[i] = exp2(p_f32x2[0]) - tSrP_f32[i+1] = exp2(p_f32x2[1]) + tSrP_f32[i + 1] = exp2(p_f32x2[1]) tSrP.store(tSrP_f32.load().to(mma_dtype)) # Store P to smem p_handle = p_producer.acquire_and_advance() - cute.autovec_copy(tSrP, tSsP[*cpy_dice, p_handle.index]) + cute.autovec_copy(tSrP, tSsP[cpy_dice + (p_handle.index,)]) cute.arch.fence_view_async_shared() p_handle.commit() # Compute correction and correct colsum if cutlass.const_expr(use_tensor_ssa_math): - correction = exp2(scale_qs_log2_e * (tSrM_prev.load() - tSrM.load())) + correction = exp2( + scale_qs_log2_e * (tSrM_prev.load() - tSrM.load()) + ) tSrL.store(tSrL.load() * correction + tSrP_f32) else: correction = cute.make_rmem_tensor_like(tSrM) for i in cutlass.range_constexpr(0, cute.size(tSrM), 2): - c_f32x2 = fadd2((tSrM_prev[i], tSrM_prev[i+1]), (-tSrM[i], -tSrM[i+1])) + c_f32x2 = fadd2( + (tSrM_prev[i], tSrM_prev[i + 1]), (-tSrM[i], -tSrM[i + 1]) + ) c_f32x2 = fmul2(c_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) c_f32x2 = (exp2(c_f32x2[0]), exp2(c_f32x2[1])) correction[i] = c_f32x2[0] - correction[i+1] = c_f32x2[1] - l_f32x2 = ffma2(c_f32x2, (tSrL[i], tSrL[i+1]), (tSrP_f32[i], tSrP_f32[i+1])) + correction[i + 1] = c_f32x2[1] + l_f32x2 = ffma2( + c_f32x2, + (tSrL[i], tSrL[i + 1]), + (tSrP_f32[i], tSrP_f32[i + 1]), + ) tSrL[i] = l_f32x2[0] - tSrL[i+1] = l_f32x2[1] + tSrL[i + 1] = l_f32x2[1] # Correct O if s > 0: @@ -1218,15 +1481,18 @@ class MixedInputFusedMultiHeadAttentionDecode: # Apply correction for dm in cutlass.range_constexpr(tiles_dm): - tSrO_dm = cute.make_rmem_tensor_like(tSsO[*cpy_dice, 0, 0]) - cute.copy(thr_load_s, tStO[*cpy_dice, dm, 0], tSrO_dm) + tSrO_dm = cute.make_rmem_tensor_like(tSsO[cpy_dice + (0, 0)]) + cute.copy(thr_load_s, tStO[cpy_dice + (dm, 0)], tSrO_dm) for i in cutlass.range_constexpr(0, cute.size(tSrO_dm), 2): - o_f32x2 = fmul2((tSrO_dm[i], tSrO_dm[i+1]), (correction[i], correction[i+1])) + o_f32x2 = fmul2( + (tSrO_dm[i], tSrO_dm[i + 1]), + (correction[i], correction[i + 1]), + ) tSrO_dm[i] = o_f32x2[0] - tSrO_dm[i+1] = o_f32x2[1] + tSrO_dm[i + 1] = o_f32x2[1] - cute.copy(thr_store_o, tSrO_dm, tStO[*cpy_dice, dm, 0]) + cute.copy(thr_store_o, tSrO_dm, tStO[cpy_dice + (dm, 0)]) # Notify MMA cute.arch.fence_view_async_tmem_store() @@ -1248,117 +1514,38 @@ class MixedInputFusedMultiHeadAttentionDecode: # Store partial colsum in smem if lane_store_max: - tSsL[*cpy_dice, warpgroup_widx][lane_idx] = tSrL_lane + tSsL[cpy_dice + (warpgroup_widx,)][lane_idx] = tSrL_lane - # Reduce cluster colmax and correct O - if do_cluster_reduction: - # Reduce cluster colmax - if warpgroup_widx == 0: - if lane_store_max: - self.dsmem_fmax( - sM_cluster.iterator + sM_layout((0, lane_idx)), - sM[(0, lane_idx)], - m_cluster_full_ptr - ) + # Wait for colsum + cute.arch.barrier( + barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads + ) - # split 0 waits for cluster colmax to finish reduction - # other splits wait for split 0 to notify cluster colmax is ready - cute.arch.mbarrier_wait(m_cluster_full_ptr, phase=0) + if warpgroup_widx == 0 and lane_store_max and inbound_hr: + # Load colsum and colmax + sL_lane_wg = sL[0, lane_idx, None] + sL_lane = ( + sL_lane_wg[0] + sL_lane_wg[1] + sL_lane_wg[2] + sL_lane_wg[3] + ) + sM_lane = sM[0, lane_idx] - if warpgroup_widx == 0: - if kv_split_in_cluster == 0: - # notify other splits that cluster colmax is ready in split 0 smem - if lane_idx > 0 and lane_idx < kv_cluster_dim: - waiting_split_in_cluster = lane_idx - cute.arch.mbarrier_arrive( - m_cluster_full_ptr, waiting_split_in_cluster, arrive_count=1 - ) - else: - # other splits copy cluster colmax into local smem - if lane_store_max: - sM_cluster[0, lane_idx] = self.dsmem_load( - sM_cluster.iterator + sM_layout((0, lane_idx)) - ) + # Scale colmax + sM_lane = sM_lane * scale_qs - # warpgroup waits for cluster colmax to load into local smem - cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + # Store colsum and colmax + gL_partial[lane_idx] = sL_lane + gM_partial[lane_idx] = sM_lane + self.gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane) - if warpgroup_widx == 0 and lane_store_max and inbound_hr: - # Load colsum and colmax - sL_lane_wg = sL[0, lane_idx, None] - sL_lane = sL_lane_wg[0] + sL_lane_wg[1] + sL_lane_wg[2] + sL_lane_wg[3] - sM_prev_lane = sM[0, lane_idx] - sM_lane = sM_cluster[0, lane_idx] - - # Correct colsum and scale colmax - correction = exp2(scale_qs_log2_e * (sM_prev_lane - sM_lane)) - sL_lane = sL_lane * correction - sM_lane = sM_lane * scale_qs - - # Store colsum and colmax - self.gmem_fadd( - gL_partial.iterator + gL_partial.layout(lane_idx), sL_lane - ) - self.gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane) - if kv_split_in_cluster == 0: - gM_partial[lane_idx] = sM_lane - - # Load cluster colmax - tSrM = cute.make_rmem_tensor_like(tSsM) - cute.autovec_copy(tSsM_cluster, tSrM) - - # Wait and load O - o_handle = o_consumer.wait_and_advance() - cute.copy(thr_load_s, tStO, tSrO) - cute.arch.fence_view_async_tmem_load() - o_handle.release() # Final release signals tmem dealloc - - # Apply cluster correction - if cutlass.const_expr(use_tensor_ssa_math): - correction = exp2(scale_qs_log2_e * (tSrM_prev.load() - tSrM.load())) - tSrO.store(tSrO.load() * correction) - else: - correction = cute.make_rmem_tensor_like(tSsM) - for i in cutlass.range_constexpr(0, cute.size(tSrM), 2): - c_f32x2 = fadd2((tSrM_prev[i], tSrM_prev[i+1]), (-tSrM[i], -tSrM[i+1])) - c_f32x2 = fmul2(c_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) - correction[i] = exp2(c_f32x2[0]) - correction[i+1] = exp2(c_f32x2[1]) - - for i in cutlass.range_constexpr(0, cute.size(tSrO), cute.size(correction)): - for j in cutlass.range_constexpr(0, cute.size(correction), 2): - o_f32x2 = fmul2((tSrO[i+j], tSrO[i+j+1]), (correction[j], correction[j+1])) - tSrO[i+j] = o_f32x2[0] - tSrO[i+j+1] = o_f32x2[1] - - # Wait and load O without cluster correction - else: - # Wait for colsum - cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) - - if warpgroup_widx == 0 and lane_store_max and inbound_hr: - # Load colsum and colmax - sL_lane_wg = sL[0, lane_idx, None] - sL_lane = sL_lane_wg[0] + sL_lane_wg[1] + sL_lane_wg[2] + sL_lane_wg[3] - sM_lane = sM[0, lane_idx] - - # Scale colmax - sM_lane = sM_lane * scale_qs - - # Store colsum and colmax - gL_partial[lane_idx] = sL_lane - gM_partial[lane_idx] = sM_lane - self.gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane) - - o_handle = o_consumer.wait_and_advance() - cute.copy(thr_load_s, tStO, tSrO) - cute.arch.fence_view_async_tmem_load() - o_handle.release() # Final release signals tmem dealloc + o_handle = o_consumer.wait_and_advance() + cute.copy(thr_load_s, tStO, tSrO) + cute.arch.fence_view_async_tmem_load() + o_handle.release() # Final release signals tmem dealloc # Store O to smem for dm in cutlass.range_constexpr(tiles_dm): - tOrO_dm = tSrO[*cpy_dice, dm, 0] - tOsO_dm = tSsO[*cpy_dice, dm, 0] + tOrO_dm = tSrO[cpy_dice + (dm, 0)] + tOsO_dm = tSsO[cpy_dice + (dm, 0)] cute.autovec_copy(tOrO_dm, tOsO_dm) cute.arch.fence_view_async_shared() @@ -1367,25 +1554,18 @@ class MixedInputFusedMultiHeadAttentionDecode: q_consumer.release() q_consumer.advance() - - # Ensure split 0 doesn't exit before all splits read cluster colmax - if do_cluster_reduction: - cute.arch.cluster_arrive_relaxed() - if kv_split_in_cluster == 0 and warp_idx == self.tma_qo_warp_id: - cute.arch.cluster_wait() - return @staticmethod @cute.kernel def reduction( - o : cute.Tensor, - m : cute.Tensor, - l : cute.Tensor, - o_partial : cute.Tensor, - m_partial : cute.Tensor, - l_partial : cute.Tensor, - scale_o : Float32, + o: cute.Tensor, + m: cute.Tensor, + l: cute.Tensor, + o_partial: cute.Tensor, + m_partial: cute.Tensor, + l_partial: cute.Tensor, + scale_o: Float32, ): d_blk_idx, coord_h, coord_b = cute.arch.block_idx() d_per_blk, _, _ = cute.arch.block_dim() @@ -1415,48 +1595,9 @@ class MixedInputFusedMultiHeadAttentionDecode: @staticmethod @cute.jit - def _mapa(ptr : Pointer, cta_rank_in_cluster : Int32 = 0): - llvm_ptr = ptr.llvm_ptr - return nvvm.mapa_shared_cluster( - llvm_ptr.type, - llvm_ptr, - Int32(cta_rank_in_cluster).ir_value(), - ) - - @cute.jit - def dsmem_load(self, val_ptr: Pointer): - val_llvm_ptr = self._mapa(val_ptr, 0) - - ret = llvm.inline_asm( - Float32.mlir_type, - [val_llvm_ptr], - "ld.relaxed.cta.shared::cluster.f32 $0, [$1];", - "=f,r", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - - return Float32(ret) - - @staticmethod - @cute.jit - def warp_fmax(val : Float32): - ret = llvm.inline_asm( - Float32.mlir_type, - [val.ir_value()], - "redux.sync.max.NaN.f32 $0, $1, 0xffffffff;", - "=f,f", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - return Float32(ret) - - @cute.jit - def smem_fmax(ptr : Pointer, val : Float32): + def smem_fmax(ptr: Pointer, val: Float32): # https://stackoverflow.com/a/72461459 - # Works with canonical NaN which warp_redux_fmax should return + # Works with canonical NaN which warp_redux_sync(kind="fmax") should return llvm.inline_asm( None, [ptr.llvm_ptr, val.ir_value()], @@ -1472,37 +1613,9 @@ class MixedInputFusedMultiHeadAttentionDecode: asm_dialect=llvm.AsmDialect.AD_ATT, ) - @cute.jit - def dsmem_fmax(self, val_ptr: Pointer, val: Float32, mbar_ptr: Pointer): - expect_tx_bytes = Int32(Float32.width // 8) - val_llvm_ptr = self._mapa(val_ptr, 0) - mbar_llvm_ptr = self._mapa(mbar_ptr, 0) - - nvvm.mbarrier_txn( - mbar_llvm_ptr, - expect_tx_bytes.ir_value(), - kind=nvvm.MBarrierTxnKind.ARRIVE_EXPECT_TX, - space=nvvm.MBarrierSpaceKind.CLUSTER, - ) - - llvm.inline_asm( - None, - [val_llvm_ptr, val.ir_value(), mbar_llvm_ptr], - """{\n\t - .reg .pred p;\n\t - setp.lt.s32 p, $1, 0x0; - @p red.async.relaxed.cluster.shared::cluster.mbarrier::complete_tx::bytes.min.u32 [$0], $1, [$2];\n\t - @!p red.async.relaxed.cluster.shared::cluster.mbarrier::complete_tx::bytes.max.s32 [$0], $1, [$2];\n\t - }\n\t""", - "r,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - @staticmethod @cute.jit - def gmem_fmax(ptr : Pointer, val : Float32): + def gmem_fmax(ptr: Pointer, val: Float32): llvm.inline_asm( None, [ptr.llvm_ptr, val.ir_value()], @@ -1518,53 +1631,26 @@ class MixedInputFusedMultiHeadAttentionDecode: asm_dialect=llvm.AsmDialect.AD_ATT, ) - @cute.jit - def smem_fadd(ptr : Pointer, val : Float32): - # Expensive - llvm.inline_asm( - None, - [ptr.llvm_ptr, val.ir_value()], - "red.relaxed.shared::cta.add.f32 [$0], $1;", - "r,f", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - - @staticmethod - @cute.jit - def gmem_fadd(ptr : Pointer, val : Float32): - llvm.inline_asm( - None, - [ptr.llvm_ptr, val.ir_value()], - "red.relaxed.cluster.global.add.f32 [$0], $1;", - "l,f", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - def run( - batches : int, - seqlen : int, - heads_q : int, - heads_k : int, - headdim : int, - block_scaledim: int, - kv_splits: int, - kv_cluster_dim: int, - q_dtype: Type[cutlass.Numeric], - kv_dtype: Type[cutlass.Numeric], - o_dtype: Type[cutlass.Numeric], - acc_dtype: Type[cutlass.Numeric], - tolerance: float, - scale_q: float, - scale_o: float, - scale_s: float, - warmup_iterations: int, - iterations: int, - skip_ref_check: bool, + batches: int = 1, + seqlen: int = 1024, + heads_q: int = 32, + heads_k: int = 4, + headdim: int = 512, + block_scaledim: int = 512, + kv_splits: int = 0, + q_dtype: Type[cutlass.Numeric] = BFloat16, + kv_dtype: Type[cutlass.Numeric] = Int8, + o_dtype: Type[cutlass.Numeric] = BFloat16, + acc_dtype: Type[cutlass.Numeric] = Float32, + tolerance: float = 0.1, + scale_q: float = 1.0, + scale_o: float = 1.0, + scale_s: float = 0.0, + warmup_iterations: int = 0, + iterations: int = 0, + skip_ref_check: bool = False, use_cold_l2: bool = False, **kwargs, ): @@ -1572,7 +1658,7 @@ def run( print(f"\tbatches: {batches}, seqlen: {seqlen}") print(f"\theads_q: {heads_q}, heads_k: {heads_k}") print(f"\theaddim: {headdim}, block_scaledim: {block_scaledim}") - print(f"\tkv_splits: {kv_splits}, kv_cluster_dim: {kv_cluster_dim}") + print(f"\tkv_splits: {kv_splits}") print(f"\tq_dtype: {q_dtype}") print(f"\tkv_dtype: {kv_dtype}") print(f"\to_dtype: {o_dtype}") @@ -1593,12 +1679,16 @@ def run( # Config Kernel # grouped_heads = heads_q // heads_k + convert_warpgroups = 1 + if headdim == 512 and grouped_heads == 8 and kv_dtype.width == 4: + convert_warpgroups = 4 + elif headdim > 128: + convert_warpgroups = 2 fmha = MixedInputFusedMultiHeadAttentionDecode( headdim=headdim, block_scaledim=block_scaledim, grouped_head_tile=min(cute.round_up(grouped_heads, 8), 32), - dual_convert=(headdim > 128), - deterministic=(kv_cluster_dim == 1), + convert_warpgroups=convert_warpgroups, ) if scale_s == 0.0: # default to 1/sqrt(d) @@ -1610,25 +1700,27 @@ def run( sm_count = hardware_info.get_device_multiprocessor_count() sm_count = 148 if sm_count <= 0 else sm_count grid_yz = batches * heads_k * math.ceil(grouped_heads / fmha.grouped_head_tile) - kv_splits = sm_count // grid_yz # 1 wave + kv_splits = sm_count // grid_yz # 1 wave kv_splits = max(1, kv_splits) if sm_count == 148 and grid_yz == 32: - kv_splits = 9 # 2 waves + kv_splits = 9 # 2 waves print(f"\tauto kv_splits: {kv_splits}") seqlen_q = 1 seqlen_k = seqlen - kv_clusters = kv_splits // kv_cluster_dim problem_shape = (batches, heads_q, heads_k, seqlen_k, headdim) - fmha.can_implement(problem_shape, kv_splits, kv_cluster_dim, q_dtype, kv_dtype, o_dtype, acc_dtype) + fmha.can_implement( + problem_shape, kv_splits, q_dtype, kv_dtype, o_dtype, acc_dtype + ) # # Allocate Tensors # torch.manual_seed(1111) - def create_tensor(shape, dtype, init = True): + + def create_tensor(shape, dtype, init=True): init_type = cutlass.torch.TensorInitType.RANDOM init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) @@ -1641,10 +1733,14 @@ def run( elif isinstance(init, tuple) or isinstance(init, list): if len(init) == 2: init_type = cutlass.torch.TensorInitType.RANDOM - init_config = cutlass.torch.RandomInitConfig(min_val=init[0], max_val=init[1]) + init_config = cutlass.torch.RandomInitConfig( + min_val=init[0], max_val=init[1] + ) if len(init) == 3: init_type = cutlass.torch.TensorInitType.GAUSSIAN - init_config = cutlass.torch.RandomInitConfig(mean=init[0], std=init[1], scale=init[2]) + init_config = cutlass.torch.RandomInitConfig( + mean=init[0], std=init[1], scale=init[2] + ) f32_torch_tensor = cutlass_torch.create_and_permute_torch_tensor( shape, @@ -1671,20 +1767,26 @@ def run( torch_tensor, ) - qo_shape = (kv_clusters, batches, heads_q, seqlen_q, headdim) - kv_shape = ( batches, heads_k, seqlen_k, headdim) - scale_shape = ( batches, heads_k, seqlen_k, fmha.scaledim) + qo_shape = (kv_splits, batches, heads_q, seqlen_q, headdim) + kv_shape = (batches, heads_k, seqlen_k, headdim) + scale_shape = (batches, heads_k, seqlen_k, fmha.scaledim) - q_ref, q_cute, q_torch = create_tensor(qo_shape[1:], q_dtype, init=[-8,7]) - k_ref, k_cute, k_torch = create_tensor(kv_shape, kv_dtype, init=[-8,7]) - v_ref, v_cute, v_torch = create_tensor(kv_shape, kv_dtype, init=[-8,7]) - k_scale_ref, k_scale_cute, k_scale_torch = create_tensor(scale_shape, q_dtype, init=[-2, 2]) - v_scale_ref, v_scale_cute, v_scale_torch = create_tensor(scale_shape, q_dtype, init=[-2, 2]) + q_ref, q_cute, q_torch = create_tensor(qo_shape[1:], q_dtype, init=[-8, 7]) + k_ref, k_cute, k_torch = create_tensor(kv_shape, kv_dtype, init=[-8, 7]) + v_ref, v_cute, v_torch = create_tensor(kv_shape, kv_dtype, init=[-8, 7]) + k_scale_ref, k_scale_cute, k_scale_torch = create_tensor( + scale_shape, q_dtype, init=[-2, 2] + ) + v_scale_ref, v_scale_cute, v_scale_torch = create_tensor( + scale_shape, q_dtype, init=[-2, 2] + ) _, o_cute, o_torch = create_tensor(qo_shape[1:], o_dtype, init=False) _, m_cute, m_torch = create_tensor(qo_shape[1:-1], acc_dtype, init=-math.inf) _, l_cute, l_torch = create_tensor(qo_shape[1:-1], acc_dtype, init=False) _, o_partial_cute, o_partial_torch = create_tensor(qo_shape, acc_dtype, init=0) - _, m_partial_cute, m_partial_torch = create_tensor(qo_shape[:-1], acc_dtype, init=-math.inf) + _, m_partial_cute, m_partial_torch = create_tensor( + qo_shape[:-1], acc_dtype, init=-math.inf + ) _, l_partial_cute, l_partial_torch = create_tensor(qo_shape[:-1], acc_dtype, init=0) # @@ -1695,7 +1797,6 @@ def run( fmha, problem_shape, kv_splits, - kv_cluster_dim, q_cute.iterator, k_cute.iterator, v_cute.iterator, @@ -1717,13 +1818,17 @@ def run( # # Refcheck # - def run_torch_fmha(q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs=1.0, scale_o=1.0): + def run_torch_fmha( + q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs=1.0, scale_o=1.0 + ): for i in range(0, headdim // block_scaledim): j = i * block_scaledim - k_ref[..., j:j+block_scaledim] *= k_scale_ref[..., i:i+1] - v_ref[..., j:j+block_scaledim] *= v_scale_ref[..., i:i+1] + k_ref[..., j : j + block_scaledim] *= k_scale_ref[..., i : i + 1] + v_ref[..., j : j + block_scaledim] *= v_scale_ref[..., i : i + 1] - with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.MATH], set_priority=True): + with sdpa_kernel( + [SDPBackend.FLASH_ATTENTION, SDPBackend.MATH], set_priority=True + ): o_ref = scaled_dot_product_attention( q_ref, k_ref, @@ -1743,7 +1848,6 @@ def run( compiled_fmha( problem_shape, kv_splits, - kv_cluster_dim, q_cute.iterator, k_cute.iterator, v_cute.iterator, @@ -1760,25 +1864,28 @@ def run( current_stream, ) print("Verifying results...") - o_ref = run_torch_fmha(q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs, scale_o) - torch.testing.assert_close(o_ref, o_torch.float().cpu(), atol=tolerance, rtol=1e-05) + o_ref = run_torch_fmha( + q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs, scale_o + ) + torch.testing.assert_close( + o_ref, o_torch.float().cpu(), atol=tolerance, rtol=1e-05 + ) def generate_tensors(): - _, q_cute, _ = create_tensor(qo_shape[1:], q_dtype, init=[-8,7]) - _, k_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8,7]) - _, v_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8,7]) - _, k_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8,7]) - _, v_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8,7]) + _, q_cute, _ = create_tensor(qo_shape[1:], q_dtype, init=[-8, 7]) + _, k_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8, 7]) + _, v_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8, 7]) + _, k_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8, 7]) + _, v_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8, 7]) _, o_cute, _ = create_tensor(qo_shape[1:], o_dtype, init=False) _, m_cute, _ = create_tensor(qo_shape[1:-1], acc_dtype, init=-math.inf) _, l_cute, _ = create_tensor(qo_shape[1:-1], acc_dtype, init=False) _, o_partial_cute, _ = create_tensor(qo_shape, acc_dtype, init=0) _, m_partial_cute, _ = create_tensor(qo_shape[:-1], acc_dtype, init=-math.inf) _, l_partial_cute, _ = create_tensor(qo_shape[:-1], acc_dtype, init=0) - return testing.JitArguments( + args = testing.JitArguments( problem_shape, kv_splits, - kv_cluster_dim, q_cute.iterator, k_cute.iterator, v_cute.iterator, @@ -1794,6 +1901,21 @@ def run( scale_o, current_stream, ) + args.add_to_scope( + [ + q_cute, + k_cute, + v_cute, + k_scale_cute, + v_scale_cute, + o_cute, + m_cute, + l_cute, + o_partial_cute, + m_partial_cute, + l_partial_cute, + ] + ) # # Profile @@ -1803,14 +1925,24 @@ def run( q_torch_effective = q_torch.values() if q_torch.is_nested else q_torch k_torch_effective = k_torch.values() if k_torch.is_nested else k_torch v_torch_effective = v_torch.values() if v_torch.is_nested else v_torch - k_scale_torch_effective = k_scale_torch.values() if k_scale_torch.is_nested else k_scale_torch - v_scale_torch_effective = v_scale_torch.values() if v_scale_torch.is_nested else v_scale_torch + k_scale_torch_effective = ( + k_scale_torch.values() if k_scale_torch.is_nested else k_scale_torch + ) + v_scale_torch_effective = ( + v_scale_torch.values() if v_scale_torch.is_nested else v_scale_torch + ) o_torch_effective = o_torch.values() if o_torch.is_nested else o_torch m_torch_effective = m_torch.values() if m_torch.is_nested else m_torch l_torch_effective = l_torch.values() if l_torch.is_nested else l_torch - o_partial_torch_effective = o_partial_torch.values() if o_partial_torch.is_nested else o_partial_torch - m_partial_torch_effective = m_partial_torch.values() if m_partial_torch.is_nested else m_partial_torch - l_partial_torch_effective = l_partial_torch.values() if l_partial_torch.is_nested else l_partial_torch + o_partial_torch_effective = ( + o_partial_torch.values() if o_partial_torch.is_nested else o_partial_torch + ) + m_partial_torch_effective = ( + m_partial_torch.values() if m_partial_torch.is_nested else m_partial_torch + ) + l_partial_torch_effective = ( + l_partial_torch.values() if l_partial_torch.is_nested else l_partial_torch + ) one_workspace_bytes = ( q_torch_effective.numel() * q_torch_effective.element_size() + k_torch_effective.numel() * k_torch_effective.element_size() @@ -1820,9 +1952,12 @@ def run( + o_torch_effective.numel() * o_torch_effective.element_size() + m_torch_effective.numel() * m_torch_effective.element_size() + l_torch_effective.numel() * l_torch_effective.element_size() - + o_partial_torch_effective.numel() * o_partial_torch_effective.element_size() - + m_partial_torch_effective.numel() * m_partial_torch_effective.element_size() - + l_partial_torch_effective.numel() * l_partial_torch_effective.element_size() + + o_partial_torch_effective.numel() + * o_partial_torch_effective.element_size() + + m_partial_torch_effective.numel() + * m_partial_torch_effective.element_size() + + l_partial_torch_effective.numel() + * l_partial_torch_effective.element_size() ) workspace_count = testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations @@ -1842,6 +1977,7 @@ def run( return exec_time # Return execution time in microseconds + if __name__ == "__main__": def parse_comma_separated_ints(s: str): @@ -1855,61 +1991,45 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="Example of FMHA on Blackwell.") parser.add_argument( - "--batches","--batch","--b", - type=int, - default=1, - help="batch size" + "--batches", "--batch", "--b", type=int, default=1, help="batch size" ) parser.add_argument( - "--seqlen","--seqlen_k","--seq","--s", + "--seqlen", + "--seqlen_k", + "--seq", + "--s", type=int, default=1024, - help="key/value sequence length" + help="key/value sequence length", + ) + + parser.add_argument("--heads_q", "--h_q", type=int, default=32, help="query heads") + + parser.add_argument( + "--heads_k", "--h_k", type=int, default=4, help="key/value heads" ) parser.add_argument( - "--heads_q","--h_q", - type=int, - default=32, - help="query heads" + "--headdim", "--d", type=int, default=512, help="head dimmension" ) parser.add_argument( - "--heads_k","--h_k", - type=int, - default=4, - help="key/value heads" - ) - - parser.add_argument( - "--headdim","--d", + "--block_scaledim", + "--bs", type=int, default=512, - help="head dimmension" + help="headdim per scale factor", ) parser.add_argument( - "--block_scaledim","--bs", - type=int, - default=512, - help="headdim per scale factor" - ) - - parser.add_argument( - "--kv_splits","--splits", + "--kv_splits", + "--splits", type=int, default=0, help="threadblocks per sequence", ) - parser.add_argument( - "--kv_cluster_dim","--cluster", - type=int, - default=1, - help="threadblocks per partial buffer", - ) - parser.add_argument( "--q_dtype", type=cutlass.dtype, @@ -1957,7 +2077,8 @@ if __name__ == "__main__": ) parser.add_argument( - "--scale_s","--scale", + "--scale_s", + "--scale", type=float, default=0.0, help="score (Q*K) scale factor; if zero, defaults to 1/sqrt(D)", @@ -1990,29 +2111,7 @@ if __name__ == "__main__": help="Use circular buffer tensor sets to ensure L2 cold cache", ) - args = parser.parse_args() - - run( - args.batches, - args.seqlen, - args.heads_q, - args.heads_k, - args.headdim, - args.block_scaledim, - args.kv_splits, - args.kv_cluster_dim, - args.q_dtype, - args.kv_dtype, - args.o_dtype, - args.acc_dtype, - args.tolerance, - args.scale_q, - args.scale_o, - args.scale_s, - args.warmup_iterations, - args.iterations, - args.skip_ref_check, - args.use_cold_l2, - ) + kwargs = vars(parser.parse_args()) + run(**kwargs) print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py new file mode 100644 index 00000000..aaaff2ac --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py @@ -0,0 +1,2030 @@ +# 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 argparse +import math +import os +import sys +from typing import Type, Tuple, Optional + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Int32, Int64, Float32 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from helpers import fmha_helpers as fmha_utils +from blackwell.mixed_input_fmha import prefill_helpers as prefill_utils + + +class MixedInputFusedMultiHeadAttentionPrefillD256: + def __init__( + self, + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + is_persistent: bool, + mask_type: fmha_utils.MaskEnum, + ): + self.qk_acc_dtype = qk_acc_dtype + self.pv_acc_dtype = pv_acc_dtype + self.cta_tiler = (128, 128, 256) + self.qk_mma_tiler = ( + self.cta_tiler[0] * 2, # default 2cta + self.cta_tiler[1], + min(self.cta_tiler[2], 128), # avoid too large GemmK + ) + self.pv_mma_tiler = self.qk_mma_tiler # keep BMM1 & BMM2 at the same pace + self.pv_block_tiler = ( + self.pv_mma_tiler[0] // 2, # default 2cta + self.pv_mma_tiler[1], + self.pv_mma_tiler[2], + ) + self.scale_granularity = scale_granularity + self.iterations_qk = self.cta_tiler[2] // self.qk_mma_tiler[2] + self.iterations_pv = self.cta_tiler[2] // self.pv_mma_tiler[1] + self.cluster_shape_mn = (2, 1) # use 2x1 cluster by default + self.tmem_warp_shape_mn = (4, 1) + self.is_persistent = is_persistent + self.mask_type = mask_type + self.transform_warp_ids = (0, 1, 2, 3, 4, 5, 6, 7) # i8 -> bf16 for kv + self.softmax_warp_ids = (8, 9, 10, 11) # softmax + self.correction_warp_ids = (12, 13, 14, 15) # correction + self.mma_warp_id = 16 # mma + self.load_warp_id = 17 # load + self.empty_warp_ids = (18, 19) # empty + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + self.tmem_alloc_sync_bar_id = 1 + self.tmem_s_offset = 256 + self.tmem_p_offset = self.tmem_s_offset + self.tmem_o_offset = 0 + self.num_regs_softmax = 256 + self.num_regs_correction = 112 + self.num_regs_other = 32 + self.num_regs_transform = 40 + self.buffer_align_bytes = 1024 + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.transform_warp_ids, + *self.softmax_warp_ids, + *self.correction_warp_ids, + self.load_warp_id, + self.mma_warp_id, + *self.empty_warp_ids, + ) + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the FMHA kernel operation. + + This method initializes and configures various attributes required for the + execution of the fused multi-head attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.q_stage = self.iterations_qk + self.kv_stage = 4 + self.scale_k_stage = self.kv_stage + self.scale_v_stage = self.kv_stage + self.qk_acc_stage = 2 + self.pv_acc_stage = 1 + self.kv_trans_stage = 2 + + @cute.jit + def __call__( + self, + q_iter: cute.Pointer, + k_iter: cute.Pointer, + v_iter: cute.Pointer, + o_iter: cute.Pointer, + scale_k_iter: cute.Pointer, + scale_v_iter: cute.Pointer, + problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32, Int32], + scale_softmax_log2: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + stream: cuda.CUstream, + ): + self._setup_attributes() + b, s_q, s_k, h_q, h_k, d = problem_shape + h_r = h_q // h_k + self.d_r = self.cta_tiler[2] // self.scale_granularity + # (s, d, ((h_r, h_k), b)) + q_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=(d, 1, ((d * s_q, d * s_q * h_r), h_r * h_k * s_q * d)), + ) + q = cute.make_tensor(q_iter, q_layout) + # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + k_layout = cute.make_layout( + (s_k, d, ((h_r, h_k), b)), + stride=(d, 1, ((0, d * s_k), h_k * s_k * d)), + ) + k = cute.make_tensor(k_iter, k_layout) + # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast + v_layout = cute.make_layout( + (d, s_k, ((h_r, h_k), b)), + stride=(1, d, ((0, d * s_k), h_k * s_k * d)), + ) + v = cute.make_tensor(v_iter, v_layout) + # (s, d, ((h_r, h_k), b)) + # set divby for better gmem store vectorization + o_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=( + cute.assume(d, divby=256), + 1, + ( + ( + cute.assume(d * s_q, divby=256), + cute.assume(d * s_q * h_r, divby=256), + ), + cute.assume(h_r * h_k * s_q * d, divby=256), + ), + ), + ) + o = cute.make_tensor(o_iter, o_layout) + # (d_r * s, ((h_r, h_k), b)) + scale_k_layout = cute.make_layout( + (s_k * self.d_r, ((h_r, h_k), b)), + stride=(1, ((0, self.d_r * s_k), s_k * self.d_r * h_k)), + ) + scale_k = cute.make_tensor(scale_k_iter, scale_k_layout) + # (d_r * s, ((h_r, h_k), b)) + scale_v_layout = cute.make_layout( + (self.d_r * s_k, ((h_r, h_k), b)), + stride=(1, ((0, self.d_r * s_k), s_k * self.d_r * h_k)), + ) + scale_v = cute.make_tensor(scale_v_iter, scale_v_layout) + + self.q_dtype = q.element_type + self.k_dtype = k.element_type + self.v_dtype = v.element_type + self.o_dtype = o.element_type + self.p_dtype = self.q_dtype # pv should has the same dtype + self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.p_dtype.width + self.scale_k_dtype = scale_k.element_type + self.scale_v_dtype = scale_v.element_type + + self.tile_sched_params, grid = fmha_utils.compute_grid( + o.shape, + self.cta_tiler, + self.is_persistent, + ) + + self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.o_layout = utils.LayoutEnum.from_tensor(o) + cta_group = tcgen05.CtaGroup.TWO + p_major_mode = tcgen05.OperandMajorMode.K + p_source = tcgen05.OperandSource.TMEM + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.qk_acc_dtype, + cta_group, + self.qk_mma_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + p_major_mode, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + self.pv_mma_tiler[:2], + p_source, + ) + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.pv_block_tiler[:2] + + q_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.q_stage, + ) + k_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.kv_stage, + ) + k_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, k_smem_layout_staged.outer + ) + k_trans_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.kv_trans_stage, + ) + p_tmem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.pv_mma_tiler, + self.p_dtype, + self.qk_acc_stage, + ) + p_tmem_layout = cute.select(p_tmem_layout_staged, mode=[0, 1, 2]) + v_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.kv_stage, + ) + v_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, v_smem_layout_staged.outer + ) + v_trans_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.kv_trans_stage, + ) + scale_k_smem_layout, self.scale_k_tiler, scale_k_s2r_view_layout = ( + prefill_utils.get_scale_smem_layout( + self.scale_granularity, + self.d_r, + self.qk_mma_tiler, + self.k_major_mode, + ) + ) + scale_k_smem_layout_staged = cute.append( + scale_k_smem_layout, + cute.make_layout( + (self.scale_k_stage), + stride=(cute.cosize(scale_k_smem_layout.outer)), + ), + ) + scale_k_s2r_view_layout_staged = cute.append( + scale_k_s2r_view_layout, + cute.make_layout( + (self.scale_k_stage), + stride=(cute.cosize(scale_k_s2r_view_layout)), + ), + ) + scale_v_smem_layout, self.scale_v_tiler, scale_v_s2r_view_layout = ( + prefill_utils.get_scale_smem_layout( + self.scale_granularity, + self.d_r, + self.pv_mma_tiler, + self.v_major_mode, + ) + ) + scale_v_smem_layout_staged = cute.append( + scale_v_smem_layout, + cute.make_layout( + (self.scale_v_stage), + stride=(cute.cosize(scale_v_smem_layout.outer)), + ), + ) + scale_v_s2r_view_layout_staged = cute.append( + scale_v_s2r_view_layout, + cute.make_layout( + (self.scale_v_stage), + stride=(cute.cosize(scale_v_s2r_view_layout)), + ), + ) + + tma_load_q_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + # For TMA Async, use one cta to sync with corresponding cta only + tma_load_kv_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp( + tcgen05.CtaGroup.ONE + ) + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_q_op, + q, + q_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for K + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_kv_op, + k, + k_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_scale_k, tma_tensor_scale_k = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_kv_op, + scale_k, + scale_k_smem_layout, + (self.scale_k_tiler[0] // 2,), + ) + + # TMA load for V + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_kv_op, + v, + v_smem_layout, + self.pv_mma_tiler, + pv_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_scale_v, tma_tensor_scale_v = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_kv_op, + scale_v, + scale_v_smem_layout, + self.scale_v_tiler, + ) + + self.tma_copy_q_bytes = cute.size_in_bytes( + self.q_dtype, q_smem_layout + ) * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_kv_bytes = cute.size_in_bytes(self.k_dtype, k_smem_layout) + self.tma_copy_scale_k_bytes = cute.size_in_bytes( + self.scale_k_dtype, scale_k_smem_layout + ) + self.tma_copy_scale_v_bytes = cute.size_in_bytes( + self.scale_v_dtype, scale_v_smem_layout + ) + + @cute.struct + class SharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[Int64, self.q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_stage * 2] + load_scale_k_mbar_ptr: cute.struct.MemRange[Int64, self.scale_k_stage * 2] + load_scale_v_mbar_ptr: cute.struct.MemRange[Int64, self.scale_v_stage * 2] + dequant_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_trans_stage * 2] + mma_s_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + s_corr_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + sum_mbar_ptr: cute.struct.MemRange[Int64, 2] + mma_o_mbar_ptr: cute.struct.MemRange[Int64, self.pv_acc_stage * 2] + tmem_dealloc_mbar_ptr: Int64 + tmem_holding_buf: Int32 + + self.shared_storage = SharedStorage + + grid = cute.round_up(grid, self.cluster_shape_mnk) + + # Launch the kernel synchronously + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_scale_k, + tma_tensor_scale_k, + tma_atom_v, + tma_tensor_v, + tma_atom_scale_v, + tma_tensor_scale_v, + o, + scale_softmax_log2, + scale_output, + window_size_left, + window_size_right, + self.cluster_layout_vmnk, + q_smem_layout_staged, + k_smem_layout_staged, + k_trans_smem_layout_staged, + scale_k_smem_layout_staged, + scale_k_s2r_view_layout_staged, + p_tmem_layout, + v_smem_layout_staged, + v_trans_smem_layout_staged, + scale_v_smem_layout_staged, + scale_v_s2r_view_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + pv_tiled_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_qdl: cute.Tensor, + tma_atom_k: cute.CopyAtom, + mK_kdl: cute.Tensor, + tma_atom_scale_k: cute.CopyAtom, + mScaleK_kdl: cute.Tensor, + tma_atom_v: cute.CopyAtom, + mV_dkl: cute.Tensor, + tma_atom_scale_v: cute.CopyAtom, + mScaleV_dkl: cute.Tensor, + mO_qdl: cute.Tensor, + scale_softmax_log2: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + cluster_layout_vmnk: cute.Layout, + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + k_trans_smem_layout_staged: cute.ComposedLayout, + scale_k_smem_layout_staged: cute.ComposedLayout, + scale_k_s2r_view_layout_staged: cute.Layout, + p_tmem_layout: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + v_trans_smem_layout_staged: cute.ComposedLayout, + scale_v_smem_layout_staged: cute.ComposedLayout, + scale_v_s2r_view_layout_staged: cute.Layout, + epi_tile: cute.Tile, + tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # Prefetch tma desc + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_scale_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_scale_v) + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(qk_tiled_mma.thr_id.shape) + 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 + ) + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.q_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + tx_count=self.tma_copy_q_bytes, + barrier_storage=storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_kv_producer, load_kv_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_kv_bytes, + barrier_storage=storage.load_kv_mbar_ptr.data_ptr(), + tidx=0, + defer_sync=True, + ).make_participants() + load_scale_k_producer, load_scale_k_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.scale_k_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_scale_k_bytes, + barrier_storage=storage.load_scale_k_mbar_ptr.data_ptr(), + tidx=0, + defer_sync=True, + ).make_participants() + load_scale_v_producer, load_scale_v_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.scale_v_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_scale_v_bytes, + barrier_storage=storage.load_scale_v_mbar_ptr.data_ptr(), + tidx=0, + defer_sync=True, + ).make_participants() + dequant_kv_producer, dequant_kv_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.kv_trans_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + barrier_storage=storage.dequant_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_s_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + p_mma_producer, p_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + barrier_storage=storage.p_mma_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + s_corr_producer, s_corr_consumer = pipeline.PipelineAsync.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.softmax_warp_ids), + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.correction_warp_ids), + ), + barrier_storage=storage.s_corr_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + sum_producer, sum_consumer = pipeline.PipelineAsync.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.softmax_warp_ids), + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.correction_warp_ids), + ), + barrier_storage=storage.sum_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + mma_o_producer, mma_o_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.pv_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.correction_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_o_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=self.threads_per_warp + * len( + (self.mma_warp_id, *self.softmax_warp_ids, *self.correction_warp_ids) + ), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.correction_warp_ids[0], + is_two_cta=True, + 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) + + sK_trans = smem.allocate_tensor( + element_type=self.q_dtype, + layout=k_trans_smem_layout_staged.outer, + swizzle=k_trans_smem_layout_staged.inner, + byte_alignment=128, + ) + sV_trans_ptr = cute.recast_ptr( + sK_trans.iterator, v_trans_smem_layout_staged.inner + ) + sV_trans = cute.make_tensor(sV_trans_ptr, v_trans_smem_layout_staged.outer) + sQ = smem.allocate_tensor( + element_type=self.q_dtype, + layout=q_smem_layout_staged.outer, + swizzle=q_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleK = smem.allocate_tensor( + element_type=self.scale_k_dtype, + layout=scale_k_smem_layout_staged.outer, + swizzle=scale_k_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleK_s2r_view = cute.make_tensor( + sScaleK.iterator, scale_k_s2r_view_layout_staged + ) + sScaleV = smem.allocate_tensor( + element_type=self.scale_v_dtype, + layout=scale_v_smem_layout_staged.outer, + swizzle=scale_v_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleV_s2r_view = cute.make_tensor( + sScaleV.iterator, scale_v_s2r_view_layout_staged + ) + sK = smem.allocate_tensor( + element_type=self.k_dtype, + layout=k_smem_layout_staged.outer, + swizzle=k_smem_layout_staged.inner, + byte_alignment=128, + ) + sV_ptr = cute.recast_ptr(sK.iterator, v_smem_layout_staged.inner) + sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer) + + sSum = smem.allocate_tensor( + element_type=self.qk_acc_dtype, + layout=cute.make_layout(len(self.softmax_warp_ids) * self.threads_per_warp), + byte_alignment=128, + ) + + qk_thr_mma = qk_tiled_mma.get_slice(mma_tile_coord_v) + pv_thr_mma = pv_tiled_mma.get_slice(mma_tile_coord_v) + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tSrK_trans = qk_thr_mma.make_fragment_B(sK_trans) + tOrV_trans = pv_thr_mma.make_fragment_B(sV_trans) + qk_acc_shape = pv_thr_mma.partition_shape_C( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + # (atomV, restM, restN, accStage) + tStS = qk_tiled_mma.make_fragment_C( + cute.append(qk_acc_shape, self.qk_acc_stage) + ) + pv_acc_shape = pv_thr_mma.partition_shape_C( + cute.select(self.pv_mma_tiler, mode=[0, 1]) + ) + # (atomV, restM, restN) + tOtO = pv_thr_mma.make_fragment_C(pv_acc_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + self.iterations_pv, + stride=self.pv_mma_tiler[1] // self.tmem_warp_shape_mn[1], + ), + ) + tStS = cute.make_tensor(tStS.iterator + self.tmem_s_offset, tStS.layout) + tOtO_staged = cute.make_tensor(tOtO.iterator + self.tmem_o_offset, tOtO_layout) + # Local_tile partition global tensors + q_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # (bM, bK, restM, restK, loopM, loopK, loopL) + gQ_qdl = cute.flat_divide(mQ_qdl, cute.select(self.qk_mma_tiler, mode=[0, 2])) + tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl) + tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, + block_in_cluster_coord_vmnk[2], + q_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ_qdl, 0, 3), + ) + kv_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # (bN, bK, loopN, loopK, loopL) + gK_kdl = cute.flat_divide(mK_kdl, cute.select(self.qk_mma_tiler, mode=[1, 2])) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + # (blk, loopBlk, loopL) + gScaleK_kdl = cute.flat_divide(mScaleK_kdl, self.scale_k_tiler) + # Deal with 2cta + gScaleK_kdl_ = cute.logical_divide(gScaleK_kdl, (self.scale_k_tiler[0] // 2,))[ + (None, mma_tile_coord_v), None, None + ] + tKsScaleK, tKgScaleK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_scale_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + sScaleK, + gScaleK_kdl_, + ) + # (bN, bK, loopN, loopK, loopL) + gV_dkl = cute.flat_divide(mV_dkl, cute.select(self.pv_mma_tiler, mode=[1, 2])) + tOgV_dkl = pv_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tOgV_dkl, 0, 3), + ) + # (bBlk, loopBlk, loopL) + gScaleV_dkl = cute.flat_divide(mScaleV_dkl, self.scale_v_tiler) + tVsScaleV, tVgScaleV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_scale_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + sScaleV, + gScaleV_dkl, + ) + # (bM, bN, loopM, loopN, loopL) + gO_qdl = cute.flat_divide(mO_qdl, cute.select(self.pv_block_tiler, mode=[0, 1])) + cO_qdl = cute.flat_divide( + cute.make_identity_tensor(mO_qdl.shape), + cute.select(self.pv_block_tiler, mode=[0, 1]), + ) + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + tile_sched = fmha_utils.create_fmha_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + # Cluster wait + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + # ((atom_v, rest_v), RestK) + tQgQ = tQgQ_qdl[None, mma_block_coord[0], None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tKgK = tKgK_kdl[None, None, None, mma_block_coord[2]] + tKgScaleK = tKgScaleK_kdl[None, None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tVgV = tVgV_dkl[None, None, None, mma_block_coord[2]] + tVgScaleV = tVgScaleV_dkl[None, None, mma_block_coord[2]] + load_kv_producer, load_scale_k_producer, load_q_producer = ( + prefill_utils.load_qk( # Q & K0 & ScaleK0 + self.iterations_qk, + kv_step=0, + k_args=(tKgK, tKsK, tma_atom_k, load_kv_producer), + scale_k_args=( + tKgScaleK, + tKsScaleK, + tma_atom_scale_k, + load_scale_k_producer, + ), + q_args=(tQgQ, tQsQ, tma_atom_q, load_q_producer), + ) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + load_kv_producer, load_scale_k_producer = ( + prefill_utils.load_qk( # Ki & ScaleKi + self.iterations_qk, + kv_step=step, + k_args=(tKgK, tKsK, tma_atom_k, load_kv_producer), + scale_k_args=( + tKgScaleK, + tKsScaleK, + tma_atom_scale_k, + load_scale_k_producer, + ), + ) + ) + load_kv_producer, load_scale_v_producer = ( + prefill_utils.load_v( # Vi-1 & ScaleVi-1 + self.iterations_pv, + kv_step=step - 1, + v_args=(tVgV, tVsV, tma_atom_v, load_kv_producer), + scale_v_args=( + tVgScaleV, + tVsScaleV, + tma_atom_scale_v, + load_scale_v_producer, + ), + ) + ) + load_kv_producer, load_scale_v_producer = ( + prefill_utils.load_v( # Vend & ScaleVend + self.iterations_pv, + kv_step=seqlen_kv_loop_steps - 1, + v_args=(tVgV, tVsV, tma_atom_v, load_kv_producer), + scale_v_args=( + tVgScaleV, + tVsScaleV, + tma_atom_scale_v, + load_scale_v_producer, + ), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_kv_producer.tail() + load_scale_k_producer.tail() + load_scale_v_producer.tail() + load_q_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + tmem.wait_for_alloc() + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + load_q_releaser = load_q_consumer.clone() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + if seqlen_kv_loop_steps > 1: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QK0 + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + load_q_consumer, + None, + dequant_kv_consumer, + ), + ) + ) + for i in cutlass.range(1, seqlen_kv_loop_steps - 1, 1, unroll=1): + mma_s_producer, _, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QKi + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + (mma_s_producer, None, None, dequant_kv_consumer), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + ) = self.mma_pv( # PVi + (pv_tiled_mma, pv_thr_mma), + (tOtO_staged, tStS, tOrV_trans, p_tmem_layout), + (p_mma_consumer, mma_o_producer, dequant_kv_consumer), + ) + mma_s_producer, _, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QKend needs to release Q + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + None, + load_q_releaser, + dequant_kv_consumer, + ), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + ) = self.mma_pv( # PVend-1 + (pv_tiled_mma, pv_thr_mma), + (tOtO_staged, tStS, tOrV_trans, p_tmem_layout), + (p_mma_consumer, mma_o_producer, dequant_kv_consumer), + ) + else: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QK0 + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + load_q_consumer, + load_q_releaser, + dequant_kv_consumer, + ), + ) + ) + pv_tiled_mma, p_mma_consumer, mma_o_producer, dequant_kv_consumer = ( + self.mma_pv( # PVend + (pv_tiled_mma, pv_thr_mma), + (tOtO_staged, tStS, tOrV_trans, p_tmem_layout), + (p_mma_consumer, mma_o_producer, dequant_kv_consumer), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + mma_s_producer.tail() + mma_o_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx < self.correction_warp_ids[0] + and warp_idx >= self.softmax_warp_ids[0] + ): + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.wait_for_alloc() + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + unmask_steps = fmha_utils.FusedMask.get_unmasked_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + cS_base = cute.make_identity_tensor( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + cS = cute.domain_offset( + (mma_block_coord[0] * self.qk_mma_tiler[0], 0), cS_base + ) + tScS = qk_thr_mma.partition_C(cS) + row_max = -Float32.inf + row_max_prev = -Float32.inf + row_sum = 0.0 + for step in cutlass.range(seqlen_kv_loop_steps, unroll=1): + cS_iter = cute.domain_offset((0, step * self.qk_mma_tiler[1]), cS) + tScS_iter = qk_thr_mma.partition_C(cS_iter) + # Si -> Pi + ( + row_max, + row_sum, + mma_s_consumer, + p_mma_producer, + s_corr_producer, + ) = self.softmax_step( + (step >= unmask_steps, window_size_left, window_size_right), + ( + row_max_prev, + row_sum, + seqlen_q, + seqlen_k, + scale_softmax_log2, + ), + (tStS, tScS_iter), + (mma_s_consumer, p_mma_producer, s_corr_producer), + ) + row_max_prev = row_max + sum_producer = self.store_sum(row_sum, sSum, sum_producer) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_mma_producer.tail() + s_corr_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Correction + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.mma_warp_id and warp_idx >= self.correction_warp_ids[0]: + cute.arch.setmaxregister_increase(self.num_regs_correction) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + gO_staged = gO_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + cO_staged = cO_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + cS = cute.make_identity_tensor( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + tScS = qk_thr_mma.partition_C(cS) + # Empty step as the first step is no need for correction + stats_handle = s_corr_consumer.wait_and_advance() + stats_handle.release() + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + # Oi-1 -> Oi + mma_o_consumer, s_corr_consumer = self.correction_rescale( + scale_softmax_log2, + (s_corr_consumer, tStS, tScS), + (mma_o_consumer, tOtO_staged, cO_staged), + epi_tile, + ) + # O_partial -> O_final + mma_o_consumer, sum_consumer = self.correction_epilog( + (seqlen_q, scale_output), + (sum_consumer, sSum), + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged), + epi_tile, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Trans + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.softmax_warp_ids[0]: + cute.arch.setmaxregister_decrease(self.num_regs_transform) + qk_thr_mma_leader_cta = qk_tiled_mma.get_slice(0) + pv_thr_mma_leader_cta = pv_tiled_mma.get_slice(0) + sScaleK_ = qk_thr_mma_leader_cta.partition_B(sScaleK_s2r_view) + sScaleV_ = pv_thr_mma_leader_cta.partition_B(sScaleV_s2r_view) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + load_kv_consumer, load_scale_k_consumer, dequant_kv_producer = ( + prefill_utils.dequant_k( # K0 + self.iterations_qk, + self.transform_warp_ids, + (self.k_dtype, self.q_dtype), + (sK, sScaleK_, sK_trans), + (load_kv_consumer, load_scale_k_consumer, dequant_kv_producer), + ) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + load_kv_consumer, load_scale_k_consumer, dequant_kv_producer = ( + prefill_utils.dequant_k( # Ki + self.iterations_qk, + self.transform_warp_ids, + (self.k_dtype, self.q_dtype), + (sK, sScaleK_, sK_trans), + ( + load_kv_consumer, + load_scale_k_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + prefill_utils.dequant_v( # Vi-1 + self.iterations_pv, + self.transform_warp_ids, + (self.v_dtype, self.q_dtype), + (sV, sScaleV_, sV_trans), + ( + load_kv_consumer, + load_scale_v_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + prefill_utils.dequant_v( # Vend + self.iterations_pv, + self.transform_warp_ids, + (self.v_dtype, self.q_dtype), + (sV, sScaleV_, sV_trans), + (load_kv_consumer, load_scale_v_consumer, dequant_kv_producer), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + dequant_kv_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Empty + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx > self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + return + + @cute.jit + def mma_pv( + self, + mma_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + ): + pv_tiled_mma, pv_thr_mma = mma_args + tOtO_staged, tStS, tOrV_trans, p_tmem_layout = tensor_args + p_mma_consumer, mma_o_producer, dequant_kv_consumer = pipeline_args + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + if is_leader_cta: + p_handle = p_mma_consumer.wait_and_advance() + o_handle = mma_o_producer.acquire_and_advance() + pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_pv, unroll=1): + v_trans_handle = dequant_kv_consumer.wait_and_advance() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + tOtO_slice = tOtO_staged[None, None, None, iter] + tStS_slice = tStS[None, None, None, p_handle.index] + tP = cute.make_tensor(tStS_slice.iterator, p_tmem_layout.outer) + tOrP = pv_thr_mma.make_fragment_A(tP) + tOrP_slice = cute.make_tensor( + cute.recast_ptr(tStS_slice.iterator, dtype=self.p_dtype), + tOrP.layout, + ) + tOrV_trans_slice = tOrV_trans[None, None, None, v_trans_handle.index] + num_kphases = cute.size(tOrV_trans_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + pv_tiled_mma, + tOtO_slice, + tOrP_slice[kphase_coord], + tOrV_trans_slice[kphase_coord], + tOtO_slice, + ) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_trans_handle.release() + o_handle.commit() + p_handle.release() + return pv_tiled_mma, p_mma_consumer, mma_o_producer, dequant_kv_consumer + + @cute.jit + def softmax_step( + self, + mask_args: Tuple, + value_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]: + need_apply_mask, window_size_left, window_size_right = mask_args + row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2 = value_args + tStS, tScS = tensor_args + mma_s_consumer, p_mma_producer, s_corr_producer = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + s_handle = mma_s_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, s_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_load.partition_S(tStS_slice) + tTMEM_LOADcS = thr_load.partition_D(tScS_slice) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + cute.copy(tmem_tiled_load, tTMEM_LOADtS, tTMEM_LOADrS) + cute.arch.fence_view_async_tmem_load() + s_handle.release() + if need_apply_mask: + fmha_utils.FusedMask.apply_mask( + self.mask_type, + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + old_row_max = row_max + row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + + stats_handle = s_corr_producer.acquire_and_advance() + stats_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], 2)) + ) + stats_c_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], 2)) + ) + tOtStats = cute.make_tensor( + tStS_slice.iterator + self.tilePlikeFP32, stats_layout + ) + tOcStats = cute.make_tensor(tScS_slice.iterator, stats_c_layout) + tmem_store_stats_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_store_stats = tcgen05.make_tmem_copy(tmem_store_stats_atom, tOtStats) + thr_tmem_store_stats = tiled_tmem_store_stats.get_slice(thread_idx) + tTMEM_STOREcStats = thr_tmem_store_stats.partition_S(tOcStats) + tTMEM_STORErStats = cute.make_rmem_tensor( + tTMEM_STOREcStats.shape, self.qk_acc_dtype + ) + tTMEM_STORErStats[0] = old_row_max + tTMEM_STORErStats[1] = row_max_safe + tTMEM_STOREtStats = thr_tmem_store_stats.partition_D(tOtStats) + cute.copy(tiled_tmem_store_stats, tTMEM_STORErStats, tTMEM_STOREtStats) + cute.arch.fence_view_async_tmem_store() + stats_handle.commit() + + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.p_dtype) + for k in cutlass.range(cute.size(tTMEM_LOADrS), vectorize=True): + tTMEM_LOADrS[k] = tTMEM_LOADrS[k] * scale + minus_row_max_scale + tTMEM_LOADrS[k] = cute.math.exp2(tTMEM_LOADrS[k], fastmath=True) + s_vec = tTMEM_LOADrS.load() + tTMEM_STORErP.store(s_vec.to(self.p_dtype)) + + p_handle = p_mma_producer.acquire_and_advance() + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tilePlikeFP32 = tStS_slice.shape[1] // Float32.width * self.p_dtype.width + tStS_P_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], tilePlikeFP32)) + ) + tStS_P = cute.make_tensor(tStS_slice.iterator, tStS_P_layout) + tScS_P_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], tilePlikeFP32)) + ) + tScS_P = cute.make_tensor(tScS_slice.iterator, tScS_P_layout) + tmem_tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) + thr_store = tmem_tiled_store.get_slice(thread_idx) + tTMEM_STOREtP = thr_store.partition_D(tStS_P) + tTMEM_STOREcS = thr_store.partition_S(tScS_P) + tTMEM_STORErP_ = cute.make_tensor( + cute.recast_ptr(tTMEM_STORErP.iterator, dtype=self.qk_acc_dtype), + tTMEM_STOREcS.shape, + ) + cute.copy(tmem_tiled_store, tTMEM_STORErP_, tTMEM_STOREtP) + cute.arch.fence_view_async_tmem_store() + + p_handle.commit() + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) * 0.5 + # TODO: calc row sum with TensorSSA + row_sum *= acc_scale + local_row_sum_0 = (row_sum, row_sum) + local_row_sum_1 = (0.0, 0.0) + local_row_sum_2 = (0.0, 0.0) + local_row_sum_3 = (0.0, 0.0) + reduction_unroll = 4 + frg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll + tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) + for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2): + local_row_sum_0 = cute.arch.add_packed_f32x2( + local_row_sum_0, (tTMEM_LOADrS_frg[j, 0], tTMEM_LOADrS_frg[j + 1, 0]) + ) + local_row_sum_1 = cute.arch.add_packed_f32x2( + local_row_sum_1, (tTMEM_LOADrS_frg[j, 1], tTMEM_LOADrS_frg[j + 1, 1]) + ) + local_row_sum_2 = cute.arch.add_packed_f32x2( + local_row_sum_2, (tTMEM_LOADrS_frg[j, 2], tTMEM_LOADrS_frg[j + 1, 2]) + ) + local_row_sum_3 = cute.arch.add_packed_f32x2( + local_row_sum_3, (tTMEM_LOADrS_frg[j, 3], tTMEM_LOADrS_frg[j + 1, 3]) + ) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) + local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) + row_sum = local_row_sum_0[0] + local_row_sum_0[1] + return row_max, row_sum, mma_s_consumer, p_mma_producer, s_corr_producer + + @cute.jit + def correction_rescale( + self, + scale_softmax_log2: Float32, + stats_args: tuple, + o_args: tuple, + epi_tile: cute.Tile, + ) -> pipeline.PipelineConsumer: + (s_corr_consumer, tStS, tScS) = stats_args + (mma_o_consumer, tOtO_staged, cO_staged) = o_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + + stats_handle = s_corr_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, stats_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + stats_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], 2)) + ) + stats_c_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], 2)) + ) + tOtStats = cute.make_tensor( + tStS_slice.iterator + self.tilePlikeFP32, stats_layout + ) + tOcStats = cute.make_tensor(tScS_slice.iterator, stats_c_layout) + tmem_load_stats_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_load_stats = tcgen05.make_tmem_copy(tmem_load_stats_atom, tOtStats) + thr_tmem_load_stats = tiled_tmem_load_stats.get_slice(thread_idx) + tTMEM_LOADtStats = thr_tmem_load_stats.partition_S(tOtStats) + tTMEM_LOADcStats = thr_tmem_load_stats.partition_D(tOcStats) + tTMEM_LOADrStats = cute.make_rmem_tensor( + tTMEM_LOADcStats.shape, self.qk_acc_dtype + ) + cute.copy(tiled_tmem_load_stats, tTMEM_LOADtStats, tTMEM_LOADrStats) + + scale = scale_softmax_log2 * (tTMEM_LOADrStats[0] - tTMEM_LOADrStats[1]) + scale = cute.math.exp2(scale, fastmath=True) + stats_handle.release() + o_handle = mma_o_consumer.wait_and_advance() + for iter in cutlass.range(self.iterations_pv, unroll_full=True): + tOtO = tOtO_staged[(None, None), 0, 0, iter] + cO = cO_staged[None, None, iter] + tOtO_epi = cute.zipped_divide(tOtO, epi_tile) + cO_epi = cute.zipped_divide(cO, epi_tile) + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(16)), + self.pv_acc_dtype, + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_epi) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(16)), + self.pv_acc_dtype, + ) + tmem_store_atom = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_epi) + thr_store = tmem_store_atom.get_slice(thread_idx) + tTMEM_LOADtO = thr_load.partition_S(tOtO_epi) + tTMEM_LOADcO = thr_load.partition_D(cO_epi) + tTMEM_STOREtO = thr_store.partition_D(tOtO_epi) + tTMrO = cute.make_rmem_tensor_like( + cute.append( + cute.make_layout(tTMEM_LOADcO[None, 0, 0].shape), + cute.make_layout( + 2, stride=cute.size(tTMEM_LOADcO[None, 0, 0].shape) + ), + ), + self.pv_acc_dtype, + ) + tTMEM_LOADtO_0 = tTMEM_LOADtO[None, 0, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_0, tTMrO[None, 0]) + iter_num = cute.size(tTMEM_LOADtO, mode=[1]) + for i in cutlass.range(1, iter_num, unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_i, tTMrO[None, i % 2]) + for j in cutlass.range( + cute.size(tTMrO, mode=[0]), unroll_full=True, vectorize=True + ): + tTMrO[j, (i - 1) % 2] = tTMrO[j, (i - 1) % 2] * scale + tTMEM_STOREtO_prev_i = tTMEM_STOREtO[None, i - 1, 0] + cute.copy( + tmem_store_atom, tTMrO[None, (i - 1) % 2], tTMEM_STOREtO_prev_i + ) + + for j in cutlass.range( + cute.size(tTMrO, mode=[0]), unroll_full=True, vectorize=True + ): + tTMrO[j, (iter_num - 1) % 2] = tTMrO[j, (iter_num - 1) % 2] * scale + cute.copy( + tmem_store_atom, + tTMrO[None, (iter_num - 1) % 2], + tTMEM_STOREtO[None, iter_num - 1, 0], + ) + cute.arch.fence_view_async_tmem_store() + o_handle.release() + return mma_o_consumer, s_corr_consumer + + @cute.jit + def correction_epilog( + self, + value_args: Tuple, + sum_args: Tuple, + o_args: Tuple, + epi_tile: cute.Tile, + ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]: + (seqlen_q, scale_output) = value_args + (sum_consumer, sSum) = sum_args + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged) = o_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + sum_handle = sum_consumer.wait_and_advance() + row_sum = sSum[thread_idx] + cute.arch.fence_view_async_shared() + sum_handle.release() + scale = scale_output / row_sum + o_handle = mma_o_consumer.wait_and_advance() + for iter in cutlass.range(self.iterations_pv): + gO = gO_staged[None, None, iter] + cO = cO_staged[None, None, iter] + tOtO = tOtO_staged[(None, None), 0, 0, iter] + tOtO_epi = cute.zipped_divide(tOtO, epi_tile) + cO_epi = cute.zipped_divide(cO, epi_tile) + gO_epi = cute.zipped_divide(gO, epi_tile) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tmem_copy_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.pv_acc_dtype + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_epi) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_epi) + tTMEM_LOADgO = thr_tmem_load.partition_D(gO_epi) + tTMEM_LOADcO = thr_tmem_load.partition_D(cO_epi) + for i in cutlass.range(cute.size(tTMEM_LOADtO, mode=[1]), unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + tTMEM_LOADgO_i = tTMEM_LOADgO[None, i, 0] + tTMEM_LOADcO_i = tTMEM_LOADcO[None, i, 0] + tTMrO = cute.make_rmem_tensor( + tTMEM_LOADcO[None, 0, i].shape, self.pv_acc_dtype + ) + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) + for j in cutlass.range( + cute.size(tTMrO), unroll_full=True, vectorize=True + ): + tTMrO[j] = tTMrO[j] * scale + tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) + o_vec = tTMrO.load() + tSMrO.store(o_vec.to(self.o_dtype)) + if cute.elem_less(tTMEM_LOADcO_i[0][0], seqlen_q): + cute.autovec_copy(tSMrO, tTMEM_LOADgO_i) + o_handle.release() + return mma_o_consumer, sum_consumer + + def store_sum(self, row_sum, sSum, sum_producer): + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + sum_handle = sum_producer.acquire_and_advance() + sSum[thread_idx] = row_sum + cute.arch.fence_view_async_shared() + sum_handle.commit() + return sum_producer + + +def run( + q_shape: Tuple[int, int, int, int], + k_shape: Tuple[int, int, int, int], + q_dtype: Type[cutlass.Numeric], + kv_dtype: Type[cutlass.Numeric], + o_dtype: Type[cutlass.Numeric], + scale_dtype: Type[cutlass.Numeric], + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + is_persistent: bool, + is_causal: bool, + scale_q: float, + scale_k: float, + scale_v: float, + inv_scale_o: float, + scale_softmax: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool = False, + **kwargs, +): + print(f"Running Blackwell SM100 Mixed Input FMHA Prefill D256 test with:") + print(f" q_shape: {q_shape}") + print(f" k_shape: {k_shape}") + print(f" q_dtype: {q_dtype}") + print(f" kv_dtype: {kv_dtype}") + print(f" o_dtype: {o_dtype}") + print(f" scale_dtype: {scale_dtype}") + print(f" scale_granularity: {scale_granularity}") + print(f" qk_acc_dtype: {qk_acc_dtype}") + print(f" pv_acc_dtype: {pv_acc_dtype}") + print(f" is_persistent: {is_persistent}") + print(f" is_causal: {is_causal}") + print(f" scale_q: {scale_q}") + print(f" scale_k: {scale_k}") + print(f" scale_v: {scale_v}") + print(f" inv_scale_o: {inv_scale_o}") + print(f" scale_softmax: {scale_softmax}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + import cutlass.torch as cutlass_torch + + # Unpack parameters + b, h_q, s_q, d = q_shape + b_, h_k, s_k, d_ = k_shape + window_size_left, window_size_right = None, None + if is_causal: + window_size_right = 0 + + if b != b_: + raise ValueError("q & k must have the same batch size") + + if d != d_: + raise ValueError("q & k must have the same head dimension") + + if d not in {256}: + raise ValueError("head dimension must be 256") + + if d % scale_granularity != 0: + raise ValueError("head dimension must be divisible by scale_granularity") + + if scale_granularity not in {128, 256}: + raise ValueError("scale_granularity must be 128, 256") + + if h_q % h_k != 0: + raise ValueError("h_q must be divisible by h_k") + + if isinstance(s_q, tuple) and len(s_q) != b: + raise ValueError("variable_seqlen s_q must have the length of batch size") + if isinstance(s_k, tuple) and len(s_k) != b: + raise ValueError("variable_seqlen s_k must have the length of batch size") + + if q_dtype not in {cutlass.BFloat16}: + raise ValueError("in_dtype must be BFloat16") + + if o_dtype not in {cutlass.BFloat16}: + raise ValueError("o_dtype must be BFloat16") + + if kv_dtype not in {cutlass.Int8}: + raise ValueError("kv_dtype must be Int8") + + if qk_acc_dtype not in {cutlass.Float32}: + raise ValueError("qk_acc_dtype must be Float32") + + if pv_acc_dtype not in {cutlass.Float32}: + raise ValueError("pv_acc_dtype must be Float32") + + h_r = h_q // h_k + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + def create_tensor(shape, dtype): + f32_torch_tensor = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch.float32, + permute_order=None, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig( + min_val=-2 if dtype.is_float or dtype.signed else 0, max_val=2 + ), + ) + + _, torch_tensor = cutlass_torch.cute_tensor_like( + f32_torch_tensor, + dtype, + is_dynamic_layout=True, + assumed_align=32, + ) + + # Create dtype cute tensor with offset (gpu) + cute_tensor = from_dlpack(torch_tensor, assumed_align=128) + cute_tensor.element_type = dtype + + return ( + f32_torch_tensor, + cute_tensor, + torch_tensor, + ) + + scale_shape = (b, h_k, s_k, d // scale_granularity) + + q_ref, q_tensor, q_torch = create_tensor(q_shape, q_dtype) + k_ref, k_tensor, k_torch = create_tensor(k_shape, kv_dtype) + v_ref, v_tensor, v_torch = create_tensor(k_shape, kv_dtype) + o_ref, o_tensor, o_torch = create_tensor(q_shape, o_dtype) + scale_k_ref, scale_k_tensor, scale_k_torch = create_tensor(scale_shape, scale_dtype) + scale_v_ref, scale_v_tensor, scale_v_torch = create_tensor(scale_shape, scale_dtype) + + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + if is_causal: + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + else: + if s_k % 128 != 0: + mask_type = fmha_utils.MaskEnum.RESIDUAL_MASK + + fmha = MixedInputFusedMultiHeadAttentionPrefillD256( + scale_granularity, + qk_acc_dtype, + pv_acc_dtype, + is_persistent, + mask_type, + ) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + if scale_softmax == 0.0: # default to 1/sqrt(d) + scale_softmax = 1.0 / math.sqrt(d) + log2_e = math.log2( + math.exp(1.0) + ) # gpu uses exp2 for perf concerns, we need an extra factor 'log2_e' here + + scale_softmax = scale_q * scale_k * scale_softmax + scale_softmax_log2 = scale_softmax * log2_e + scale_output = scale_v * inv_scale_o + problem_size = (b, s_q, s_k, h_q, h_k, d) + compiled_fmha = cute.compile( + fmha, + q_tensor.iterator, + k_tensor.iterator, + v_tensor.iterator, + o_tensor.iterator, + scale_k_tensor.iterator, + scale_v_tensor.iterator, + problem_size, + scale_softmax_log2, + scale_output, + window_size_left if window_size_left is None else Int32(window_size_left), + window_size_right if window_size_right is None else Int32(window_size_right), + current_stream, + options=f"--opt-level 2", + ) + + def run_torch_fmha( + q, k, v, scale_k, scale_v, scale_softmax=1.0, scale_output=1.0, is_causal=False + ): + h_q = q.shape[1] + h_k = k.shape[1] + if not h_q == h_k: + repeat_factor = h_q // h_k + k = k.repeat_interleave(repeat_factor, dim=1) + v = v.repeat_interleave(repeat_factor, dim=1) + scale_k = scale_k.repeat_interleave(repeat_factor, dim=1) + scale_v = scale_v.repeat_interleave(repeat_factor, dim=1) + scale_k = ( + scale_k.unsqueeze(-1) + .repeat(1, 1, 1, 1, k.shape[3] // scale_k.shape[3]) + .reshape(k.shape) + ) + scale_v = ( + scale_v.unsqueeze(-1) + .repeat(1, 1, 1, 1, v.shape[3] // scale_v.shape[3]) + .reshape(v.shape) + ) + batch = q.shape[0] + ref_list = [] + for batch_idx in range(batch): + q_i = q[batch_idx] + k_i = k[batch_idx] + v_i = v[batch_idx] + scale_k_i = scale_k[batch_idx] + scale_v_i = scale_v[batch_idx] + s_i = torch.einsum("hqd,hkd->hqk", q_i, k_i * scale_k_i) * scale_softmax + s_q = q_i.shape[1] + s_k = k_i.shape[1] + if is_causal: + q_coords = torch.arange(0, s_q).view(-1, 1) + k_coords = torch.arange(0, s_k).view(1, -1) + _mask = k_coords > q_coords + s_k - s_q + s_i = s_i.masked_fill(_mask, -torch.inf) + p_i = s_i.softmax(dim=-1) + ref_i = torch.einsum("hqk,hkd->hqd", p_i, v_i * scale_v_i) * scale_output + ref_list.append(ref_i) + ref = torch.stack(ref_list) + return ref + + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_fmha( + q_tensor.iterator, + k_tensor.iterator, + v_tensor.iterator, + o_tensor.iterator, + scale_k_tensor.iterator, + scale_v_tensor.iterator, + problem_size, + scale_softmax_log2, + scale_output, + window_size_left if window_size_left is None else Int32(window_size_left), + ( + window_size_right + if window_size_right is None + else Int32(window_size_right) + ), + current_stream, + ) + print("Verifying results...") + o_ref = run_torch_fmha( + q_ref, + k_ref, + v_ref, + scale_k_ref, + scale_v_ref, + scale_softmax, + scale_output, + is_causal, + ) + + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o_tensor, o_fp32) + o_result = o_fp32_torch.cpu() + torch.testing.assert_close(o_ref, o_result, atol=tolerance, rtol=1e-05) + + print("Results verified successfully!") + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str): + 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 FMHA on Blackwell.") + + parser.add_argument( + "--q_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Input data type", + ) + + parser.add_argument( + "--kv_dtype", + type=cutlass.dtype, + default=cutlass.Int8, + help="Input data type", + ) + + parser.add_argument( + "--o_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Output data type", + ) + + parser.add_argument( + "--scale_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Scale data type", + ) + + parser.add_argument( + "--scale_granularity", + type=int, + default=256, + help="Scale granularity", + ) + + parser.add_argument( + "--qk_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="QK accumulator data type", + ) + + parser.add_argument( + "--pv_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="PV accumulator data type", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--is_causal", + action="store_true", + help="Whether to use casual mask", + ) + + parser.add_argument( + "--q_shape", + type=parse_comma_separated_ints, + default=(1, 8, 256, 256), + help="Shape of Q (B, H, S_q, D)", + ) + + parser.add_argument( + "--k_shape", + type=parse_comma_separated_ints, + default=(1, 8, 256, 256), + help="Shape of K (B, H_k, S_k, D)", + ) + + parser.add_argument( + "--scale_q", + type=float, + default=1.0, + help="Scaling factors to dequantize Q", + ) + + parser.add_argument( + "--scale_k", + type=float, + default=1.0, + help="Scaling factors to dequantize K", + ) + + parser.add_argument( + "--scale_v", + type=float, + default=1.0, + help="Scaling factors to dequantize V", + ) + + parser.add_argument( + "--inv_scale_o", + type=float, + default=1.0, + help="Scaling factor to quantize O", + ) + + parser.add_argument( + "--scale_softmax", + type=float, + default=0.0, + help="Scaling factor to scale S (i.e. Q*K); if zero, defaults to 1/sqrt(D)", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + 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.q_shape) != 4: + parser.error("--q_shape must contain exactly 4 values") + + if len(args.k_shape) != 4: + parser.error("--k_shape must contain exactly 4 values") + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + run( + args.q_shape, + args.k_shape, + args.q_dtype, + args.kv_dtype, + args.o_dtype, + args.scale_dtype, + args.scale_granularity, + args.qk_acc_dtype, + args.pv_acc_dtype, + args.is_persistent, + args.is_causal, + args.scale_q, + args.scale_k, + args.scale_v, + args.inv_scale_o, + args.scale_softmax, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py new file mode 100644 index 00000000..874ffc49 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py @@ -0,0 +1,2172 @@ +# 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 argparse +import math +import os +import sys + +from typing import Type, Tuple, Optional + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Int32, Int64, Float32 + + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from helpers import fmha_helpers as fmha_utils +from blackwell.mixed_input_fmha import prefill_helpers as prefill_utils + + +class MixedInputFusedMultiHeadAttentionPrefillD512: + def __init__( + self, + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + is_persistent: bool, + mask_type: fmha_utils.MaskEnum, + ): + self.qk_acc_dtype = qk_acc_dtype + self.pv_acc_dtype = pv_acc_dtype + self.cta_tiler = (128, 128, 512) # seq_q, seq_k, d + self.qk_mma_tiler = ( + self.cta_tiler[0] * 2, # default 2cta + self.cta_tiler[1], + min(self.cta_tiler[2], 128), + ) + self.pv_mma_tiler = self.qk_mma_tiler # keep BMM1 & BMM2 at the same pace + self.pv_block_tiler = ( + self.pv_mma_tiler[0] // 2, # default 2cta + self.pv_mma_tiler[1], + self.pv_mma_tiler[2], + ) + self.scale_granularity = scale_granularity + self.iterations_qk = self.cta_tiler[2] // self.qk_mma_tiler[2] + self.iterations_pv = self.cta_tiler[2] // self.pv_mma_tiler[1] + self.cluster_shape_mn = (2, 1) # use 2x1 cluster by default + self.tmem_warp_shape_mn = (4, 1) + self.is_persistent = is_persistent + self.mask_type = mask_type + self.transform_warp_ids = (0, 1, 2, 3, 4, 5, 6, 7) # i8 -> bf16 for kv + self.softmax_warp_ids = (8, 9, 10, 11) # softmax + correction + self.mma_warp_id = 12 # mma + self.load_warp_id = 13 # load + self.empty_warp_ids = (14, 15) # empty + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + self.tmem_alloc_sync_bar_id = 1 + self.tmem_s_offset = 128 + self.tmem_o_offset = 0 + self.num_regs_softmax = 256 + self.num_regs_other = 32 + self.num_regs_transform = 112 + self.buffer_align_bytes = 1024 + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.transform_warp_ids, + *self.softmax_warp_ids, + self.load_warp_id, + self.mma_warp_id, + *self.empty_warp_ids, + ) + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the FMHA kernel operation. + + This method initializes and configures various attributes required for the + execution of the fused multi-head attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.q_stage = self.iterations_qk + self.kv_stage = 4 + self.scale_k_stage = 1 + self.scale_v_stage = 1 + self.qk_acc_stage = 1 + self.pv_acc_stage = 1 + self.swap_stage = 1 + self.kv_trans_stage = 2 + + @cute.jit + def __call__( + self, + q_iter: cute.Pointer, + k_iter: cute.Pointer, + v_iter: cute.Pointer, + o_iter: cute.Pointer, + scale_k_iter: cute.Pointer, + scale_v_iter: cute.Pointer, + problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32, Int32], + scale_softmax_log2: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + stream: cuda.CUstream, + ): + self._setup_attributes() + b, s_q, s_k, h_q, h_k, d = problem_shape + h_r = h_q // h_k + self.d_r = self.cta_tiler[2] // self.scale_granularity + # (s, d, ((h_r, h_k), b)) + q_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=(d, 1, ((d * s_q, d * s_q * h_r), h_r * h_k * s_q * d)), + ) + q = cute.make_tensor(q_iter, q_layout) + # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + k_layout = cute.make_layout( + (s_k, d, ((h_r, h_k), b)), + stride=(d, 1, ((0, d * s_k), h_k * s_k * d)), + ) + k = cute.make_tensor(k_iter, k_layout) + # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast + v_layout = cute.make_layout( + (d, s_k, ((h_r, h_k), b)), + stride=(1, d, ((0, d * s_k), h_k * s_k * d)), + ) + v = cute.make_tensor(v_iter, v_layout) + # (s, d, ((h_r, h_k), b)) + # set divby for better gmem store vectorization + o_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=( + cute.assume(d, divby=256), + 1, + ( + ( + cute.assume(d * s_q, divby=256), + cute.assume(d * s_q * h_r, divby=256), + ), + cute.assume(h_r * h_k * s_q * d, divby=256), + ), + ), + ) + o = cute.make_tensor(o_iter, o_layout) + # (d_r * s, ((h_r, h_k), b)) + scale_k_layout = cute.make_layout( + (s_k * self.d_r, ((h_r, h_k), b)), + stride=(1, ((0, self.d_r * s_k), s_k * self.d_r * h_k)), + ) + scale_k = cute.make_tensor(scale_k_iter, scale_k_layout) + # (d_r * s, ((h_r, h_k), b)) + scale_v_layout = cute.make_layout( + (self.d_r * s_k, ((h_r, h_k), b)), + stride=(1, ((0, self.d_r * s_k), s_k * self.d_r * h_k)), + ) + scale_v = cute.make_tensor(scale_v_iter, scale_v_layout) + self.q_dtype = q.element_type + self.k_dtype = k.element_type + self.v_dtype = v.element_type + self.o_dtype = o.element_type + self.p_dtype = self.q_dtype # pv should has the same dtype + self.scale_k_dtype = scale_k.element_type + self.scale_v_dtype = scale_v.element_type + self.tile_sched_params, grid = fmha_utils.compute_grid( + o.shape, + self.cta_tiler, + self.is_persistent, + ) + self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.o_layout = utils.LayoutEnum.from_tensor(o) + cta_group = tcgen05.CtaGroup.TWO + p_major_mode = tcgen05.OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.qk_acc_dtype, + cta_group, + self.qk_mma_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + p_major_mode, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + self.pv_mma_tiler[:2], + ) + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.pv_block_tiler[:2] + q_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.q_stage, + ) + k_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.kv_stage, + ) + k_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, k_smem_layout_staged.outer + ) + k_trans_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.kv_trans_stage, + ) + p_smem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.pv_mma_tiler, + self.p_dtype, + self.qk_acc_stage, + ) + v_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.kv_stage, + ) + v_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, v_smem_layout_staged.outer + ) + v_trans_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.kv_trans_stage, + ) + scale_k_smem_layout, self.scale_k_tiler, scale_k_s2r_view_layout = ( + prefill_utils.get_scale_smem_layout( + self.scale_granularity, + self.d_r, + self.qk_mma_tiler, + self.k_major_mode, + ) + ) + scale_k_smem_layout_staged = cute.append( + scale_k_smem_layout, + cute.make_layout( + (self.scale_k_stage), + stride=(cute.cosize(scale_k_smem_layout.outer)), + ), + ) + scale_k_s2r_view_layout_staged = cute.append( + scale_k_s2r_view_layout, + cute.make_layout( + (self.scale_k_stage), + stride=(cute.cosize(scale_k_s2r_view_layout)), + ), + ) + scale_v_smem_layout, self.scale_v_tiler, scale_v_s2r_view_layout = ( + prefill_utils.get_scale_smem_layout( + self.scale_granularity, + self.d_r, + self.pv_mma_tiler, + self.v_major_mode, + ) + ) + scale_v_smem_layout_staged = cute.append( + scale_v_smem_layout, + cute.make_layout( + (self.scale_v_stage), + stride=(cute.cosize(scale_v_smem_layout.outer)), + ), + ) + scale_v_s2r_view_layout_staged = cute.append( + scale_v_s2r_view_layout, + cute.make_layout( + (self.scale_v_stage), + stride=(cute.cosize(scale_v_s2r_view_layout)), + ), + ) + tma_load_q_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + # For TMA Async, use one cta to sync with corresponding cta only + tma_load_kv_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp( + tcgen05.CtaGroup.ONE + ) + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_q_op, + q, + q_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for K + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_kv_op, + k, + k_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_scale_k, tma_tensor_scale_k = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_kv_op, + scale_k, + scale_k_smem_layout, + (self.scale_k_tiler[0] // 2,), + ) + # TMA load for V + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_kv_op, + v, + v_smem_layout, + self.pv_mma_tiler, + pv_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_scale_v, tma_tensor_scale_v = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_kv_op, + scale_v, + scale_v_smem_layout, + self.scale_v_tiler, + ) + self.tma_copy_q_bytes = cute.size_in_bytes( + self.q_dtype, q_smem_layout + ) * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_kv_bytes = cute.size_in_bytes(self.k_dtype, k_smem_layout) + self.tma_copy_scale_k_bytes = cute.size_in_bytes( + self.scale_k_dtype, scale_k_smem_layout + ) + self.tma_copy_scale_v_bytes = cute.size_in_bytes( + self.scale_v_dtype, scale_v_smem_layout + ) + + @cute.struct + class SharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[Int64, self.q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_stage * 2] + load_scale_k_mbar_ptr: cute.struct.MemRange[Int64, self.scale_k_stage * 2] + load_scale_v_mbar_ptr: cute.struct.MemRange[Int64, self.scale_v_stage * 2] + dequant_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_trans_stage * 2] + mma_s_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + mma_o_mbar_ptr: cute.struct.MemRange[Int64, self.pv_acc_stage * 2] + swap_mbar_ptr: cute.struct.MemRange[Int64, self.swap_stage * 2] + tmem_dealloc_mbar_ptr: Int64 + tmem_holding_buf: Int32 + + self.shared_storage = SharedStorage + + grid = cute.round_up(grid, self.cluster_shape_mnk) + + # Launch the kernel synchronously + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_scale_k, + tma_tensor_scale_k, + tma_atom_v, + tma_tensor_v, + tma_atom_scale_v, + tma_tensor_scale_v, + o, + scale_softmax_log2, + scale_output, + window_size_left, + window_size_right, + self.cluster_layout_vmnk, + q_smem_layout_staged, + k_smem_layout_staged, + k_trans_smem_layout_staged, + scale_k_smem_layout_staged, + scale_k_s2r_view_layout_staged, + p_smem_layout_staged, + v_smem_layout_staged, + v_trans_smem_layout_staged, + scale_v_smem_layout_staged, + scale_v_s2r_view_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + pv_tiled_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_qdl: cute.Tensor, + tma_atom_k: cute.CopyAtom, + mK_kdl: cute.Tensor, + tma_atom_scale_k: cute.CopyAtom, + mScaleK_kdl: cute.Tensor, + tma_atom_v: cute.CopyAtom, + mV_dkl: cute.Tensor, + tma_atom_scale_v: cute.CopyAtom, + mScaleV_dkl: cute.Tensor, + mO_qdl: cute.Tensor, + scale_softmax_log2: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + cluster_layout_vmnk: cute.Layout, + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + k_trans_smem_layout_staged: cute.ComposedLayout, + scale_k_smem_layout_staged: cute.ComposedLayout, + scale_k_s2r_view_layout_staged: cute.Layout, + p_smem_layout_staged: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + v_trans_smem_layout_staged: cute.ComposedLayout, + scale_v_smem_layout_staged: cute.ComposedLayout, + scale_v_s2r_view_layout_staged: cute.Layout, + epi_tile: cute.Tile, + tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # Prefetch tma desc + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_scale_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_scale_v) + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(qk_tiled_mma.thr_id.shape) + 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 + ) + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.q_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + tx_count=self.tma_copy_q_bytes, + barrier_storage=storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_kv_producer, load_kv_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_kv_bytes, + barrier_storage=storage.load_kv_mbar_ptr.data_ptr(), + tidx=0, + defer_sync=True, + ).make_participants() + load_scale_k_producer, load_scale_k_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.scale_k_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_scale_k_bytes, + barrier_storage=storage.load_scale_k_mbar_ptr.data_ptr(), + defer_sync=True, + tidx=0, + ).make_participants() + load_scale_v_producer, load_scale_v_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.scale_v_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_scale_v_bytes, + barrier_storage=storage.load_scale_v_mbar_ptr.data_ptr(), + defer_sync=True, + tidx=0, + ).make_participants() + dequant_kv_producer, dequant_kv_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.kv_trans_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + barrier_storage=storage.dequant_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_s_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + p_mma_producer, p_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + barrier_storage=storage.p_mma_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + mma_o_producer, mma_o_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.pv_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_o_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + swap_producer, swap_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.swap_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.swap_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=self.threads_per_warp + * len((self.mma_warp_id, *self.softmax_warp_ids)), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.softmax_warp_ids[0], + is_two_cta=True, + 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) + + sK_trans = smem.allocate_tensor( + element_type=self.q_dtype, + layout=k_trans_smem_layout_staged.outer, + swizzle=k_trans_smem_layout_staged.inner, + byte_alignment=128, + ) + sV_trans_ptr = cute.recast_ptr( + sK_trans.iterator, v_trans_smem_layout_staged.inner + ) + sV_trans = cute.make_tensor(sV_trans_ptr, v_trans_smem_layout_staged.outer) + sQ = smem.allocate_tensor( + element_type=self.q_dtype, + layout=q_smem_layout_staged.outer, + swizzle=q_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleK = smem.allocate_tensor( + element_type=self.scale_k_dtype, + layout=scale_k_smem_layout_staged.outer, + swizzle=scale_k_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleK_s2r_view = cute.make_tensor( + sScaleK.iterator, scale_k_s2r_view_layout_staged + ) + sScaleV = smem.allocate_tensor( + element_type=self.scale_v_dtype, + layout=scale_v_smem_layout_staged.outer, + swizzle=scale_v_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleV_s2r_view = cute.make_tensor( + sScaleV.iterator, scale_v_s2r_view_layout_staged + ) + sP = smem.allocate_tensor( + element_type=self.p_dtype, + layout=p_smem_layout_staged.outer, + swizzle=p_smem_layout_staged.inner, + byte_alignment=128, + ) + sK = smem.allocate_tensor( + element_type=self.k_dtype, + layout=k_smem_layout_staged.outer, + swizzle=k_smem_layout_staged.inner, + byte_alignment=128, + ) + sV_ptr = cute.recast_ptr(sK.iterator, v_smem_layout_staged.inner) + sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer) + qk_thr_mma = qk_tiled_mma.get_slice(mma_tile_coord_v) + pv_thr_mma = pv_tiled_mma.get_slice(mma_tile_coord_v) + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tOrP = qk_thr_mma.make_fragment_A(sP) + tSrK_trans = qk_thr_mma.make_fragment_B(sK_trans) + tOrV_trans = pv_thr_mma.make_fragment_B(sV_trans) + qk_acc_shape = pv_thr_mma.partition_shape_C( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + # (atomV, restM, restN, accStage) + tStS = qk_tiled_mma.make_fragment_C( + cute.append(qk_acc_shape, self.qk_acc_stage) + ) + pv_acc_shape = pv_thr_mma.partition_shape_C( + cute.select(self.pv_mma_tiler, mode=[0, 1]) + ) + # (atomV, restM, restN) + tOtO = pv_thr_mma.make_fragment_C(pv_acc_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + self.iterations_pv, + stride=self.pv_mma_tiler[1] // self.tmem_warp_shape_mn[1], + ), + ) + tStS = cute.make_tensor(tStS.iterator + self.tmem_s_offset, tStS.layout) + tOtO_staged = cute.make_tensor(tOtO.iterator + self.tmem_o_offset, tOtO_layout) + # Local_tile partition global tensors + q_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # (bM, bK, restM, restK, loopM, loopK, loopL) + gQ_qdl = cute.flat_divide(mQ_qdl, cute.select(self.qk_mma_tiler, mode=[0, 2])) + tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl) + tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, + block_in_cluster_coord_vmnk[2], + q_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ_qdl, 0, 3), + ) + kv_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # (bN, bK, loopN, loopK, loopL) + gK_kdl = cute.flat_divide(mK_kdl, cute.select(self.qk_mma_tiler, mode=[1, 2])) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + # (blk, loopBlk, loopL) + gScaleK_kdl = cute.flat_divide(mScaleK_kdl, self.scale_k_tiler) + # Deal with 2cta + gScaleK_kdl_ = cute.logical_divide(gScaleK_kdl, (self.scale_k_tiler[0] // 2,))[ + (None, mma_tile_coord_v), None, None + ] + tKsScaleK, tKgScaleK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_scale_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + sScaleK, + gScaleK_kdl_, + ) + # (bN, bK, loopN, loopK, loopL) + gV_dkl = cute.flat_divide(mV_dkl, cute.select(self.pv_mma_tiler, mode=[1, 2])) + tOgV_dkl = pv_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tOgV_dkl, 0, 3), + ) + # (bBlk, loopBlk, loopL) + gScaleV_dkl = cute.flat_divide(mScaleV_dkl, self.scale_v_tiler) + tVsScaleV, tVgScaleV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_scale_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + sScaleV, + gScaleV_dkl, + ) + # (bM, bN, loopM, loopN, loopL) + gO_qdl = cute.flat_divide(mO_qdl, cute.select(self.pv_block_tiler, mode=[0, 1])) + cO_qdl = cute.flat_divide( + cute.make_identity_tensor(mO_qdl.shape), + cute.select(self.pv_block_tiler, mode=[0, 1]), + ) + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + tile_sched = fmha_utils.create_fmha_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + # Cluster wait + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + # ((atom_v, rest_v), RestK) + tQgQ = tQgQ_qdl[None, mma_block_coord[0], None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tKgK = tKgK_kdl[None, None, None, mma_block_coord[2]] + tKgScaleK = tKgScaleK_kdl[None, None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tVgV = tVgV_dkl[None, None, None, mma_block_coord[2]] + tVgScaleV = tVgScaleV_dkl[None, None, mma_block_coord[2]] + load_kv_producer, load_scale_k_producer, load_q_producer = ( + prefill_utils.load_qk( # Q & K0 & ScaleK0 + iterations=self.iterations_qk, + kv_step=0, + k_args=(tKgK, tKsK, tma_atom_k, load_kv_producer), + scale_k_args=( + tKgScaleK, + tKsScaleK, + tma_atom_scale_k, + load_scale_k_producer, + ), + q_args=(tQgQ, tQsQ, tma_atom_q, load_q_producer), + ) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + load_kv_producer, load_scale_k_producer = ( + prefill_utils.load_qk( # Ki & ScaleKi + iterations=self.iterations_qk, + kv_step=step, + k_args=(tKgK, tKsK, tma_atom_k, load_kv_producer), + scale_k_args=( + tKgScaleK, + tKsScaleK, + tma_atom_scale_k, + load_scale_k_producer, + ), + ) + ) + load_kv_producer, load_scale_v_producer = ( + prefill_utils.load_v( # Vi-1 & ScaleVi-1 + iterations=self.iterations_pv, + kv_step=step - 1, + v_args=(tVgV, tVsV, tma_atom_v, load_kv_producer), + scale_v_args=( + tVgScaleV, + tVsScaleV, + tma_atom_scale_v, + load_scale_v_producer, + ), + ) + ) + load_kv_producer, load_scale_v_producer = ( + prefill_utils.load_v( # Vend & ScaleVend + iterations=self.iterations_pv, + kv_step=seqlen_kv_loop_steps - 1, + v_args=(tVgV, tVsV, tma_atom_v, load_kv_producer), + scale_v_args=( + tVgScaleV, + tVsScaleV, + tma_atom_scale_v, + load_scale_v_producer, + ), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_kv_producer.tail() + load_scale_k_producer.tail() + load_scale_v_producer.tail() + load_q_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + tmem.wait_for_alloc() + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + load_q_releaser = load_q_consumer.clone() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + if seqlen_kv_loop_steps > 1: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QK0 + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + load_q_consumer, + None, + dequant_kv_consumer, + ), + ) + ) + for i in cutlass.range(1, seqlen_kv_loop_steps - 1, 1, unroll=1): + mma_s_producer, _, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QKi + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + (mma_s_producer, None, None, dequant_kv_consumer), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ) = self.mma_pv( # PVi + pv_tiled_mma, + (tOtO_staged, tOrP, tOrV_trans), + ( + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ), + ) + mma_s_producer, _, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QKend needs to release Q + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + None, + load_q_releaser, + dequant_kv_consumer, + ), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ) = self.mma_pv( # PVend-1 + pv_tiled_mma, + (tOtO_staged, tOrP, tOrV_trans), + ( + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ), + ) + else: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QK0 + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + load_q_consumer, + load_q_releaser, + dequant_kv_consumer, + ), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ) = self.mma_pv( # PVend + pv_tiled_mma, + (tOtO_staged, tOrP, tOrV_trans), + ( + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ), + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + mma_s_producer.tail() + mma_o_producer.tail() + swap_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.mma_warp_id and warp_idx >= self.softmax_warp_ids[0]: + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + unmask_steps = fmha_utils.FusedMask.get_unmasked_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + gO_staged = gO_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + cO_staged = cO_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + cS_base = cute.make_identity_tensor( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + cS = cute.domain_offset( + (mma_block_coord[0] * self.qk_mma_tiler[0], 0), cS_base + ) + tScS = qk_thr_mma.partition_C(cS) + row_max = -Float32.inf + row_max_prev = -Float32.inf + row_sum = 0.0 + # S0 -> P0 + row_max, row_sum, mma_s_consumer, p_mma_producer = self.softmax_step( + (unmask_steps == 0, window_size_left, window_size_right), + (row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2), + (tStS, tScS, sP), + (mma_s_consumer, p_mma_producer), + ) + row_max_prev = row_max + # Use stage1 to do S/O swap + tmem_tiled_load, tmem_tiled_store, tSWAPtO, tSWAPrO = ( + self.get_swap_o_partition(tOtO_staged, cO_staged) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + cS_iter = cute.domain_offset((0, step * self.qk_mma_tiler[1]), cS) + tScS_iter = qk_thr_mma.partition_C(cS_iter) + # Si -> Pi & Oi-1 -> Oi + ( + tSWAPrO, + row_max, + row_sum, + mma_s_consumer, + p_mma_producer, + mma_o_consumer, + swap_consumer, + ) = self.softmax_correction_step( + ( + step >= unmask_steps, + step > 1, + step == seqlen_kv_loop_steps - 1, + window_size_left, + window_size_right, + ), + ( + row_max_prev, + row_sum, + seqlen_q, + seqlen_k, + scale_softmax_log2, + ), + (tStS, tScS_iter, sP, tOtO_staged, cO_staged), + (mma_s_consumer, p_mma_producer, mma_o_consumer, swap_consumer), + (tmem_tiled_load, tmem_tiled_store, tSWAPtO, tSWAPrO), + epi_tile, + ) + row_max_prev = row_max + # O_partial -> O_final + mma_o_consumer, swap_consumer = self.correction_epilog( + (tmem_tiled_store, tSWAPrO, tSWAPtO, swap_consumer), + (row_sum, seqlen_q, scale_output), + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged), + epi_tile, + ) + # Make sure we start the next wave's S=QK after correction epilog + mma_s_consumer.release() + mma_s_consumer.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_mma_producer.tail() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Trans + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.softmax_warp_ids[0]: + cute.arch.setmaxregister_decrease(self.num_regs_transform) + qk_thr_mma_leader_cta = qk_tiled_mma.get_slice(0) + pv_thr_mma_leader_cta = pv_tiled_mma.get_slice(0) + sScaleK_ = qk_thr_mma_leader_cta.partition_B(sScaleK_s2r_view) + sScaleV_ = pv_thr_mma_leader_cta.partition_B(sScaleV_s2r_view) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + load_kv_consumer, load_scale_k_consumer, dequant_kv_producer = ( + prefill_utils.dequant_k( # K0 + self.iterations_qk, + self.transform_warp_ids, + (self.k_dtype, self.q_dtype), + (sK, sScaleK_, sK_trans), + (load_kv_consumer, load_scale_k_consumer, dequant_kv_producer), + ) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + load_kv_consumer, load_scale_k_consumer, dequant_kv_producer = ( + prefill_utils.dequant_k( # Ki + self.iterations_qk, + self.transform_warp_ids, + (self.k_dtype, self.q_dtype), + (sK, sScaleK_, sK_trans), + ( + load_kv_consumer, + load_scale_k_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + prefill_utils.dequant_v( # Vi-1 + self.iterations_pv, + self.transform_warp_ids, + (self.v_dtype, self.q_dtype), + (sV, sScaleV_, sV_trans), + ( + load_kv_consumer, + load_scale_v_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + prefill_utils.dequant_v( # Vend + self.iterations_pv, + self.transform_warp_ids, + (self.v_dtype, self.q_dtype), + (sV, sScaleV_, sV_trans), + (load_kv_consumer, load_scale_v_consumer, dequant_kv_producer), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + dequant_kv_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Empty + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx > self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + return + + @cute.jit + def get_swap_o_partition( + self, + tOtO_staged: cute.Tensor, + cO_staged: cute.Tensor, + ) -> cute.Tensor: + # Swap S & O on stage 1 + tOtO_stage1 = tOtO_staged[(None, None), 0, 0, 1] + cO_stage1 = cO_staged[None, None, 1] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.pv_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_stage1) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tSWAPtO = thr_load.partition_D(tOtO_stage1) + tSWAPcO = thr_load.partition_D(cO_stage1) + tSWAPrO = cute.make_rmem_tensor(tSWAPcO.shape, self.pv_acc_dtype) + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), self.pv_acc_dtype + ) + tmem_tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_staged) + return tmem_tiled_load, tmem_tiled_store, tSWAPtO, tSWAPrO + + @cute.jit + def mma_pv( + self, + pv_tiled_mma: cute.TiledMma, + tensor_args: Tuple, + pipeline_args: Tuple, + ): + tOtO_staged, tOrP, tOrV_trans = tensor_args + p_mma_consumer, mma_o_producer, dequant_kv_consumer, swap_producer = ( + pipeline_args + ) + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + if is_leader_cta: + p_handle = p_mma_consumer.wait_and_advance() + o_handle = mma_o_producer.acquire_and_advance() + pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter_n in cutlass.range(self.iterations_pv, unroll=1): + v_trans_handle = dequant_kv_consumer.wait_and_advance() + if iter_n == 1: + swap_producer.acquire() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + tOtO_slice = tOtO_staged[None, None, None, iter_n] + tOrP_slice = tOrP[None, None, None, p_handle.index] + tOrV_trans_slice = tOrV_trans[None, None, None, v_trans_handle.index] + num_kphases = cute.size(tOrV_trans_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + pv_tiled_mma, + tOtO_slice, + tOrP_slice[kphase_coord], + tOrV_trans_slice[kphase_coord], + tOtO_slice, + ) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + if iter_n == 1: + swap_producer.commit() + swap_producer.advance() + v_trans_handle.release() + o_handle.commit() + p_handle.release() + return ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ) + + @cute.jit + def softmax_step( + self, + mask_args: Tuple, + value_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]: + need_apply_mask, window_size_left, window_size_right = mask_args + row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2 = value_args + tStS, tScS, sP = tensor_args + mma_s_consumer, p_mma_producer = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + s_handle = mma_s_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, s_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_load.partition_S(tStS_slice) + tTMEM_LOADcS = thr_load.partition_D(tScS_slice) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + cute.copy(tmem_tiled_load, tTMEM_LOADtS, tTMEM_LOADrS) + cute.arch.fence_view_async_tmem_load() + s_handle.release() + if need_apply_mask: + fmha_utils.FusedMask.apply_mask( + self.mask_type, + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + old_row_max = row_max + row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.p_dtype) + for k in cutlass.range(cute.size(tTMEM_LOADrS), vectorize=True): + tTMEM_LOADrS[k] = tTMEM_LOADrS[k] * scale + minus_row_max_scale + tTMEM_LOADrS[k] = cute.math.exp2(tTMEM_LOADrS[k], fastmath=True) + s_vec = tTMEM_LOADrS.load() + tTMEM_STORErP.store(s_vec.to(self.p_dtype)) + + p_handle = p_mma_producer.acquire_and_advance() + sP_slice = sP[None, None, None, p_handle.index] + sP_mk_view = cute.make_tensor( + sP_slice.iterator, + cute.make_layout( + ( + (sP_slice.shape[0][0], sP_slice.shape[1]), + (sP_slice.shape[0][1], sP_slice.shape[2]), + ), + stride=( + (sP_slice.stride[0][0], sP_slice.stride[1]), + (sP_slice.stride[0][1], sP_slice.stride[2]), + ), + ), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_load) + smem_thr_copy = smem_tiled_copy.get_slice(thread_idx) + rP_copy_view = smem_thr_copy.retile(tTMEM_STORErP) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + cute.arch.fence_view_async_shared() + p_handle.commit() + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) * 0.5 + # TODO: calc row sum with TensorSSA + row_sum *= acc_scale + local_row_sum_0 = (row_sum, row_sum) + local_row_sum_1 = (0.0, 0.0) + local_row_sum_2 = (0.0, 0.0) + local_row_sum_3 = (0.0, 0.0) + reduction_unroll = 4 + frg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll + tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) + for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2): + local_row_sum_0 = cute.arch.add_packed_f32x2( + local_row_sum_0, (tTMEM_LOADrS_frg[j, 0], tTMEM_LOADrS_frg[j + 1, 0]) + ) + local_row_sum_1 = cute.arch.add_packed_f32x2( + local_row_sum_1, (tTMEM_LOADrS_frg[j, 1], tTMEM_LOADrS_frg[j + 1, 1]) + ) + local_row_sum_2 = cute.arch.add_packed_f32x2( + local_row_sum_2, (tTMEM_LOADrS_frg[j, 2], tTMEM_LOADrS_frg[j + 1, 2]) + ) + local_row_sum_3 = cute.arch.add_packed_f32x2( + local_row_sum_3, (tTMEM_LOADrS_frg[j, 3], tTMEM_LOADrS_frg[j + 1, 3]) + ) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) + local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) + row_sum = local_row_sum_0[0] + local_row_sum_0[1] + return row_max, row_sum, mma_s_consumer, p_mma_producer + + @cute.jit + def correction_rescale( + self, + tOtO: cute.Tensor, + cO: cute.Tensor, + epi_tile: cute.Tile, + scale: Float32, + ) -> cute.Tensor: + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tOtO_epi = cute.zipped_divide(tOtO, cute.make_layout(epi_tile)) + cO_epi = cute.zipped_divide(cO, cute.make_layout(epi_tile)) + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), + self.pv_acc_dtype, + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_epi) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), + self.pv_acc_dtype, + ) + tmem_store_atom = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_epi) + thr_store = tmem_store_atom.get_slice(thread_idx) + tTMEM_LOADtO = thr_load.partition_S(tOtO_epi) + tTMEM_LOADcO = thr_load.partition_D(cO_epi) + tTMEM_STOREtO = thr_store.partition_D(tOtO_epi) + + tTMrO = cute.make_rmem_tensor_like( + cute.append( + cute.make_layout(tTMEM_LOADcO[None, 0, 0].shape), + cute.make_layout(2, stride=cute.size(tTMEM_LOADcO[None, 0, 0].shape)), + ), + self.pv_acc_dtype, + ) + tTMEM_LOADtO_0 = tTMEM_LOADtO[None, 0, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_0, tTMrO[None, 0]) + iter_num = cute.size(tTMEM_LOADtO, mode=[1]) + for i in cutlass.range(1, iter_num, unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_i, tTMrO[None, i % 2]) + for j in cutlass.range( + cute.size(tTMrO, mode=[0]), unroll_full=True, vectorize=True + ): + tTMrO[j, (i - 1) % 2] = tTMrO[j, (i - 1) % 2] * scale + tTMEM_STOREtO_prev_i = tTMEM_STOREtO[None, i - 1, 0] + cute.copy(tmem_store_atom, tTMrO[None, (i - 1) % 2], tTMEM_STOREtO_prev_i) + for j in cutlass.range( + cute.size(tTMrO, mode=[0]), unroll_full=True, vectorize=True + ): + tTMrO[j, (iter_num - 1) % 2] = tTMrO[j, (iter_num - 1) % 2] * scale + cute.copy( + tmem_store_atom, + tTMrO[None, (iter_num - 1) % 2], + tTMEM_STOREtO[None, iter_num - 1, 0], + ) + + @cute.jit + def sum_reduction( + self, + tensor: cute.Tensor, + ) -> Float32: + local_sum_0 = (0.0, 0.0) + local_sum_1 = (0.0, 0.0) + local_sum_2 = (0.0, 0.0) + local_sum_3 = (0.0, 0.0) + reduction_unroll = 4 + frg_tile = cute.size(tensor) // reduction_unroll + tensor_frg = cute.logical_divide(tensor, cute.make_layout(frg_tile)) + for i in cutlass.range(0, cute.size(tensor_frg, mode=[0]), 2, unroll_full=True): + local_sum_0 = cute.arch.add_packed_f32x2( + local_sum_0, (tensor_frg[i, 0], tensor_frg[i + 1, 0]) + ) + local_sum_1 = cute.arch.add_packed_f32x2( + local_sum_1, (tensor_frg[i, 1], tensor_frg[i + 1, 1]) + ) + local_sum_2 = cute.arch.add_packed_f32x2( + local_sum_2, (tensor_frg[i, 2], tensor_frg[i + 1, 2]) + ) + local_sum_3 = cute.arch.add_packed_f32x2( + local_sum_3, (tensor_frg[i, 3], tensor_frg[i + 1, 3]) + ) + local_sum_0 = cute.arch.add_packed_f32x2(local_sum_0, local_sum_1) + local_sum_2 = cute.arch.add_packed_f32x2(local_sum_2, local_sum_3) + local_sum_0 = cute.arch.add_packed_f32x2(local_sum_0, local_sum_2) + return local_sum_0[0] + local_sum_0[1] + + @cute.jit + def softmax_correction_step( + self, + mask_args: Tuple, + value_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + swap_args: Tuple, + epi_tile: cute.Tile, + ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]: + ( + need_apply_mask, + need_store_o, + last_iteration, + window_size_left, + window_size_right, + ) = mask_args + row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2 = value_args + tStS, tScS_iter, sP, tOtO_staged, cO_staged = tensor_args + mma_s_consumer, p_mma_producer, mma_o_consumer, swap_consumer = pipeline_args + swap_tmem_tiled_load, swap_tmem_tiled_store, tSWAPtO, tSWAPrO = swap_args + + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tStS_slice = tStS[(None, None), 0, 0, None] + tScS_slice = tScS_iter[(None, None), 0, 0] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy( + tmem_load_atom, tStS_slice[None, None, 0] + ) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_load.partition_S(tStS_slice) + tTMEM_LOADcS = thr_load.partition_D(tScS_slice) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + mma_s_consumer.wait() + if not need_store_o: + cute.copy( + tmem_tiled_load, + tTMEM_LOADtS[None, None, None, 0], + tTMEM_LOADrS, + ) + else: + iter_num = cute.size(tTMEM_LOADrS, mode=[2]) + for i in cutlass.range(iter_num, unroll_full=True): + cute.copy( + swap_tmem_tiled_load, + tTMEM_LOADtS[None, 0, iter_num - 1 - i, 0], + tTMEM_LOADrS[None, 0, iter_num - 1 - i], + ) + cute.copy( + swap_tmem_tiled_store, + tSWAPrO[None, 0, iter_num - 1 - i], + tSWAPtO[None, 0, iter_num - 1 - i], + ) + cute.arch.fence_view_async_tmem_load() + cute.arch.fence_view_async_tmem_store() + swap_consumer.release() + swap_consumer.advance() + + if need_apply_mask: + fmha_utils.FusedMask.apply_mask( + self.mask_type, + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + old_row_max = row_max + row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) + row_sum *= acc_scale + + tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.p_dtype) + subtile_cnt = 2 + subtile_size = cute.size(tTMEM_LOADrS) // subtile_cnt + tTMEM_LOADrS_subtile = cute.logical_divide( + tTMEM_LOADrS, cute.make_layout(subtile_size) + ) + tTMEM_STORErP_subtile = cute.logical_divide( + tTMEM_STORErP, cute.make_layout(subtile_size) + ) + + tTMEM_LOADrS_0 = tTMEM_LOADrS_subtile[None, 0] + tTMEM_STORErP_0 = tTMEM_STORErP_subtile[None, 0] + for k in cutlass.range(cute.size(tTMEM_LOADrS_0), vectorize=True): + tTMEM_LOADrS_0[k] = tTMEM_LOADrS_0[k] * scale + minus_row_max_scale + tTMEM_LOADrS_0[k] = cute.math.exp2(tTMEM_LOADrS_0[k], fastmath=True) + row_sum += self.sum_reduction(tTMEM_LOADrS_0) + s_vec = tTMEM_LOADrS_0.load() + tTMEM_STORErP_0.store(s_vec.to(self.p_dtype)) + + swap_consumer.wait() + cute.copy(swap_tmem_tiled_load, tSWAPtO, tSWAPrO) + cute.arch.fence_view_async_tmem_load() + if not last_iteration: + mma_s_consumer.release() + mma_s_consumer.advance() + + tTMEM_LOADrS_1 = tTMEM_LOADrS_subtile[None, 1] + tTMEM_STORErP_1 = tTMEM_STORErP_subtile[None, 1] + for k in cutlass.range(cute.size(tTMEM_LOADrS_1), vectorize=True): + tTMEM_LOADrS_1[k] = tTMEM_LOADrS_1[k] * scale + minus_row_max_scale + tTMEM_LOADrS_1[k] = cute.math.exp2(tTMEM_LOADrS_1[k], fastmath=True) + row_sum += self.sum_reduction(tTMEM_LOADrS_1) + s_vec = tTMEM_LOADrS_1.load() + tTMEM_STORErP_1.store(s_vec.to(self.p_dtype)) + + p_handle = p_mma_producer.acquire_and_advance() + sP_slice = sP[None, None, None, p_handle.index] + sP_mk_view = cute.make_tensor( + sP_slice.iterator, + cute.make_layout( + ( + (sP_slice.shape[0][0], sP_slice.shape[1]), + (sP_slice.shape[0][1], sP_slice.shape[2]), + ), + stride=( + (sP_slice.stride[0][0], sP_slice.stride[1]), + (sP_slice.stride[0][1], sP_slice.stride[2]), + ), + ), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_load) + smem_thr_copy = smem_tiled_copy.get_slice(thread_idx) + rP_copy_view = smem_thr_copy.retile(tTMEM_STORErP) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + for i in cutlass.range(0, cute.size(tSWAPrO), 2, unroll_full=True): + tSWAPrO[i], tSWAPrO[i + 1] = cute.arch.mul_packed_f32x2( + (tSWAPrO[i], tSWAPrO[i + 1]), + (acc_scale, acc_scale), + ) + cute.arch.fence_view_async_shared() + p_handle.commit() + o_handle = mma_o_consumer.wait_and_advance() + self.correction_rescale( + tOtO_staged[(None, None), 0, 0, 0], + cO_staged[None, None, 0], + epi_tile, + acc_scale, + ) + # Skip stage1; loop starts at 2 + for iter_n in cutlass.range(2, self.iterations_pv, unroll_full=True): + self.correction_rescale( + tOtO_staged[(None, None), 0, 0, iter_n], + cO_staged[None, None, iter_n], + epi_tile, + acc_scale, + ) + cute.arch.fence_view_async_tmem_store() + o_handle.release() + + return ( + tSWAPrO, + row_max, + row_sum, + mma_s_consumer, + p_mma_producer, + mma_o_consumer, + swap_consumer, + ) + + @cute.jit + def correction_epilog( + self, + swap_args: Tuple, + value_args: Tuple, + o_args: Tuple, + epi_tile: cute.Tile, + ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]: + (swap_tmem_tiled_store, tSWAPrO, tSWAPtO, swap_consumer) = swap_args + (row_sum, seqlen_q, scale_output) = value_args + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged) = o_args + + cute.copy(swap_tmem_tiled_store, tSWAPrO, tSWAPtO) + cute.arch.fence_view_async_tmem_store() + swap_consumer.release() + swap_consumer.advance() + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + scale = scale_output / row_sum + o_handle = mma_o_consumer.wait_and_advance() + # empty step as we access tOtO_stage1 by the normal way + swap_consumer.wait() + swap_consumer.release() + swap_consumer.advance() + for iter_n in cutlass.range(self.iterations_pv): + gO = gO_staged[None, None, iter_n] + cO = cO_staged[None, None, iter_n] + tOtO = tOtO_staged[(None, None), 0, 0, iter_n] + tOtO_epi = cute.zipped_divide(tOtO, epi_tile) + cO_epi = cute.zipped_divide(cO, epi_tile) + gO_epi = cute.zipped_divide(gO, epi_tile) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tmem_copy_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.pv_acc_dtype + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_epi) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_epi) + tTMEM_LOADgO = thr_tmem_load.partition_D(gO_epi) + tTMEM_LOADcO = thr_tmem_load.partition_D(cO_epi) + for i in cutlass.range(cute.size(tTMEM_LOADtO, mode=[1]), unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + tTMEM_LOADgO_i = tTMEM_LOADgO[None, i, 0] + tTMEM_LOADcO_i = tTMEM_LOADcO[None, i, 0] + tTMrO = cute.make_rmem_tensor( + tTMEM_LOADcO[None, 0, i].shape, self.pv_acc_dtype + ) + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) + for j in cutlass.range(0, cute.size(tTMrO), 2, unroll_full=True): + tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( + (tTMrO[j], tTMrO[j + 1]), + (scale, scale), + ) + tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) + o_vec = tTMrO.load() + tSMrO.store(o_vec.to(self.o_dtype)) + if cute.elem_less(tTMEM_LOADcO_i[0][0], seqlen_q): + cute.autovec_copy(tSMrO, tTMEM_LOADgO_i) + o_handle.release() + return mma_o_consumer, swap_consumer + + +def run( + q_shape: Tuple[int, int, int, int], + k_shape: Tuple[int, int, int, int], + q_dtype: Type[cutlass.Numeric], + kv_dtype: Type[cutlass.Numeric], + o_dtype: Type[cutlass.Numeric], + scale_dtype: Type[cutlass.Numeric], + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + is_persistent: bool, + is_causal: bool, + scale_q: float, + scale_k: float, + scale_v: float, + inv_scale_o: float, + scale_softmax: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool = False, + **kwargs, +): + print(f"Running Blackwell SM100 Mixed Input FMHA Prefill D 512 test with:") + print(f" q_shape: {q_shape}") + print(f" k_shape: {k_shape}") + print(f" q_dtype: {q_dtype}") + print(f" kv_dtype: {kv_dtype}") + print(f" o_dtype: {o_dtype}") + print(f" scale_dtype: {scale_dtype}") + print(f" scale_granularity: {scale_granularity}") + print(f" qk_acc_dtype: {qk_acc_dtype}") + print(f" pv_acc_dtype: {pv_acc_dtype}") + print(f" is_persistent: {is_persistent}") + print(f" is_causal: {is_causal}") + print(f" scale_q: {scale_q}") + print(f" scale_k: {scale_k}") + print(f" scale_v: {scale_v}") + print(f" inv_scale_o: {inv_scale_o}") + print(f" scale_softmax: {scale_softmax}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + import cutlass.torch as cutlass_torch + + # Unpack parameters + b, h_q, s_q, d = q_shape + b_, h_k, s_k, d_ = k_shape + window_size_left, window_size_right = None, None + if is_causal: + window_size_right = 0 + + if b != b_: + raise ValueError("q & k must have the same batch size") + + if d != d_: + raise ValueError("q & k must have the same head dimension") + + if d not in {512}: + raise ValueError("head dimension must be 512") + + if d % scale_granularity != 0: + raise ValueError("head dimension must be divisible by scale_granularity") + + if scale_granularity not in {128, 256, 512}: + raise ValueError("scale_granularity must be 128, 256, or 512") + + if h_q % h_k != 0: + raise ValueError("h_q must be divisible by h_k") + + if isinstance(s_q, tuple) and len(s_q) != b: + raise ValueError("variable_seqlen s_q must have the length of batch size") + if isinstance(s_k, tuple) and len(s_k) != b: + raise ValueError("variable_seqlen s_k must have the length of batch size") + + if q_dtype not in {cutlass.BFloat16}: + raise ValueError("in_dtype must be BFloat16") + + if o_dtype not in {cutlass.BFloat16}: + raise ValueError("o_dtype must be BFloat16") + + if kv_dtype not in {cutlass.Int8}: + raise ValueError("kv_dtype must be Int8") + + if qk_acc_dtype not in {cutlass.Float32}: + raise ValueError("qk_acc_dtype must be Float32") + + if pv_acc_dtype not in {cutlass.Float32}: + raise ValueError("pv_acc_dtype must be Float32") + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + def create_tensor(shape, dtype): + f32_torch_tensor = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch.float32, + permute_order=None, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig( + min_val=-2 if dtype.is_float or dtype.signed else 0, max_val=2 + ), + ) + + _, torch_tensor = cutlass_torch.cute_tensor_like( + f32_torch_tensor, + dtype, + is_dynamic_layout=True, + assumed_align=32, + ) + + # Create dtype cute tensor with offset (gpu) + cute_tensor = from_dlpack(torch_tensor, assumed_align=128) + cute_tensor.element_type = dtype + + return ( + f32_torch_tensor, + cute_tensor, + torch_tensor, + ) + + scale_shape = (b, h_k, s_k, d // scale_granularity) + + q_ref, q_tensor, q_torch = create_tensor(q_shape, q_dtype) + k_ref, k_tensor, k_torch = create_tensor(k_shape, kv_dtype) + v_ref, v_tensor, v_torch = create_tensor(k_shape, kv_dtype) + o_ref, o_tensor, o_torch = create_tensor(q_shape, o_dtype) + scale_k_ref, scale_k_tensor, scale_k_torch = create_tensor(scale_shape, scale_dtype) + scale_v_ref, scale_v_tensor, scale_v_torch = create_tensor(scale_shape, scale_dtype) + + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + if is_causal: + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + else: + if s_k % 128 != 0: + mask_type = fmha_utils.MaskEnum.RESIDUAL_MASK + + fmha = MixedInputFusedMultiHeadAttentionPrefillD512( + scale_granularity, + qk_acc_dtype, + pv_acc_dtype, + is_persistent, + mask_type, + ) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + if scale_softmax == 0.0: # default to 1/sqrt(d) + scale_softmax = 1.0 / math.sqrt(d) + log2_e = math.log2( + math.exp(1.0) + ) # gpu uses exp2 for perf concerns, we need an extra factor 'log2_e' here + + scale_softmax = scale_q * scale_k * scale_softmax + scale_softmax_log2 = scale_softmax * log2_e + scale_output = scale_v * inv_scale_o + problem_size = (b, s_q, s_k, h_q, h_k, d) + compiled_fmha = cute.compile( + fmha, + q_tensor.iterator, + k_tensor.iterator, + v_tensor.iterator, + o_tensor.iterator, + scale_k_tensor.iterator, + scale_v_tensor.iterator, + problem_size, + scale_softmax_log2, + scale_output, + window_size_left if window_size_left is None else Int32(window_size_left), + window_size_right if window_size_right is None else Int32(window_size_right), + current_stream, + options=f"--opt-level 2", + ) + + def run_torch_fmha( + q, k, v, scale_k, scale_v, scale_softmax=1.0, scale_output=1.0, is_causal=False + ): + h_q = q.shape[1] + h_k = k.shape[1] + if not h_q == h_k: + repeat_factor = h_q // h_k + k = k.repeat_interleave(repeat_factor, dim=1) + v = v.repeat_interleave(repeat_factor, dim=1) + scale_k = scale_k.repeat_interleave(repeat_factor, dim=1) + scale_v = scale_v.repeat_interleave(repeat_factor, dim=1) + scale_k = ( + scale_k.unsqueeze(-1) + .repeat(1, 1, 1, 1, k.shape[3] // scale_k.shape[3]) + .reshape(k.shape) + ) + scale_v = ( + scale_v.unsqueeze(-1) + .repeat(1, 1, 1, 1, v.shape[3] // scale_v.shape[3]) + .reshape(v.shape) + ) + batch = q.shape[0] + ref_list = [] + for batch_idx in range(batch): + q_i = q[batch_idx] + k_i = k[batch_idx] + v_i = v[batch_idx] + scale_k_i = scale_k[batch_idx] + scale_v_i = scale_v[batch_idx] + s_i = torch.einsum("hqd,hkd->hqk", q_i, k_i * scale_k_i) * scale_softmax + s_q = q_i.shape[1] + s_k = k_i.shape[1] + if is_causal: + q_coords = torch.arange(0, s_q).view(-1, 1) + k_coords = torch.arange(0, s_k).view(1, -1) + _mask = k_coords > q_coords + s_k - s_q + s_i = s_i.masked_fill(_mask, -torch.inf) + p_i = s_i.softmax(dim=-1) + ref_i = torch.einsum("hqk,hkd->hqd", p_i, v_i * scale_v_i) * scale_output + ref_list.append(ref_i) + ref = torch.stack(ref_list) + return ref + + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_fmha( + q_tensor.iterator, + k_tensor.iterator, + v_tensor.iterator, + o_tensor.iterator, + scale_k_tensor.iterator, + scale_v_tensor.iterator, + problem_size, + scale_softmax_log2, + scale_output, + window_size_left if window_size_left is None else Int32(window_size_left), + ( + window_size_right + if window_size_right is None + else Int32(window_size_right) + ), + current_stream, + ) + print("Verifying results...") + o_ref = run_torch_fmha( + q_ref, + k_ref, + v_ref, + scale_k_ref, + scale_v_ref, + scale_softmax, + scale_output, + is_causal, + ) + + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o_tensor, o_fp32) + o_result = o_fp32_torch.cpu() + torch.testing.assert_close(o_ref, o_result, atol=tolerance, rtol=1e-05) + + print("Results verified successfully!") + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str): + 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 FMHA on Blackwell.") + + parser.add_argument( + "--q_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Input data type", + ) + + parser.add_argument( + "--kv_dtype", + type=cutlass.dtype, + default=cutlass.Int8, + help="Input data type", + ) + + parser.add_argument( + "--o_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Output data type", + ) + + parser.add_argument( + "--scale_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Scale data type", + ) + + parser.add_argument( + "--scale_granularity", + type=int, + default=512, + help="Scale granularity", + ) + + parser.add_argument( + "--qk_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="QK accumulator data type", + ) + + parser.add_argument( + "--pv_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="PV accumulator data type", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--is_causal", + action="store_true", + help="Whether to use casual mask", + ) + + parser.add_argument( + "--q_shape", + type=parse_comma_separated_ints, + default=(1, 8, 256, 512), + help="Shape of Q (B, H, S_q, D)", + ) + + parser.add_argument( + "--k_shape", + type=parse_comma_separated_ints, + default=(1, 8, 256, 512), + help="Shape of K (B, H_k, S_k, D)", + ) + + parser.add_argument( + "--scale_q", + type=float, + default=1.0, + help="Scaling factors to dequantize Q", + ) + + parser.add_argument( + "--scale_k", + type=float, + default=1.0, + help="Scaling factors to dequantize K", + ) + + parser.add_argument( + "--scale_v", + type=float, + default=1.0, + help="Scaling factors to dequantize V", + ) + + parser.add_argument( + "--inv_scale_o", + type=float, + default=1.0, + help="Scaling factor to quantize O", + ) + + parser.add_argument( + "--scale_softmax", + type=float, + default=0.0, + help="Scaling factor to scale S (i.e. Q*K); if zero, defaults to 1/sqrt(D)", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + 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.q_shape) != 4: + parser.error("--q_shape must contain exactly 4 values") + + if len(args.k_shape) != 4: + parser.error("--k_shape must contain exactly 4 values") + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + run( + args.q_shape, + args.k_shape, + args.q_dtype, + args.kv_dtype, + args.o_dtype, + args.scale_dtype, + args.scale_granularity, + args.qk_acc_dtype, + args.pv_acc_dtype, + args.is_persistent, + args.is_causal, + args.scale_q, + args.scale_k, + args.scale_v, + args.inv_scale_o, + args.scale_softmax, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/prefill_helpers.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/prefill_helpers.py new file mode 100644 index 00000000..2a833cce --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/prefill_helpers.py @@ -0,0 +1,400 @@ +# 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. + +from typing import Tuple, Optional + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.pipeline as pipeline + + +@cute.jit +def load_qk( + iterations: int, + kv_step: cutlass.Int32, + k_args: Tuple, + scale_k_args: Optional[Tuple] = None, + q_args: Optional[Tuple] = None, +) -> Tuple[pipeline.PipelineProducer, pipeline.PipelineProducer]: + if cutlass.const_expr(q_args is not None): + tQgQ, tQsQ, tma_atom_q, load_q_producer = q_args + else: + tQgQ, tQsQ, tma_atom_q, load_q_producer = None, None, None, None + tKgK, tKsK, tma_atom_k, load_k_producer = k_args + tKgScaleK, tKsScaleK, tma_atom_scale_k, load_scale_k_producer = scale_k_args + + scale_k_handle = load_scale_k_producer.acquire_and_advance() + cute.copy( + tma_atom_scale_k, + tKgScaleK[None, kv_step], + tKsScaleK[None, scale_k_handle.index], + tma_bar_ptr=scale_k_handle.barrier, + ) + for iter in cutlass.range(iterations, unroll=1): + if cutlass.const_expr(q_args is not None): + q_handle = load_q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tQgQ[None, iter], + tQsQ[None, q_handle.index], + tma_bar_ptr=q_handle.barrier, + ) + k_handle = load_k_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_step, iter], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + if cutlass.const_expr(q_args is not None): + return load_k_producer, load_scale_k_producer, load_q_producer + else: + return load_k_producer, load_scale_k_producer + + +@cute.jit +def load_v( + iterations: int, + kv_step: cutlass.Int32, + v_args: Tuple, + scale_v_args: Tuple, +) -> pipeline.PipelineProducer: + tVgV, tVsV, tma_atom_v, load_v_producer = v_args + tScaleVgV, tScaleVsV, tma_atom_scale_v, load_scale_v_producer = scale_v_args + scale_v_handle = load_scale_v_producer.acquire_and_advance() + cute.copy( + tma_atom_scale_v, + tScaleVgV[None, kv_step], + tScaleVsV[None, scale_v_handle.index], + tma_bar_ptr=scale_v_handle.barrier, + ) + for iter in cutlass.range(iterations, unroll=1): + v_handle = load_v_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tVgV[None, iter, kv_step], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + return load_v_producer, load_scale_v_producer + + +@cute.jit +def get_scale_smem_layout( + scale_granularity: int, + d_r: int, + mma_tiler: cute.Tile, + major_mode: tcgen05.OperandMajorMode, +) -> Tuple[cute.Layout, cute.Tile]: + size_mn = mma_tiler[1] // 2 # 2cta by default + if cutlass.const_expr(major_mode == tcgen05.OperandMajorMode.MN): # v + scale_tiler = (mma_tiler[2] * d_r,) + tma_view_layout = cute.make_layout( + (mma_tiler[2] * d_r), + ) + assert scale_granularity % mma_tiler[1] == 0, ( + "scale_granularity must be divisible by mma_tiler[1]" + ) + rest_l = scale_granularity // mma_tiler[1] + s2r_view_layout = cute.make_layout( + (size_mn, mma_tiler[2], (rest_l, d_r)), + stride=(0, d_r, (0, 1)), + ) + else: # k + scale_tiler = (mma_tiler[1] * d_r,) + tma_view_layout = cute.make_layout((size_mn * d_r)) + assert scale_granularity % mma_tiler[2] == 0, ( + "scale_granularity must be divisible by mma_tiler[2]" + ) + rest_l = scale_granularity // mma_tiler[2] + s2r_view_layout = cute.make_layout( + (size_mn, mma_tiler[2], (rest_l, d_r)), + stride=(d_r, 0, (0, 1)), + ) + # Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom + tma_view_smem_layout = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, tma_view_layout + ) + return tma_view_smem_layout, scale_tiler, s2r_view_layout + + +@cute.jit +def mma_qk( + iterations: int, + qk_tiled_mma: cute.TiledMma, + tensor_args: Tuple, + pipeline_args: Tuple, +): + tStS, tSrQ, tSrK_trans = tensor_args + mma_s_producer, load_q_consumer, load_q_releaser, dequant_kv_consumer = ( + pipeline_args + ) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + if is_leader_cta: + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(iterations, unroll=1): + if cutlass.const_expr(load_q_consumer is not None): + load_q_consumer.wait_and_advance() + tSrQ_slice = tSrQ[None, None, None, iter] + k_trans_handle = dequant_kv_consumer.wait_and_advance() + tSrK_trans_slice = tSrK_trans[None, None, None, k_trans_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_trans_handle.release() + if cutlass.const_expr(load_q_releaser is not None): + load_q_releaser.release() + load_q_releaser.advance() + s_handle.commit() + return mma_s_producer, load_q_consumer, dequant_kv_consumer + + +@cute.jit +def dequant_k( + iterations: int, + transform_warp_ids: Tuple, + dtype_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, +): + (k_dtype, q_dtype) = dtype_args + (sOrig, sScale, sTrans) = tensor_args + (load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + THREADS_PER_WARP = 32 + thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids)) + r2s_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), k_dtype, num_bits_per_copy=32 + ) + # Construct tiled_copy satisfying 16 contiguous elts per copy atom + r2s_tiled_copy = cute.make_cotiled_copy( + r2s_copy_atom, + cute.make_layout((256, 16), stride=(16, 1)), + sTrans[(None, None, None, 0)].layout, + ) + thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx) + tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig) + tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans) + tOrOrig = cute.make_rmem_tensor_like( + cute.append( + tOsOrig[None, None, None, None, 0].layout, + cute.make_layout( + 2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout) + ), + ), + k_dtype, + ) + tTrTrans = cute.make_rmem_tensor_like( + cute.append( + tTsTrans[None, None, None, None, 0].layout, + cute.make_layout( + 2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout) + ), + ), + q_dtype, + ) + tSsScale = thr_r2s_tiled_copy.partition_S(sScale) + tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0]) + scale_handle = load_scale_consumer.wait_and_advance() + cute.autovec_copy( + tSsScale[None, None, None, None, None, scale_handle.index], tSrScale + ) + cute.arch.fence_view_async_shared() + scale_handle.release() + # prefetch iter = 0 + kv_handle = load_kv_consumer.wait_and_advance() + cute.autovec_copy( + tOsOrig[None, None, None, None, kv_handle.index], + tOrOrig[None, None, None, None, 0], + ) + transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, 0].load(), + transformed_tensor.shape, + q_dtype, + ) + transformed_tensor = transformed_tensor * scale + tTrTrans[None, None, None, None, 0].store(transformed_tensor) + cute.arch.fence_view_async_shared() + kv_handle.release() + for iter in cutlass.range(1, iterations, unroll_full=True): + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (iter - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + kv_handle = load_kv_consumer.wait_and_advance() + cute.autovec_copy( + tOsOrig[None, None, None, None, kv_handle.index], + tOrOrig[None, None, None, None, iter % 2], + ) + transformed_tensor = ( + tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype) + ) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, iter].load(), + transformed_tensor.shape, + q_dtype, + ) + transformed_tensor = transformed_tensor * scale + tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor) + cute.arch.fence_view_async_shared() + kv_handle.release() + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (iterations - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + return load_kv_consumer, load_scale_consumer, dequant_kv_producer + + +@cute.jit +def dequant_v( + iterations: int, + transform_warp_ids: Tuple, + dtype_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, +): + (v_dtype, q_dtype) = dtype_args + (sOrig, sScale, sTrans) = tensor_args + (load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + THREADS_PER_WARP = 32 + thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids)) + r2s_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), v_dtype, num_bits_per_copy=32 + ) + # Construct tiled_copy satisfying 16 contiguous elts per copy atom + r2s_tiled_copy = cute.make_cotiled_copy( + r2s_copy_atom, + cute.make_layout((256, 16), stride=(16, 1)), + sTrans[(None, None, None, 0)].layout, + ) + thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx) + tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig) + tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans) + # double buffer for better perf + tOrOrig = cute.make_rmem_tensor_like( + cute.append( + tOsOrig[None, None, None, None, 0].layout, + cute.make_layout( + 2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout) + ), + ), + v_dtype, + ) + tTrTrans = cute.make_rmem_tensor_like( + cute.append( + tTsTrans[None, None, None, None, 0].layout, + cute.make_layout( + 2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout) + ), + ), + q_dtype, + ) + tSsScale = thr_r2s_tiled_copy.partition_S(sScale) + tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0]) + scale_v_handle = load_scale_consumer.wait_and_advance() + cute.autovec_copy( + tSsScale[None, None, None, None, None, scale_v_handle.index], + tSrScale, + ) + cute.arch.fence_view_async_shared() + scale_v_handle.release() + # prefetch iter = 0 + kv_handle = load_kv_consumer.wait_and_advance() + cute.autovec_copy( + tOsOrig[None, None, None, None, kv_handle.index], + tOrOrig[None, None, None, None, 0], + ) + transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, 0].load(), + transformed_tensor.shape, + q_dtype, + ) + transformed_tensor = transformed_tensor * scale + tTrTrans[None, None, None, None, 0].store(transformed_tensor) + cute.arch.fence_view_async_shared() + kv_handle.release() + for iter in cutlass.range(1, iterations, unroll_full=True): + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (iter - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + kv_handle = load_kv_consumer.wait_and_advance() + cute.autovec_copy( + tOsOrig[None, None, None, None, kv_handle.index], + tOrOrig[None, None, None, None, iter % 2], + ) + transformed_tensor = ( + tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype) + ) + scale = cute.TensorSSA( + tSrScale[ + None, + None, + None, + None, + iter, + ].load(), + transformed_tensor.shape, + q_dtype, + ) + transformed_tensor = transformed_tensor * scale + tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor) + cute.arch.fence_view_async_shared() + kv_handle.release() + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (iterations - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + return load_kv_consumer, load_scale_consumer, dequant_kv_producer diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py new file mode 100644 index 00000000..a81d1e39 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py @@ -0,0 +1,2526 @@ +# 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 argparse +from math import log2, ceil +from typing import Optional +import os +import sys + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.utils.mixed_input_helpers import TransformMode +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import cpasync, tcgen05 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mixed_input_gemm.mixed_input_host_utils import ( + create_tensors_for_contiguous_grouped_mixed_input_gemm as create_tensors, + run_contiguous_grouped_ref_and_compare as run_ref_and_compare, +) +from cutlass.cute.runtime import from_dlpack + +""" +A mixed-input grouped GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL. + +This example demonstrates an implementation of mixed-input grouped GEMM using a TMA plus Blackwell +SM100 TensorCore warp-specialized persistent kernel. It can be viewed as an extension of the batched +mixed-input GEMM example to support a specific grouped GEMM pattern: grouped GEMM with contiguous offsets. + +Specifically, the input A tensor is still in the shape of (M, K, L), and L is the number of groups. The +input B tensor is in the shape of (N, K) and the result C tensor is in the shape of (M, N). Tensor B +and tensor C are not divided into groups explicitly and there is an extra input tensor cumsum defining +the mapping between the N mode to groups. The cumsum tensor is in the shape of (N+1) and cumsum[i] +defines the accumulated size along N mode for groups up to i (not including i): + + ``` + Group 0 Group 1 Group 2 ..... Group L-1 + -+--------+--------+--------+.....+----------------+ + | | | | | + |<- N0 ->|<- N1 ->|<- N2 ->|.....|<-- NL-1 -->| + | | | | | + -+--------+--------+--------+.....+-------------------+ +cumsum: | 0 | N0 | N0+N1 |.....| sum(N0,N1,...NL-2) | sum(N0,N1,...NL-1) + ``` + +The computation flow is the same as the batched mixed-input GEMM example. A is the narrow-precision tensor +and B holds data with a wider precision. MMA will work in the wide precision of tensor B and tensor A +will be transformed to the wide precision of tensor B following 1 of the 2 possible modes as follows: + +1. convert-only mode: + C = type_convert(A) x B + +In convert-only mode, tensor A is directly converted to the wide precision of tensor B. + +2. convert-scale mode: + C = (type_convert(A) * scale) x B + +In convert-scale mode, tensor A is first converted to the wide precision of tensor B and then scaled by the scale tensor. +The scale tensor is in the same precision as tensor B. +The mode is determined by tensor A's data type as follows: +- if tensor A is in int8 or uint8, convert-only mode is used. +- if tensor A is in int4, convert-scale mode is used. + +The output tensor C could have the same precision as tensor B or fp32. + +To run this example: + +.. code-block:: bash + + python examples/blackwell/grouped_mixed_input_gemm.py \ + --a_dtype Int8 --b_dtype BFloat16 \ + --scale_granularity_m 0 --scale_granularity_k 0 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \ + --mnkl 256,512,8192,1 + +Input A and B have int8 and bf16 data types, respectively. The Blackwell tcgen05 MMA tile shape +is specified as (128,128,64) and the cluster shape is (1,1). The MMA accumulator and output data type +are set as fp32 and bf16, respectively. As tensor A is int8, convert-only mode is used. +scale_granularity_m and scale_granularity_k are set as 0 for convert-only mode. + +Here is an example of running convert-scale mode: + +.. code-block:: bash + + python examples/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py \ + --a_dtype Int4 --b_dtype BFloat16 \ + --scale_granularity_m 1 --scale_granularity_k 256 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 256,128,128 --cluster_shape_mn 2,1 \ + --use_2cta_instrs --mnkl 1024,8192,6144,16 \ + +Input A and B have int4 and bf16 data types, respectively. The scale granularity is set as (1,256), +which means each element along the m mode of tensor A has its own scale element and 256 contiguous elements +along the k mode share the same scale element. There is no scale reuse along the L mode. If the GEMM shape is +(M, N, K, L), then the scale tensor shape is (M // scale_granularity_m, K // scale_granularity_k, L), +which is (1024, 6144/256, 16) in this example. +The Blackwell tcgen05 MMA tile shape is specified as (256,128,128) and tcgen05 2CTA feature is enabled. +The cluster shape is (2,1). The MMA accumulator and output data type are set as fp32 and bf16, respectively. +As tensor A is int4, the convert-scale mode is used. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py \ + --a_dtype Int8 --b_dtype BFloat16 \ + --scale_granularity_m 0 --scale_granularity_k 0 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \ + --mnkl 256,512,8192,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + +Besides the requirements from the batched mixed-input GEMM example, there are some constraints for this example: +* --use_tma_store option is removed as no alignment assumption is made for each group. +""" + + +class GroupedMixedInputGemmKernel: + """ + Mixed-input grouped GEMM kernel for NVIDIA Blackwell SM100 architecture. + + This kernel supports GEMM operations where input tensors A and B have different + data types, with tensor A being transformed to the precision of tensor B before + matrix multiplication. + Tensor A is in shape of [M, K, L] with L being the number of groups. Tensor B is in shape of [N, K] and a group search algorithm + is applied along the N mode to find the group index for each CTA tile. A cumsum tensor provides the offsets of each group along the N mode. + + :param scale_granularity_m: Number of elements sharing the same scale factor along the M mode + :type scale_granularity_m: int + :param scale_granularity_k: Number of elements sharing the same scale factor along the K mode + :type scale_granularity_k: int + :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_mnk: Shape of the Matrix Multiply-Accumulate (MMA) tile (M, N, K) + :type mma_tiler_mnk: tuple[int, int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: tuple[int, int] + :param group_count: The total number of groups + :type group_count: int + :param shuffle_a: Whether to use shuffle intrinsic for int4-to-bf16 conversion + :type shuffle_a: bool + """ + + def __init__( + self, + scale_granularity_m: int, + scale_granularity_k: int, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + group_count: int, + shuffle_a: bool, + ): + """ + Initializes the mixed-input GEMM kernel with a specified configuration. + """ + # Scale granularity defines how many elements share the same scale factor + # along the M and K modes. + self.scale_granularity_m = scale_granularity_m + self.scale_granularity_k = scale_granularity_k + # Set transform mode + if cutlass.const_expr( + self.scale_granularity_m == 0 and self.scale_granularity_k == 0 + ): + self.scale_mode = TransformMode.ConvertOnly + else: + self.scale_mode = TransformMode.ConvertScale + self.group_count = group_count + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = mma_tiler_mnk + self.shuffle_a = shuffle_a + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.scale_tma_warp_id = 6 + # Schedule warp to do the group search + self.schedule_warp_id = 7 + self.transform_warp_id = ( + 8, + 9, + 10, + 11, + ) + # Define expected register count for different warps + self.num_regs_epilogue_warps = 192 + self.num_regs_mma_warp = 96 + self.num_regs_tma_warps = 80 + self.num_regs_transform_warps = 208 + self.num_regs_schedule_warp = 64 + self.threads_per_cta = 32 * ( + max( + ( + self.mma_warp_id, + self.tma_warp_id, + self.scale_tma_warp_id, + *self.epilog_warp_id, + *self.transform_warp_id, + ) + ) + + 1 + ) + + # Set barrier id for cta sync, epilogue sync, tmem ptr sync, and transform sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + 1, 32 * len(self.epilog_warp_id) + ) + self.tmem_ptr_sync_barrier = pipeline.NamedBarrier(2, self.threads_per_cta) + self.transform_sync_barrier = pipeline.NamedBarrier( + 3, 32 * len(self.transform_warp_id) + ) + self.cta_sync_barrier = pipeline.NamedBarrier(4, self.threads_per_cta) + self.sched_sync_barrier = pipeline.NamedBarrier(5, 32) + + self.smem_buffer_align_bytes = 1024 + + 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: + - Deduce where the transformed A tensor is stored + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue sub-tile + - Setting up A/scale/B/C stage counts in shared memory + - Setting up transformed A stage count in shared memory or tensor memory + - Computing A/transformed A/scale/B/C memory layout + - Computing tensor memory allocation columns + """ + # Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM) + self.transform_a_source = mixed_input_utils.get_transform_a_source( + self.a_major_mode + ) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_tile_shape_mnk = ( + self.cluster_shape_mn[0] * self.cta_tile_shape_mnk[0], + self.cluster_shape_mn[1] * self.cta_tile_shape_mnk[1], + self.cta_tile_shape_mnk[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + 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 + + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + # Compute tensor memory(TMEM) columns and stages for each pipeline + ( + self.num_load2trans_stage, + self.num_scale_load2trans_stage, + self.num_trans2mma_stage, + self.num_acc_stage, + self.num_c_stage, + self.num_tile_info_stage, + self.num_acc_tmem_cols, + self.num_a_tmem_cols, + ) = self._compute_stages_and_tmem_cols( + tiled_mma, + self.mma_tiler, + self.cta_tile_shape_mnk, + self.epi_tile, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.c_layout, + self.transform_a_source, + self.scale_granularity_m, + self.scale_granularity_k, + self.smem_buffer_align_bytes, + self.scale_mode, + ) + + # Align TMEM columns for allocation + # TMEM allocation requires power-of-2 column alignment + # and must meet minimum allocation requirements + self.num_tmem_alloc_cols = cute.round_up( + self.num_acc_tmem_cols + self.num_a_tmem_cols, + cute.arch.get_min_tmem_alloc_cols("sm_100"), + ) + self.num_tmem_alloc_cols = 2 ** (ceil(log2(self.num_tmem_alloc_cols))) + # Get smem layout for C tensor + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + # Get smem layout for A, transformed A, and B + ( + self.smem_layout_a, + self.smem_layout_a_transform, + self.smem_layout_b, + ) = mixed_input_utils.compute_smem_layout( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.num_load2trans_stage, + self.num_trans2mma_stage, + ) + # Get smem layout for scale tensor + self.smem_layout_scale_per_stage = None + self.smem_layout_scale = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Get scale tile shape and smem layout for scale tensor + ( + self.scale_tile_shape, + self.smem_layout_scale_per_stage, + self.smem_layout_scale, + ) = mixed_input_utils.get_smem_layout_scale( + self.mma_tiler, + self.use_2cta_instrs, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + self.a_scale_dtype, + self.num_scale_load2trans_stage, + ) + + def _validate_inputs( + self, + a: cute.Tensor, + a_scale: Optional[cute.Tensor], + b: cute.Tensor, + c: cute.Tensor, + ) -> None: + """ + Validates input tensors and their properties. + + :param a: Input tensor A. + :type a: cute.Tensor + :param a_scale: Scale tensor for tensor A (None for ConvertOnly mode). + :type a_scale: Optional[cute.Tensor] + :param b: Input tensor B. + :type b: cute.Tensor + :param c: Output tensor C. + :type c: cute.Tensor + :raises ValueError: If inputs don't meet kernel requirements. + """ + # Validate scale tensor major mode + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + and utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + != tcgen05.OperandMajorMode.MN + ): + raise ValueError("scale_major_mode must be M-major") + + @cute.jit + def __call__( + self, + a: cute.Tensor, + a_scale: Optional[cute.Tensor], # None for ConvertOnly mode + b: cute.Tensor, + cumsum: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """ + Executes the Mixed Input Grouped GEMM operation. + + This method sets up the kernel parameters, computes the grid size, + defines the shared storage, and launches the kernel. + + The execution steps are as follows: + - 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 a_scale: Scale tensor for tensor A (None for ConvertOnly mode). + :type a_scale: Optional[cute.Tensor] + :param b: Input tensor B. + :type b: cute.Tensor + :param cumsum: tensor containing the cumulative size of each group along the search mode(aka, N mode in this example). + :type cumsum: cute.Tensor + :param c: Output tensor C. + :type c: cute.Tensor + :param max_active_clusters: Maximum number of active clusters to launch. + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream to launch the kernel on. + :type stream: cuda.CUstream + """ + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.a_scale_dtype: type[cutlass.Numeric] = ( + a_scale.element_type + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.mma_dtype = self.b_dtype + + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.scale_major_mode = ( + utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Get gmem layout for scale tensor + self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale( + a.shape, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + ) + + # Validate inputs + self._validate_inputs(a, a_scale, b, c) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + # Set up gmem copy atoms for A, scale, and B + a_op = mixed_input_utils.get_tma_atom_kind( + self.is_a_mcast, self.use_2cta_instrs, False + ) + b_op = mixed_input_utils.get_tma_atom_kind( + self.is_b_mcast, self.use_2cta_instrs, True + ) + a_scale_op = a_op + # Deduce TMA copy atom and TMA tensor for A, scale, and B + smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + smem_layout_a_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + tma_atom_scale, tma_tensor_scale = None, None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Partition smem layout for scale tensor to make it compatible with TMA atom + smem_layout_for_tma_atom = cute.get( + tiled_mma._thrfrg_A(self.smem_layout_scale_per_stage.outer), mode=[1] + ) + # ((MMA_M, MMA_K), REST_M, REST_K) + smem_layout_for_tma_atom = cute.dice( + smem_layout_for_tma_atom, + (1, (1,) * cute.rank(self.smem_layout_scale_per_stage.outer)), + ) + tma_atom_scale, tma_tensor_scale = cute.nvgpu.make_tiled_tma_atom_A( + a_scale_op, + cute.make_tensor(a_scale.iterator, self.gmem_layout_scale), + smem_layout_for_tma_atom, + # (SCALE_M, 1, SCALE_K) + (self.scale_tile_shape[0], 1, self.scale_tile_shape[1]), + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 + if a_scale.element_type is cutlass.Float32 + else None + ), + ) + + smem_layout_b_per_stage = cute.slice_(self.smem_layout_b, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + smem_layout_b_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + # Calculate copy size for tensor A, B, and scale + a_copy_size = cute.size_in_bytes(self.a_dtype, smem_layout_a_per_stage) + b_copy_size = cute.size_in_bytes(self.b_dtype, smem_layout_b_per_stage) + a_scale_copy_size = ( + cute.size_in_bytes(self.a_scale_dtype, self.smem_layout_scale_per_stage) + if self.scale_mode is TransformMode.ConvertScale + else 0 + ) + + self.num_tma_load_bytes_a = a_copy_size + self.num_tma_load_bytes_b = b_copy_size * cute.size(tiled_mma.thr_id.shape) + self.num_tma_load_bytes_scale = a_scale_copy_size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + 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, + ) + + @cute.struct + class SharedStorage: + # buffer holding group search results + tile_info: cute.struct.MemRange[cutlass.Int32, 4 * self.num_tile_info_stage] + a_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_scale_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_scale_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_trans2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + a_trans2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + b_load2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + b_load2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_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] + tile_info_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_tile_info_stage + ] + tile_info_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_tile_info_stage + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + self.shared_storage = SharedStorage + + # Launch kernel + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_scale, + tma_tensor_scale, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + c, + cumsum, + self.group_count, + self.cluster_layout_vmnk, + self.smem_layout_a, + self.smem_layout_scale, + self.smem_layout_a_transform, + self.smem_layout_b, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + min_blocks_per_mp=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_s: Optional[cute.CopyAtom], + mS_mkl: Optional[cute.Tensor], + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tensor_c: cute.Tensor, + cumsum: cute.Tensor, + group_count: cutlass.Constexpr[int], + cluster_layout_vmnk: cute.Layout, + a_smem_layout: cute.ComposedLayout, + scale_smem_layout: cute.ComposedLayout, + a_smem_layout_transform: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + c_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """ + GPU device kernel performing the Persistent Mixed-Input Grouped GEMM computation. + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + # Prefetch TMA descriptors + if warp_idx == self.epilog_warp_id[0]: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + cpasync.prefetch_descriptor(tma_atom_s) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + bidx, bidy, bidz = cute.arch.block_idx() + # Compute how many k_tiles share the same scale + num_k_tiles_per_scale = self.scale_granularity_k // self.cta_tile_shape_mnk[2] + + 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 + ) + tidx, _, _ = cute.arch.thread_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize load2transform pipeline, which tracks the dependencies between TMA's loading + # of A and B, and the transformation of A and MMA's consumption + transform_thread_idx = ( + tidx - 32 * self.transform_warp_id[0] + if tidx >= 32 * self.transform_warp_id[0] + else tidx + ) + a_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mcast_ctas_a * len(self.transform_warp_id), + ), + tx_count=self.num_tma_load_bytes_a, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=(1, 0), # multicast for A will only happen on the M-mode + defer_sync=True, + ) + # Initialize scale_load2trans pipeline, which tracks the dependencies between TMA's loading + # of scale, and the transformation of A + scale_load2trans_pipeline = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + num_producers_a_scale = self.num_mcast_ctas_a + scale_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_scale_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_scale_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_producers_a_scale + * len(self.transform_warp_id) + * num_k_tiles_per_scale, + ), + tx_count=self.num_tma_load_bytes_scale, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=( + 1, + 0, + ), # multicast for scale_a will only happen on the M-mode + defer_sync=True, + ) + # Initialize transform2mma pipeline, which tracks the dependencies between the transformation + # of A and MMA's consumption of transformed A + cta_v_size = cute.size(cluster_layout_vmnk, mode=[0]) + trans2mma_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.a_trans2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_trans2mma_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.transform_warp_id) * cta_v_size, + ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize pipeline for tensor B load to MMA + # MMA warp informs TMA warp to proceed to load next tile of B tensor + b_load2mma_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_load2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ), + tx_count=self.num_tma_load_bytes_b, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), # multicast for B will only happen on the N-mode + defer_sync=True, + ) + # Initialize accumulator pipeline, which tracks the dependencies between + # MMA's computation of accumulators and epilogue warps' consumption of accumulators + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize tile info pipeline, which tracks the dependencies between + # tile scheduling warp and other warps + # Skip scheduler warp and TMA scale load warp when scale_mode is ConvertOnly + # when computing consumer thread count + num_tile_info_pipeline_consumer_threads = ( + self.threads_per_cta + - 32 + - (32 if self.scale_mode is TransformMode.ConvertOnly else 0) + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.tile_info_full_mbar_ptr.data_ptr(), + num_stages=self.num_tile_info_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * 1), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_tile_info_pipeline_consumer_threads, + ), + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_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=self.cluster_shape_mn, is_relaxed=True) + + # Setup smem tensor A/scale/B/C + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=c_smem_layout_staged.inner, + ) + sA_input = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=a_smem_layout.inner, + ) + sS_input = ( + smem.allocate_tensor( + element_type=self.mma_dtype, + layout=scale_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=scale_smem_layout.inner, + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + sB_input = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=b_smem_layout.inner, + ) + sA_transform = None + # Get smem tensor for transformed A when transform_a_source is SMEM + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.SMEM): + sA_transform = smem.allocate_tensor( + element_type=self.mma_dtype, + layout=a_smem_layout_transform.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=a_smem_layout_transform.inner, + ) + sTile_info = storage.tile_info.get_tensor( + cute.make_layout((4, self.num_tile_info_stage), stride=(1, 4)) + ) + + # Compute multicast mask for A/B buffer full + a_full_mcast_mask = None + b_full_mcast_mask = None + s_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 + ) + # Scale tensor shares the same multicast mask as the A tensor + s_full_mcast_mask = a_full_mcast_mask + 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, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bM, bK, loopM, loopK, loopL) + gS_mkl = ( + cute.local_tile( + mS_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + gC_mnl_simt = cute.local_tile( + tensor_c, 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, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgS = ( + thr_mma.partition_A(gS_mkl) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + tCgC_simt = thr_mma.partition_C(gC_mnl_simt) + + # Setup copy atom to load A from shared memory for further transformation + copy_atom_a_input = ( + cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.a_dtype, num_bits_per_copy=32 + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + a_smem_shape = tiled_mma.partition_shape_A( + cute.dice(self.mma_tiler, (1, None, 1)) + ) + # Setup copy atom to store transformed A into tensor memory or shared memory + copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform( + self.mma_dtype, + self.use_2cta_instrs, + self.transform_a_source, + a_smem_shape, + self.a_dtype, + ) + + # 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), loopM, loopK, loopL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA_input, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + + tCsS = None + tSsS = None + tSgS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + thr_mma_leader_cta = tiled_mma.get_slice(0) + # (MMA, MMA_M, MMA_K, STAGE) + tCsS = thr_mma_leader_cta.partition_A(sS_input) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS, tSgS = mixed_input_utils.scale_tma_partition( + tCsS, + tCgS, + tma_atom_s, + block_in_cluster_coord_vmnk, + a_cta_layout, + ) + + # 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), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB_input, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB_input) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc and ensure pipelines are ready + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # TMEM allocation + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + + # Schedule warp + if warp_idx == self.schedule_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_schedule_warp) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentRuntimeTileScheduler.create( + tile_sched_params, + (bidx, bidy, bidz), + cute.arch.grid_dim(), + inner_mode=0, + ) + work_tile = tile_sched.initial_work_tile_info() + tile_info_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_tile_info_stage + ) + # Create initial group search state + search_state = ( + mixed_input_utils.create_initial_contiguous_group_search_state() + ) + not_last_tile = cutlass.Boolean(1) + while not_last_tile: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + cluster_tile_coord_mnl = work_tile.tile_idx + cta_tile_coord_m = ( + cluster_tile_coord_mnl[0] * self.cluster_shape_mn[0] + + block_in_cluster_coord_vmnk[1] * cute.size(tiled_mma.thr_id.shape) + + block_in_cluster_coord_vmnk[0] + ) + cta_tile_offset_n = block_in_cluster_coord_vmnk[2] + search_state = mixed_input_utils.contiguous_group_search( + self.cluster_tile_shape_mnk, + group_count, + cluster_tile_coord_mnl[1], + search_state, + cumsum, + 1, # mode index to perform the search. 0 for M and 1 for N + ) + cur_sTile_info = sTile_info[(None, tile_info_producer_state.index)] + not_last_tile = search_state.cur_group_idx <= group_count + # Store tile info into shared memory buffer + with cute.arch.elect_one(): + cur_sTile_info[0] = cta_tile_coord_m + cur_sTile_info[1] = ( + search_state.cur_start + + cta_tile_offset_n * self.cta_tile_shape_mnk[1] + ) + cur_sTile_info[2] = search_state.cur_group_idx - 1 + cur_sTile_info[3] = ( + search_state.cur_boundary + - search_state.cur_start + - (cta_tile_offset_n * self.cta_tile_shape_mnk[1]) + ) + # Fence and barrier to ensure tile info store has finished + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.sched_sync_barrier.arrive_and_wait() + # Commit tile info pipeline + tile_info_pipeline.producer_commit(tile_info_producer_state) + # Advance to next tile + tile_info_producer_state.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # Specialized TMA load warp for A/B tensor + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + a_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + b_load2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + + while work_tile.is_valid_tile: + tAgA_slice = tAgA[ + ( + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + None, + work_tile.group_idx, + ) + ] + # Apply offset to B tensor based on group search result + coord_n_offset = ( + (work_tile.coord_n, 0, 0) + if cutlass.const_expr( + self.b_major_mode == tcgen05.OperandMajorMode.MN + ) + else (0, work_tile.coord_n, 0) + ) + tBgB_slice = cute.make_tensor( + ( + tBgB.iterator[0] + coord_n_offset[0], + coord_n_offset[1] + tBgB.iterator[1], + coord_n_offset[2] + tBgB.iterator[2], + ), + cute.slice_(tBgB.layout, (None, 0, None, 0)), + ) + + a_load2trans_producer_state.reset_count() + peek_load2trans_empty_status = cutlass.Boolean(1) + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + b_load2mma_producer_state.reset_count() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.producer_acquire( + a_load2trans_producer_state, peek_load2trans_empty_status + ) + b_load2mma_pipeline.producer_acquire(b_load2mma_producer_state) + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, a_load2trans_producer_state.count)], + tAsA[(None, a_load2trans_producer_state.index)], + tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier( + a_load2trans_producer_state + ), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, b_load2mma_producer_state.count)], + tBsB[(None, b_load2mma_producer_state.index)], + tma_bar_ptr=b_load2mma_pipeline.producer_get_barrier( + b_load2mma_producer_state + ), + mcast_mask=b_full_mcast_mask, + ) + a_load2trans_pipeline.producer_commit(a_load2trans_producer_state) + b_load2mma_pipeline.producer_commit(b_load2mma_producer_state) + a_load2trans_producer_state.advance() + b_load2mma_producer_state.advance() + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait A/B buffer empty + a_load2trans_pipeline.producer_tail(a_load2trans_producer_state) + b_load2mma_pipeline.producer_tail(b_load2mma_producer_state) + + # Specialized TMA load for scale tensor + if warp_idx == self.scale_tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + scale_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scale_load2trans_stage + ) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + + while work_tile.is_valid_tile: + # ((atom_v, rest_v), RestK) + tSgS_slice = tSgS[ + ( + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + None, + work_tile.group_idx, + ) + ] + # Filter zeros in rest mode + rest_filtered = cute.filter_zeros(tSgS_slice[(0, None)].layout) + tSgS_slice_filtered = cute.make_tensor( + tSgS_slice.iterator, + cute.make_layout( + (tSgS_slice.layout[0].shape, rest_filtered.shape), + stride=(tSgS_slice.layout[0].stride, rest_filtered.stride), + ), + ) + + scale_load2trans_producer_state.reset_count() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + scale_load2trans_pipeline.producer_acquire( + scale_load2trans_producer_state, + peek_scale_load2trans_empty_status, + ) + # TMA load scale + cute.copy( + tma_atom_s, + tSgS_slice_filtered[ + (None, scale_load2trans_producer_state.count) + ], + tSsS[(None, scale_load2trans_producer_state.index)], + tma_bar_ptr=scale_load2trans_pipeline.producer_get_barrier( + scale_load2trans_producer_state + ), + mcast_mask=s_full_mcast_mask, + ) + + scale_load2trans_producer_state.advance() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait scale buffer empty + scale_load2trans_pipeline.producer_tail(scale_load2trans_producer_state) + + # Specialized transform warps + if warp_idx >= self.transform_warp_id[0]: + cute.arch.setmaxregister_increase(self.num_regs_transform_warps) + transform_local_tidx = tidx - 32 * self.transform_warp_id[0] + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + # Partition tensors for transform input and output and set up the copy atom + # used for loading and storing transformed A tensor + src_copy_a, dst_copy_a, tAsA_input, tAsA_transform = ( + mixed_input_utils.transform_partition( + self.transform_a_source, + self.scale_mode, + copy_atom_a_input, + copy_atom_a_transform, + sA_input, + ( + tCrA + if self.transform_a_source == tcgen05.OperandSource.TMEM + else sA_transform + ), + transform_local_tidx, + ) + ) + # make fragment for input A and transformed A + tArA = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, tAsA_input.element_type + ) + tArA_transform = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, self.mma_dtype + ) + # Partition scale tensor + smem_thr_copy_S = None + tSsS_trans = None + tSrS_copy = None + tSrS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = ( + mixed_input_utils.scale_partition( + src_copy_a, tCsS, transform_local_tidx, self.mma_dtype + ) + ) + assert cute.size(tSrS, mode=[0]) == cute.size(tArA, mode=[0]), ( + "tSrS and tArA have different leading dimension" + ) + assert cute.size(tSrS) == cute.size(tArA), ( + "tSrS and tArA have different shape" + ) + # Deduce a sub-tile size and tile tensors + transform_tiler_size = min( + cute.size(cute.coalesce(tAsA_input.layout), mode=[0]), 64 + ) + transform_tiler = cute.make_layout(transform_tiler_size) + tArA_load = cute.flat_divide(tArA, transform_tiler) + tArA_load = cute.group_modes(tArA_load, 1, cute.rank(tArA_load)) + tSrS_load = ( + cute.flat_divide(tSrS, transform_tiler) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tSrS_load = ( + cute.group_modes(tSrS_load, 1, cute.rank(tSrS_load)) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tArA_transform_store = cute.flat_divide(tArA_transform, transform_tiler) + tArA_transform_store = cute.group_modes( + tArA_transform_store, 1, cute.rank(tArA_transform_store) + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + a_load2trans_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_load2trans_stage, + ) + scale_load2trans_consumer_state = ( + pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_scale_load2trans_stage, + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + trans2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_trans2mma_stage, + ) + while work_tile.is_valid_tile: + a_load2trans_consumer_state.reset_count() + peek_load2trans_full_status = cutlass.Boolean(1) + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + peek_scale_load2trans_full_status = cutlass.Boolean(1) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + scale_load2trans_consumer_state.reset_count() + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + trans2mma_producer_state.reset_count() + peek_trans2mma_empty_status = cutlass.Boolean(1) + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.consumer_wait( + a_load2trans_consumer_state, peek_load2trans_full_status + ) + tAsA_input_slice = tAsA_input[ + (None, None, None, None, a_load2trans_consumer_state.index) + ] + tAsA_input_slice = cute.flat_divide( + tAsA_input_slice, transform_tiler + ) + tAsA_input_slice = cute.group_modes( + tAsA_input_slice, 1, cute.rank(tAsA_input_slice) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_wait( + scale_load2trans_consumer_state, + peek_scale_load2trans_full_status, + ) + trans2mma_pipeline.producer_acquire( + trans2mma_producer_state, peek_trans2mma_empty_status + ) + # load scale tensor when needed + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + if k_tile % num_k_tiles_per_scale == 0: + tSsS_slice = tSsS_trans[ + ( + None, + None, + None, + None, + scale_load2trans_consumer_state.index, + ) + ] + tSsS_slice_filtered = cute.make_tensor( + tSsS_slice.iterator, + cute.filter_zeros(tSsS_slice.layout), + ) + cute.autovec_copy(tSsS_slice_filtered, tSrS_copy) + cur_scale_load2trans_consumer_state = ( + scale_load2trans_consumer_state.clone() + ) + if (k_tile + 1) % num_k_tiles_per_scale == 0: + scale_load2trans_consumer_state.advance() + + cur_a_load2trans_consumer_state = ( + a_load2trans_consumer_state.clone() + ) + for idx in cutlass.range_constexpr(cute.size(tArA_load, mode=[1])): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice[(None, idx)], + tArA_load[(None, idx)], + ) + if cutlass.const_expr( + idx == cute.size(tArA_load, mode=[1]) - 1 + ): + a_load2trans_consumer_state.advance() + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load[(None, idx)], self.mma_dtype, self.shuffle_a + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale = cute.TensorSSA( + tSrS_load[(None, idx)].load(), + tensor_transformed.shape, + self.mma_dtype, + ) + # Apply scale + tensor_transformed = tensor_transformed * scale + tArA_transform_store[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + dst_copy_a, + ) + # Ensure all transform threads have finished the copy and reached the fence + self.transform_sync_barrier.arrive_and_wait() + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_release( + cur_scale_load2trans_consumer_state + ) + + a_load2trans_pipeline.consumer_release( + cur_a_load2trans_consumer_state + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state) + trans2mma_producer_state.advance() + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait a_transform buffer empty + trans2mma_pipeline.producer_tail(trans2mma_producer_state) + + # Specialized MMA warp + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_mma_warp) + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + tCtAcc_base = accumulators + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + trans2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_trans2mma_stage + ) + b_load2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_load2trans_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + b_load2mma_consumer_state.reset_count() + trans2mma_consumer_state.reset_count() + peek_trans2mma_full_status = cutlass.Boolean(1) + if is_leader_cta: + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Mma mainloop + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + trans2mma_pipeline.consumer_wait( + trans2mma_consumer_state, peek_trans2mma_full_status + ) + b_load2mma_pipeline.consumer_wait(b_load2mma_consumer_state) + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord_a = ( + None, + None, + kblock_idx, + trans2mma_consumer_state.index, + ) + kblock_coord_b = ( + None, + None, + kblock_idx, + b_load2mma_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord_a], + tCrB[kblock_coord_b], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + trans2mma_pipeline.consumer_release(trans2mma_consumer_state) + b_load2mma_pipeline.consumer_release(b_load2mma_consumer_state) + trans2mma_consumer_state.advance() + b_load2mma_consumer_state.advance() + peek_trans2mma_full_status = cutlass.Boolean(1) + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + # Async arrive accumulator buffer full + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait for accumulator buffer empty + acc_pipeline.producer_tail(acc_producer_state) + + # Specialized epilogue warps + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + epi_tidx = tidx + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCtAcc_base = accumulators + # Partition for epilogue + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = ( + mixed_input_utils.epilog_tmem_copy_and_partition( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + 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 = ( + mixed_input_utils.epilog_smem_copy_and_partition( + self.c_layout, + self.c_dtype, + self.acc_dtype, + tiled_copy_t2r, + tTR_rC, + epi_tidx, + sC, + ) + ) + (tma_atom_c, bSG_sC, bSG_gC_partitioned, simt_atom, tTR_gC_partitioned) = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, + epi_tidx, + tma_atom_c, + tiled_copy_t2r, + tCgC, + tCgC_simt, + epi_tile, + sC, + ) + ) + + # Predicates + thr_mapping = cute.make_identity_tensor( + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1]) + ) + thr_mapping_mn = cute.flat_divide(thr_mapping, epi_tile) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + m_thr_offset = thr_copy_t2r.partition_D(thr_mapping_mn) + m_thr_offset = cute.group_modes(m_thr_offset, 3, cute.rank(m_thr_offset)) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + 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, + ) + + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + num_prev_subtiles = cutlass.Int32(0) + while work_tile.is_valid_tile: + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + 0, + 0, + ) + ] + tma_store_offset_coord = ( + (work_tile.coord_n, 0, 0) + if cutlass.const_expr(self.c_layout.is_n_major_c()) + else (0, work_tile.coord_n, 0) + ) + bSG_gC = cute.make_tensor( + ( + tma_store_offset_coord[0] + bSG_gC.iterator[0], + tma_store_offset_coord[1] + bSG_gC.iterator[1], + tma_store_offset_coord[2] + bSG_gC.iterator[2], + ), + bSG_gC.layout, + ) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + 0, + 0, + ) + ] + tTR_gC = cute.make_tensor( + tTR_gC.iterator + (work_tile.coord_n * tensor_c.layout.stride[1]), + tTR_gC.layout, + ) + + 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)) + 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]) + 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 work_tile.distance_to_boundary >= self.cta_tile_shape_mnk[1]: + # Convert to C type + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + num_prev_subtiles += 1 + c_buffer = num_prev_subtiles % self.num_c_stage + # Store C to shared memory + 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( + "async.shared", + space="cta", + ) + self.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)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # Convert to C type + acc_vec = tTR_rAcc.load() + acc_vec = acc_vec.to(self.c_dtype) + tTR_rC.store(acc_vec) + # Compute predicate for SIMT store + tCpC = cute.make_rmem_tensor( + cute.make_layout(tTR_rC.shape), + cutlass.Boolean, + ) + m_thr_slice = m_thr_offset[(None, None, None, subtile_idx)] + for i in cutlass.range(cute.size(tCpC), unroll_full=True): + tCpC[i] = ( + m_thr_slice[(i)][0] + + work_tile.cta_coord_m * self.cta_tile_shape_mnk[0] + < tensor_c.shape[0] + ) and (m_thr_slice[(i)][1] < work_tile.distance_to_boundary) + # Store C to global memory + cute.copy( + simt_atom, + cute.flatten(tTR_rC), + cute.flatten(tTR_gC[(None, None, None, subtile_idx)]), + pred=cute.flatten(tCpC), + ) + # 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_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + c_pipeline.producer_tail() + + @staticmethod + def _compute_stages_and_tmem_cols( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + cta_tile_shape_mnk: tuple[int, int, int], + epi_tile: cute.Tile, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + transform_a_source: tcgen05.OperandSource, + scale_granularity_m: int, + scale_granularity_k: int, + smem_buffer_align_bytes: int, + scale_mode: TransformMode, + ) -> tuple[int, int, int, int, int, int, int, int]: + """ + Compute pipeline stages and TMEM column allocation configurations. + + This method calculates the number of pipeline stages for different operations + (tile_info, load2trans, trans2mma, accumulator, etc.) and determines TMEM column allocation + based on available memory resources and tile configuration. + + :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 cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :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. + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param transform_a_source: The source of the transformed A tensor. + :type transform_a_source: tcgen05.OperandSource + :param scale_granularity_m: The granularity of the scale tensor along the M mode. + :type scale_granularity_m: int + :param scale_granularity_k: The granularity of the scale tensor along the K mode. + :type scale_granularity_k: int + :param smem_buffer_align_bytes: The alignment of the shared memory buffer. + :type smem_buffer_align_bytes: int + :param scale_mode: The transform mode. + :type scale_mode: TransformMode + + :return: A tuple containing the number of stages for: + (load2trans, scale_load2trans, transform2mma, accumulator, c, tile_info, tmem_acc_cols, tmem_a_cols) + :rtype: tuple[int, int, int, int, int, int, int] + - num_load2trans_stage: Stages for load-to-transform A and B tensors pipeline + - num_scale_load2trans_stage: Stages for scale load-to-transform A tensor pipeline + - num_trans2mma_stage: Stages for transform-to-MMA pipeline + - num_acc_stage: Stages for accumulator-to-epilogue pipeline + - num_c_stage: Stages for epilogue-to-output C pipeline + - num_tile_info_stage: Stages for buffers storing tile info + - num_acc_tmem_cols: TMEM columns for accumulator + - num_a_tmem_cols: TMEM columns for transformed A tensor + """ + # Compute tmem columns required for accumulator + acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) + tCtAcc_stage1 = tiled_mma.make_fragment_C(cute.append(acc_shape, 1)) + num_tmem_acc_col_per_stage = utils.get_num_tmem_alloc_cols(tCtAcc_stage1, True) + # Heuristic to decide the number of stages for accumulator + sm100_tmem_columns = cute.arch.get_max_tmem_alloc_cols("sm_100") + accumulator_stage_count = sm100_tmem_columns // num_tmem_acc_col_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM: + if num_tmem_acc_col_per_stage < 128: + accumulator_stage_count = 3 + elif num_tmem_acc_col_per_stage < 256: + accumulator_stage_count = 2 + else: + accumulator_stage_count = 1 + # transformed A in 16bit, thus 1 tmem column could hold 2 elements + num_elts_per_tmem_col = 32 // tiled_mma.op.a_dtype.width + num_tmem_cols_a_per_stage = cute.round_up( + ( + cta_tile_shape_mnk[2] // num_elts_per_tmem_col + if transform_a_source == tcgen05.OperandSource.TMEM + else 0 + ), + 4, + ) + + bytes_per_pipeline_stage = 16 + # By default, we use 2 stages for tile info + num_tile_info_stage = 2 + tile_info_bytes = ( + cute.size_in_bytes(cute.Int32, cute.make_layout((4, num_tile_info_stage))) + + bytes_per_pipeline_stage * num_tile_info_stage + ) + + c_stage_count = 2 + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * c_stage_count + + smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + if scale_mode == TransformMode.ConvertOnly: + scale_load2trans_stage_count = 0 + a_scale_bytes_per_stage = 0 + else: + # Ensure we have 4 buffers for scale tiles needed for 1 CTA tile + a_scale_k_mode = max(cta_tile_shape_mnk[2] // scale_granularity_k, 1) + a_scale_m_mode = max(cta_tile_shape_mnk[0] // scale_granularity_m, 1) + scale_load2trans_stage_count = 4 + a_scale_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((a_scale_m_mode, a_scale_k_mode)), + ), + smem_buffer_align_bytes, + ) + a_scale_bytes = ( + a_scale_bytes_per_stage + bytes_per_pipeline_stage + ) * scale_load2trans_stage_count + carveout_smem_bytes = ( + bytes_per_pipeline_stage * accumulator_stage_count + + a_scale_bytes + + c_bytes + + tile_info_bytes + ) + + # Compute transform stages if A is in TMEM + num_tmem_acc_cols = cute.round_up( + accumulator_stage_count * num_tmem_acc_col_per_stage, 4 + ) + + transform2mma_stage_count_a_source_tmem_potential = ( + (sm100_tmem_columns - num_tmem_acc_cols) // num_tmem_cols_a_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM + else -1 + ) + if ( + transform_a_source == tcgen05.OperandSource.TMEM + and transform2mma_stage_count_a_source_tmem_potential <= 0 + ): + raise ValueError("Not enough TMEM capacity for selected tile size") + a_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + b_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + b_dtype, + cute.make_layout( + ( + cta_tile_shape_mnk[1] // cute.size(tiled_mma.thr_id), + cta_tile_shape_mnk[2], + ) + ), + ), + smem_buffer_align_bytes, + ) + ab_load_bytes_per_stage = ( + a_load_bytes_per_stage + + b_load_bytes_per_stage + + 2 * bytes_per_pipeline_stage + ) + a_transform_bytes_per_stage = ( + cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + if transform_a_source == tcgen05.OperandSource.SMEM + else 0 + ) + + a_transform_bytes_per_stage = ( + a_transform_bytes_per_stage + bytes_per_pipeline_stage + ) + transform2mma_stage_count_a_source_smem_potential = ( + smem_capacity - carveout_smem_bytes + ) // (ab_load_bytes_per_stage + a_transform_bytes_per_stage) + transform2mma_stage_count = ( + min( + transform2mma_stage_count_a_source_tmem_potential, + transform2mma_stage_count_a_source_smem_potential, + ) + if transform_a_source == tcgen05.OperandSource.TMEM + else transform2mma_stage_count_a_source_smem_potential + ) + load2transform_stage_count = ( + smem_capacity + - carveout_smem_bytes + - (transform2mma_stage_count * a_transform_bytes_per_stage) + ) // ab_load_bytes_per_stage + if ( + load2transform_stage_count < 2 + or transform2mma_stage_count < 2 + or accumulator_stage_count < 1 + ): + raise ValueError("Not enough SMEM or TMEM capacity for selected tile size") + num_tmem_a_cols = transform2mma_stage_count * num_tmem_cols_a_per_stage + # Check if we can increase c_stage_count with leftover smem + c_stage_count += ( + smem_capacity + - load2transform_stage_count * ab_load_bytes_per_stage + - transform2mma_stage_count * a_transform_bytes_per_stage + - scale_load2trans_stage_count * a_scale_bytes_per_stage + - c_bytes + ) // c_bytes_per_stage + + return ( + load2transform_stage_count, + scale_load2trans_stage_count, + transform2mma_stage_count, + accumulator_stage_count, + c_stage_count, + num_tile_info_stage, + num_tmem_acc_cols, + num_tmem_a_cols, + ) + + @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. + """ + 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 = (cluster_shape_mn[0], cluster_shape_mn[1], max_active_clusters) + + return tile_sched_params, grid + + def can_implement( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + scale_granularity_m: int, + scale_granularity_k: int, + mma_tiler: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + ) -> bool: + """ + Check if the kernel can be implemented for the given tensor shapes and data types. + """ + m, n, k, l = mnkl + + if not mixed_input_utils.is_valid_mma_tiler_and_cluster_shape( + mma_tiler, cluster_shape_mn, use_2cta_instrs + ): + return False + if not mixed_input_utils.is_valid_scale_granularity( + scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2] + ): + return False + if not mixed_input_utils.is_valid_tensor_alignment( + m, + n, + k, + a_dtype, + b_dtype, + c_dtype, + b_dtype, + a_major, + b_major, + c_major, + mma_tiler, + use_2cta_instrs, + cluster_shape_mn, + scale_granularity_m, + scale_granularity_k, + ): + return False + return True + + +def get_advanced_compiler_control_path(): + """ + Return the path to the advanced compiler control file of this example. If not found, return None. + """ + import os + + need_advanced_compiler_control = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major == 13 and CUDA_VERSION.minor == 1: + need_advanced_compiler_control = True + except ImportError: + pass + + if not need_advanced_compiler_control: + return None + # Get the path to the advanced compiler control file + current_dir = os.path.dirname(os.path.abspath(__file__)) + target_path = os.path.join(current_dir, "../../advanced_compiler_control/gemm0.bin") + if os.path.exists(target_path): + print(f"Found advanced compiler control file at {target_path}") + return target_path + else: + return None + + +def run( + mnkl: tuple[int, int, int, int], + scale_granularity_m: int, + scale_granularity_k: int, + a_dtype: type[cutlass.Numeric], + b_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_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + uniform_group_sizes: bool = False, + use_cold_l2: bool = False, + **kwargs, +) -> None: + """ + Run the mixed-input GEMM kernel with specified parameters. + + This function creates tensors, validates parameters, executes the kernel, + optionally compares results with a reference implementation and reports + kernel execution time. + """ + m, n, k, l = mnkl + import torch + + if not torch.cuda.is_available(): + raise ValueError("CUDA is not available") + + # Check if given configuration is supported + if not GroupedMixedInputGemmKernel.can_implement( + mnkl, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + scale_granularity_m, + scale_granularity_k, + mma_tiler_mnk, + cluster_shape_mn, + use_2cta_instrs, + ): + raise ValueError("GEMM configuration not supported") + + # 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) + + group_count = l + shuffle_a = mixed_input_utils.is_shuffle_a( + a_major, k, a_dtype, b_dtype, scale_granularity_k + ) + # shuffle is supported since CUDA 13.1 + shuffle_supported = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major > 13 or ( + CUDA_VERSION.major == 13 and CUDA_VERSION.minor >= 1 + ): + shuffle_supported = True + except ImportError: + pass + + shuffle_a = shuffle_a and shuffle_supported + mixed_input_gemm = GroupedMixedInputGemmKernel( + scale_granularity_m, + scale_granularity_k, + acc_dtype, + use_2cta_instrs, + mma_tiler_mnk, + cluster_shape_mn, + group_count, + shuffle_a, + ) + torch.manual_seed(2025) + ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + uniform_group_sizes, + ) + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1], + ) + advanced_compiler_options = None + advanced_compiler_control_path = get_advanced_compiler_control_path() + if advanced_compiler_control_path: + advanced_compiler_options = ( + f"--ptxas-options '--apply-controls={advanced_compiler_control_path}'" + ) + + compiled_kernel = cute.compile( + mixed_input_gemm, + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + max_active_clusters, + current_stream, + options=advanced_compiler_options, + ) + + if not skip_ref_check: + compiled_kernel( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + current_stream, + ) + run_ref_and_compare( + a_torch_cpu, + b_torch_cpu, + a_scale_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + c_dtype, + tolerance, + ) + + # Early return if no performance measurement is needed + if iterations <= 0: + return + + def generate_tensors(): + ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + uniform_group_sizes, + ) + return testing.JitArguments( + a_tensor, a_scale_tensor, b_tensor, cumsum_tensor, c_tensor, current_stream + ) + + 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_gpu.numel() * c_torch_gpu.element_size() + + a_scale_torch_cpu.numel() * a_scale_torch_cpu.element_size() + if a_scale_torch_cpu is not None + else 0 + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_kernel, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +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() + parser.add_argument( + "--mnkl", type=parse_comma_separated_ints, default=(128, 128, 128, 1) + ) + parser.add_argument( + "--mma_tiler_mnk", type=parse_comma_separated_ints, default=(128, 128, 128) + ) + parser.add_argument( + "--cluster_shape_mn", type=parse_comma_separated_ints, default=(1, 1) + ) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument( + "--a_dtype", + type=cutlass.dtype, + default=cutlass.Int4, + choices=[cutlass.Int8, cutlass.Uint8, cutlass.Int4], + ) + parser.add_argument( + "--b_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + choices=[cutlass.BFloat16, cutlass.Float16], + ) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="m") + 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( + "--scale_granularity_m", + type=int, + default=1, + help="Scale granularity along M dimension.", + ) + parser.add_argument( + "--scale_granularity_k", + type=int, + default=128, + help="Scale granularity along K dimension.", + ) + 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( + "--uniform_group_sizes", action="store_true", help="Use uniform group sizes" + ) + args = parser.parse_args() + run( + args.mnkl, + args.scale_granularity_m, + args.scale_granularity_k, + args.a_dtype, + args.b_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mnk, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.uniform_group_sizes, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py new file mode 100644 index 00000000..472816b6 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py @@ -0,0 +1,2502 @@ +# 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 argparse +from math import log2, ceil +from typing import Union +import os +import sys + +import torch +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.utils.mixed_input_helpers import TransformMode +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import cpasync, tcgen05 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mixed_input_gemm.mixed_input_host_utils import ( + create_tensors_for_contiguous_grouped_mixed_input_gemm as create_tensors, + run_contiguous_grouped_ref_and_compare as run_ref_and_compare, +) + +""" +A mixed-input grouped GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL. + +Compared to the grouped_mixed_input_gemm.py example, this acc_scale example demonstrates a different implementation by +performing the scaling step on the accumulator instead of the input A tensor. The original computation is +``` +C = (type_convert(A) * scale) x B. +``` +In this example, we swap the scaling step and MMA and reformulate the computation as follows: +``` +C = scale * (type_convert(A) x B). +``` +The reformulation is valid when `scale_granularity_k` is an exact multiple of `mma_tiler_mnk[2]`. With this change, +the number of operations in the scaling step changes from `CTA_TILE_M * CTA_TILE_K` to `CTA_TILE_M * CTA_TILE_N`. +When `CTA_TILE_N` is smaller than `CTA_TILE_K` (common in decoding), the acc_scale implementation can +be more efficient due to fewer scaling operations. +Only convert-scale mode is supported in this example; convert-only mode has no scaling. +Other than the above changes, the computation flow is the same as the grouped_mixed_input_gemm.py example. + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py \ + --a_dtype Int4 --b_dtype BFloat16 \ + --scale_granularity_m 1 --scale_granularity_k 256 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 256,16,128 --cluster_shape_mn 2,1 \ + --use_2cta_instrs --mnkl 1024,8192,6144,16 \ + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py \ + --a_dtype Int4 --b_dtype BFloat16 \ + --scale_granularity_m 1 --scale_granularity_k 256 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,8,256 --cluster_shape_mn 1,1 \ + --mnkl 4096,8,8192,32 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check +""" + + +class GroupedMixedInputGemmAccScaleKernel: + """ + Mixed-input grouped GEMM kernel with scaling applied on accumulator for NVIDIA Blackwell SM100 architecture. + + This kernel supports GEMM operations where input tensors A and B have different data types, with tensor A + being converted to the precision of tensor B before matrix multiplication. The scaling is applied on + the accumulator instead of the input A tensor to reduce the number of operations in the scaling step. + Tensor A is in shape of [M, K, L] with L being the number of groups. Tensor B is in shape of [N, K] and a group search algorithm + is applied along the N mode to find the group index for each CTA tile. A cumsum tensor provides the offsets of each group along the N mode. + + :param scale_granularity_m: Number of elements sharing the same scale factor along the M mode + :type scale_granularity_m: int + :param scale_granularity_k: Number of elements sharing the same scale factor along the K mode + :type scale_granularity_k: int + :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_mnk: Shape of the Matrix Multiply-Accumulate (MMA) tile (M, N, K) + :type mma_tiler_mnk: tuple[int, int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: tuple[int, int] + :param group_count: The total number of groups + :type group_count: int + :param shuffle_a: Whether to use shuffle intrinsic for int4-to-bf16 conversion + :type shuffle_a: bool + """ + + def __init__( + self, + scale_granularity_m: int, + scale_granularity_k: int, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + group_count: int, + shuffle_a: bool, + ): + """ + Initializes the mixed-input GEMM kernel with a specified configuration. + """ + # Scale granularity defines how many elements share the same scale factor + # along the M and K modes. + self.scale_granularity_m = scale_granularity_m + self.scale_granularity_k = scale_granularity_k + # Set transform mode + if cutlass.const_expr( + self.scale_granularity_m == 0 and self.scale_granularity_k == 0 + ): + # Acc-update kernel variant is only for convert-scale mode. + raise ValueError("convert-only mode is not supported for acc-scale kernel") + self.scale_mode = TransformMode.ConvertScale + # scale_granularity_k must be exactly multiple of CTA tile shape K to allow acc-update recipe. + if cutlass.const_expr(self.scale_granularity_k % mma_tiler_mnk[2] != 0): + raise ValueError( + "scale_granularity_k must be exactly multiple of CTA tile shape K" + ) + + self.group_count = group_count + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = mma_tiler_mnk + self.shuffle_a = shuffle_a + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + # transformation ktile loop unrolling factor + self.transform_k_tile_unroll_factor = 2 + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.scale_tma_warp_id = 6 + # Schedule warp to do the group search + self.schedule_warp_id = 7 + self.transform_warp_id = ( + 8, + 9, + 10, + 11, + ) + # Define expected register count for different warps + # Reserve more registers for transformation and epilogue warps + self.num_regs_epilogue_warps = 168 + self.num_regs_mma_warp = 80 + self.num_regs_tma_warps = 72 + self.num_regs_transform_warps = 240 + self.num_regs_schedule_warp = 64 + self.threads_per_cta = 32 * ( + max( + ( + self.mma_warp_id, + self.tma_warp_id, + self.scale_tma_warp_id, + *self.epilog_warp_id, + *self.transform_warp_id, + ) + ) + + 1 + ) + + # Set barrier id for cta sync, epilogue sync, tmem ptr sync, and transform sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + 1, 32 * len(self.epilog_warp_id) + ) + self.tmem_ptr_sync_barrier = pipeline.NamedBarrier(2, self.threads_per_cta) + self.cta_sync_barrier = pipeline.NamedBarrier(3, self.threads_per_cta) + self.sched_sync_barrier = pipeline.NamedBarrier(4, 32) + + self.smem_buffer_align_bytes = 1024 + + def _setup_attributes(self): + """ + Set up configurations that are dependent on GEMM inputs + """ + # Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM) + self.transform_a_source = mixed_input_utils.get_transform_a_source( + self.a_major_mode + ) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_tile_shape_mnk = ( + self.cluster_shape_mn[0] * self.cta_tile_shape_mnk[0], + self.cluster_shape_mn[1] * self.cta_tile_shape_mnk[1], + self.cta_tile_shape_mnk[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + 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 + + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + # Compute tensor memory(TMEM) columns and stages for each pipeline + ( + self.num_load2trans_stage, + self.num_scale_load2accu_stage, + self.num_trans2mma_stage, + self.num_acc_stage, + self.num_c_stage, + self.num_tile_info_stage, + self.num_acc_tmem_cols, + self.num_a_tmem_cols, + ) = self._compute_stages_and_tmem_cols( + tiled_mma, + self.mma_tiler, + self.cta_tile_shape_mnk, + self.epi_tile, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.c_layout, + self.transform_a_source, + self.scale_granularity_m, + self.scale_granularity_k, + self.smem_buffer_align_bytes, + ) + + # Align TMEM columns for allocation + # TMEM allocation requires power-of-2 column alignment + # and must meet minimum allocation requirements + self.num_tmem_alloc_cols = cute.round_up( + self.num_acc_tmem_cols + self.num_a_tmem_cols, + cute.arch.get_min_tmem_alloc_cols("sm_100"), + ) + self.num_tmem_alloc_cols = 2 ** (ceil(log2(self.num_tmem_alloc_cols))) + # Get smem layout for C tensor + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + # Get smem layout for A, transformed A, and B + ( + self.smem_layout_a, + self.smem_layout_a_transform, + self.smem_layout_b, + ) = mixed_input_utils.compute_smem_layout( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.num_load2trans_stage, + self.num_trans2mma_stage, + ) + # Check if stages match the requirements for unrolling + if ( + self.num_scale_load2accu_stage < self.transform_k_tile_unroll_factor + or self.num_trans2mma_stage < self.transform_k_tile_unroll_factor + ): + raise ValueError("Not enough SMEM capacity for selected tile size") + # Get scale tile shape and smem layout for scale tensor + # ((M_SHARING_SCALE, NUM_SCALES_M),(K_SHARING_SCALE, NUM_SCALES_K), STAGES) + ( + self.scale_tile_shape, + self.smem_layout_scale_per_stage, + self.smem_layout_scale, + ) = mixed_input_utils.get_smem_layout_scale( + self.mma_tiler, + self.use_2cta_instrs, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + self.a_scale_dtype, + self.num_scale_load2accu_stage, + ) + + def _validate_inputs( + self, + a: cute.Tensor, + a_scale: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + ) -> None: + """ + Validates input tensors and their properties. + """ + # Validate scale tensor major mode + if cutlass.const_expr( + utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + != tcgen05.OperandMajorMode.MN + ): + raise ValueError("scale_major_mode must be M-major") + + @cute.jit + def __call__( + self, + a: cute.Tensor, + a_scale: cute.Tensor, + b: cute.Tensor, + cumsum: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """ + Executes the Mixed Input Grouped GEMM operation. + """ + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.a_scale_dtype: type[cutlass.Numeric] = a_scale.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.mma_dtype = self.b_dtype + + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.scale_major_mode = utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + # Get gmem layout for scale tensor + self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale( + a.shape, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + ) + + # Validate inputs + self._validate_inputs(a, a_scale, b, c) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + # Set up gmem copy atoms for A, scale, and B + a_op = mixed_input_utils.get_tma_atom_kind( + self.is_a_mcast, self.use_2cta_instrs, is_b=False + ) + b_op = mixed_input_utils.get_tma_atom_kind( + self.is_b_mcast, self.use_2cta_instrs, is_b=True + ) + a_scale_op = a_op + # Deduce TMA copy atom and TMA tensor for A, scale, and B + smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + smem_layout_a_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + # Partition smem layout for scale tensor to make it compatible with TMA atom + smem_layout_for_tma_atom = cute.get( + tiled_mma._thrfrg_A(self.smem_layout_scale_per_stage.outer), mode=[1] + ) + # ((MMA_M, MMA_K), REST_M, REST_K) + smem_layout_for_tma_atom = cute.dice( + smem_layout_for_tma_atom, + (1, (1,) * cute.rank(self.smem_layout_scale_per_stage.outer)), + ) + tma_atom_scale, tma_tensor_scale = cute.nvgpu.make_tiled_tma_atom_A( + a_scale_op, + cute.make_tensor(a_scale.iterator, self.gmem_layout_scale), + smem_layout_for_tma_atom, + # (SCALE_M, 1, SCALE_K) + (self.scale_tile_shape[0], 1, self.scale_tile_shape[1]), + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a_scale.element_type is cutlass.Float32 else None + ), + ) + + smem_layout_b_per_stage = cute.slice_(self.smem_layout_b, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + smem_layout_b_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + # Calculate copy size for tensor A, B, and scale + a_copy_size = cute.size_in_bytes(self.a_dtype, smem_layout_a_per_stage) + b_copy_size = cute.size_in_bytes(self.b_dtype, smem_layout_b_per_stage) + a_scale_copy_size = cute.size_in_bytes( + self.a_scale_dtype, self.smem_layout_scale_per_stage + ) + + self.num_tma_load_bytes_a = a_copy_size + self.num_tma_load_bytes_b = b_copy_size * cute.size(tiled_mma.thr_id.shape) + self.num_tma_load_bytes_scale = a_scale_copy_size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + 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, + ) + + # Shared memory structure + @cute.struct + class SharedStorage: + # buffer holding group search results + tile_info: cute.struct.MemRange[cutlass.Int32, 4 * self.num_tile_info_stage] + a_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_scale_load2accu_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2accu_stage + ] + a_scale_load2accu_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2accu_stage + ] + a_trans2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + a_trans2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + b_load2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + b_load2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_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] + tile_info_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_tile_info_stage + ] + tile_info_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_tile_info_stage + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + self.shared_storage = SharedStorage + + # Launch kernel + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_scale, + tma_tensor_scale, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + c, + cumsum, + self.group_count, + self.cluster_layout_vmnk, + self.smem_layout_a, + self.smem_layout_scale, + self.smem_layout_a_transform, + self.smem_layout_b, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + min_blocks_per_mp=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_s: cute.CopyAtom, + mS_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tensor_c: cute.Tensor, + cumsum: cute.Tensor, + group_count: cutlass.Constexpr[int], + cluster_layout_vmnk: cute.Layout, + a_smem_layout: cute.ComposedLayout, + scale_smem_layout: cute.ComposedLayout, + a_smem_layout_transform: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + c_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """ + GPU device kernel performing the Persistent Mixed-Input Grouped GEMM computation. + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + # Prefetch TMA descriptors + if warp_idx == self.epilog_warp_id[0]: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_s) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + bidx, bidy, bidz = cute.arch.block_idx() + # Compute how many k_tiles share the same scale + num_k_tiles_per_scale = self.scale_granularity_k // self.cta_tile_shape_mnk[2] + + 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 + ) + tidx, _, _ = cute.arch.thread_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize load2transform pipeline, which tracks the dependencies between TMA's loading + # of A and B, and the transformation of A and MMA's consumption + transform_thread_idx = ( + tidx - 32 * self.transform_warp_id[0] + if tidx >= 32 * self.transform_warp_id[0] + else tidx + ) + a_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mcast_ctas_a * len(self.transform_warp_id), + ), + tx_count=self.num_tma_load_bytes_a, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=(1, 0), # multicast for A will only happen on the M-mode + defer_sync=True, + ) + # Initialize scale_load2accu pipeline, which tracks the dependencies between TMA's loading + # of scale, and the accumulator update + num_producers_a_scale = self.num_mcast_ctas_a + scale_load2accu_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_scale_load2accu_full_mbar_ptr.data_ptr(), + num_stages=self.num_scale_load2accu_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_producers_a_scale * len(self.epilog_warp_id), + ), + tx_count=self.num_tma_load_bytes_scale, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + mcast_mode_mn=( + 1, + 0, + ), # multicast for scale_a will only happen on the M-mode + defer_sync=True, + ) + # Initialize transform2mma pipeline, which tracks the dependencies between the type conversion + # of A and MMA's consumption on converted A + cta_v_size = cute.size(cluster_layout_vmnk, mode=[0]) + trans2mma_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.a_trans2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_trans2mma_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.transform_warp_id) * cta_v_size, + ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize pipeline for tensor B load to MMA + # MMA warp informs TMA warp to proceed to load next tile of B tensor + b_load2mma_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_load2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ), + tx_count=self.num_tma_load_bytes_b, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), # multicast for B will only happen on the N-mode + defer_sync=True, + ) + # Initialize accumulator pipeline, which tracks the dependencies between + # MMA's computation of accumulators and epilogue warps' consumption of accumulators + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize tile info pipeline, which tracks the dependencies between + # tile scheduling warp and other warps + # Skip scheduler warp when computing consumer thread count + num_tile_info_pipeline_consumer_threads = self.threads_per_cta - 32 + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.tile_info_full_mbar_ptr.data_ptr(), + num_stages=self.num_tile_info_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_tile_info_pipeline_consumer_threads, + ), + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_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=self.cluster_shape_mn, is_relaxed=True) + + # Setup smem tensor A/scale/B/C + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=c_smem_layout_staged.inner, + ) + sA_input = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=a_smem_layout.inner, + ) + sS_input = smem.allocate_tensor( + element_type=self.mma_dtype, + layout=scale_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=scale_smem_layout.inner, + ) + sB_input = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=b_smem_layout.inner, + ) + sA_transform = None + # Get smem tensor for transformed A when transform_a_source is SMEM + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.SMEM): + sA_transform = smem.allocate_tensor( + element_type=self.mma_dtype, + layout=a_smem_layout_transform.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=a_smem_layout_transform.inner, + ) + sTile_info = storage.tile_info.get_tensor( + cute.make_layout((4, self.num_tile_info_stage), stride=(1, 4)) + ) + + # Compute multicast mask for A/B buffer full + a_full_mcast_mask = None + b_full_mcast_mask = None + s_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 + ) + # Scale tensor shares the same multicast mask as the A tensor + s_full_mcast_mask = a_full_mcast_mask + 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, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bM, bK, loopM, loopK, loopL) + gS_mkl = cute.local_tile( + mS_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + gC_mnl_simt = cute.local_tile( + tensor_c, 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, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgS = thr_mma.partition_A(gS_mkl) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + tCgC_simt = thr_mma.partition_C(gC_mnl_simt) + + # Setup copy atom to load A from shared memory for further transformation + copy_atom_a_input = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.a_dtype, num_bits_per_copy=32 + ) + a_smem_shape = tiled_mma.partition_shape_A( + cute.dice(self.mma_tiler, (1, None, 1)) + ) + # Setup copy atom to store transformed A into tensor memory or shared memory + copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform( + self.mma_dtype, + self.use_2cta_instrs, + self.transform_a_source, + a_smem_shape, + self.a_dtype, + ) + + # 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), loopM, loopK, loopL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA_input, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + + thr_mma_leader_cta = tiled_mma.get_slice(0) + # (MMA, MMA_M, MMA_K, STAGE) + tCsS = thr_mma_leader_cta.partition_A(sS_input) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS, tSgS = mixed_input_utils.scale_tma_partition( + tCsS, + tCgS, + tma_atom_s, + block_in_cluster_coord_vmnk, + a_cta_layout, + ) + + # 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), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB_input, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB_input) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc and ensure pipelines are ready + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # TMEM allocation + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + + # Schedule warp + if warp_idx == self.schedule_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_schedule_warp) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentRuntimeTileScheduler.create( + tile_sched_params, + (bidx, bidy, bidz), + cute.arch.grid_dim(), + inner_mode=0, + ) + work_tile = tile_sched.initial_work_tile_info() + tile_info_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_tile_info_stage + ) + # Create initial group search state + search_state = ( + mixed_input_utils.create_initial_contiguous_group_search_state() + ) + not_last_tile = cutlass.Boolean(1) + while not_last_tile: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + cluster_tile_coord_mnl = work_tile.tile_idx + cta_tile_coord_m = ( + cluster_tile_coord_mnl[0] * self.cluster_shape_mn[0] + + block_in_cluster_coord_vmnk[1] * cute.size(tiled_mma.thr_id.shape) + + block_in_cluster_coord_vmnk[0] + ) + cta_tile_offset_n = block_in_cluster_coord_vmnk[2] + search_state = mixed_input_utils.contiguous_group_search( + self.cluster_tile_shape_mnk, + group_count, + cluster_tile_coord_mnl[1], + search_state, + cumsum, + 1, # mode index to perform the search. 0 for M and 1 for N + ) + cur_sTile_info = sTile_info[(None, tile_info_producer_state.index)] + not_last_tile = search_state.cur_group_idx <= group_count + # Store tile info into shared memory buffer + with cute.arch.elect_one(): + cur_sTile_info[0] = cta_tile_coord_m + cur_sTile_info[1] = ( + search_state.cur_start + + cta_tile_offset_n * self.cta_tile_shape_mnk[1] + ) + cur_sTile_info[2] = search_state.cur_group_idx - 1 + cur_sTile_info[3] = ( + search_state.cur_boundary + - search_state.cur_start + - (cta_tile_offset_n * self.cta_tile_shape_mnk[1]) + ) + # Fence and barrier to ensure tile info store has finished + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.sched_sync_barrier.arrive_and_wait() + # Commit tile info pipeline + tile_info_pipeline.producer_commit(tile_info_producer_state) + # Advance to next tile + tile_info_producer_state.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # Specialized TMA load warp for A/B tensor + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + a_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + b_load2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + + while work_tile.is_valid_tile: + tAgA_slice = tAgA[ + ( + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + None, + work_tile.group_idx, + ) + ] + # Apply offset to B tensor based on group search result + coord_n_offset = ( + (work_tile.coord_n, 0, 0) + if cutlass.const_expr( + self.b_major_mode == tcgen05.OperandMajorMode.MN + ) + else (0, work_tile.coord_n, 0) + ) + tBgB_slice = cute.make_tensor( + ( + tBgB.iterator[0] + coord_n_offset[0], + coord_n_offset[1] + tBgB.iterator[1], + coord_n_offset[2] + tBgB.iterator[2], + ), + cute.slice_(tBgB.layout, (None, 0, None, 0)), + ) + + a_load2trans_producer_state.reset_count() + peek_load2trans_empty_status = cutlass.Boolean(1) + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + b_load2mma_producer_state.reset_count() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.producer_acquire( + a_load2trans_producer_state, peek_load2trans_empty_status + ) + b_load2mma_pipeline.producer_acquire(b_load2mma_producer_state) + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, a_load2trans_producer_state.count)], + tAsA[(None, a_load2trans_producer_state.index)], + tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier( + a_load2trans_producer_state + ), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, b_load2mma_producer_state.count)], + tBsB[(None, b_load2mma_producer_state.index)], + tma_bar_ptr=b_load2mma_pipeline.producer_get_barrier( + b_load2mma_producer_state + ), + mcast_mask=b_full_mcast_mask, + ) + a_load2trans_pipeline.producer_commit(a_load2trans_producer_state) + b_load2mma_pipeline.producer_commit(b_load2mma_producer_state) + a_load2trans_producer_state.advance() + b_load2mma_producer_state.advance() + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait A/B buffer empty + a_load2trans_pipeline.producer_tail(a_load2trans_producer_state) + b_load2mma_pipeline.producer_tail(b_load2mma_producer_state) + + # Specialized TMA load for scale tensor + if warp_idx == self.scale_tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + scale_load2accu_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scale_load2accu_stage + ) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + + while work_tile.is_valid_tile: + # ((atom_v, rest_v), RestK) + tSgS_slice = tSgS[ + ( + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + None, + work_tile.group_idx, + ) + ] + # Filter zeros in rest mode + rest_filtered = cute.filter_zeros(tSgS_slice[(0, None)].layout) + tSgS_slice_filtered = cute.make_tensor( + tSgS_slice.iterator, + cute.make_layout( + (tSgS_slice.layout[0].shape, rest_filtered.shape), + stride=(tSgS_slice.layout[0].stride, rest_filtered.stride), + ), + ) + + scale_load2accu_producer_state.reset_count() + peek_scale_load2accu_empty_status = cutlass.Boolean(1) + if scale_load2accu_producer_state.count < scale_k_tile_cnt: + peek_scale_load2accu_empty_status = ( + scale_load2accu_pipeline.producer_try_acquire( + scale_load2accu_producer_state + ) + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + scale_load2accu_pipeline.producer_acquire( + scale_load2accu_producer_state, + peek_scale_load2accu_empty_status, + ) + # TMA load scale + cute.copy( + tma_atom_s, + tSgS_slice_filtered[ + (None, scale_load2accu_producer_state.count) + ], + tSsS[(None, scale_load2accu_producer_state.index)], + tma_bar_ptr=scale_load2accu_pipeline.producer_get_barrier( + scale_load2accu_producer_state + ), + mcast_mask=s_full_mcast_mask, + ) + + scale_load2accu_producer_state.advance() + peek_scale_load2accu_empty_status = cutlass.Boolean(1) + if scale_load2accu_producer_state.count < scale_k_tile_cnt: + peek_scale_load2accu_empty_status = ( + scale_load2accu_pipeline.producer_try_acquire( + scale_load2accu_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait scale buffer empty + scale_load2accu_pipeline.producer_tail(scale_load2accu_producer_state) + + # Specialized transform warps + if warp_idx >= self.transform_warp_id[0]: + cute.arch.setmaxregister_increase(self.num_regs_transform_warps) + transform_local_tidx = tidx - 32 * self.transform_warp_id[0] + # Partition tensors for transform input and output and set up the copy atom + # used for loading and storing transformed A tensor + src_copy_a, dst_copy_a, tAsA_input, tAsA_transform = ( + mixed_input_utils.transform_partition( + self.transform_a_source, + self.scale_mode, + copy_atom_a_input, + copy_atom_a_transform, + sA_input, + ( + tCrA + if self.transform_a_source == tcgen05.OperandSource.TMEM + else sA_transform + ), + transform_local_tidx, + ) + ) + # make fragment for input A and transformed A + tArA_load = cute.make_rmem_tensor( + cute.append( + tAsA_input[(None, None, None, None, 0)].shape, + self.transform_k_tile_unroll_factor, + ), + tAsA_input.element_type, + ) + + tArA_transform = cute.make_rmem_tensor( + cute.append( + tAsA_input[(None, None, None, None, 0)].shape, + self.transform_k_tile_unroll_factor, + ), + self.mma_dtype, + ) + # Deduce a sub-tile size and tile tensors + transform_tiler_size = min( + cute.size(cute.coalesce(tAsA_input.layout), mode=[0]), 32 + ) + transform_tiler = cute.make_layout(transform_tiler_size) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + a_load2trans_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_load2trans_stage, + ) + a_load2trans_consumer_state0 = a_load2trans_consumer_state.clone() + trans2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_trans2mma_stage, + ) + trans2mma_producer_state0 = trans2mma_producer_state.clone() + k_tile_cnt_unrolled2 = k_tile_cnt // self.transform_k_tile_unroll_factor + is_tile_cnt_odd = k_tile_cnt % self.transform_k_tile_unroll_factor == 1 + while work_tile.is_valid_tile: + a_load2trans_consumer_state.reset_count() + a_load2trans_consumer_state0, a_load2trans_consumer_state = ( + self.pipeline_state_clone_and_advance(a_load2trans_consumer_state) + ) + trans2mma_producer_state.reset_count() + trans2mma_producer_state0, trans2mma_producer_state = ( + self.pipeline_state_clone_and_advance(trans2mma_producer_state) + ) + for k_tile in cutlass.range(0, k_tile_cnt_unrolled2, 1, unroll=1): + tAsA_input_slice0, tAsA_input_slice1 = ( + self.slice_and_divide_with_index_pair( + tAsA_input, + ( + a_load2trans_consumer_state0.index, + a_load2trans_consumer_state.index, + ), + transform_tiler, + ) + ) + # reg buffer0 and buffer1 for A_load + tArA_load_slice0, tArA_load_slice1 = ( + self.slice_and_divide_with_index_pair( + tArA_load, + (0, 1), + transform_tiler, + ) + ) + # reg buffer0 and buffer1 for A_transform + tArA_transform_buffer0 = tArA_transform[(None, None, None, None, 0)] + tArA_transform_buffer1 = tArA_transform[(None, None, None, None, 1)] + tArA_transform_slice0 = self.divide_tensor_by_tiler( + tArA_transform_buffer0, + transform_tiler, + ) + tArA_transform_slice1 = self.divide_tensor_by_tiler( + tArA_transform_buffer1, + transform_tiler, + ) + # Check if input A data are ready + a_load2trans_pipeline.consumer_wait(a_load2trans_consumer_state0) + a_load2trans_pipeline.consumer_wait(a_load2trans_consumer_state) + trans2mma_pipeline.producer_acquire(trans2mma_producer_state) + # Transformation in buffer0 + for idx in cutlass.range_constexpr( + cute.size(tArA_load_slice0, mode=[1]) + ): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice0[(None, idx)], + tArA_load_slice0[(None, idx)], + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load_slice0[(None, idx)], + self.mma_dtype, + self.shuffle_a, + ) + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice1[(None, idx)], + tArA_load_slice1[(None, idx)], + ) + tArA_transform_slice0[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform_buffer0, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state0.index) + ], + dst_copy_a, + ) + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation in buffer0 + trans2mma_pipeline.producer_commit(trans2mma_producer_state0) + a_load2trans_pipeline.consumer_release(a_load2trans_consumer_state0) + for idx in cutlass.range_constexpr( + cute.size(tArA_load_slice1, mode=[1]) + ): + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load_slice1[(None, idx)], + self.mma_dtype, + self.shuffle_a, + ) + tArA_transform_slice1[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform_buffer1, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + dst_copy_a, + ) + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state) + trans2mma_producer_state.advance() + trans2mma_producer_state0, trans2mma_producer_state = ( + self.pipeline_state_clone_and_advance(trans2mma_producer_state) + ) + + a_load2trans_pipeline.consumer_release(a_load2trans_consumer_state) + a_load2trans_consumer_state.advance() + a_load2trans_consumer_state0, a_load2trans_consumer_state = ( + self.pipeline_state_clone_and_advance( + a_load2trans_consumer_state + ) + ) + # Handle the last tile if needed + if is_tile_cnt_odd: + tAsA_input_slice = tAsA_input[ + (None, None, None, None, a_load2trans_consumer_state0.index) + ] + tAsA_input_slice = self.divide_tensor_by_tiler( + tAsA_input_slice, transform_tiler + ) + tArA_load_slice = tArA_load[(None, None, None, None, 0)] + tArA_load_slice = self.divide_tensor_by_tiler( + tArA_load_slice, transform_tiler + ) + tArA_transform_buffer = tArA_transform[(None, None, None, None, 0)] + tArA_transform_slice = self.divide_tensor_by_tiler( + tArA_transform_buffer, transform_tiler + ) + a_load2trans_pipeline.consumer_wait(a_load2trans_consumer_state0) + trans2mma_pipeline.producer_acquire(trans2mma_producer_state0) + for idx in cutlass.range_constexpr( + cute.size(tArA_load_slice, mode=[1]) + ): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice[(None, idx)], + tArA_load_slice[(None, idx)], + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load_slice[(None, idx)], + self.mma_dtype, + self.shuffle_a, + ) + tArA_transform_slice[(None, idx)].store(tensor_transformed) + a_load2trans_pipeline.consumer_release(a_load2trans_consumer_state0) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform_buffer, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state0.index) + ], + dst_copy_a, + ) + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state0) + trans2mma_producer_state0.advance() + a_load2trans_consumer_state0.advance() + # Keep pipeline state ready for next available buffer + trans2mma_producer_state = trans2mma_producer_state0.clone() + a_load2trans_consumer_state = a_load2trans_consumer_state0.clone() + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait a_transform buffer empty + trans2mma_pipeline.producer_tail(trans2mma_producer_state0) + + # Specialized MMA warp + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_mma_warp) + tCtAcc_base = accumulators + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + trans2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_trans2mma_stage + ) + b_load2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_load2trans_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + b_load2mma_consumer_state.reset_count() + trans2mma_consumer_state.reset_count() + peek_trans2mma_full_status = cutlass.Boolean(1) + if is_leader_cta: + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + k_block_cnt = cute.ceil_div(k_tile_cnt, num_k_tiles_per_scale) + num_k_tiles_executed = 0 + # Loop over K blocks with different scales and commit + # before starting the next tile requiring new scales + for k_block in cutlass.range(0, k_block_cnt, 1, unroll=1): + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[ + (None, None, None, acc_producer_state.index) + ] + acc_pipeline.producer_acquire(acc_producer_state) + cur_num_k_tiles = min( + num_k_tiles_per_scale, k_tile_cnt - num_k_tiles_executed + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, cur_num_k_tiles, 1, unroll=1): + trans2mma_pipeline.consumer_wait( + trans2mma_consumer_state, peek_trans2mma_full_status + ) + b_load2mma_pipeline.consumer_wait(b_load2mma_consumer_state) + num_kslices = cute.size(tCrA, mode=[2]) + for kslice_idx in cutlass.range( + num_kslices, unroll_full=True + ): + kblock_coord_a = ( + None, + None, + kslice_idx, + trans2mma_consumer_state.index, + ) + kblock_coord_b = ( + None, + None, + kslice_idx, + b_load2mma_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord_a], + tCrB[kblock_coord_b], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + trans2mma_pipeline.consumer_release( + trans2mma_consumer_state + ) + b_load2mma_pipeline.consumer_release( + b_load2mma_consumer_state + ) + trans2mma_consumer_state.advance() + b_load2mma_consumer_state.advance() + peek_trans2mma_full_status = cutlass.Boolean(1) + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + num_k_tiles_executed += cur_num_k_tiles + # Async arrive accumulator buffer full + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait for accumulator buffer empty + acc_pipeline.producer_tail(acc_producer_state) + + # Specialized acc update and epilogue warps + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + epi_tidx = tidx + tCtAcc_base = accumulators + # Construct scale tensor view as C + scale_view_as_C_layout = cute.make_layout( + ( + scale_smem_layout.outer[0].shape, + self.cta_tile_shape_mnk[1], + scale_smem_layout.outer[2].shape, + ), + stride=( + scale_smem_layout.outer[0].stride, + 0, + scale_smem_layout.outer[2].stride, + ), + ) + scale_view_as_C = cute.make_tensor( + sS_input.iterator, + scale_view_as_C_layout, + ) + # Partition for epilogue and accumulator update + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc, tTR_rAcc_final, tTR_sScale = ( + self.epilog_and_acc_update_tmem_copy_and_partition( + epi_tidx, + tCtAcc_base, + tCgC, + scale_view_as_C, + epi_tile, + use_2cta_instrs, + ) + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = ( + mixed_input_utils.epilog_smem_copy_and_partition( + self.c_layout, + self.c_dtype, + self.acc_dtype, + tiled_copy_t2r, + tTR_rC, + epi_tidx, + sC, + ) + ) + (tma_atom_c, bSG_sC, bSG_gC_partitioned, simt_atom, tTR_gC_partitioned) = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, + epi_tidx, + tma_atom_c, + tiled_copy_t2r, + tCgC, + tCgC_simt, + epi_tile, + sC, + ) + ) + + # Predicates + thr_mapping = cute.make_identity_tensor( + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1]) + ) + thr_mapping_mn = cute.flat_divide(thr_mapping, epi_tile) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + m_thr_offset = thr_copy_t2r.partition_D(thr_mapping_mn) + m_thr_offset = cute.group_modes(m_thr_offset, 3, cute.rank(m_thr_offset)) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + scale_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_scale_load2accu_stage, + ) + + 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, + ) + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + num_prev_subtiles = cutlass.Int32(0) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + while work_tile.is_valid_tile: + # perform accumulator update with scales + tTR_rAcc_final.fill(0.0) + tTR_rScale = cute.make_rmem_tensor( + cute.slice_(tTR_sScale, (None, None, None, 0, None, 0)).shape, + self.a_scale_dtype, + ) + scale_consumer_state.reset_count() + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < scale_k_tile_cnt: + peek_scale_full_status = scale_load2accu_pipeline.consumer_try_wait( + scale_consumer_state + ) + acc_consumer_state.reset_count() + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < scale_k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait( + acc_consumer_state + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + tTR_sScale_slice = cute.slice_( + tTR_sScale, + (None, None, None, 0, None, scale_consumer_state.index), + ) + scale_load2accu_pipeline.consumer_wait( + scale_consumer_state, peek_scale_full_status + ) + cute.autovec_copy(tTR_sScale_slice, tTR_rScale) + acc_pipeline.consumer_wait(acc_consumer_state, peek_acc_full_status) + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range_constexpr(subtile_cnt): + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + tTR_rAcc_subtile = tTR_rAcc_final[ + (None, None, None, subtile_idx) + ] + tTR_rScale_subtile = tTR_rScale[(None, None, None, subtile_idx)] + acc_vec = tTR_rAcc.load() + final_vec = tTR_rAcc_subtile.load() + scale = tTR_rScale_subtile.load().to(self.acc_dtype) + final_vec = acc_vec * scale + final_vec + tTR_rAcc_subtile.store(final_vec) + scale_load2accu_pipeline.consumer_release(scale_consumer_state) + scale_consumer_state.advance() + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < scale_k_tile_cnt: + peek_scale_full_status = ( + scale_load2accu_pipeline.consumer_try_wait( + scale_consumer_state + ) + ) + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < scale_k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait( + acc_consumer_state + ) + # epilogue partition + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + 0, + 0, + ) + ] + tma_store_offset_coord = ( + (work_tile.coord_n, 0, 0) + if cutlass.const_expr(self.c_layout.is_n_major_c()) + else (0, work_tile.coord_n, 0) + ) + bSG_gC = cute.make_tensor( + ( + tma_store_offset_coord[0] + bSG_gC.iterator[0], + tma_store_offset_coord[1] + bSG_gC.iterator[1], + tma_store_offset_coord[2] + bSG_gC.iterator[2], + ), + bSG_gC.layout, + ) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + 0, + 0, + ) + ] + tTR_gC = cute.make_tensor( + tTR_gC.iterator + (work_tile.coord_n * tensor_c.layout.stride[1]), + tTR_gC.layout, + ) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # Store accumulator to global memory in subtiles + subtile_cnt = cute.size(tTR_rAcc_final.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + tTR_rAcc_subtile = tTR_rAcc_final[(None, None, None, subtile_idx)] + if work_tile.distance_to_boundary >= self.cta_tile_shape_mnk[1]: + # Convert to C type + acc_vec = tiled_copy_r2s.retile(tTR_rAcc_subtile).load() + acc_vec = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + num_prev_subtiles += 1 + c_buffer = num_prev_subtiles % self.num_c_stage + # Store C to shared memory + 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( + "async.shared", + space="cta", + ) + self.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)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # Convert to C type + acc_vec = tTR_rAcc_subtile.load() + acc_vec = acc_vec.to(self.c_dtype) + tTR_rC.store(acc_vec) + # Compute predicate for SIMT store + tCpC = cute.make_rmem_tensor( + cute.make_layout(tTR_rC.shape), + cutlass.Boolean, + ) + m_thr_slice = m_thr_offset[(None, None, None, subtile_idx)] + for i in cutlass.range(cute.size(tCpC), unroll_full=True): + tCpC[i] = ( + m_thr_slice[(i)][0] + + work_tile.cta_coord_m * self.cta_tile_shape_mnk[0] + < tensor_c.shape[0] + ) and (m_thr_slice[(i)][1] < work_tile.distance_to_boundary) + # Store C to global memory + cute.copy( + simt_atom, + cute.flatten(tTR_rC), + cute.flatten(tTR_gC[(None, None, None, subtile_idx)]), + pred=cute.flatten(tCpC), + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + c_pipeline.producer_tail() + + def divide_tensor_by_tiler( + self, tensor: cute.Tensor, transform_tiler: cute.Layout + ) -> cute.Tensor: + """ + Divide the input tensor by given tiler and organize the resulting layout to 2 modes. + The first mode is the tile mode and the second mode is the rest mode. + """ + divided_tensor = cute.flat_divide(tensor, transform_tiler) + divided_tensor = cute.group_modes(divided_tensor, 1, cute.rank(divided_tensor)) + return divided_tensor + + def slice_and_divide_with_index_pair( + self, + tensor: cute.Tensor, + index_pair: tuple[cutlass.Int32, cutlass.Int32], + tiler: cute.Layout, + slice_mode=4, + ) -> tuple[cute.Tensor, cute.Tensor]: + """ + Perform the slice and divide_tensor_by_tiler operation on the sliced tensor with index pair. Coords used for slice are + """ + # pad None before the slice_mode + tensor0_slice, tensor1_slice = ( + tensor[(None,) * slice_mode + (index_pair[0],)], + tensor[(None,) * slice_mode + (index_pair[1],)], + ) + return self.divide_tensor_by_tiler( + tensor0_slice, tiler + ), self.divide_tensor_by_tiler(tensor1_slice, tiler) + + def pipeline_state_clone_and_advance( + self, pipeline_state: pipeline.PipelineState + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """ + Clones the pipeline state and advances it. + """ + pipeline_state_clone = pipeline_state.clone() + pipeline_state.advance() + return pipeline_state_clone, pipeline_state + + def epilog_and_acc_update_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + scale_tensor: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor]: + """ + Partitions source and destination tensors for a tensor memory load together with + the scale and accumulator tensors for the accumulator update. + """ + # 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, loopM, loopN, loopL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + sScale_epi = cute.flat_divide(scale_tensor, epi_tile) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + tTR_sScale = thr_copy_t2r.partition_D(sScale_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 + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_rAcc_final_ = cute.make_rmem_tensor( + tTR_gC[(None, None, None, None, None, 0, 0, 0)].shape, self.acc_dtype + ) + tTR_rAcc_final = cute.group_modes( + tTR_rAcc_final_, 3, cute.rank(tTR_rAcc_final_) + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_rAcc_final, tTR_sScale + + @staticmethod + def _compute_stages_and_tmem_cols( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + cta_tile_shape_mnk: tuple[int, int, int], + epi_tile: cute.Tile, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + transform_a_source: tcgen05.OperandSource, + scale_granularity_m: int, + scale_granularity_k: int, + smem_buffer_align_bytes: int, + ) -> tuple[int, int, int, int, int, int, int, int]: + """ + Compute pipeline stages and TMEM column allocation configurations. + """ + # Compute tmem columns required for accumulator + acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) + tCtAcc_stage1 = tiled_mma.make_fragment_C(cute.append(acc_shape, 1)) + num_tmem_acc_col_per_stage = cute.round_up( + tcgen05.find_tmem_tensor_col_offset(tCtAcc_stage1), 2 + ) + # Heuristic to decide the number of stages for accumulator + sm100_tmem_columns = cute.arch.get_max_tmem_alloc_cols("sm_100") + accumulator_stage_count = sm100_tmem_columns // num_tmem_acc_col_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM: + if num_tmem_acc_col_per_stage < 128: + accumulator_stage_count = 3 + elif num_tmem_acc_col_per_stage < 256: + accumulator_stage_count = 2 + else: + accumulator_stage_count = 1 + # transformed A in 16bit, thus 1 tmem column could hold 2 elements + num_elts_per_tmem_col = 32 // tiled_mma.op.a_dtype.width + num_tmem_cols_a_per_stage = cute.round_up( + ( + cta_tile_shape_mnk[2] // num_elts_per_tmem_col + if transform_a_source == tcgen05.OperandSource.TMEM + else 0 + ), + 4, + ) + + bytes_per_pipeline_stage = 16 + # By default, we use 2 stages for tile info + num_tile_info_stage = 2 + tile_info_bytes = ( + cute.size_in_bytes(cute.Int32, cute.make_layout((4, num_tile_info_stage))) + + bytes_per_pipeline_stage * num_tile_info_stage + ) + # Just keep 1 stage as tileN is small for decoding cases + c_stage_count = 1 + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * c_stage_count + + smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + # Ensure we have 4 buffers for scale tiles needed for 1 CTA tile + a_scale_k_mode = max(cta_tile_shape_mnk[2] // scale_granularity_k, 1) + a_scale_m_mode = max(cta_tile_shape_mnk[0] // scale_granularity_m, 1) + scale_load2accu_stage_count = 4 + a_scale_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((a_scale_m_mode, a_scale_k_mode)), + ), + smem_buffer_align_bytes, + ) + a_scale_bytes = ( + a_scale_bytes_per_stage + bytes_per_pipeline_stage + ) * scale_load2accu_stage_count + carveout_smem_bytes = ( + bytes_per_pipeline_stage * accumulator_stage_count + + a_scale_bytes + + c_bytes + + tile_info_bytes + ) + + # Compute transform stages if A is in TMEM + num_tmem_acc_cols = cute.round_up( + accumulator_stage_count * num_tmem_acc_col_per_stage, 4 + ) + + transform2mma_stage_count_a_source_tmem_potential = ( + (sm100_tmem_columns - num_tmem_acc_cols) // num_tmem_cols_a_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM + else -1 + ) + if ( + transform_a_source == tcgen05.OperandSource.TMEM + and transform2mma_stage_count_a_source_tmem_potential <= 0 + ): + raise ValueError("Not enough TMEM capacity for selected tile size") + a_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + b_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + b_dtype, + cute.make_layout( + ( + cta_tile_shape_mnk[1] // cute.size(tiled_mma.thr_id), + cta_tile_shape_mnk[2], + ) + ), + ), + smem_buffer_align_bytes, + ) + ab_load_bytes_per_stage = ( + a_load_bytes_per_stage + + b_load_bytes_per_stage + + 2 * bytes_per_pipeline_stage + ) + a_transform_bytes_per_stage = ( + cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + if transform_a_source == tcgen05.OperandSource.SMEM + else 0 + ) + + a_transform_bytes_per_stage = ( + a_transform_bytes_per_stage + bytes_per_pipeline_stage + ) + transform2mma_stage_count_a_source_smem_potential = ( + smem_capacity - carveout_smem_bytes + ) // (ab_load_bytes_per_stage + a_transform_bytes_per_stage) + transform2mma_stage_count = ( + min( + transform2mma_stage_count_a_source_tmem_potential, + transform2mma_stage_count_a_source_smem_potential, + ) + if transform_a_source == tcgen05.OperandSource.TMEM + else transform2mma_stage_count_a_source_smem_potential + ) + load2transform_stage_count = ( + smem_capacity + - carveout_smem_bytes + - (transform2mma_stage_count * a_transform_bytes_per_stage) + ) // ab_load_bytes_per_stage + if ( + load2transform_stage_count < 2 + or transform2mma_stage_count < 2 + or accumulator_stage_count < 1 + ): + raise ValueError("Not enough SMEM or TMEM capacity for selected tile size") + num_tmem_a_cols = transform2mma_stage_count * num_tmem_cols_a_per_stage + # Check if we can increase c_stage_count with leftover smem + c_stage_count += ( + smem_capacity + - load2transform_stage_count * ab_load_bytes_per_stage + - transform2mma_stage_count * a_transform_bytes_per_stage + - scale_load2accu_stage_count * a_scale_bytes_per_stage + - c_bytes + ) // c_bytes_per_stage + + return ( + load2transform_stage_count, + scale_load2accu_stage_count, + transform2mma_stage_count, + accumulator_stage_count, + c_stage_count, + num_tile_info_stage, + num_tmem_acc_cols, + num_tmem_a_cols, + ) + + @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. + """ + 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 = (cluster_shape_mn[0], cluster_shape_mn[1], max_active_clusters) + + return tile_sched_params, grid + + def can_implement( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + scale_granularity_m: int, + scale_granularity_k: int, + mma_tiler: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + ) -> bool: + """ + Check if the kernel can be implemented for the given tensor shapes and data types. + """ + m, n, k, l = mnkl + + if not mixed_input_utils.is_valid_mma_tiler_and_cluster_shape( + mma_tiler, cluster_shape_mn, use_2cta_instrs + ): + return False + if not mixed_input_utils.is_valid_scale_granularity( + scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2] + ): + return False + + # Check tensor alignment + def check_contiguous_NB_alignment( + dtype, contiguous_dim_size, expected_align_bytes + ): + expected_alignment = expected_align_bytes * 8 // dtype.width + return contiguous_dim_size % expected_alignment == 0 + + if not ( + check_contiguous_NB_alignment(a_dtype, m if a_major == "m" else k, 16) + and check_contiguous_NB_alignment(b_dtype, n if b_major == "n" else k, 16) + ): + return False + return True + + +def get_advanced_compiler_control_path(): + """ + Return the path to the advanced compiler control file of this example. If not found, return None. + """ + import os + + need_advanced_compiler_control = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major == 13 and CUDA_VERSION.minor == 1: + need_advanced_compiler_control = True + except ImportError: + pass + + if not need_advanced_compiler_control: + return None + # Get the path to the advanced compiler control file + current_dir = os.path.dirname(os.path.abspath(__file__)) + target_path = os.path.join(current_dir, "../../advanced_compiler_control/gemm0.bin") + if os.path.exists(target_path): + print(f"Found advanced compiler control file at {target_path}") + return target_path + else: + return None + + +def run( + mnkl: tuple[int, int, int, int], + scale_granularity_m: int, + scale_granularity_k: int, + a_dtype: type[cutlass.Numeric], + b_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_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + uniform_group_sizes: bool = False, + use_cold_l2: bool = False, + **kwargs, +) -> None: + """ + Run the mixed-input GEMM kernel with specified parameters. + + This function creates tensors, validates parameters, executes the kernel, + optionally compares results with a reference implementation and reports + kernel execution time. + """ + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise ValueError("CUDA is not available") + + # Check if given configuration is supported + if not GroupedMixedInputGemmAccScaleKernel.can_implement( + mnkl, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + scale_granularity_m, + scale_granularity_k, + mma_tiler_mnk, + cluster_shape_mn, + use_2cta_instrs, + ): + raise ValueError("GEMM configuration not supported") + + # 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) + + group_count = l + shuffle_a = mixed_input_utils.is_shuffle_a( + a_major, k, a_dtype, b_dtype, scale_granularity_k + ) + # shuffle is supported since CUDA 13.1 + shuffle_supported = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major > 13 or ( + CUDA_VERSION.major == 13 and CUDA_VERSION.minor >= 1 + ): + shuffle_supported = True + except ImportError: + pass + + shuffle_a = shuffle_a and shuffle_supported + mixed_input_gemm = GroupedMixedInputGemmAccScaleKernel( + scale_granularity_m, + scale_granularity_k, + acc_dtype, + use_2cta_instrs, + mma_tiler_mnk, + cluster_shape_mn, + group_count, + shuffle_a, + ) + torch.manual_seed(2025) + ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + uniform_group_sizes, + ) + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1], + ) + advanced_compiler_options = None + advanced_compiler_control_path = get_advanced_compiler_control_path() + if advanced_compiler_control_path: + advanced_compiler_options = ( + f"--ptxas-options '--apply-controls={advanced_compiler_control_path}'" + ) + + compiled_kernel = cute.compile( + mixed_input_gemm, + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + max_active_clusters, + current_stream, + options=advanced_compiler_options, + ) + + if not skip_ref_check: + compiled_kernel( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + current_stream, + ) + run_ref_and_compare( + a_torch_cpu, + b_torch_cpu, + a_scale_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + c_dtype, + tolerance, + ) + + # Early return if no performance measurement is needed + if iterations <= 0: + return + + def generate_tensors(): + ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + uniform_group_sizes, + ) + return testing.JitArguments( + a_tensor, a_scale_tensor, b_tensor, cumsum_tensor, c_tensor, current_stream + ) + + 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_gpu.numel() * c_torch_gpu.element_size() + + a_scale_torch_cpu.numel() * a_scale_torch_cpu.element_size() + if a_scale_torch_cpu is not None + else 0 + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_kernel, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +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() + parser.add_argument( + "--mnkl", type=parse_comma_separated_ints, default=(128, 128, 128, 1) + ) + parser.add_argument( + "--mma_tiler_mnk", type=parse_comma_separated_ints, default=(128, 128, 128) + ) + parser.add_argument( + "--cluster_shape_mn", type=parse_comma_separated_ints, default=(1, 1) + ) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument( + "--a_dtype", + type=cutlass.dtype, + default=cutlass.Int4, + choices=[cutlass.Int8, cutlass.Uint8, cutlass.Int4], + ) + parser.add_argument( + "--b_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + choices=[cutlass.BFloat16, cutlass.Float16], + ) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="m") + 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( + "--scale_granularity_m", + type=int, + default=1, + help="Scale granularity along M dimension.", + ) + parser.add_argument( + "--scale_granularity_k", + type=int, + default=128, + help="Scale granularity along K dimension.", + ) + 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( + "--uniform_group_sizes", action="store_true", help="Use uniform group sizes" + ) + args = parser.parse_args() + print(f"skip_ref_check={args.skip_ref_check}") + run( + args.mnkl, + args.scale_granularity_m, + args.scale_granularity_k, + args.a_dtype, + args.b_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mnk, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.uniform_group_sizes, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py new file mode 100644 index 00000000..2517e69a --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py @@ -0,0 +1,2321 @@ +# 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 argparse +from math import log2, ceil +from typing import Optional, Union +import os +import sys + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.utils.mixed_input_helpers import TransformMode +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import cpasync, tcgen05 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mixed_input_gemm.mixed_input_host_utils import ( + create_tensors_for_batched_mixed_input_gemm as create_tensors, + run_batched_mixed_input_ref_and_compare as run_ref_and_compare, +) + +""" +A mixed-input GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL. + +This example demonstrates an implementation of mixed-input GEMM using a TMA plus Blackwell SM100 TensorCore +warp-specialized persistent kernel. + +The inputs A and B have different data types. In this example, it's assumed that A is the narrow-precision tensor +and B holds data with a wider precision. +MMA will work in the wide precision of tensor B and tensor A will be transformed to the wide precision of tensor B +following 1 of the 2 possible modes as follows: + +1. convert-only mode: + C = type_convert(A) x B + +In convert-only mode, tensor A is directly converted to the wide precision of tensor B. + +2. convert-scale mode: + C = (type_convert(A) * scale) x B + +In convert-scale mode, tensor A is first converted to the wide precision of tensor B and then scaled by the scale tensor. +The scale tensor is in the same precision as tensor B. +The mode is determined by tensor A's data type as follows: +- if tensor A is in int8 or uint8, convert-only mode is used. +- if tensor A is in int4, convert-scale mode is used. + +The output tensor C could have the same precision as tensor B or fp32. + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mixed_input_gemm/mixed_input_gemm.py \ + --a_dtype Int8 --b_dtype BFloat16 \ + --scale_granularity_m 0 --scale_granularity_k 0 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \ + --mnkl 256,512,8192,1 + +Input A and B have int8 and bf16 data types, respectively. The Blackwell tcgen05 MMA tile shape +is specified as (128,128,64) and the cluster shape is (1,1). The MMA accumulator and output data type +are set as fp32 and bf16, respectively. As tensor A is int8, convert-only mode is used. +scale_granularity_m and scale_granularity_k are set as 0 for convert-only mode. + +Here is an example of running convert-scale mode: + +.. code-block:: bash + + python examples/blackwell/mixed_input_gemm/mixed_input_gemm.py \ + --a_dtype Int4 --b_dtype BFloat16 \ + --scale_granularity_m 1 --scale_granularity_k 256 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 256,128,128 --cluster_shape_mn 2,1 \ + --use_2cta_instrs --use_tma_store \ + --mnkl 1024,8192,6144,16 + +Input A and B have int4 and bf16 data types, respectively. The scale granularity is set as (1,256), +which means each element along the m mode of tensor A has its own scale element and 256 contiguous elements +along the k mode share the same scale element. There is no scale reuse along the L mode. If the GEMM shape is +(M, N, K, L), then the scale tensor shape is (M // scale_granularity_m, K // scale_granularity_k, L), +which is (1024, 6144/256, 16) in this example. +The Blackwell tcgen05 MMA tile shape is specified as (256,128,128) and tcgen05 2CTA feature is enabled. +The cluster shape is (2,1). The MMA accumulator and output data type are set as fp32 and bf16, respectively. +As tensor A is int4, the convert-scale mode is used. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mixed_input_gemm/mixed_input_gemm.py \ + --a_dtype Int8 --b_dtype BFloat16 \ + --scale_granularity_m 0 --scale_granularity_k 0 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \ + --mnkl 256,512,8192,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + +Besides the requirements from the Blackwell dense GEMM example, there are some constraints for this example: +* The narrow-precision is constrained to be int8, uint8, or int4 and the other data type is bf16 or f16. +* Output data types could only be fp16, bf16, or fp32. +* The scale_granularity_m must be 1 currently. +* The scale_granularity_k must be a multiple of mma_tiler_k and also be divisible by gemm_k. +* The scale tensor must be in M-major mode. +* OOB tiles are not allowed when TMA store is disabled +""" + + +class MixedInputGemmKernel: + """ + Mixed-input GEMM kernel for NVIDIA Blackwell SM100 architecture. + + This kernel supports GEMM operations where input tensors A and B have different + data types, with tensor A being transformed to the precision of tensor B before + matrix multiplication. + + :param scale_granularity_m: Number of elements sharing the same scale factor along the M mode + :type scale_granularity_m: int + :param scale_granularity_k: Number of elements sharing the same scale factor along the K mode + :type scale_granularity_k: int + :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_mnk: Shape of the Matrix Multiply-Accumulate (MMA) tile (M, N, K) + :type mma_tiler_mnk: tuple[int, 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 shuffle_a: Whether to use shuffle intrinsic for int4-to-bf16 conversion + :type shuffle_a: bool + """ + + def __init__( + self, + scale_granularity_m: int, + scale_granularity_k: int, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_tma_store: bool, + shuffle_a: bool, + ): + """ + Initializes the mixed-input GEMM kernel with a specified configuration. + """ + # Scale granularity defines how many elements share the same scale factor + # along the M and K modes. + self.scale_granularity_m = scale_granularity_m + self.scale_granularity_k = scale_granularity_k + # Set transform mode + if cutlass.const_expr( + self.scale_granularity_m == 0 and self.scale_granularity_k == 0 + ): + self.scale_mode = TransformMode.ConvertOnly + else: + self.scale_mode = TransformMode.ConvertScale + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = mma_tiler_mnk + self.use_tma_store = use_tma_store + self.shuffle_a = shuffle_a + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.scale_tma_warp_id = 6 + self.idle_warp_id = 7 + # 4 warps to do the transformation + self.transform_warp_id = ( + 8, + 9, + 10, + 11, + ) + self.num_regs_epilogue_warps = 192 + self.num_regs_mma_warp = 96 + self.num_regs_tma_warps = 96 + self.num_regs_transform_warps = 208 + self.num_regs_idle_warp = 24 + self.threads_per_cta = 32 * ( + max( + ( + self.mma_warp_id, + self.tma_warp_id, + self.scale_tma_warp_id, + *self.epilog_warp_id, + *self.transform_warp_id, + ) + ) + + 1 + ) + + # Set barrier id for epilogue sync, tmem ptr sync, and transform sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + 1, 32 * len(self.epilog_warp_id) + ) + self.tmem_ptr_sync_barrier = pipeline.NamedBarrier(2, self.threads_per_cta) + self.transform_sync_barrier = pipeline.NamedBarrier( + 3, 32 * len(self.transform_warp_id) + ) + + self.smem_buffer_align_bytes = 1024 + + 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: + - Deduce where the transformed A tensor is stored + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue sub-tile + - Setting up A/scale/B/C stage counts in shared memory + - Setting up transformed A stage count in shared memory or tensor memory + - Computing A/transformed A/scale/B/C memory layout + - Computing tensor memory allocation columns + """ + # Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM) + self.transform_a_source = mixed_input_utils.get_transform_a_source( + self.a_major_mode + ) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + 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 + + 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] + + # Compute tensor memory(TMEM) columns and stages for each pipeline + ( + self.num_load2trans_stage, + self.num_scale_load2trans_stage, + self.num_trans2mma_stage, + self.num_acc_stage, + self.num_c_stage, + self.num_acc_tmem_cols, + self.num_a_tmem_cols, + ) = self._compute_stages_and_tmem_cols( + tiled_mma, + self.mma_tiler, + self.cta_tile_shape_mnk, + self.epi_tile, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.c_layout, + self.transform_a_source, + self.scale_granularity_m, + self.scale_granularity_k, + self.smem_buffer_align_bytes, + self.use_tma_store, + self.scale_mode, + ) + + # Align TMEM columns for allocation + # TMEM allocation requires power-of-2 column alignment + # and must meet minimum allocation requirements + self.num_tmem_alloc_cols = cute.round_up( + self.num_acc_tmem_cols + self.num_a_tmem_cols, + cute.arch.get_min_tmem_alloc_cols("sm_100"), + ) + self.num_tmem_alloc_cols = 2 ** (ceil(log2(self.num_tmem_alloc_cols))) + # Get smem layout for C tensor when TMA store is enabled + 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 + ) + # Get smem layout for A, transformed A, and B + ( + self.smem_layout_a, + self.smem_layout_a_transform, + self.smem_layout_b, + ) = mixed_input_utils.compute_smem_layout( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.num_load2trans_stage, + self.num_trans2mma_stage, + ) + # Get smem layout for scale tensor + self.smem_layout_scale_per_stage = None + self.smem_layout_scale = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Get scale tile shape and smem layout for scale tensor + ( + self.scale_tile_shape, + self.smem_layout_scale_per_stage, + self.smem_layout_scale, + ) = mixed_input_utils.get_smem_layout_scale( + self.mma_tiler, + self.use_2cta_instrs, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + self.a_scale_dtype, + self.num_scale_load2trans_stage, + ) + + def _validate_inputs( + self, + a: cute.Tensor, + a_scale: Optional[cute.Tensor], + b: cute.Tensor, + c: cute.Tensor, + ) -> None: + """ + Validates input tensors and their properties. + + :param a: Input tensor A. + :type a: cute.Tensor + :param a_scale: Scale tensor for tensor A (None for ConvertOnly mode). + :type a_scale: Optional[cute.Tensor] + :param b: Input tensor B. + :type b: cute.Tensor + :param c: Output tensor C. + :type c: cute.Tensor + :raises ValueError: If inputs don't meet kernel requirements. + """ + # Validate scale tensor major mode + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + and utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + != tcgen05.OperandMajorMode.MN + ): + raise ValueError("scale_major_mode must be M-major") + + @cute.jit + def __call__( + self, + a: cute.Tensor, + a_scale: Optional[cute.Tensor], # None for ConvertOnly mode + b: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """ + Executes the Mixed Input GEMM operation. + + This method sets up the kernel parameters, computes the grid size, + defines the shared storage, and launches the kernel. + + The execution steps are as follows: + - 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 a_scale: Scale tensor for tensor A (None for ConvertOnly mode). + :type a_scale: Optional[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 to launch. + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream to launch the kernel on. + :type stream: cuda.CUstream + """ + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.a_scale_dtype: type[cutlass.Numeric] = ( + a_scale.element_type + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.mma_dtype = self.b_dtype + + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.scale_major_mode = ( + utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Get gmem layout for scale tensor + self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale( + a.shape, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + ) + + # Validate inputs + self._validate_inputs(a, a_scale, b, c) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + # Set up gmem copy atoms for A, scale, and B + a_op = mixed_input_utils.get_tma_atom_kind( + self.is_a_mcast, self.use_2cta_instrs, False + ) + b_op = mixed_input_utils.get_tma_atom_kind( + self.is_b_mcast, self.use_2cta_instrs, True + ) + a_scale_op = a_op + # Deduce TMA copy atom and TMA tensor for A, scale, and B + smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + smem_layout_a_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + tma_atom_scale, tma_tensor_scale = None, None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Partition smem layout for scale tensor to make it compatible with TMA atom + smem_layout_for_tma_atom = cute.get( + tiled_mma._thrfrg_A(self.smem_layout_scale_per_stage.outer), mode=[1] + ) + # ((MMA_M, MMA_K), REST_M, REST_K) + smem_layout_for_tma_atom = cute.dice( + smem_layout_for_tma_atom, + (1, (1,) * cute.rank(self.smem_layout_scale_per_stage.outer)), + ) + tma_atom_scale, tma_tensor_scale = cute.nvgpu.make_tiled_tma_atom_A( + a_scale_op, + cute.make_tensor(a_scale.iterator, self.gmem_layout_scale), + smem_layout_for_tma_atom, + # (SCALE_M, 1, SCALE_K) + (self.scale_tile_shape[0], 1, self.scale_tile_shape[1]), + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 + if a_scale.element_type is cutlass.Float32 + else None + ), + ) + + smem_layout_b_per_stage = cute.slice_(self.smem_layout_b, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + smem_layout_b_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + # Calculate copy size for tensor A, B, and scale + a_copy_size = cute.size_in_bytes(self.a_dtype, smem_layout_a_per_stage) + b_copy_size = cute.size_in_bytes(self.b_dtype, smem_layout_b_per_stage) + a_scale_copy_size = ( + cute.size_in_bytes(self.a_scale_dtype, self.smem_layout_scale_per_stage) + if self.scale_mode is TransformMode.ConvertScale + else 0 + ) + + self.num_tma_load_bytes_a = a_copy_size + self.num_tma_load_bytes_b = b_copy_size * cute.size(tiled_mma.thr_id.shape) + self.num_tma_load_bytes_scale = a_scale_copy_size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + tma_atom_c = None + tma_tensor_c = None + c_smem_size = 0 + 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, + ) + c_smem_size = cute.cosize(self.c_smem_layout_staged.outer) + + # Shared memory structure + a_smem_size = cute.cosize(self.smem_layout_a.outer) + b_smem_size = cute.cosize(self.smem_layout_b.outer) + a_transform_smem_size = ( + cute.cosize(self.smem_layout_a_transform.outer) + if self.transform_a_source == tcgen05.OperandSource.SMEM + else 0 + ) + a_scale_smem_size = ( + cute.cosize(self.smem_layout_scale.outer) + if self.scale_mode is TransformMode.ConvertScale + else 0 + ) + + @cute.struct + class SharedStorage: + a_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_scale_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_scale_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_trans2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + a_trans2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + b_load2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + b_load2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_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 + # Tensor buffers + # (EPI_TILE_M, EPI_TILE_N, STAGE) + smem_C: cute.struct.Align[ + cute.struct.MemRange[self.c_dtype, c_smem_size], + self.smem_buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + smem_A: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, a_smem_size], + self.smem_buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + smem_B: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, b_smem_size], + self.smem_buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + smem_A_transform: cute.struct.Align[ + cute.struct.MemRange[self.mma_dtype, a_transform_smem_size], + self.smem_buffer_align_bytes, + ] + # (MMA, MMA_M_SCALE, MMA_K_SCALE, STAGE) + smem_A_scale: cute.struct.Align[ + cute.struct.MemRange[self.mma_dtype, a_scale_smem_size], + self.smem_buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch kernel + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_scale, + tma_tensor_scale, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + self.cluster_layout_vmnk, + self.smem_layout_a, + self.smem_layout_scale, + self.smem_layout_a_transform, + self.smem_layout_b, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + min_blocks_per_mp=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_s: Optional[cute.CopyAtom], + mS_mkl: Optional[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: cute.ComposedLayout, + scale_smem_layout: cute.ComposedLayout, + a_smem_layout_transform: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + c_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """ + GPU device kernel performing the Persistent Mixed-Input GEMM computation. + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + # Prefetch TMA descriptors + if warp_idx == self.epilog_warp_id[0]: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + cpasync.prefetch_descriptor(tma_atom_s) + 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 + bidx, bidy, bidz = cute.arch.block_idx() + # Compute how many k_tiles share the same scale + num_k_tiles_per_scale = self.scale_granularity_k // self.cta_tile_shape_mnk[2] + + 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 + ) + tidx, _, _ = cute.arch.thread_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize load2transform pipeline, which tracks the dependencies between TMA's loading + # of A and B, and the transformation of A and MMA's consumption + transform_thread_idx = ( + tidx - 32 * self.transform_warp_id[0] + if tidx >= 32 * self.transform_warp_id[0] + else tidx + ) + a_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mcast_ctas_a * len(self.transform_warp_id), + ), + tx_count=self.num_tma_load_bytes_a, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=(1, 0), # multicast for A will only happen on the M-mode + defer_sync=True, + ) + # Initialize scale_load2trans pipeline, which tracks the dependencies between TMA's loading + # of scale, and the transformation of A + scale_load2trans_pipeline = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + num_producers_a_scale = self.num_mcast_ctas_a + scale_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_scale_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_scale_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_producers_a_scale + * len(self.transform_warp_id) + * num_k_tiles_per_scale, + ), + tx_count=self.num_tma_load_bytes_scale, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=( + 1, + 0, + ), # multicast for scale_a will only happen on the M-mode + defer_sync=True, + ) + # Initialize transform2mma pipeline, which tracks the dependencies between the transformation + # of A and MMA's consumption of transformed A + cta_v_size = cute.size(cluster_layout_vmnk, mode=[0]) + trans2mma_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.a_trans2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_trans2mma_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.transform_warp_id) * cta_v_size, + ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize pipeline for tensor B load to MMA + # MMA warp informs TMA warp to proceed to load next tile of B tensor + b_load2mma_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_load2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ), + tx_count=self.num_tma_load_bytes_b, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), # multicast for B will only happen on the N-mode + defer_sync=True, + ) + # Initialize accumulator pipeline, which tracks the dependencies between + # MMA's computation of accumulators and epilogue warps' consumption of accumulators + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_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=self.cluster_shape_mn, is_relaxed=True) + + # Setup smem tensor A/scale/B/C + sC = ( + storage.smem_C.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + if self.use_tma_store + else None + ) + sA_input = storage.smem_A.get_tensor( + a_smem_layout.outer, swizzle=a_smem_layout.inner + ) + sS_input = ( + storage.smem_A_scale.get_tensor( + scale_smem_layout.outer, swizzle=scale_smem_layout.inner + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + sB_input = storage.smem_B.get_tensor( + b_smem_layout.outer, swizzle=b_smem_layout.inner + ) + sA_transform = None + # Get smem tensor for transformed A when transform_a_source is SMEM + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.SMEM): + sA_transform = storage.smem_A_transform.get_tensor( + a_smem_layout_transform.outer, swizzle=a_smem_layout_transform.inner + ) + + # Compute multicast mask for A/B buffer full + a_full_mcast_mask = None + b_full_mcast_mask = None + s_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 + ) + # scale tensor share the same multicast mask with A tensor + s_full_mcast_mask = a_full_mcast_mask + 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, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bM, bK, loopM, loopK, loopL) + gS_mkl = ( + cute.local_tile( + mS_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + 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, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgS = ( + thr_mma.partition_A(gS_mkl) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + + # Setup copy atom to load A from shared memory for further transformation + copy_atom_a_input = ( + cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.a_dtype, num_bits_per_copy=32 + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + a_smem_shape = tiled_mma.partition_shape_A( + cute.dice(self.mma_tiler, (1, None, 1)) + ) + # Setup copy atom to store transformed A into tensor memory or shared memory + copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform( + self.mma_dtype, + self.use_2cta_instrs, + self.transform_a_source, + a_smem_shape, + self.a_dtype, + ) + + # 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), loopM, loopK, loopL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA_input, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + + tCsS = None + tSsS = None + tSgS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + thr_mma_leader_cta = tiled_mma.get_slice(0) + # (MMA, MMA_M, MMA_K, STAGE) + tCsS = thr_mma_leader_cta.partition_A(sS_input) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS, tSgS = mixed_input_utils.scale_tma_partition( + tCsS, + tCgS, + tma_atom_s, + block_in_cluster_coord_vmnk, + a_cta_layout, + ) + + # 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), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB_input, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB_input) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc and ensure pipelines are ready + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # TMEM allocation + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + + # Specialized TMA load warp for A/B tensor + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + a_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + b_load2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_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], + ) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + a_load2trans_producer_state.reset_count() + peek_load2trans_empty_status = cutlass.Boolean(1) + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + b_load2mma_producer_state.reset_count() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.producer_acquire( + a_load2trans_producer_state, peek_load2trans_empty_status + ) + b_load2mma_pipeline.producer_acquire(b_load2mma_producer_state) + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, a_load2trans_producer_state.count)], + tAsA[(None, a_load2trans_producer_state.index)], + tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier( + a_load2trans_producer_state + ), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, b_load2mma_producer_state.count)], + tBsB[(None, b_load2mma_producer_state.index)], + tma_bar_ptr=b_load2mma_pipeline.producer_get_barrier( + b_load2mma_producer_state + ), + mcast_mask=b_full_mcast_mask, + ) + a_load2trans_pipeline.producer_commit(a_load2trans_producer_state) + b_load2mma_pipeline.producer_commit(b_load2mma_producer_state) + a_load2trans_producer_state.advance() + b_load2mma_producer_state.advance() + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait A/B buffer empty + a_load2trans_pipeline.producer_tail(a_load2trans_producer_state) + b_load2mma_pipeline.producer_tail(b_load2mma_producer_state) + + # Specialized TMA load for scale tensor + if warp_idx == self.scale_tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + scale_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scale_load2trans_stage + ) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + + 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), + cur_tile_coord[1], + cur_tile_coord[2], + ) + # ((atom_v, rest_v), RestK) + tSgS_slice = tSgS[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # Filter zeros in rest mode + rest_filtered = cute.filter_zeros(tSgS_slice[(0, None)].layout) + tSgS_slice_filtered = cute.make_tensor( + tSgS_slice.iterator, + cute.make_layout( + (tSgS_slice.layout[0].shape, rest_filtered.shape), + stride=(tSgS_slice.layout[0].stride, rest_filtered.stride), + ), + ) + + scale_load2trans_producer_state.reset_count() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + scale_load2trans_pipeline.producer_acquire( + scale_load2trans_producer_state, + peek_scale_load2trans_empty_status, + ) + # TMA load scale + cute.copy( + tma_atom_s, + tSgS_slice_filtered[ + (None, scale_load2trans_producer_state.count) + ], + tSsS[(None, scale_load2trans_producer_state.index)], + tma_bar_ptr=scale_load2trans_pipeline.producer_get_barrier( + scale_load2trans_producer_state + ), + mcast_mask=s_full_mcast_mask, + ) + + scale_load2trans_producer_state.advance() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait scale buffer empty + scale_load2trans_pipeline.producer_tail(scale_load2trans_producer_state) + + # Specialized transform warps + if warp_idx >= self.transform_warp_id[0]: + cute.arch.setmaxregister_increase(self.num_regs_transform_warps) + transform_local_tidx = tidx - 32 * self.transform_warp_id[0] + # Partition tensors for transform input and output and set up the copy atom + # used for loading and storing transformed A tensor + ( + src_copy_a, + dst_copy_a, + tAsA_input, + tAsA_transform, + ) = mixed_input_utils.transform_partition( + self.transform_a_source, + self.scale_mode, + copy_atom_a_input, + copy_atom_a_transform, + sA_input, + ( + tCrA + if self.transform_a_source == tcgen05.OperandSource.TMEM + else sA_transform + ), + transform_local_tidx, + ) + # make rmem tensor for input A and transformed A + tArA = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, tAsA_input.element_type + ) + tArA_transform = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, self.mma_dtype + ) + # Partition scale tensor + smem_thr_copy_S = None + tSsS_trans = None + tSrS_copy = None + tSrS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = ( + mixed_input_utils.scale_partition( + src_copy_a, tCsS, transform_local_tidx, self.mma_dtype + ) + ) + assert cute.size(tSrS, mode=[0]) == cute.size(tArA, mode=[0]), ( + "tSrS and tArA have different leading dimension" + ) + assert cute.size(tSrS) == cute.size(tArA), ( + "tSrS and tArA have different shape" + ) + # Deduce a sub-tile size and tile tensors + transform_tiler_size = min( + cute.size(cute.coalesce(tAsA_input.layout), mode=[0]), 64 + ) + transform_tiler = cute.make_layout(transform_tiler_size) + tArA_load = cute.flat_divide(tArA, transform_tiler) + tArA_load = cute.group_modes(tArA_load, 1, cute.rank(tArA_load)) + tSrS_load = ( + cute.flat_divide(tSrS, transform_tiler) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tSrS_load = ( + cute.group_modes(tSrS_load, 1, cute.rank(tSrS_load)) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tArA_transform_store = cute.flat_divide(tArA_transform, transform_tiler) + tArA_transform_store = cute.group_modes( + tArA_transform_store, 1, cute.rank(tArA_transform_store) + ) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + a_load2trans_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_load2trans_stage, + ) + scale_load2trans_consumer_state = ( + pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_scale_load2trans_stage, + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + trans2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_trans2mma_stage, + ) + while work_tile.is_valid_tile: + a_load2trans_consumer_state.reset_count() + peek_load2trans_full_status = cutlass.Boolean(1) + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + peek_scale_load2trans_full_status = cutlass.Boolean(1) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + scale_load2trans_consumer_state.reset_count() + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + trans2mma_producer_state.reset_count() + peek_trans2mma_empty_status = cutlass.Boolean(1) + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.consumer_wait( + a_load2trans_consumer_state, peek_load2trans_full_status + ) + # Load A from shared memory + tAsA_input_slice = tAsA_input[ + (None, None, None, None, a_load2trans_consumer_state.index) + ] + tAsA_input_slice = cute.flat_divide( + tAsA_input_slice, transform_tiler + ) + tAsA_input_slice = cute.group_modes( + tAsA_input_slice, 1, cute.rank(tAsA_input_slice) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_wait( + scale_load2trans_consumer_state, + peek_scale_load2trans_full_status, + ) + trans2mma_pipeline.producer_acquire( + trans2mma_producer_state, peek_trans2mma_empty_status + ) + # load scale tensor when needed + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + if k_tile % num_k_tiles_per_scale == 0: + tSsS_slice = tSsS_trans[ + ( + None, + None, + None, + None, + scale_load2trans_consumer_state.index, + ) + ] + tSsS_slice_filtered = cute.make_tensor( + tSsS_slice.iterator, + cute.filter_zeros(tSsS_slice.layout), + ) + cute.autovec_copy(tSsS_slice_filtered, tSrS_copy) + cur_scale_load2trans_consumer_state = ( + scale_load2trans_consumer_state.clone() + ) + if (k_tile + 1) % num_k_tiles_per_scale == 0: + scale_load2trans_consumer_state.advance() + + cur_a_load2trans_consumer_state = ( + a_load2trans_consumer_state.clone() + ) + for idx in cutlass.range_constexpr(cute.size(tArA_load, mode=[1])): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice[(None, idx)], + tArA_load[(None, idx)], + ) + if cutlass.const_expr( + idx == cute.size(tArA_load, mode=[1]) - 1 + ): + a_load2trans_consumer_state.advance() + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load[(None, idx)], self.mma_dtype, self.shuffle_a + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale = cute.TensorSSA( + tSrS_load[(None, idx)].load(), + tensor_transformed.shape, + self.mma_dtype, + ) + # Apply scale + tensor_transformed = tensor_transformed * scale + tArA_transform_store[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + dst_copy_a, + ) + # Ensure all transform threads have finished the copy and reached the fence + self.transform_sync_barrier.arrive_and_wait() + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_release( + cur_scale_load2trans_consumer_state + ) + a_load2trans_pipeline.consumer_release( + cur_a_load2trans_consumer_state + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state) + trans2mma_producer_state.advance() + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait a_transform buffer empty + trans2mma_pipeline.producer_tail(trans2mma_producer_state) + + # Specialized MMA warp + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_mma_warp) + tCtAcc_base = accumulators + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + trans2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_trans2mma_stage + ) + b_load2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_load2trans_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + b_load2mma_consumer_state.reset_count() + trans2mma_consumer_state.reset_count() + peek_trans2mma_full_status = cutlass.Boolean(1) + if is_leader_cta: + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Mma mainloop + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + trans2mma_pipeline.consumer_wait( + trans2mma_consumer_state, peek_trans2mma_full_status + ) + b_load2mma_pipeline.consumer_wait(b_load2mma_consumer_state) + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord_a = ( + None, + None, + kblock_idx, + trans2mma_consumer_state.index, + ) + kblock_coord_b = ( + None, + None, + kblock_idx, + b_load2mma_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord_a], + tCrB[kblock_coord_b], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + trans2mma_pipeline.consumer_release(trans2mma_consumer_state) + b_load2mma_pipeline.consumer_release(b_load2mma_consumer_state) + trans2mma_consumer_state.advance() + b_load2mma_consumer_state.advance() + peek_trans2mma_full_status = cutlass.Boolean(1) + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + # Async arrive accumulator buffer full + 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: + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + epi_tidx = tidx + tCtAcc_base = accumulators + # Partition for epilogue + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = mixed_input_utils.epilog_tmem_copy_and_partition( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tidx, + tCtAcc_base, + tCgC, + epi_tile, + self.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_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = ( + mixed_input_utils.epilog_smem_copy_and_partition( + self.c_layout, + self.c_dtype, + self.acc_dtype, + 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, (bidx, bidy, bidz), 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): + 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: + 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], + ) + + bSG_gC = None + tTR_gC = None + if cutlass.const_expr(self.use_tma_store): + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + else: + tTR_gC = tTR_gC_partitioned[ + (None, None, None, None, None, *mma_tile_coord_mnl) + ] + + 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 = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + # Store C to shared memory + 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( + "async.shared", + space="cta", + ) + self.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)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # Convert to C type + acc_vec = tTR_rAcc.load() + acc_vec = acc_vec.to(self.c_dtype) + tTR_rC.store(acc_vec) + # Store C to global memory + cute.autovec_copy( + 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() + + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + if cutlass.const_expr(self.use_tma_store): + c_pipeline.producer_tail() + + # Idle warp + if warp_idx == self.idle_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_idle_warp) + + 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]: + """ + Partitions source and destination tensors for a TMA store or SIMT store. + """ + if self.use_tma_store: + tma_atom_c, bSG_sC, bSG_gC, _, _ = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, tidx, atom, None, gC_mnl, None, epi_tile, sC + ) + ) + return tma_atom_c, bSG_sC, bSG_gC + else: + _, _, _, simt_atom, tTR_gC = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, tidx, None, atom, None, gC_mnl, epi_tile, sC + ) + ) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_rmem_tensor( + 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_and_tmem_cols( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + cta_tile_shape_mnk: tuple[int, int, int], + epi_tile: cute.Tile, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + transform_a_source: tcgen05.OperandSource, + scale_granularity_m: int, + scale_granularity_k: int, + smem_buffer_align_bytes: int, + use_tma_store: bool, + scale_mode: TransformMode, + ) -> tuple[int, int, int, int, int, int, int]: + """ + Compute pipeline stages and TMEM column allocation configurations. + + This method calculates the number of pipeline stages for different operations + (load2trans, trans2mma, accumulator, etc.) and determines TMEM column allocation + based on available memory resources and tile configuration. + + :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 cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :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. + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param transform_a_source: The source of the transformed A tensor. + :type transform_a_source: tcgen05.OperandSource + :param scale_granularity_m: The granularity of the scale tensor along the M mode. + :type scale_granularity_m: int + :param scale_granularity_k: The granularity of the scale tensor along the K mode. + :type scale_granularity_k: int + :param smem_buffer_align_bytes: The alignment of the shared memory buffer. + :type smem_buffer_align_bytes: int + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + :param scale_mode: The transform mode. + :type scale_mode: TransformMode + + :return: A tuple containing the number of stages for: + (load2trans, scale_load2trans, transform2mma, accumulator, c, tmem_acc_cols, tmem_a_cols) + :rtype: tuple[int, int, int, int, int, int, int] + - num_load2trans_stage: Stages for load-to-transform A and B tensors pipeline + - num_scale_load2trans_stage: Stages for scale load-to-transform A tensor pipeline + - num_trans2mma_stage: Stages for transform-to-MMA pipeline + - num_acc_stage: Stages for accumulator-to-epilogue pipeline + - num_c_stage: Stages for epilogue-to-output C pipeline + - num_acc_tmem_cols: TMEM columns for accumulator + - num_a_tmem_cols: TMEM columns for transformed A tensor + """ + # Compute tmem columns required for accumulator + acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) + tCtAcc_stage1 = tiled_mma.make_fragment_C(cute.append(acc_shape, 1)) + num_tmem_acc_col_per_stage = utils.get_num_tmem_alloc_cols(tCtAcc_stage1, True) + # Heuristic to decide the number of stages for accumulator + sm100_tmem_columns = cute.arch.get_max_tmem_alloc_cols("sm_100") + accumulator_stage_count = sm100_tmem_columns // num_tmem_acc_col_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM: + if num_tmem_acc_col_per_stage < 128: + accumulator_stage_count = 3 + elif num_tmem_acc_col_per_stage < 256: + accumulator_stage_count = 2 + else: + accumulator_stage_count = 1 + # transformed A in 16bit, thus 1 tmem column could hold 2 elements + num_elts_per_tmem_col = 32 // tiled_mma.op.a_dtype.width + num_tmem_cols_a_per_stage = cute.round_up( + ( + cta_tile_shape_mnk[2] // num_elts_per_tmem_col + if transform_a_source == tcgen05.OperandSource.TMEM + else 0 + ), + 4, + ) + + c_stage_count = 2 if use_tma_store else 0 + c_smem_layout_staged_one = ( + sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + if use_tma_store + else None + ) + 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 * c_stage_count + + smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + bytes_per_pipeline_stage = 16 + if scale_mode == TransformMode.ConvertOnly: + scale_load2trans_stage_count = 0 + a_scale_bytes_per_stage = 0 + else: + # Ensure we have 2 buffers for scale tiles needed for 1 CTA tile + a_scale_k_mode = max(cta_tile_shape_mnk[2] // scale_granularity_k, 1) + a_scale_m_mode = max(cta_tile_shape_mnk[0] // scale_granularity_m, 1) + scale_load2trans_stage_count = 4 + a_scale_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((a_scale_m_mode, a_scale_k_mode)), + ), + smem_buffer_align_bytes, + ) + a_scale_bytes = ( + a_scale_bytes_per_stage + bytes_per_pipeline_stage + ) * scale_load2trans_stage_count + carveout_smem_bytes = ( + bytes_per_pipeline_stage * accumulator_stage_count + a_scale_bytes + c_bytes + ) + + # Compute transform stages if A is in TMEM + num_tmem_acc_cols = cute.round_up( + accumulator_stage_count * num_tmem_acc_col_per_stage, 4 + ) + + transform2mma_stage_count_a_source_tmem_potential = ( + (sm100_tmem_columns - num_tmem_acc_cols) // num_tmem_cols_a_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM + else -1 + ) + if ( + transform_a_source == tcgen05.OperandSource.TMEM + and transform2mma_stage_count_a_source_tmem_potential <= 0 + ): + raise ValueError("Not enough TMEM capacity for selected tile size") + a_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + b_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + b_dtype, + cute.make_layout( + ( + cta_tile_shape_mnk[1] // cute.size(tiled_mma.thr_id), + cta_tile_shape_mnk[2], + ) + ), + ), + smem_buffer_align_bytes, + ) + ab_load_bytes_per_stage = ( + a_load_bytes_per_stage + + b_load_bytes_per_stage + + 2 * bytes_per_pipeline_stage + ) + a_transform_bytes_per_stage = ( + cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + if transform_a_source == tcgen05.OperandSource.SMEM + else 0 + ) + + a_transform_bytes_per_stage = ( + a_transform_bytes_per_stage + bytes_per_pipeline_stage + ) + transform2mma_stage_count_a_source_smem_potential = ( + smem_capacity - carveout_smem_bytes + ) // (ab_load_bytes_per_stage + a_transform_bytes_per_stage) + transform2mma_stage_count = ( + min( + transform2mma_stage_count_a_source_tmem_potential, + transform2mma_stage_count_a_source_smem_potential, + ) + if transform_a_source == tcgen05.OperandSource.TMEM + else transform2mma_stage_count_a_source_smem_potential + ) + load2transform_stage_count = ( + smem_capacity + - carveout_smem_bytes + - (transform2mma_stage_count * a_transform_bytes_per_stage) + ) // ab_load_bytes_per_stage + if ( + load2transform_stage_count < 2 + or transform2mma_stage_count < 2 + or accumulator_stage_count < 1 + ): + raise ValueError("Not enough SMEM or TMEM capacity for selected tile size") + num_tmem_a_cols = transform2mma_stage_count * num_tmem_cols_a_per_stage + # Check if we can increase c_stage_count with leftover smem + if use_tma_store: + c_stage_count += ( + smem_capacity + - load2transform_stage_count * ab_load_bytes_per_stage + - transform2mma_stage_count * a_transform_bytes_per_stage + - scale_load2trans_stage_count * a_scale_bytes_per_stage + - c_bytes + ) // c_bytes_per_stage + + return ( + load2transform_stage_count, + scale_load2trans_stage_count, + transform2mma_stage_count, + accumulator_stage_count, + c_stage_count, + num_tmem_acc_cols, + num_tmem_a_cols, + ) + + @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. + """ + 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 + + def is_valid_epilog_store_option( + m: int, + n: int, + mma_tiler_mn: tuple[int, int], + use_tma_store: bool, + use_2cta_instrs: bool, + ) -> bool: + """ + Check if the epilogue store option is valid for the given problem size. + """ + cta_tile_shape_mn = ( + mma_tiler_mn[0] // (2 if use_2cta_instrs else 1), + mma_tiler_mn[1], + ) + # No OOB tile support when TMA store is disabled + if not use_tma_store: + if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): + return False + return True + + def can_implement( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + scale_granularity_m: int, + scale_granularity_k: int, + mma_tiler: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + use_tma_store: bool, + ) -> bool: + """ + Check if the kernel can be implemented for the given tensor shapes and data types. + """ + m, n, k, l = mnkl + + if not mixed_input_utils.is_valid_mma_tiler_and_cluster_shape( + mma_tiler, cluster_shape_mn, use_2cta_instrs + ): + return False + if not mixed_input_utils.is_valid_scale_granularity( + scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2] + ): + return False + if not mixed_input_utils.is_valid_tensor_alignment( + m, + n, + k, + a_dtype, + b_dtype, + c_dtype, + b_dtype, + a_major, + b_major, + c_major, + mma_tiler, + use_2cta_instrs, + cluster_shape_mn, + scale_granularity_m, + scale_granularity_k, + ): + return False + if not MixedInputGemmKernel.is_valid_epilog_store_option( + m, n, mma_tiler[:2], use_tma_store, use_2cta_instrs + ): + return False + return True + + +def run( + mnkl: tuple[int, int, int, int], + scale_granularity_m: int, + scale_granularity_k: int, + a_dtype: type[cutlass.Numeric], + b_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_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + use_tma_store: bool, + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +) -> None: + """ + Run the mixed-input GEMM kernel with specified parameters. + + This function creates tensors, validates parameters, executes the kernel, + optionally compares results with a reference implementation and reports + kernel execution time. + """ + m, n, k, l = mnkl + import torch + import cutlass.torch as cutlass_torch + + if not torch.cuda.is_available(): + raise ValueError("CUDA is not available") + + # Check if given configuration is supported + if not MixedInputGemmKernel.can_implement( + mnkl, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + scale_granularity_m, + scale_granularity_k, + mma_tiler_mnk, + cluster_shape_mn, + use_2cta_instrs, + use_tma_store, + ): + raise ValueError("GEMM configuration not supported") + + # 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) + shuffle_a = mixed_input_utils.is_shuffle_a( + a_major, k, a_dtype, b_dtype, scale_granularity_k + ) + # shuffle is supported since CUDA 13.1 + shuffle_supported = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major > 13 or ( + CUDA_VERSION.major == 13 and CUDA_VERSION.minor >= 1 + ): + shuffle_supported = True + except ImportError: + pass + + shuffle_a = shuffle_a and shuffle_supported + mixed_input_gemm = MixedInputGemmKernel( + scale_granularity_m, + scale_granularity_k, + acc_dtype, + use_2cta_instrs, + mma_tiler_mnk, + cluster_shape_mn, + use_tma_store, + shuffle_a, + ) + ( + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + ) + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1], + ) + # try to check CUDA version to decide the opt level + try: + from cutlass import CUDA_VERSION + + opt_level = 3 if (CUDA_VERSION.major == 13 and CUDA_VERSION.minor < 1) else 2 + except ImportError: + opt_level = 3 + compiled_kernel = cute.compile( + mixed_input_gemm, + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + max_active_clusters, + current_stream, + options=f"--opt-level {opt_level}", + ) + + if not skip_ref_check: + compiled_kernel( + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + current_stream, + ) + run_ref_and_compare( + a_torch_cpu, b_torch_cpu, a_scale_torch_cpu, c_torch_gpu, c_dtype, tolerance + ) + + # Early return if no performance measurement is needed + if iterations <= 0: + return + + def generate_tensors(): + ( + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + ) + return testing.JitArguments( + a_tensor, a_scale_tensor, b_tensor, c_tensor, current_stream + ) + + 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_gpu.numel() * c_torch_gpu.element_size() + + a_scale_torch_cpu.numel() * a_scale_torch_cpu.element_size() + if a_scale_torch_cpu is not None + else 0 + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_kernel, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +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() + parser.add_argument( + "--mnkl", type=parse_comma_separated_ints, default=(128, 128, 128, 1) + ) + parser.add_argument( + "--mma_tiler_mnk", type=parse_comma_separated_ints, default=(128, 128, 128) + ) + parser.add_argument( + "--cluster_shape_mn", type=parse_comma_separated_ints, default=(1, 1) + ) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument( + "--a_dtype", + type=cutlass.dtype, + default=cutlass.Int4, + choices=[cutlass.Int8, cutlass.Uint8, cutlass.Int4], + ) + parser.add_argument( + "--b_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + choices=[cutlass.BFloat16, cutlass.Float16], + ) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="m") + 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( + "--scale_granularity_m", + type=int, + default=1, + help="Scale granularity along M dimension.", + ) + parser.add_argument( + "--scale_granularity_k", + type=int, + default=128, + help="Scale granularity along K dimension.", + ) + 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" + ) + args = parser.parse_args() + + run( + args.mnkl, + args.scale_granularity_m, + args.scale_granularity_k, + args.a_dtype, + args.b_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mnk, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.use_tma_store, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_host_utils.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_host_utils.py new file mode 100644 index 00000000..107dc5a7 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_host_utils.py @@ -0,0 +1,506 @@ +# 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. + +from typing import Optional + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.cute.runtime import from_dlpack + +""" +This file contains common host-side utilities for mixed-input GEMM. +""" + + +def create_cumsum_tensor( + num_groups: int, + fused_n: int, + alignment: int, + uniform_distribution: bool = False, +) -> tuple[cute.Tensor, torch.Tensor]: + """ + Create a tensor of shape (num_groups + 1) recording the cumulative sum of the elements in each group. + """ + assert fused_n % alignment == 0, "fused_n must be divisible by alignment" + if uniform_distribution: + # keep a uniform distribution for debug and performance collection + group_counts = torch.tensor([fused_n // num_groups] * num_groups) + else: + # sample group sizes with equal probability for each group + probs = torch.ones(num_groups) / num_groups + group_sizes = torch.multinomial(probs, fused_n // alignment, replacement=True) + group_counts = torch.bincount(group_sizes, minlength=num_groups) * alignment + print(group_counts.tolist()) + + # Create cumulative sum + cumsum_torch = torch.cat([torch.tensor([0]), group_counts.cumsum(0)]) + print(cumsum_torch.tolist()) + + cumsum_tensor, _ = cutlass_torch.cute_tensor_like( + cumsum_torch, cutlass.Int32, is_dynamic_layout=False + ) + + return cumsum_tensor, cumsum_torch.to("cpu") + + +def create_i4_tensor_and_scale( + l: int, + m: int, + k: int, + is_m_major: bool, + dtype: type[cutlass.Numeric], + shuffle_a: bool, + scale_granularity_m: int, + scale_granularity_k: int, + is_dynamic_layout: bool = True, + init_config: tuple = ( + cutlass_torch.TensorInitType.RANDOM, + cutlass_torch.RandomInitConfig(min_val=-7, max_val=6), + ), + divisibility: int = 16, + transformed_dtype: Optional[type[cutlass.Numeric]] = None, +) -> tuple[ + cute.Tensor, + torch.Tensor, + torch.Tensor, + cute.Tensor, + torch.Tensor, + torch.Tensor, +]: + """ + Create quantized 4-bit tensor and corresponding scale tensor. + """ + lb_4b = -8 if dtype == cutlass.Int4 else 0 + up_4b = 7 if dtype == cutlass.Int4 else 15 + if not ( + init_config[0] == cutlass_torch.TensorInitType.RANDOM + or init_config[0] == cutlass_torch.TensorInitType.SCALAR + ): + raise ValueError( + "Only random and scalar initialization is supported for 4bit data type" + ) + + # Construct reference tensor in f32 + ref_fp32 = cutlass_torch.matrix(l, m, k, is_m_major, cutlass.Float32, *init_config) + # Generate scale data and perform quantization + num_scales = k // scale_granularity_k + ref = ref_fp32.to(dtype=cutlass_torch.dtype(transformed_dtype)).reshape( + m, num_scales, scale_granularity_k, l + ) + # Get elements with maximum absolute value to compute scaling factors + a_max = ( + torch.maximum(ref / up_4b, ref / lb_4b) + if dtype == cutlass.Int4 + else ref / up_4b + ) + a_scales, _ = torch.max(a_max, dim=2, keepdim=True) + a_scale_inv = torch.where(a_scales == 0, 0, 1 / a_scales) + a_quant = ref * a_scale_inv + # Convert values to integer to avoid computation errors + a_quant = a_quant.to(dtype=torch.int32).reshape((m, k, l)).to(dtype=torch.float32) + # Construct cute scale tensor + a_scales = a_scales.random_(-3, 3).reshape((m, num_scales, l)) + # Scale tensor is always m-major + a_scales = a_scales.permute(2, 1, 0).contiguous().permute(2, 1, 0).to(device="cuda") + if shuffle_a: + # shuffle within each group of 8 elements + perm = torch.tensor([0, 2, 1, 3, 4, 6, 5, 7], device=a_quant.device) + a_shuffled = ( + a_quant.view(m, k // 8, 8, l)[:, :, perm, :] + .reshape(a_quant.shape) + .permute(2, 0, 1) + .contiguous() + .permute(1, 2, 0) + ) + # Construct A quantized tensor + cute_a_quant_tensor, torch_a_quant_tensor = cutlass_torch.cute_tensor_like( + a_shuffled, + dtype, + is_dynamic_layout=is_dynamic_layout, + assumed_align=divisibility, + ) + else: + # Construct A quantized tensor + cute_a_quant_tensor, torch_a_quant_tensor = cutlass_torch.cute_tensor_like( + a_quant, + dtype, + is_dynamic_layout=is_dynamic_layout, + assumed_align=divisibility, + ) + cute_scale_tensor = from_dlpack(a_scales, assumed_align=divisibility) + for i, stride in enumerate(a_scales.stride()): + if stride == 1: + leading_dim = i + break + if is_dynamic_layout: + cute_scale_tensor = cute_scale_tensor.mark_layout_dynamic( + leading_dim=leading_dim + ) + + return ( + cute_a_quant_tensor, + torch_a_quant_tensor, + a_quant.to("cpu"), + cute_scale_tensor, + a_scales, + a_scales.to("cpu"), + ) + + +def create_tensor_a( + l: int, + m: int, + k: int, + a_major: str, + a_dtype: type[cutlass.Numeric], + shuffle_a: bool, + scale_granularity_m: int = 0, + scale_granularity_k: int = 0, + transformed_dtype: Optional[type[cutlass.Numeric]] = None, +) -> tuple[cute.Tensor, Optional[cute.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """ + Create tensor A and scale tensor. + """ + a_scale_tensor = None + a_scale_torch_cpu = None + if a_dtype in (cutlass.Int4,): + ( + a_tensor, + a_torch_gpu, + a_torch_cpu, + a_scale_tensor, + a_scale_torch_gpu, + a_scale_torch_cpu, + ) = create_i4_tensor_and_scale( + l, + m, + k, + a_major == "m", + a_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + divisibility=mixed_input_utils.get_divisibility(m if a_major == "m" else k), + transformed_dtype=transformed_dtype, + ) + else: + a_torch_cpu = cutlass_torch.matrix( + l, + m, + k, + a_major == "m", + a_dtype, + ) + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_torch_cpu, + a_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility( + m if a_major == "m" else k + ), + ) + return a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu + + +def create_tensors_for_contiguous_grouped_mixed_input_gemm( + l: int, + m: int, + n: int, + k: int, + a_major: str, + b_major: str, + c_major: str, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + shuffle_a: bool = False, + scale_granularity_m: int = 0, + scale_granularity_k: int = 0, + uniform_group_sizes: bool = False, +) -> tuple: + """ + Create all input and output tensors for the contiguous grouped mixed-input GEMM. + """ + a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a( + l, + m, + k, + a_major, + a_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + b_dtype, + ) + + # In GROUP mode, l specifies the number of groups. We'll fuse group into the n mode for tensor B and C. + # Batch mode will be set to 1. + num_groups = l + fused_n = n * num_groups + b_torch_cpu = cutlass_torch.matrix( + 1, # batch=1 + fused_n, + k, + b_major == "n", + b_dtype, + cutlass_torch.TensorInitType.RANDOM, + cutlass_torch.RandomInitConfig(min_val=-10, max_val=10), + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, + b_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility(n if b_major == "n" else k), + ) + + c_torch_cpu = cutlass_torch.matrix( + 1, # batch=1 + m, + fused_n, + c_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=mixed_input_utils.get_divisibility(m if c_major == "m" else n), + ) + c_tensor = c_tensor.mark_compact_shape_dynamic( + mode=(0 if c_major == "m" else 1), + stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1), + divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n), + ) + # We need to ensure mode N satisfies 16B alignment for each group + alignment_n = 16 * 8 // b_dtype.width + cumsum_tensor, cumsum_torch = create_cumsum_tensor( + num_groups, fused_n, alignment_n, uniform_distribution=uniform_group_sizes + ) + + return ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch, + c_torch_gpu, + ) + + +def create_tensors_for_batched_mixed_input_gemm( + l: int, + m: int, + n: int, + k: int, + a_major: str, + b_major: str, + c_major: str, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + shuffle_a: bool = False, + scale_granularity_m: int = 0, + scale_granularity_k: int = 0, +) -> tuple: + """ + Create all input and output tensors for the batched mixed-input GEMM. + """ + torch.manual_seed(2025) + + a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a( + l, + m, + k, + a_major, + a_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + b_dtype, + ) + + b_torch_cpu = cutlass_torch.matrix( + l, + n, + k, + b_major == "n", + b_dtype, + cutlass_torch.TensorInitType.RANDOM, + cutlass_torch.RandomInitConfig(min_val=-10, max_val=10), + ) + c_torch_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + c_dtype, + ) + + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, + b_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility(n if b_major == "n" else k), + ) + c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( + c_torch_cpu, + c_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility(m if c_major == "m" else n), + ) + c_tensor = c_tensor.mark_compact_shape_dynamic( + mode=(0 if c_major == "m" else 1), + stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1), + divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n), + ) + + return ( + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + c_torch_gpu, + ) + + +def run_contiguous_grouped_ref_and_compare( + a_torch_cpu: torch.Tensor, + b_torch_cpu: torch.Tensor, + a_scale_torch_cpu: Optional[torch.Tensor], + cumsum_torch_cpu: torch.Tensor, + c_torch_gpu: torch.Tensor, + c_dtype: type[cutlass.Numeric], + tolerance: float, +) -> None: + """ + Compare kernel result with reference computation. + """ + kernel_result = c_torch_gpu.cpu() + assert kernel_result.shape[2] == 1, "batch mode must be 1" + kernel_result = kernel_result.reshape( + kernel_result.shape[0], kernel_result.shape[1] + ) + # Compute reference result + a_for_gemm = a_torch_cpu + if a_scale_torch_cpu is not None: + scale_shape = a_scale_torch_cpu.shape + a_shape = a_torch_cpu.shape + a_scale_torch_cpu = a_scale_torch_cpu.to(dtype=torch.float32).reshape( + scale_shape[0], scale_shape[1], 1, scale_shape[2] + ) + a_torch_cpu = a_torch_cpu.to(dtype=torch.float32).reshape( + a_torch_cpu.shape[0], scale_shape[1], -1, a_torch_cpu.shape[2] + ) + a_for_gemm = (a_torch_cpu * a_scale_torch_cpu).reshape(a_shape) + # A in (m, k, l), b in (n, k), c in (m, n) + assert cumsum_torch_cpu.shape[0] == a_for_gemm.shape[-1] + 1, ( + "cumsum tensor must have one more element than a_for_gemm" + ) + assert b_torch_cpu.shape[2] == 1, ( + "b_torch_cpu must have a singleton dimension in the last position" + ) + prev_idx = 0 + ref = torch.zeros((a_for_gemm.shape[0], b_torch_cpu.shape[0]), dtype=torch.float32) + for group_idx in range(1, cumsum_torch_cpu.shape[0]): + # No computation for current group + if cumsum_torch_cpu[group_idx] == prev_idx: + continue + # Get A slice for current group + sliced_a = a_for_gemm[:, :, group_idx - 1] + # Get B slice for current group + sliced_b = b_torch_cpu[prev_idx : cumsum_torch_cpu[group_idx], :, 0] + sliced_ref = torch.einsum( + "mk,nk->mn", + sliced_a.to(dtype=torch.float32), + sliced_b.to(dtype=torch.float32), + ) + ref[:, prev_idx : cumsum_torch_cpu[group_idx]] = sliced_ref + prev_idx = cumsum_torch_cpu[group_idx] + # 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_batched_mixed_input_ref_and_compare( + a_torch_cpu: torch.Tensor, + b_torch_cpu: torch.Tensor, + a_scale_torch_cpu: Optional[torch.Tensor], + c_torch_gpu: torch.Tensor, + c_dtype: type[cutlass.Numeric], + tolerance: float, +) -> None: + """ + Compare kernel result with reference computation. + """ + kernel_result = c_torch_gpu.cpu() + # Compute reference result + if a_scale_torch_cpu is not None: + scale_shape = a_scale_torch_cpu.shape + a_shape = a_torch_cpu.shape + a_scale_torch_cpu = a_scale_torch_cpu.to(dtype=torch.float32).reshape( + scale_shape[0], scale_shape[1], 1, scale_shape[2] + ) + a_torch_cpu = a_torch_cpu.to(dtype=torch.float32).reshape( + a_torch_cpu.shape[0], scale_shape[1], -1, a_torch_cpu.shape[2] + ) + a_dequant = a_torch_cpu * a_scale_torch_cpu + ref = torch.einsum( + "mkl,nkl->mnl", + a_dequant.reshape(a_shape), + b_torch_cpu.to(dtype=torch.float32), + ) + else: + ref = torch.einsum( + "mkl,nkl->mnl", + a_torch_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) diff --git a/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py b/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py new file mode 100644 index 00000000..7eba992b --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py @@ -0,0 +1,4373 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import os +import sys +import argparse +import math +from typing import Type, Tuple, Optional +from types import SimpleNamespace + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode +import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.base_dsl.arch import Arch +from cutlass.cutlass_dsl import BaseDSL + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mla.mla_helpers import ( + ceil_div, + MAX_SPLITS, + LOG2_E, + MLAStaticTileScheduler, + MLAStaticTileSchedulerParams, + create_mla_static_tile_scheduler, + create_mla_static_tile_scheduler_params, +) + +""" +A Multi-Head Latent Attention (MLA) example with FP16 data type for the NVIDIA Blackwell SM100 architecture using CUTE DSL + +This example demonstrates an implementation of inference of multi-head latent attention using a TMA + Blackwell +SM100 TensorCore warp-specialized persistent kernel. The implementation integrates the (Qc + Qr)*(Kc + Kr)^T +matrix multiplication, softmax normalization, and softmax((Qc + Qr)*(Kc + Kr)^T)*Vc into a single kernel. +The kernel provides support for page table storage and variable-length KV cache sequences. It implements KV splitting +functionality to minimize latency when processing long KV sequences. + +The kernel implements key optimizations including: +- Warp specialization for different computation phases (load, MMA, softmax, correction, epilogue) +- Pipeline stages between different warps for overlapping computation and memory access +- Support for different precision data types +- Two sub-kernels (split KV kernel and reduction kernel) that enable split KV processing + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mla_fp16.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float16 --out_dtype Float16 \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent + +The above example runs Multi-Head Latent Attention (MLA) with the following configuration: +- Batch size: 4 +- Sequence length of Q: 1 +- Sequence length of K: 1024 +- Latent dimension: 512 +- RoPE dimension: 64 +- Number of heads: 128 +- Data types: Float16 (input), Float16 (output), Float32 (accumulation and LSE) + +It utilizes page table storage for the KV cache and enables both variable-length KV cache sequences +and variable split KV processing with persistent scheduling. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mla_fp16.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float16 --out_dtype Float16 \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent --warmup_iterations 3 \ + --iterations 10 --skip_ref_check + +Constraints for this example: +* Data type requirements: + - Input/output: Float16 + - Accumulation and LSE: Float32 +* Fixed architecture parameters: + - Number of attention heads: 128 + - Latent dimension: 512 + - RoPE dimension: 64 +* Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) +* Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) +* Query sequence length must be 1-4 +* Only supports 2-CTA instructions +* Variable sequence length requires page table storage enabled +""" + + +class BlackwellMultiHeadLatentAttentionForwardFP16: + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + max_active_clusters: int, + page_size: int, + skip_correction_threshold: float, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + ): + """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. + + :param acc_dtype: Data type for accumulation S and O + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Data type for output LSE + :type lse_dtype: Type[cutlass.Numeric] + :param mma_s_tiler: The (H, K) tile shape of the MMA instruction for S + :type mma_s_tiler: Tuple[int, int] + :param mma_p_tiler: The (H, D) tile shape of the MMA instruction for P + :type mma_p_tiler: Tuple[int, int] + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: int + :param page_size: The page size of the page table + :type page_size: int + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split KV + :type is_var_split_kv: bool + """ + + self.latent_dim = 512 + self.rope_dim = 64 + self.acc_dtype = acc_dtype + self.lse_dtype = lse_dtype + self.mma_qk_tiler_mn = mma_qk_tiler_mn + self.mma_pv_tiler_mn = mma_pv_tiler_mn + self.max_active_clusters = max_active_clusters + self.skip_correction_threshold = skip_correction_threshold + self.is_persistent = is_persistent + self.page_size = page_size + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv + self.cluster_shape_mnk = (2, 1, 1) + self.use_2cta_instrs = True + # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), + # while warps 2-3 handle accumulation for second half [n/2, n) + self.warps_in_n = 2 + self.num_compute_warps = 4 + self.threads_per_warp = 32 + mma_qk_tiler_k = self.rope_dim + self.mma_qk_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + mma_qk_tiler_k, + ) + self.mma_qk_rope_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + self.rope_dim, + ) + self.mma_pv_tiler = ( + self.mma_pv_tiler_mn[0], + self.mma_pv_tiler_mn[1], + self.mma_qk_tiler[1] * self.mma_qk_tiler[2] // self.mma_pv_tiler_mn[1], + ) + self.iterations_qk_latent = self.latent_dim // self.mma_qk_tiler[2] + self.iterations_qk_rope = mma_qk_tiler_k // self.mma_qk_tiler[2] + self.iterations_qk = self.iterations_qk_latent + self.iterations_qk_rope + self.iterations_pv_k = self.mma_qk_tiler[1] // self.mma_pv_tiler[2] + self.iterations_pv_n = self.latent_dim // self.mma_pv_tiler[1] + + # Set specialized warp ids + self.compute_warp_ids = (0, 1, 2, 3) + self.correction_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + + self.load_tma_warp_id = 9 + self.load_pt_warp_id = 10 + self.empty_warp_ids = (11,) + self.threads_per_cta = self.threads_per_warp * len( + ( + self.mma_warp_id, + self.load_tma_warp_id, + self.load_pt_warp_id, + *self.compute_warp_ids, + *self.correction_warp_ids, + *self.empty_warp_ids, + ) + ) + + # register settings + self.softmax_reg_num = 192 + self.correction_reg_num = 208 + self.other_reg_num = 96 + # Named barriers + self.tmem_ptr_sync_bar = pipeline.NamedBarrier( + barrier_id=1, + num_threads=( + self.threads_per_warp + + self.threads_per_warp * self.num_compute_warps * 2 + ), + ) + self.softmax_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=2, num_threads=(self.threads_per_warp * self.num_compute_warps) + ) + self.epilogue_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=3, num_threads=(self.threads_per_warp * self.num_compute_warps) + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the MLA kernel operation. + + This method initializes and configures various attributes required for the + execution of the multi-head latent attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.load_q_stage = 1 + self.load_kv_stage = 15 + self.mma_s_stage = 2 + self.p_mma_stage = 2 + self.p_cor_stage = 2 + self.mma_o_stage = 1 + self.load_pt_stage = 4 + + self.tmem_o_offset = self.mma_s_stage * self.mma_qk_tiler[1] // self.warps_in_n + self.correction_factor_offset = ( + self.tmem_o_offset + self.latent_dim // self.warps_in_n + ) + + @cute.jit + def __call__( + self, + q_latent: cute.Tensor, + q_rope: cute.Tensor, + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_table: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: Optional[cute.Tensor], + block_split_kvs: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + stream: cuda.CUstream, + ): + """Execute the Multi-Head Latent Attention operation on the provided tensors. + + The method handles: + 1. Initialization of workspace for temporary split KV buffers + 2. Validation of tensor data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch(split KV kernel and reduction kernel) with appropriate parameters + + :param q_latent: The query tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type q_latent: cute.Tensor + :param q_rope: The query RoPE tensor with shape [num_head, rope_dim, seq_len_q, batch_size] + :type q_rope: cute.Tensor + :param c_latent: The key tensor with shape [seq_len_k, latent_dim, batch_size] + :type c_latent: cute.Tensor + :param c_rope: The key RoPE tensor with shape [seq_len_k, rope_dim, batch_size] + :type c_rope: cute.Tensor + :param page_table: The page table tensor with shape [page_count, batch_size] + :type page_table: cute.Tensor + :param o: The output tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type o: cute.Tensor + :param lse: The LSE tensor with shape [num_head, seq_len_q, batch_size] + :type lse: cute.Tensor + :param workspace: The workspace tensor with 1-d shape prepared for acc_o and acc_lse + :type workspace: cute.Tensor + :param split_kv: The scalar factor for split KV + :type split_kv: cutlass.Int32 + :param cache_seqs: The cache sequences tensor with shape [batch_size] + :type cache_seqs: cute.Tensor + :param block_split_kvs: The block split KV tensor with shape [batch_size] + :type block_split_kvs: cute.Tensor + :param softmax_scale: The scale factor for softmax + :type softmax_scale: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param stream: The CUDA stream to execute the kernel on + :type stream: cuda.CUstream + + :raises TypeError: If tensor data types don't match or aren't supported + """ + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q_latent.element_type + self.k_dtype = c_latent.element_type + self.v_dtype = c_latent.element_type + self.o_dtype = o.element_type + + # check type consistency + if cutlass.const_expr( + self.q_dtype != self.k_dtype or self.q_dtype != self.v_dtype + ): + raise TypeError( + f"Type mismatch: {self.q_dtype} != {self.k_dtype} or {self.q_dtype} != {self.v_dtype}" + ) + # check leading dimensions of input/output + if cutlass.const_expr(q_latent.stride[1] != 1 or q_rope.stride[1] != 1): + raise ValueError("q_latent and q_rope must have leading dimension 1") + if cutlass.const_expr(c_latent.stride[1] != 1 or c_rope.stride[1] != 1): + raise ValueError("c_latent and c_rope must have leading dimension 1") + if cutlass.const_expr(o.stride[1] != 1): + raise ValueError("o must have leading dimension 1") + if cutlass.const_expr(lse.stride[0] != 1): + raise ValueError("lse must have leading dimension 0") + + acc_o, acc_lse = self.initialize_workspace( + q_latent.shape[0], + q_latent.shape[1], + q_latent.shape[2], + q_latent.shape[3], + split_kv, + self.acc_dtype, + workspace, + ) + + c_latent_tranpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor( + c_latent.iterator, c_latent_tranpose_layout + ) + + self.q_major_mode = tcgen05.OperandMajorMode.K + self.k_major_mode = tcgen05.OperandMajorMode.K + self.v_major_mode = tcgen05.OperandMajorMode.MN + + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + # the intermediate tensor p is from smem & k-major + p_major_mode = tcgen05.OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.acc_dtype, + cta_group, + self.mma_qk_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.acc_dtype, + cta_group, + self.mma_pv_tiler[:2], + ) + + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + + self.epi_tile = self.mma_pv_tiler[:2] + + q_latent_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_tiler, + self.q_dtype, + (self.iterations_qk_latent * self.load_q_stage), + ) + q_latent_smem_layout_staged = cute.logical_divide( + q_latent_smem_layout_staged, (None, None, None, self.iterations_qk_latent) + ) + q_rope_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.q_dtype, + self.load_q_stage, + ) + + # rope reuse the same smem layout as latent + kc_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_tiler, + self.k_dtype, + self.load_kv_stage, + ) + kc_page_tile_size = min( + self.page_size, qk_tiled_mma.op.shape_mnk[0] // qk_tiled_mma.thr_id.shape + ) + + kc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + (self.mma_qk_tiler[0] // qk_tiled_mma.thr_id.shape, self.mma_qk_tiler[2]), + self.k_dtype, + self.load_kv_stage, + ) + kc_smem_layout_for_tma = cute.tiled_divide( + kc_smem_layout_for_tma, (kc_page_tile_size, self.mma_qk_tiler[2]) + ) + + p_smem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.mma_pv_tiler, + self.q_dtype, + (self.iterations_pv_k * self.p_mma_stage), + ) + p_smem_layout_staged = cute.logical_divide( + p_smem_layout_staged, (None, None, None, self.iterations_pv_k) + ) + + vc_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.mma_pv_tiler, + self.v_dtype, + self.load_kv_stage, + ) + vc_page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + vc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.MN, + (self.mma_pv_tiler[1] // pv_tiled_mma.thr_id.shape, self.mma_pv_tiler[2]), + self.v_dtype, + self.load_kv_stage, + ) + vc_smem_layout_for_tma = cute.tiled_divide( + vc_smem_layout_for_tma, + ( + pv_tiled_mma.op.shape_mnk[1] // pv_tiled_mma.thr_id.shape, + vc_page_tile_size, + ), + ) + # TMA load for Q latent and rope + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + + q_latent_smem_layout = cute.select(q_latent_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q_latent, tma_tensor_q_latent = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_latent, + q_latent_smem_layout, + self.mma_qk_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + q_rope_smem_layout = cute.select(q_rope_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q_rope, tma_tensor_q_rope = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_rope, + q_rope_smem_layout, + self.mma_qk_rope_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + # TMA load for c latent and k rope + kc_smem_layout = cute.select(kc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent, tma_tensor_c_latent = self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + tma_atom_c_rope, tma_tensor_c_rope = self.make_paged_tiled_tma_atom( + tma_load_op, + c_rope, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + # TMA load for c latent transpose + vc_smem_layout = cute.select(vc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent_transpose, tma_tensor_c_latent_transpose = ( + self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent_transpose, + vc_smem_layout, + (self.mma_pv_tiler[1], self.mma_pv_tiler[2]), + pv_tiled_mma, + is_k_load=False, + ) + ) + + q_latent_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_latent_smem_layout) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_latent + ) + q_rope_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_rope_smem_layout) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_rope + ) + q_copy_size = q_latent_copy_size + q_rope_copy_size + kc_copy_size = cute.size_in_bytes( + self.k_dtype, cute.select(kc_smem_layout_staged, mode=[0, 1, 2]) + ) * cute.size(qk_tiled_mma.thr_id.shape) + vc_copy_size = cute.size_in_bytes( + self.v_dtype, cute.select(vc_smem_layout_staged, mode=[0, 1, 2]) + ) * cute.size(pv_tiled_mma.thr_id.shape) + assert ( + kc_copy_size == vc_copy_size + ), "kc_copy_size and vc_copy_size must be the same" + + self.tma_copy_q_bytes = q_copy_size + self.tma_copy_kc_bytes = kc_copy_size + + tile_sched_params, grid = self._compute_grid( + o, + split_kv, + self.cluster_shape_mnk, + self.max_active_clusters, + self.is_persistent, + ) + + @cute.struct + class SplitKVKernelSharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.load_kv_stage * 2 + ] + mma_s_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_s_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_mma_stage * 2] + p_cor_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_cor_stage * 2] + mma_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_o_stage * 2] + load_pt_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.load_pt_stage * 2 + ] + # Tmem dealloc cluster barrier + tmem_dealloc_mbar_ptr: cutlass.Int64 + + # Tmem holding buffer + tmem_holding_buf: cutlass.Int32 + # Smem tensors + softmax_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp + ] + epilogue_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp + ] + smem_q_latent: cute.struct.Align[ + cute.struct.MemRange[ + self.q_dtype, cute.cosize(q_latent_smem_layout_staged) + ], + 1024, + ] + smem_q_rope: cute.struct.Align[ + cute.struct.MemRange[ + self.q_dtype, cute.cosize(q_rope_smem_layout_staged) + ], + 1024, + ] + smem_kc: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, cute.cosize(kc_smem_layout_staged)], + 1024, + ] + smem_p: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, cute.cosize(p_smem_layout_staged)], + 1024, + ] + smem_page_table: cute.struct.MemRange[ + cutlass.Int32, self.load_pt_stage * self.mma_qk_tiler[1] // 2 + ] + + softmax_scale_log2 = softmax_scale * LOG2_E + self.split_kv_kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q_latent, + tma_tensor_q_latent, + tma_atom_q_rope, + tma_tensor_q_rope, + tma_atom_c_latent, + tma_tensor_c_latent, + tma_atom_c_rope, + tma_tensor_c_rope, + tma_atom_c_latent_transpose, + tma_tensor_c_latent_transpose, + page_table, + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale_log2, + output_scale, + q_latent_smem_layout_staged, + q_rope_smem_layout_staged, + kc_smem_layout_staged, + p_smem_layout_staged, + vc_smem_layout_staged, + kc_smem_layout_for_tma, + vc_smem_layout_for_tma, + cta_layout_vmnk, + tile_sched_params, + SplitKVKernelSharedStorage, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + smem=SplitKVKernelSharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + if cutlass.const_expr(acc_o is not None): + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + ).launch( + grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), + block=[self.threads_per_warp * self.num_compute_warps, 1, 1], + smem=MAX_SPLITS * self.acc_dtype.width // 8, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.jit + def make_paged_tiled_tma_atom( + self, + tma_load_op: cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp, + gmem: cute.Tensor, + smem_layout: cute.Layout, + mma_tiler, + tiled_mma: cute.TiledMma, + is_k_load: bool, + ): + ident = cute.make_identity_layout(gmem.shape) + g_tile = cute.composition(ident, mma_tiler) + cta_mn = mma_tiler[0] // tiled_mma.thr_id.shape + cta_v_map = cute.flat_divide(g_tile, (cta_mn,)) + cta_v_map = cute.select(cta_v_map, mode=[0, 2]) + page_tile_size = ( + min(self.page_size, cta_mn) + if is_k_load + else min(self.page_size, mma_tiler[1]) + ) + cta_v_map = cute.zipped_divide( + cta_v_map, + (page_tile_size, mma_tiler[1]) if is_k_load else (cta_mn, page_tile_size), + ) + cta_v_map = cute.select(cta_v_map, mode=[0]) + from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir + + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( + gmem.value, + smem_layout.value, + cta_v_map, + tma_load_op._to_ir(), + num_multicast=1, + ) + return ( + cute.CopyAtom( + tma_load_op, cpasync.CopyBulkTensorTileG2SNonExecTrait(res[0]) + ), + res[1], + ) + + @cute.kernel + def split_kv_kernel( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tma_atom_q_latent: Optional[cute.CopyAtom], + mQL: cute.Tensor, + tma_atom_q_rope: Optional[cute.CopyAtom], + mQR: cute.Tensor, + tma_atom_c_latent: Optional[cute.CopyAtom], + mCL: cute.Tensor, + tma_atom_c_rope: Optional[cute.CopyAtom], + mKR: cute.Tensor, + tma_atom_c_latent_transpose: Optional[cute.CopyAtom], + mCLT: cute.Tensor, + mPT: cute.Tensor, + mO: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + mAccO: Optional[cute.Tensor], + mAccLSE: Optional[cute.Tensor], + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + q_latent_smem_layout_staged: cute.ComposedLayout, + q_rope_smem_layout_staged: cute.ComposedLayout, + kc_smem_layout_staged: cute.ComposedLayout, + p_smem_layout_staged: cute.ComposedLayout, + vc_smem_layout_staged: cute.ComposedLayout, + kc_smem_layout_for_tma: cute.ComposedLayout, + vc_smem_layout_for_tma: cute.ComposedLayout, + cta_layout_vmnk: cute.Layout, + tile_sched_params: MLAStaticTileSchedulerParams, + SharedStorage: cutlass.Constexpr, + ): + """The device split_kv kernel implementation of the Multi-Head Latent Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the MLA computation: + 1. Load warp: Loads Q/C latent/rope data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Compute warps: Compute softmax and do rescaling on accumulators, and store the intermediate/final results + to global memory + + The kernel produces either intermediate or final results of the MLA computation based on the split_kv parameter. + When split_kv is 1, the kernel generates the final results directly. Otherwise, it produces intermediate results + that will later be combined by a reduction kernel. + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases. + + :param tiled_mma_qk: Tiled MMA for Q*K^T + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: Tiled MMA for P*V + :type tiled_mma_pv: cute.TiledMma + :param tma_atom_q_latent: TMA copy atom for query latent tensor + :type tma_atom_q_latent: cute.CopyAtom + :param mQL: query latent tensor + :type mQL: cute.Tensor + :param tma_atom_q_rope: TMA copy atom for query rope tensor + :type tma_atom_q_rope: cute.CopyAtom + :param mKR: Compressed rope tensor + :type mKR: cute.Tensor + :param tma_atom_c_latent: TMA copy atom for c latent tensor + :type tma_atom_c_latent: cute.CopyAtom + :param mCL: Compressed latent tensor + :type mCL: cute.Tensor + :param tma_atom_c_rope: TMA copy atom for c rope tensor + :type tma_atom_c_rope: cute.CopyAtom + :param mCLT: Compressed latent transpose tensor + :type mCLT: cute.Tensor + :param mPT: Page table tensor + :type mPT: cute.Tensor + :param mO: Output tensor + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor + :type mLSE: cute.Tensor + :param mAccO: Intermediate accumulator output tensor + :type mAccO: cute.Tensor + :param mAccLSE: Intermediate accumulator log-sum-exp tensor + :type mAccLSE: cute.Tensor + :param split_kv: The split_kv parameter + :type split_kv: cutlass.Int32 + :param cache_seqs: The variable sequence length tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: The per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param softmax_scale_log2: The log2 scale factor for softmax + :type softmax_scale_log2: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param q_latent_smem_layout_staged: Shared memory layout for query latent tensor + :type q_latent_smem_layout_staged: cute.ComposedLayout + :param q_rope_smem_layout_staged: Shared memory layout for query rope tensor + :type q_rope_smem_layout_staged: cute.ComposedLayout + :param kc_smem_layout_staged: Shared memory layout for key/value latent/rope tensor + :type kc_smem_layout_staged: cute.ComposedLayout + :param p_smem_layout_staged: Shared memory layout for probability matrix + :type p_smem_layout_staged: cute.ComposedLayout + :param vc_smem_layout_staged: Shared memory layout for value tensor + :type vc_smem_layout_staged: cute.ComposedLayout + :param kc_smem_layout_for_tma: Shared memory layout for key/value latent tensor for TMA + :type kc_smem_layout_for_tma: cute.ComposedLayout + :param vc_smem_layout_for_tma: Shared memory layout for value tensor for TMA + :type vc_smem_layout_for_tma: cute.ComposedLayout + :param cta_layout_vmnk: Layout for compute threads + :type cta_layout_vmnk: cute.Layout + :param tile_sched_params: Scheduling parameters for work distribution + :type tile_sched_params: MLAStaticTileSchedulerParams + :param SharedStorage: Shared storage for the kernel + :type SharedStorage: cutlass.Constexpr + """ + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + + # Prefetch tma descriptor + if warp_idx == self.mma_warp_id: + cpasync.prefetch_descriptor(tma_atom_q_latent) + cpasync.prefetch_descriptor(tma_atom_q_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent) + cpasync.prefetch_descriptor(tma_atom_c_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent_transpose) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_bar, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + load_q_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_q_stage, + self.tma_copy_q_bytes, + ) + load_kv_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_kv_stage, + self.tma_copy_kc_bytes, + ) + mma_s_pipeline = self.make_and_init_mma_s_pipeline( + storage.mma_s_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + p_mma_pipeline = self.make_and_init_p_mma_pipeline( + storage.p_mma_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + p_cor_pipeline = self.make_and_init_p_cor_pipeline( + storage.p_cor_mbar_ptr.data_ptr() + ) + mma_o_pipeline = self.make_and_init_mma_o_pipeline( + storage.mma_o_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + load_pt_pipeline = self.make_and_init_load_pt_pipeline( + storage.load_pt_mbar_ptr.data_ptr() + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) + + # Generate smem tensor Q/KC/VC/exchange + # (MMA, MMA_H, MMA_R, PIPE) + sQ = storage.smem_q_latent.get_tensor( + q_latent_smem_layout_staged.outer, swizzle=q_latent_smem_layout_staged.inner + ) + sQ_rope = storage.smem_q_rope.get_tensor( + q_rope_smem_layout_staged.outer, swizzle=q_rope_smem_layout_staged.inner + ) + # (MMA, MMA_K, MMA_R, PIPE) + sKC = storage.smem_kc.get_tensor( + kc_smem_layout_staged.outer, swizzle=kc_smem_layout_staged.inner + ) + sKC_for_tma = storage.smem_kc.get_tensor( + kc_smem_layout_for_tma.outer, + swizzle=kc_smem_layout_for_tma.inner, + ) + # (MMA, MMA_D, MMA_K, PIPE) + # reuse smem + sVC_ptr = cute.recast_ptr(sKC.iterator, vc_smem_layout_staged.inner) + sVC = cute.make_tensor(sVC_ptr, vc_smem_layout_staged.outer) + sVC_for_tma = cute.make_tensor(sVC_ptr, vc_smem_layout_for_tma.outer) + # (MMA, MMA_H, MMA_K) + sP = storage.smem_p.get_tensor( + p_smem_layout_staged.outer, swizzle=p_smem_layout_staged.inner + ) + sPT = storage.smem_page_table.get_tensor( + cute.make_layout((self.mma_qk_tiler[1] // 2, self.load_pt_stage)) + ) + # (compute_threads,) + softmax_smem_exchange = storage.softmax_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp) + ) + epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load warps, including page table and data tensors + # /////////////////////////////////////////////////////////////////////////////// + + if warp_idx >= self.empty_warp_ids[0] and warp_idx <= self.empty_warp_ids[-1]: + cute.arch.setmaxregister_decrease(self.other_reg_num) + if warp_idx == self.load_pt_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_pt_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_pt_stage + ) + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + load_pt_common_params = SimpleNamespace( + blk_coord=blk_coord, + load_pt_pipeline=load_pt_pipeline, + mPT=mPT, + sPT=sPT, + tidx=tidx, + page_size=mCL.shape[0], + ) + load_pt_producer_state = self.load_page_table( + load_pt_common_params, + k_index, + k_tile_count, + load_pt_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_pt_pipeline.producer_tail(load_pt_producer_state) + if warp_idx == self.load_tma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_q_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_q_stage + ) + load_kv_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_kv_stage + ) + load_pt_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_pt_stage + ) + load_pt_release_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_pt_stage + ) + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_kv_pipeline=load_kv_pipeline, + mPT=mPT, + sPT=sPT, + load_pt_pipeline=load_pt_pipeline, + ) + tma_qk_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + tma_atom_q_latent=tma_atom_q_latent, + tma_atom_q_rope=tma_atom_q_rope, + tma_atom_c_latent=tma_atom_c_latent, + tma_atom_c_rope=tma_atom_c_rope, + mQL=mQL, + mQR=mQR, + mCL=mCL, + mKR=mKR, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC_for_tma, + ) + tma_pv_params = SimpleNamespace( + tiled_mma_pv=tiled_mma_pv, + tma_atom_c_latent_transpose=tma_atom_c_latent_transpose, + mCL=mCL, + mKR=mKR, + mCLT=mCLT, + sVC=sVC_for_tma, + ) + # Load tma + ( + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) = self.load_tma( + tma_common_params, + tma_qk_params, + tma_pv_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + load_q_pipeline.producer_tail(load_q_producer_state) + load_kv_pipeline.producer_tail(load_kv_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA warp + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + # Alloc tensor memory buffer + tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100")) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + load_q_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_q_stage + ) + load_kv_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_kv_stage + ) + mma_s_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_s_stage + ) + p_mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_mma_stage + ) + mma_o_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_o_stage + ) + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + mma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_kv_pipeline=load_kv_pipeline, + tmem_ptr=tmem_ptr, + is_leader_cta=is_leader_cta, + L=mCL.shape[1], + ) + mma_qk_params = SimpleNamespace( + mma_s_pipeline=mma_s_pipeline, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC, + ) + mma_pv_params = SimpleNamespace( + p_mma_pipeline=p_mma_pipeline, + mma_o_pipeline=mma_o_pipeline, + sP=sP, + sVC=sVC, + ) + ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma( + mma_common_params, + mma_qk_params, + mma_pv_params, + k_tile_count, + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + mma_s_pipeline.producer_tail(mma_s_producer_state) + mma_o_pipeline.producer_tail(mma_o_producer_state) + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Compute warp + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx >= self.compute_warp_ids[0] + and warp_idx <= self.compute_warp_ids[-1] + ): + cute.arch.setmaxregister_increase(self.softmax_reg_num) + mma_s_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_s_stage + ) + p_mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_mma_stage + ) + p_cor_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_cor_stage + ) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage + ) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=softmax_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, + p_cor_pipeline=p_cor_pipeline, + ) + compute_softmax_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + sP=sP, + mma_s_pipeline=mma_s_pipeline, + p_mma_pipeline=p_mma_pipeline, + softmax_scale_log2=softmax_scale_log2, + ) + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = ( + self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_cor_pipeline.producer_tail(p_cor_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # Correction warp + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx >= self.correction_warp_ids[0] + and warp_idx <= self.correction_warp_ids[-1] + ): + cute.arch.setmaxregister_increase(self.correction_reg_num) + p_cor_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_cor_stage + ) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage + ) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=epilogue_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + H=mQL.shape[0], + tmem_ptr=tmem_ptr, + tidx=tidx, + tiled_mma_pv=tiled_mma_pv, + p_cor_pipeline=p_cor_pipeline, + mma_o_pipeline=mma_o_pipeline, + ) + compute_epilogue_params = SimpleNamespace( + output_scale=output_scale, + softmax_scale_log2=softmax_scale_log2, + mAccLSE=mAccLSE, + mLSE=mLSE, + ) + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + return + + @cute.kernel + def reduction_kernel( + self, + mO: cute.Tensor, + mLSE: cute.Tensor, + mAccO: cute.Tensor, + mAccLSE: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + ): + """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results + from multiple split_kv blocks into final outputs. + + :param mO: Output tensor for storing final results + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor for storing final LSE values + :type mLSE: cute.Tensor + :param mAccO: Accumulated output tensor from split_kv blocks + :type mAccO: cute.Tensor + :param mAccLSE: Accumulated LSE tensor from split_kv blocks + :type mAccLSE: cute.Tensor + :param split_kv: Number of split_kv blocks + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor (for variable split_kv) + :type block_split_kvs: cute.Tensor + """ + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + blk_coord = (bidx, bidy, bidz) + local_split_kv = ( + block_split_kvs[blk_coord[2]] if self.is_var_split_kv else split_kv + ) + k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) + local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) + + # Alloc shared memory + smem = utils.SmemAllocator() + storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) + lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) + smem_lse_scale = cute.make_tensor(lse_scale_ptr, cute.make_layout(MAX_SPLITS)) + + gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) + lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) + + local_lse = cute.make_rmem_tensor( + cute.make_layout(lse_per_thread), self.lse_dtype + ) + lse_max = -self.lse_dtype.inf + # find the max lse + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + local_lse[i] = ( + gLSE[split_kv_idx] + if cute.elem_less(split_kv_idx, local_split_kv) + else -self.lse_dtype.inf + ) + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) + lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 0.0 + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if not sum_lse == self.lse_dtype(0.0) or sum_lse != sum_lse + else self.lse_dtype.inf + ) + if tidx == 0: + mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[split_kv_idx] = cute.math.exp2( + local_lse[i] - global_lse, fastmath=True + ) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( + self.latent_dim, self.threads_per_warp * self.num_compute_warps + ) + gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] + rAccO = cute.make_rmem_tensor( + cute.make_layout(elements_per_thread), self.acc_dtype + ) + rO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), self.o_dtype) + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] + return + + @staticmethod + def get_split_kv( + B: int, S: int, K: int, mma_qk_tiler_mn: tuple, max_active_blocks: int + ) -> int: + """Get the proper split_kv value for the MLA kernel based on parameters. + + :param B: Batch size + :type B: int + :param S: Sequence length + :type S: int + :param K: Sequence length + :type K: int + :param mma_qk_tiler_mn: MLA tiling parameters + :type mma_qk_tiler_mn: tuple + :param max_active_blocks: Maximum number of active blocks + :type max_active_blocks: int + :return: Split_kv value + :rtype: int + """ + max_splits = ceil_div(K, mma_qk_tiler_mn[1]) + blocks_per_batch = max(1, max_active_blocks // B // (S * 2)) + split_heur = min(max_splits, blocks_per_batch) + k_waves = ceil_div(max_splits, split_heur) + split_wave_aware = ceil_div(max_splits, k_waves) + max_split_kv = 32 + return min(split_wave_aware, max_split_kv) + + @cute.jit + def get_k_tile_count( + self, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + blk_coord: cute.Coord, + ) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Get the current k_index, k_tile_count, and local split_kv value for the MLA kernel. + + :param split_kv: Split_kv value + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param blk_coord: Block coordinate + :type blk_coord: cute.Coord + :return: k_index, k_tile_count, split_kv + :rtype: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32] + """ + K = cache_seqs[blk_coord[2]] + if cutlass.const_expr(self.is_var_split_kv): + split_kv = block_split_kvs[blk_coord[2]] + + k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv) + k_index = blk_coord[3] * k_tile_per_cta + k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index) + return k_index, k_tile_count, split_kv + + @cute.jit + def load_page_table( + self, + common_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_pt_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load warp to load page table. Updates the load pt producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_pt_producer_state: The load pt producer state + :type load_pt_producer_state: pipeline.PipelineState + + :return: The load pt producer state + :rtype: pipeline.PipelineState + """ + mPT = common_params.mPT[None, common_params.blk_coord[2]] + page_per_tile = self.mma_qk_tiler[1] // self.page_size + tidx = common_params.tidx % self.threads_per_warp + + load_pt_pipeline = common_params.load_pt_pipeline + while k_tile_count > 0: + load_pt_pipeline.producer_acquire(load_pt_producer_state) + + elem_per_thread = cute.ceil_div(page_per_tile, self.threads_per_warp) + + # atom_async_copy: async copy atom for page table load + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS), + cutlass.Int32, + num_bits_per_copy=cutlass.Int32.width, + ) + mPT_for_copy = cute.flat_divide(mPT, (1,)) + sPT_for_copy = cute.flat_divide(common_params.sPT, (1,)) + # elem_per_thread is a dynamic value depends on the page_size setting. + for i in range(elem_per_thread): + idx = i * self.threads_per_warp + tidx + if cute.elem_less( + k_index * page_per_tile + idx, mPT.shape[0] + ) and cute.elem_less(idx, page_per_tile): + cute.copy( + atom_async_copy, + mPT_for_copy[None, k_index * page_per_tile + idx], + sPT_for_copy[None, idx, load_pt_producer_state.index], + ) + else: + sPT_for_copy[None, idx, load_pt_producer_state.index].fill(0) + mbar_ptr = load_pt_pipeline.producer_get_barrier(load_pt_producer_state) + load_pt_pipeline.producer_commit(load_pt_producer_state) + load_pt_producer_state.advance() + k_index += 1 + k_tile_count -= 1 + + return load_pt_producer_state + + @cute.jit + def load_tma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_kv_producer_state: pipeline.PipelineState, + load_pt_consumer_state: pipeline.PipelineState, + load_pt_release_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Load wrap to load Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param v_params: The v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_kv_producer_state: The load kv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_consumer_state: The load pt consumer state + :type load_pt_consumer_state: pipeline.PipelineState + :param load_pt_release_state: The load pt release state + :type load_pt_release_state: pipeline.PipelineState + + :return: The load q producer state, load kv producer state, load pt consumer state, and load pt release state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for QK TMA load + # (bM, bK, rM, rK, rL) + mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) + gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) + mma_qk_tiler_mk_rope = cute.select(self.mma_qk_rope_tiler, mode=[0, 2]) + gQR = cute.flat_divide(qk_params.mQR, mma_qk_tiler_mk_rope) + + thr_mma_qk = qk_params.tiled_mma_qk.get_slice( + common_params.blk_coord[0] % cute.size(qk_params.tiled_mma_qk.thr_id) + ) + tSgQL = thr_mma_qk.partition_A(gQL) + tSgQR = thr_mma_qk.partition_A(gQR) + + cta_m = min( + qk_params.tiled_mma_qk.op.shape_mnk[0] + // qk_params.tiled_mma_qk.thr_id.shape, + self.page_size, + ) + page_tile_size = min(self.page_size, cta_m) + gCL = cute.tiled_divide(qk_params.mCL, (page_tile_size, self.mma_qk_tiler[2])) + tSgCL = ( + gCL[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] + if cta_m < self.page_size + else gCL[None, 0, None, None] + ) + gKR = cute.tiled_divide(qk_params.mKR, (page_tile_size, self.mma_qk_tiler[2])) + tSgKR = ( + gKR[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] + if cta_m < self.page_size + else gKR[None, 0, None, None] + ) + + # tma partition for q, k latent/rope + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tQsQ, tQLgQL_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_latent, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ, 0, 3), + cute.group_modes(tSgQL, 0, 3), + ) + + tQsQ_rope, tQRgQR_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_rope, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ_rope, 0, 3), + cute.group_modes(tSgQR, 0, 3), + ) + + tKCsKC, tCLgCL = cpasync.tma_partition( + qk_params.tma_atom_c_latent, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgCL, + ) + + _, tKRgKR = cpasync.tma_partition( + qk_params.tma_atom_c_rope, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgKR, + ) + + tQLgQL = tQLgQL_mkl[ + None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] + ] + tQRgQR = tQRgQR_mkl[ + None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] + ] + + # Flatten divide and partition global tensors for V TMA load + page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + gCLT = cute.flat_divide(v_params.mCLT, (self.mma_pv_tiler[1], page_tile_size)) + cta_n = self.mma_pv_tiler[1] // v_params.tiled_mma_pv.thr_id.shape + gCLT = cute.logical_divide(gCLT, (cta_n,))[ + (None, common_params.blk_coord[0]), None, None, None, None + ] + tOgCLT = cute.tiled_divide(gCLT, (cta_n, page_tile_size)) + tOgCLT = tOgCLT[None, 0, 0, None, None, None] + + # tma partition for vc + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tVCsVC, tCLTgCLT = cpasync.tma_partition( + v_params.tma_atom_c_latent_transpose, + 0, + cute.make_layout(1), + v_params.sVC, + tOgCLT, + ) + + # set extra params + common_params.mPT = mPT + qk_params.tQLgQL = tQLgQL + qk_params.tQRgQR = tQRgQR + qk_params.tCLgCL = tCLgCL + qk_params.tKRgKR = tKRgKR + qk_params.tQsQ = tQsQ + qk_params.tQsQ_rope = tQsQ_rope + qk_params.tKCsKC = tKCsKC + v_params.tCLTgCLT = tCLTgCLT + v_params.tVCsVC = tVCsVC + + load_q_producer_state, load_kv_producer_state, load_pt_consumer_state = ( + self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_q=True, + ) + ) + k_index += 1 + k_tile_count -= 1 + while k_tile_count > 0: + load_q_producer_state, load_kv_producer_state, load_pt_consumer_state = ( + self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_q=False, + ) + ) + load_kv_producer_state, load_pt_release_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index - 1, + load_kv_producer_state, + load_pt_release_state, + ) + k_index += 1 + k_tile_count -= 1 + + # load last v tile + load_kv_producer_state, load_pt_release_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index - 1, + load_kv_producer_state, + load_pt_release_state, + ) + return ( + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) + + @cute.jit + def load_tma_qk_one_k_tile( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_kv_producer_state: pipeline.PipelineState, + load_pt_consumer_state: pipeline.PipelineState, + load_q: bool, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]: + """Load one k-tile of Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_kv_producer_state: The load kv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_consumer_state: The load pt consumer state + :type load_pt_consumer_state: pipeline.PipelineState + :param load_q: Whether to load q + :type load_q: bool + + :return: The load q producer state, load kv producer state, and load pt consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = ceil_div( + self.mma_qk_tiler[1] // self.page_size, qk_params.tiled_mma_qk.thr_id.shape + ) + common_params.load_pt_pipeline.consumer_wait(load_pt_consumer_state) + page_table_stage = load_pt_consumer_state.index + load_pt_consumer_state.advance() + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = ( + common_params.sPT[0, page_table_stage] + if self.mma_qk_tiler[1] // self.page_size == 1 + else common_params.sPT[ + i + common_params.blk_coord[0] * page_per_tile, page_table_stage + ] + ) + # load q once at first iteration + if cutlass.const_expr(load_q): + common_params.load_q_pipeline.producer_acquire(load_q_producer_state) + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_q_pipeline.producer_get_barrier( + load_q_producer_state + ) + for i in cutlass.range(self.iterations_qk_latent): + # load q latent + cute.copy( + qk_params.tma_atom_q_latent, + qk_params.tQLgQL[None, 0, i], + qk_params.tQsQ[None, (i, 0)], + tma_bar_ptr=tma_bar_ptr, + ) + for i in cutlass.range(self.iterations_qk_rope): + # load q rope + cute.copy( + qk_params.tma_atom_q_rope, + qk_params.tQRgQR[None, 0, i], + qk_params.tQsQ_rope[None, i], + tma_bar_ptr=tma_bar_ptr, + ) + load_q_producer_state.advance() + load_kv_pipeline = common_params.load_kv_pipeline + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) + for i in cutlass.range(self.iterations_qk_latent): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_tile): + # load k latent + cute.copy( + qk_params.tma_atom_c_latent, + qk_params.tCLgCL[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_kv_producer_state.advance() + + for i in cutlass.range(self.iterations_qk_rope): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_tile): + # load k rope + cute.copy( + qk_params.tma_atom_c_rope, + qk_params.tKRgKR[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_kv_producer_state.advance() + + return load_q_producer_state, load_kv_producer_state, load_pt_consumer_state + + @cute.jit + def load_tma_v_one_k_tile( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + load_kv_producer_state: pipeline.PipelineState, + load_pt_release_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load one k-tile of compressed latent transpose tensor(v). Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The load tma v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param load_kv_producer_state: The load qkv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_release_state: The load pt release state + :type load_pt_release_state: pipeline.PipelineState + + :return: The load kv producer state and load pt release state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = self.mma_pv_tiler[2] * self.iterations_pv_k // self.page_size + page_per_subtile = ceil_div(page_per_tile, self.iterations_pv_k) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) + page_table_stage = load_pt_release_state.index + for i in cutlass.range(page_per_tile): + k_idx[i] = ( + common_params.sPT[0, page_table_stage] + if page_per_tile == 1 + else common_params.sPT[i, page_table_stage] + ) + common_params.load_pt_pipeline.consumer_release(load_pt_release_state) + load_pt_release_state.advance() + load_kv_pipeline = common_params.load_kv_pipeline + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) + for i in cutlass.range(self.iterations_pv_k): + for j in cutlass.range(self.iterations_pv_n): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier( + load_kv_producer_state + ) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_subtile): + k_idx_i = k_idx[ + k + + i + // ceil_div(self.iterations_pv_k, page_per_tile) + * page_per_subtile + ] + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[ + None, + j, + i % ceil_div(self.iterations_pv_k, page_per_tile), + k_idx_i, + ], + v_params.tVCsVC[None, 0, k, load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + + load_kv_producer_state.advance() + return load_kv_producer_state, load_pt_release_state + + @cute.jit + def mma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + pv_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_kv_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """MMA warp to compute the result of Q*K^T and P*V. Updates the tiled mma and pipeline states. + + :param common_params: The common parameters for mma qk and pv + :type common_params: SimpleNamespace + :param qk_params: The mma qk parameters + :type qk_params: SimpleNamespace + :param pv_params: The mma pv parameters + :type pv_params: SimpleNamespace + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + :param p_mma_consumer_state: The p mma consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The mma o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the tiled mma pv, the load q consumer state, the load kv consumer state, the mma s producer state, the p mma consumer state, and the mma o producer state + :rtype: tuple[cute.TiledMma, cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + tSrQ = tiled_mma_qk.make_fragment_A(qk_params.sQ) + tSrQ_rope = tiled_mma_qk.make_fragment_A(qk_params.sQ_rope) + tSrKC = tiled_mma_qk.make_fragment_B(qk_params.sKC) + tOrP = tiled_mma_pv.make_fragment_A(pv_params.sP) + tOrVC = tiled_mma_pv.make_fragment_B(pv_params.sVC) + + tStS_shape = tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1]) + ) + tStS_staged_fake = tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage) + ) + # use real tmem ptr for tStS + tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout) + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1]) + ) + # mma O has 1 stage. + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO_staged = cute.make_tensor( + tStS_staged.iterator + self.tmem_o_offset, tOtO_layout + ) + + # set more parameters + qk_params.tSrQ = tSrQ + qk_params.tSrQ_rope = tSrQ_rope + qk_params.tSrKC = tSrKC + qk_params.tStS_staged = tStS_staged + pv_params.tOrP = tOrP + pv_params.tOrVC = tOrVC + pv_params.tOtO_staged = tOtO_staged + + # mma O accumulates on K, so the accumlate flag is set to False once before all K blocks. + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + if common_params.is_leader_cta: + load_q_release_state = load_q_consumer_state.clone() + + ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + wait_q=True, + ) + k_tile_count -= 1 + while k_tile_count > 0: + ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + wait_q=False, + ) + ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + k_tile_count -= 1 + + # release q consumer states + load_q_pipeline.consumer_release(load_q_release_state) + load_q_release_state.advance() + ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + return ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def mma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + tiled_mma_qk: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_kv_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + wait_q: bool, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for Q*K^T. Updates the tiled MMA QK and pipeline states. + + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the load q consumer state, the load kv consumer state, and the mma s producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + tStS = qk_params.tStS_staged[None, None, None, mma_s_producer_state.index] + + qk_params.mma_s_pipeline.producer_acquire(mma_s_producer_state) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + load_kv_pipeline = common_params.load_kv_pipeline + if cutlass.const_expr(wait_q): + load_q_pipeline.consumer_wait(load_q_consumer_state) + load_q_consumer_state.advance() + for q_stage in range(self.iterations_qk_latent): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + kc_stage = load_kv_consumer_state.index + for k_block in cutlass.range(cute.size(qk_params.tSrQ.shape[2])): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ[None, None, k_block, q_stage], + qk_params.tSrKC[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + for q_stage in range(self.iterations_qk_rope): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + kc_stage = load_kv_consumer_state.index + for k_block in cutlass.range(self.rope_dim // tiled_mma_qk.shape_mnk[2]): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ_rope[None, None, k_block, q_stage], + qk_params.tSrKC[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + + qk_params.mma_s_pipeline.producer_commit(mma_s_producer_state) + mma_s_producer_state.advance() + return ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) + + @cute.jit + def mma_pv( + self, + common_params: SimpleNamespace, + pv_params: SimpleNamespace, + tiled_mma_pv: cute.TiledMma, + load_kv_consumer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for P*V. Updates the tiled mma pv and pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param pv_params: The pv parameters + :type pv_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param p_mma_consumer_state: The P MMA consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The MMA o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma pv, the load qkv consumer state, the P MMA consumer state, and the MMA o producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + pv_params.mma_o_pipeline.producer_acquire(mma_o_producer_state) + pv_params.p_mma_pipeline.consumer_wait(p_mma_consumer_state) + load_kv_pipeline = common_params.load_kv_pipeline + for p_stage in range(self.iterations_pv_k): + accumulate_flag = tiled_mma_pv.get(tcgen05.Field.ACCUMULATE) + for acc_stage in range(self.iterations_pv_n): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, accumulate_flag) + vc_stage = load_kv_consumer_state.index + tOtO = pv_params.tOtO_staged[None, None, None, acc_stage] + for k_block in cutlass.range(pv_params.tOrP.shape[2]): + cute.gemm( + tiled_mma_pv, + tOtO, + pv_params.tOrP[ + None, + None, + k_block, + (p_stage, p_mma_consumer_state.index), + ], + pv_params.tOrVC[None, None, k_block, vc_stage], + tOtO, + ) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + pv_params.p_mma_pipeline.consumer_release(p_mma_consumer_state) + p_mma_consumer_state.advance() + pv_params.mma_o_pipeline.producer_commit(mma_o_producer_state) + mma_o_producer_state.advance() + + return ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def compute( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + + :return: The MMA s consumer state, the P MMA producer state, and the P correction producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_total = cute.ceil_div(common_params.K, self.mma_qk_tiler[1]) + + row_max = -self.acc_dtype.inf + row_sum = self.acc_dtype(0) + correction_factor = self.acc_dtype(1) + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # no mask applied + while k_tile_count > 1: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + False, + False, + ) + k_index = k_index + 1 + k_tile_count = k_tile_count - 1 + + # mask applied + if cutlass.const_expr(common_params.mAccO is not None): + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + k_index == k_tile_total - 1, + True, + ) + else: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + True, + True, + ) + + return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state + + @cute.jit + def correction( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + p_cor_consumer_state: pipeline.PipelineState, + mma_o_consumer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + :param mma_o_consumer_state: The MMA o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, and the MMA o consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction = ( + self.get_correction_factor(common_params, p_cor_consumer_state) + ) + if k_tile_count_init != k_tile_count: + mma_o_consumer_state = self.rescale( + common_params, + mma_o_consumer_state, + correction_factor, + no_correction, + ) + k_tile_count = k_tile_count - 1 + if k_tile_count == 0: + mma_o_consumer_state = self.epilogue( + common_params, + epilogue_params, + mma_o_consumer_state, + row_sum, + row_max, + ) + + return p_cor_consumer_state, mma_o_consumer_state + + @cute.jit + def exchange_p_cor_metadata( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + correction_factor: cutlass.Float32, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + row_max_new: cutlass.Float32, + tAcc: cute.Tensor, + tidx: cutlass.Int32, + p_cor_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Compute the correction factor for the last k tile.""" + no_correction = 0 + if ( + row_max_new - row_max + ) * softmax_params.softmax_scale_log2 <= self.skip_correction_threshold: + no_correction = 1 + row_max_new = row_max + + # pad for 4x32b + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.mma_s_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, + corr_layout, + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype + ) + corr_tmem_store_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_store_atom, tCor) + corr_tmem_store_thr_copy = corr_tmem_store_tiled_copy.get_slice(tidx) + cCor_for_copy = corr_tmem_store_thr_copy.partition_S(cCor) + tCor_for_copy = corr_tmem_store_thr_copy.partition_D(tCor) + rCor = cute.make_fragment_like( + cCor_for_copy[None, None, None, 0], self.acc_dtype + ) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout + ) + rCor[0] = row_sum + rCor[1] = row_max_new + rCor[2] = correction_factor + rCor_int[3] = no_correction + + cute.copy( + corr_tmem_store_tiled_copy, + rCor, + tCor_for_copy[None, None, None, p_cor_producer_state.index], + ) + # fence between tmem store and correction warp + cute.arch.fence_view_async_tmem_store() + common_params.p_cor_pipeline.producer_commit(p_cor_producer_state) + p_cor_producer_state.advance() + return p_cor_producer_state, row_max_new + + @cute.jit + def softmax( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + row_max: cutlass.Float32, + row_sum: cutlass.Float32, + correction_factor: cutlass.Float32, + is_last_tile: bool, + is_local_last_tile: cutlass.Boolean, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + ]: + """Softmax for one k-tile. Updates the related pipeline states and returns the computed results. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + :param row_max: The row max + :type row_max: cutlass.Float32 + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param is_last_tile: Whether the last tile + :type is_last_tile: bool + :param is_local_last_tile: Whether the last tile is local + :type is_local_last_tile: cutlass.Boolean + + :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32] + """ + + softmax_params.p_mma_pipeline.producer_acquire(p_mma_producer_state) + softmax_params.mma_s_pipeline.consumer_wait(mma_s_consumer_state) + + # load S from tmem + tStS_shape = softmax_params.tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1]) + ) + tStS_staged_fake = softmax_params.tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage) + ) + tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout) + tStS = tStS_staged[None, None, None, mma_s_consumer_state.index] + + tAcc = tStS[(None, None), 0, 0] + cta_qk_tiler = ( + self.mma_qk_tiler[0] // self.cluster_shape_mnk[0], + self.mma_qk_tiler[1], + self.mma_qk_tiler[2], + ) + cS = cute.make_identity_tensor(cute.select(cta_qk_tiler, mode=[0, 1])) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + + tmem_thr_copy = tmem_tiled_copy.get_slice(tidx) + tTR_tAcc = tmem_thr_copy.partition_S(tAcc) + tTR_tS = tmem_thr_copy.partition_D(cS) + + tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) + + row_max_new = row_max + arch = BaseDSL._get_dsl().get_arch_enum() + if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): + cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + if is_last_tile: + tTR_rAcc[i] = ( + tTR_rAcc[i] + if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) + else -self.acc_dtype.inf + ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) + + elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + tmem_load_red_atom = cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp( + tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX + ), + self.acc_dtype, + ) + tmem_red_tiled_copy = tcgen05.make_tmem_copy(tmem_load_red_atom, tAcc) + tmem_red_thr_copy = tmem_red_tiled_copy.get_slice(tidx) + tTR_tAcc_red = tmem_red_thr_copy.partition_S(tAcc) + tTR_tS_red = tmem_red_thr_copy.partition_D(cS) + tTR_rAcc_red = cute.make_fragment_like(tTR_tS_red, self.acc_dtype) + tTR_rMax = cute.make_rmem_tensor( + cute.make_layout((1, tTR_tS_red.shape[1], tTR_tS_red.shape[2])), + self.acc_dtype, + ) + cute.copy( + tmem_red_tiled_copy, + tTR_tAcc_red, + (tTR_rAcc_red, tTR_rMax), + ) + tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) + if is_last_tile: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[i] = ( + tTR_rAcc[i] + if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) + else -self.acc_dtype.inf + ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce( + cute.ReductionOp.MAX, row_max_new, 0 + ) + else: + row_max_new = cute.arch.fmax(row_max_new, tTR_rMax[0]) + + # if warps in N is 2, reduce row_max across warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_max_new + self.softmax_exchange_sync_bar.wait() + row_max_new = cute.arch.fmax( + row_max_new, + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp) + ], + ) + + # find correction factor + correction_factor = cute.math.exp2( + (row_max - row_max_new) * softmax_params.softmax_scale_log2, fastmath=True + ) + # split kv case + if cutlass.const_expr(not is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # softmax + fma_b = softmax_params.softmax_scale_log2 + fma_c = (0.0 - row_max_new) * softmax_params.softmax_scale_log2 + + for i in cutlass.range(cute.size(tTR_rAcc), vectorize=True, unroll_full=True): + tTR_rAcc[i] = tTR_rAcc[i] * fma_b + fma_c + tTR_rAcc[i] = cute.math.exp2(tTR_rAcc[i], fastmath=True) + + tTR_rS = cute.make_fragment_like(tTR_tS, self.q_dtype) + + # quantize + tTR_rS.store(tTR_rAcc.load().to(self.q_dtype)) + + # create sP + sP = softmax_params.sP[None, None, None, (None, p_mma_producer_state.index)] + sP_mk_view = cute.make_tensor( + sP.iterator, + cute.make_layout( + ( + (sP.shape[0][0], sP.shape[1]), + (sP.shape[0][1], sP.shape[2], sP.shape[3]), + ), + stride=( + (sP.stride[0][0], sP.stride[1]), + (sP.stride[0][1], sP.stride[2], sP.stride[3]), + ), + ), + ) + # change to PISL + sP_wo_swizzle_iter = cute.recast_ptr(sP.iterator, swizzle_=None) + swizzle_bits = ( + int(math.log2(self.mma_pv_tiler[2] * self.q_dtype.width // 8 // 32)) + 1 + ) + swizzle_base = 3 if self.q_dtype.width == 16 else 4 + sP_swizzle = cute.make_swizzle(swizzle_bits, swizzle_base, 3) + sP_mk_view = cute.make_tensor( + sP_wo_swizzle_iter, + cute.make_composed_layout(sP_swizzle, 0, sP_mk_view.layout), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_copy) + smem_thr_copy = smem_tiled_copy.get_slice(tidx) + rP_copy_view = smem_thr_copy.retile(tTR_rS) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + + # fence between smem store and mma o + cute.arch.fence_view_async_shared() + softmax_params.p_mma_pipeline.producer_commit(p_mma_producer_state) + p_mma_producer_state.advance() + + # row_sum, using `add_packed_f32x2` to reduce the number of instructions + row_sum = row_sum * correction_factor + row_sum_vec = (0.0, 0.0) + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + row_sum_vec = cute.arch.add_packed_f32x2( + row_sum_vec, (tTR_rAcc[i], tTR_rAcc[i + 1]) + ) + row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum + + # split kv case + if cutlass.const_expr(is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # store correction factor/row_sum/row_max to tmem for correction warp + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # fence between tmem load and mma s + cute.arch.fence_view_async_tmem_load() + + softmax_params.mma_s_pipeline.consumer_release(mma_s_consumer_state) + mma_s_consumer_state.advance() + + return ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max_new, + row_sum, + correction_factor, + ) + + @cute.jit + def _tmem_load_partition( + self, common_params: SimpleNamespace, tiled_mma_pv: cute.TiledMma, iter_n: int + ) -> tuple[ + cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma + ]: + """Tensor memory load partition for rescale and epilogue. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param iter_n: The iteration number + :type iter_n: int + + :return: The tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv + :rtype: tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma] + """ + + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1]) + ) + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO = cute.make_tensor( + common_params.tmem_ptr + self.tmem_o_offset, tOtO_layout + ) + tOtO = tOtO[None, None, None, iter_n] + + tAcc = tOtO[(None, None), 0, 0] + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_load_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + tmem_load_thr_copy = tmem_load_tiled_copy.get_slice( + common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + ) + + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + # Flatten divide and partition global tensors for O + cta_pv_tiler_mn = cute.select(cta_pv_tiler, mode=[0, 1]) + + gO = None + if cutlass.const_expr(common_params.mAccO is not None): + gO = cute.local_tile( + common_params.mAccO[None, common_params.blk_coord[3], None, None, None], + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor( + common_params.mAccO[ + None, common_params.blk_coord[3], None, None, None + ].shape + ), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + else: + gO = cute.local_tile( + common_params.mO, + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor(common_params.mO.shape), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + tTR_tAcc = tmem_load_thr_copy.partition_S(tAcc) + tTR_gO = tmem_load_thr_copy.partition_D(gO) + tTR_cO = tmem_load_thr_copy.partition_D(cO) + tTR_rAcc = cute.make_fragment_like(tTR_gO, self.acc_dtype) + return tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc + + def get_correction_factor( + self, + common_params: SimpleNamespace, + p_cor_consumer_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + cutlass.Int32, + ]: + """Get the correction factor from the P correction consumer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, the row_sum, the row_max, and the correction factor + :rtype: tuple[pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Int32] + """ + common_params.p_cor_pipeline.consumer_wait(p_cor_consumer_state) + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + # load correction factor + _, tAcc, _, _, _, _ = self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, 0 + ) + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.p_cor_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, corr_layout + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype + ) + corr_tmem_load_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_load_atom, tCor) + corr_tmem_load_thr_copy = corr_tmem_load_tiled_copy.get_slice(tidx) + tCor_for_copy = corr_tmem_load_thr_copy.partition_S(tCor) + cCor_for_copy = corr_tmem_load_thr_copy.partition_D(cCor) + rCor = cute.make_fragment_like( + cCor_for_copy[None, None, None, 0], self.acc_dtype + ) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout + ) + cute.copy( + corr_tmem_load_tiled_copy, + tCor_for_copy[None, None, None, p_cor_consumer_state.index], + rCor, + ) + row_sum = rCor[0] + row_max = rCor[1] + correction_factor = rCor[2] + no_correction = rCor_int[3] + + common_params.p_cor_pipeline.consumer_release(p_cor_consumer_state) + p_cor_consumer_state.advance() + return p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction + + @cute.jit + def rescale( + self, + common_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + correction_factor: cutlass.Float32, + no_correction: cutlass.Int32, + ) -> pipeline.PipelineState: + """Rescale for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param no_correction: Whether to apply correction factor + :type no_correction: cutlass.Int32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + skip_correction = cute.arch.vote_all_sync(no_correction == 1) + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + if not skip_correction: + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, iter_n + ) + ) + + # tmem store tiled copy + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_store_tiled_copy = tcgen05.make_tmem_copy(tmem_store_atom, tAcc) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + # rescale, using `mul_packed_f32x2` to reduce the number of instructions + for i in cutlass.range( + cute.size(tTR_rAcc), vectorize=True, unroll_full=True + ): + tTR_rAcc[i] = tTR_rAcc[i] * correction_factor + + # store o to tensor memory for next k tile + cute.copy(tmem_store_tiled_copy, tTR_rAcc, tTR_tAcc) + + cute.arch.fence_view_async_tmem_store() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + @cute.jit + def epilogue( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + ) -> pipeline.PipelineState: + """Epilogue for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param row_max: The row max + :type row_max: cutlass.Float32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + + # exchange row_sum between warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_sum + self.epilogue_exchange_sync_bar.wait() + # (64, 2) + row_sum = ( + row_sum + + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp) + ] + ) + # mma_o pipeline consumer wait + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, iter_n + ) + ) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + + # apply output scale and normalize by row_sum + for i in cutlass.range( + cute.size(tTR_rAcc), vectorize=True, unroll_full=True + ): + tTR_rAcc[i] = ( + tTR_rAcc[i] + * epilogue_params.output_scale + * cute.arch.rcp_approx(row_sum) + ) + + # store o to global memory + tR2G_rO_src = None + tR2G_rO_dst = tTR_gO + if cutlass.const_expr(common_params.mAccO is None): + tR2G_rO_src = cute.make_fragment_like(tTR_gO, self.o_dtype) + # using final output dtype for o + tR2G_rO_src.store(tTR_rAcc.load().to(self.o_dtype)) + else: + # using accumulate dtype for o + tR2G_rO_src = tTR_rAcc + + if cute.elem_less(tTR_cO[0][0], common_params.H): + cute.autovec_copy( + tR2G_rO_src, + tR2G_rO_dst, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, + ) + + # store the lse to global memory + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + gLSE = None + cLSE = None + if cutlass.const_expr(epilogue_params.mAccLSE is None): + gLSE = cute.local_tile( + epilogue_params.mLSE, + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor(epilogue_params.mLSE.shape), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + + else: + gLSE = cute.local_tile( + epilogue_params.mAccLSE[ + None, common_params.blk_coord[3], None, None + ], + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor( + epilogue_params.mAccLSE[ + None, common_params.blk_coord[3], None, None + ].shape + ), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + lse = ( + cute.math.log2(row_sum, fastmath=True) + + epilogue_params.softmax_scale_log2 * row_max + ) + if cutlass.const_expr(self.warps_in_n == 2): + if cute.elem_less(cLSE[tidx][0], common_params.H): + gLSE[tidx] = lse + + cute.arch.fence_view_async_tmem_load() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + def make_and_init_load_pt_pipeline(self, load_pt_mbar_ptr): + """Create and initialize the load page table pipeline. + + :param load_pt_mbar_ptr: The load page table mbar pointer + :type load_pt_mbar_ptr: cute.Tensor + + :return: The load page table pipeline + :rtype: pipeline.PipelineAsync + """ + load_pt_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len([self.load_pt_warp_id]), + ) + load_pt_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len([self.load_tma_warp_id]), + ) + return pipeline.PipelineCpAsync.create( + barrier_storage=load_pt_mbar_ptr, + num_stages=self.load_pt_stage, + producer_group=load_pt_producer_group, + consumer_group=load_pt_consumer_group, + defer_sync=True, + ) + + def make_and_init_load_qkv_pipeline( + self, load_qkv_mbar_ptr, cta_layout_vmnk, load_stages, tx_count + ) -> pipeline.PipelineTmaUmma: + """Create and initialize the tma load qkv pipeline. + + :param load_qkv_mbar_ptr: The load qkv mbar pointer + :type load_qkv_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + :param load_stages: The load stages + :type load_stages: list[int] + :param tx_count: The tx count + :type tx_count: int + + :return: The tma load qkv pipeline + :rtype: pipeline.PipelineTmaUmma + """ + load_qkv_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_tma_warp_id]) + ) + load_qkv_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_qkv_mbar_ptr, + num_stages=load_stages, + producer_group=load_qkv_producer_group, + consumer_group=load_qkv_consumer_group, + tx_count=tx_count, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_mma_s_pipeline( + self, mma_s_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma s pipeline. + + :param mma_s_mbar_ptr: The mma s mbar pointer + :type mma_s_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma s pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_s_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + consumer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + mma_s_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_s_mbar_ptr, + num_stages=self.mma_s_stage, + producer_group=mma_s_producer_group, + consumer_group=mma_s_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_mma_pipeline( + self, p_mma_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p mma pipeline. + + :param p_mma_mbar_ptr: The p mma mbar pointer + :type p_mma_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The p mma pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + p_mma_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_mma_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineAsyncUmma.create( + barrier_storage=p_mma_mbar_ptr, + num_stages=self.p_mma_stage, + producer_group=p_mma_producer_group, + consumer_group=p_mma_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_cor_pipeline( + self, p_cor_mbar_ptr + ) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p correction pipeline. + + :param p_cor_mbar_ptr: The p correction mbar pointer + :type p_cor_mbar_ptr: cute.Tensor + + :return: The p correction pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = self.threads_per_warp * len(self.compute_warp_ids) + p_cor_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_cor_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + return pipeline.PipelineAsync.create( + barrier_storage=p_cor_mbar_ptr, + num_stages=self.p_cor_stage, + producer_group=p_cor_producer_group, + consumer_group=p_cor_consumer_group, + defer_sync=True, + ) + + def make_and_init_mma_o_pipeline( + self, mma_o_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma o pipeline. + + :param mma_o_mbar_ptr: The mma o mbar pointer + :type mma_o_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma o pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_o_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + consumer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + mma_o_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_o_mbar_ptr, + num_stages=self.mma_o_stage, + producer_group=mma_o_producer_group, + consumer_group=mma_o_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + @staticmethod + def _compute_grid( + o: cute.Tensor, + split_kv: cutlass.Int32, + cluster_shape_mnk: Tuple[int, int, int], + max_active_clusters: int, + is_persistent: bool, + ) -> Tuple[MLAStaticTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid shape 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] + + :return: Tile scheduler parameters and grid shape. + :rtype: tuple[MLAStaticTileSchedulerParams, tuple[int, int, int]] + """ + o_shape = o.shape + tile_sched_params = create_mla_static_tile_scheduler_params( + is_persistent, + cute.size(o_shape[3]), + cute.size(o_shape[2]), + cluster_shape_mnk, + split_kv, + ) + grid = MLAStaticTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def get_workspace_size( + H: int, + S: int, + D: int, + B: int, + split_kv: int, + acc_dtype: Type[cutlass.Numeric], + ) -> int: + """Get the extra workspace(device memory) size for the MLA kernel when split_kv is not 1. + + :param H: The height of the output tensor C + :type H: int + :param S: The sequence length of the output tensor C + :type S: int + :param D: The depth of the output tensor C + :type D: int + :param B: The batch size of the output tensor C + :type B: int + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + + :return: The workspace size for the MLA kernel + :rtype: int + """ + if split_kv == 1: + return 0 + return B * H * S * split_kv * (D + 1) * acc_dtype.width // 8 + + @cute.jit + def initialize_workspace( + self, + H: cutlass.Int32, + D: cutlass.Int32, + S: cutlass.Int32, + B: cutlass.Int32, + split_kv: cutlass.Int32, + acc_dtype: Type[cutlass.Numeric], + workspace: cute.Tensor, + ) -> tuple[cute.Tensor, cute.Tensor]: + """Initialize the workspace for the MLA kernel. Construct the intermediate tensors + acc_o and acc_lse. + + :param H: The height of the output tensor C + :type H: cutlass.Int32 + :param D: The depth of the output tensor C + :type D: cutlass.Int32 + :param S: The sequence length of the output tensor C + :type S: cutlass.Int32 + :param B: The batch size of the output tensor C + :type B: cutlass.Int32 + :param split_kv: The split key-value of the output tensor C + :type split_kv: cutlass.Int32 + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + :param workspace: The workspace tensor + :type workspace: cute.Tensor + + :return: The output tensor C and the workspace tensor + :rtype: tuple[cute.Tensor, cute.Tensor] + """ + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): + align = 256 // self.q_dtype.width + acc_o_layout = cute.make_layout( + (H, split_kv, D, S, B), + stride=( + cute.assume(split_kv * D, align), + cute.assume(D, align), + 1, + cute.assume(split_kv * H * D, align), + cute.assume(H * split_kv * S * D, align), + ), + ) + acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype) + acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) + acc_lse_layout = cute.make_layout( + (H, split_kv, S, B), + stride=(split_kv, 1, H * split_kv, H * split_kv * S), + ) + acc_lse_iter = cute.recast_ptr( + workspace.iterator + cute.cosize(acc_o_layout) * acc_dtype.width // 8, + dtype=acc_dtype, + ) + acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout) + return acc_o, acc_lse + + @staticmethod + def can_implement( + B: int, + S: int, + K: int, + H: int, + L: int, + R: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + ) -> bool: + """Check if the MLA kernel can be implemented. + + :param B: The batch size of the output tensor C + :type B: int + :param S: The sequence length of the output tensor C + :type S: int + :param K: The width of the output tensor KV + :type K: int + :param H: The number of heads of the output tensor C + :type H: int + :param L: The number of latent dimensions of the tensor KV + :type L: int + :param R: The number of rope dimensions of the tensor C_rope + :type R: int + :param in_dtype: The data type of the input tensor + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: The data type of the output tensor + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: The data type of the log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: The tile shape of the query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: The tile shape of the probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: The page size of the page table + :type page_size: int + + :return: Whether the MLA kernel can be implemented + :rtype: bool + """ + if L != 512 or R != 64: + return False + if in_dtype not in [cutlass.Float16]: + return False + if out_dtype not in [cutlass.Float16]: + return False + if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32: + return False + # page size equals 1 is prohibited by tma specification, not 128B aligned. + if mma_qk_tiler_mn[1] % page_size != 0 or page_size == 1: + return False + if mma_qk_tiler_mn[0] != mma_pv_tiler_mn[0] or mma_qk_tiler_mn[0] != 128: + return False + if is_var_split_kv and not is_var_seq: + return False + if H > 128 or (H < 128 and split_kv != 1): + return False + if S < 1 or S > 4: + return False + if K <= 0: + return False + return True + + +def run( + batch_size: int, + seq_len_q: int, + seq_len_k: int, + num_heads: int, + latent_dim: int, + rope_dim: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + softmax_scale: float, + output_scale: float, + skip_correction_threshold: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool, + **kwargs, +): + """Execute Multi-Head Latent Attention (MLA) on Blackwell architecture and validate results. + + This function creates random input tensors for query latent/rope, compressed latent/rope, and value, + then performs the complete MLA computation pipeline. It supports configurable data types, tiling parameters, + page table, variable sequence length, and variable split_kv. Results can be validated against a PyTorch reference + implementation or run multiple times for performance measurement. + + :param batch_size: Batch size + :type batch_size: int + :param seq_len_q: Sequence length of Q + :type seq_len_q: int + :param seq_len_k: Sequence length of K + :type seq_len_k: int + :param num_heads: Number of heads + :type num_heads: int + :param latent_dim: dimension of query/compressed latent + :type latent_dim: int + :param rope_dim: dimension of query/compressed rope + :type rope_dim: int + :param in_dtype: Input data type for query/compressed latent/rope tensors + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: Output data type for attention output + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type for query-key matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Accumulator data type for log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: Matrix multiply accumulate tile shape (M, N) for query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: Matrix multiply accumulate tile shape (M, N) for probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: Split key-value + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: Page size of the page table + :type page_size: int + :param softmax_scale: Attention score scaling factor + :type softmax_scale: float + :param output_scale: Output scaling factor + :type output_scale: float + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param tolerance: Maximum acceptable error for validation + :type tolerance: float + :param warmup_iterations: Number of warmup iterations + :type warmup_iterations: int + :param iterations: Number of iterations to run for performance testing + :type iterations: int + :param skip_ref_check: Skip validation against reference implementation + :type skip_ref_check: bool + :param use_cold_l2: Whether to use cold L2 cache + :type use_cold_l2: bool + + :raises ValueError: If input shapes are incompatible or head dimension is unsupported + :raises RuntimeError: If GPU is unavailable for computation + """ + + print("Running Blackwell MLA test with:") + print(f" batch_size: {batch_size}") + print(f" seq_len_q: {seq_len_q}") + print(f" seq_len_k: {seq_len_k}") + print(f" num_heads: {num_heads}") + print(f" latent_dim: {latent_dim}") + print(f" rope_dim: {rope_dim}") + print(f" in_dtype: {in_dtype}") + print(f" out_dtype: {out_dtype}") + print(f" acc_dtype: {acc_dtype}") + print(f" mma_qk_tiler_mn: {mma_qk_tiler_mn}") + print(f" mma_pv_tiler_mn: {mma_pv_tiler_mn}") + print(f" split_kv: {split_kv}") + print(f" is_persistent: {is_persistent}") + print(f" is_var_seq: {is_var_seq}") + print(f" is_var_split_kv: {is_var_split_kv}") + print(f" page_size: {page_size}") + print(f" softmax_scale: {softmax_scale}") + print(f" output_scale: {output_scale}") + print(f" skip_correction_threshold: {skip_correction_threshold}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + + import torch + import cutlass.torch as cutlass_torch + + # Prepare pytorch tensors: Q, K, V (random from 0 to 2) and O (all zero) + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + if not BlackwellMultiHeadLatentAttentionForwardFP16.can_implement( + batch_size, + seq_len_q, + seq_len_k, + num_heads, + latent_dim, + rope_dim, + in_dtype, + out_dtype, + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + split_kv, + is_persistent, + is_var_seq, + is_var_split_kv, + page_size, + ): + raise TypeError( + f"Unsupported testcase {batch_size}, {seq_len_q}, {seq_len_k}, {num_heads}, {latent_dim}, {rope_dim}, {in_dtype}, {out_dtype}, {acc_dtype}, {lse_dtype}, {mma_qk_tiler_mn}, {mma_pv_tiler_mn}, {split_kv}, {is_persistent}, {is_var_seq}, {is_var_split_kv}, {page_size}" + ) + + torch.manual_seed(1111) + + def create_data_tensor( + B, + HK, + D, + dtype, + is_dynamic_layout=True, + page_table=None, + cache_seqs=None, + is_lse=False, + seq_len_q=None, + ): + shape = (B, HK, D) + if page_table is not None: + if cache_seqs is not None: + max_seq_len = torch.max(cache_seqs) + shape = (B * ceil_div(max_seq_len, page_size), page_size, D) + else: + shape = (B * ceil_div(HK, page_size), page_size, D) + + if seq_len_q is not None: + shape = (B, seq_len_q, HK, D) + + permute_order = (1, 2, 0) + stride_order = (2, 0, 1) + leading_dim = 1 + if is_lse: + shape = (B, seq_len_q, HK) + permute_order = (2, 1, 0) + stride_order = (2, 1, 0) + leading_dim = 0 + elif seq_len_q is not None: + permute_order = (2, 3, 1, 0) + stride_order = (3, 2, 0, 1) + leading_dim = 1 + + init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) + + torch_dtype = ( + cutlass_torch.dtype(dtype) if dtype != cutlass.Float8E4M3FN else torch.int8 + ) + + # Create dtype torch tensor (cpu) + torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch_dtype, + permute_order=permute_order, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=init_config, + ) + + # Create dtype torch tensor (gpu) + torch_tensor_gpu = torch_tensor_cpu.cuda() + + # Create f32 torch tensor (cpu) + f32_torch_tensor = torch_tensor_cpu.to(dtype=torch.float32) + + # Create dtype cute tensor (gpu) + cute_tensor = from_dlpack(torch_tensor_gpu, assumed_align=16) + cute_tensor.element_type = dtype + if is_dynamic_layout: + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + if not is_lse: + cute_tensor = cute_tensor.mark_compact_shape_dynamic( + mode=leading_dim, + stride_order=stride_order, + divisibility=(128 // dtype.width), + ) + + cute_tensor = cutlass_torch.convert_cute_tensor( + f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=is_dynamic_layout, + ) + + return f32_torch_tensor, cute_tensor, torch_tensor_gpu + + def create_cache_seqs(batch_size, seq_len_k, is_var_seq): + cache_seqs_ref = torch.ones(batch_size, dtype=torch.int32) * seq_len_k + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack(cache_seqs_gpu, assumed_align=16).mark_layout_dynamic() + if is_var_seq: + max_seq_len = seq_len_k + min_seq_len = int(seq_len_k * 0.8) + cache_seqs_ref = cutlass_torch.create_and_permute_torch_tensor( + (batch_size,), + torch.int32, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig( + min_val=min_seq_len, max_val=max_seq_len + 1 + ), + ) + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack( + cache_seqs_gpu, + assumed_align=16, + ).mark_layout_dynamic() + return cache_seqs_ref, cache_seqs, cache_seqs_gpu + + def create_page_table(batch_size, seq_len_k, is_var_seq, page_size): + max_seq_len = seq_len_k if not is_var_seq else torch.max(cache_seqs_ref) + page_count = ceil_div(max_seq_len, page_size) + page_table_ref = torch.empty([batch_size, page_count], dtype=torch.int32) + # use transposed index for page table to make sure the value is in bound of `batch_size * seq_len_block`. In practice, the value could be any positive values. This setting is only for testing purpose. + for b in range(batch_size): + for j in range(page_count): + page_table_ref[b, j] = b + j * batch_size + page_table_gpu = page_table_ref.permute(1, 0).cuda() + page_table = from_dlpack(page_table_gpu, assumed_align=16).mark_layout_dynamic( + leading_dim=0 + ) + return page_table_ref, page_table, page_table_gpu + + def create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ): + block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu = None, None, None + # check if split_kv is valid otherwise do auto setting of split_kv + if is_var_split_kv: + block_split_kvs_ref = torch.zeros([batch_size], dtype=torch.int32) + for b in range(batch_size): + block_split_kvs_ref[b] = ( + BlackwellMultiHeadLatentAttentionForwardFP16.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[b].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + ) + split_kv = torch.max(block_split_kvs_ref).item() + block_split_kvs_gpu = block_split_kvs_ref.cuda() + block_split_kvs = from_dlpack( + block_split_kvs_gpu, assumed_align=16 + ).mark_layout_dynamic() + elif split_kv <= 0: + split_kv = BlackwellMultiHeadLatentAttentionForwardFP16.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[0].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + return split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu + + def create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, split_kv, acc_dtype + ): + workspace_size = ( + BlackwellMultiHeadLatentAttentionForwardFP16.get_workspace_size( + num_heads, + seq_len_q, + latent_dim, + batch_size, + split_kv, + acc_dtype, + ) + ) + + workspace, workspace_torch = None, None + if workspace_size > 0: + workspace_torch = torch.empty([workspace_size], dtype=torch.int8).cuda() + workspace = from_dlpack(workspace_torch, assumed_align=32) + return workspace, workspace_torch + + cache_seqs_ref, cache_seqs, cache_seqs_torch = create_cache_seqs( + batch_size, seq_len_k, is_var_seq + ) + page_table_ref, page_table, page_table_torch = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size + ) + cluster_shape_mnk = (2, 1, 1) + hardware_info = utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mnk[0] * cluster_shape_mnk[1] + ) + split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_torch = ( + create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + ) + + q_latent_ref, q_latent, q_latent_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + q_rope_ref, q_rope, q_rope_torch = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + c_latent_ref, c_latent, c_latent_torch = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + c_rope_ref, c_rope, c_rope_torch = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + o_ref, o, o_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + lse_ref, lse, lse_torch = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, split_kv, acc_dtype + ) + + mla = BlackwellMultiHeadLatentAttentionForwardFP16( + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + max_active_clusters, + page_size, + skip_correction_threshold, + is_persistent, + is_var_seq, + is_var_split_kv, + ) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + stream = cuda.CUstream(torch_stream.cuda_stream) + + # compile mla kernel + compiled_mla = cute.compile( + mla, + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + options="--opt-level 2", + ) + + def torch_reference_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + softmax_scale=1.0, + output_scale=1.0, + ): + # expand and concat q_latent and q_rope to have the dimension of sequence length for q + q_ref = torch.cat([q_latent, q_rope], dim=1).permute(3, 2, 0, 1) + # expand and concat c_latent and c_rope to have the dimension of num_heads for k and v + page_count = page_table_ref.shape[1] + k_ref_paged = ( + torch.cat([c_latent, c_rope], dim=1) + .permute(2, 0, 1) + .reshape(batch_size * page_count, page_size, latent_dim + rope_dim) + ) + v_ref_paged = c_latent.permute(2, 0, 1).reshape( + batch_size * page_count, page_size, latent_dim + ) + + if is_var_seq: + max_seq_len = torch.max(cache_seqs_ref) + else: + max_seq_len = seq_len_k + + k_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim + rope_dim]) + v_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim]) + k_ref = torch.index_select( + k_ref_paged, 0, torch.flatten(page_table_ref) + ).reshape(batch_size, 1, -1, latent_dim + rope_dim)[:, :, :max_seq_len, :] + v_ref = torch.index_select( + v_ref_paged, 0, torch.flatten(page_table_ref) + ).reshape(batch_size, 1, -1, latent_dim)[:, :, :max_seq_len, :] + for b in range(batch_size): + k_ref[b, :, cache_seqs_ref[b] :, :] = 0 + v_ref[b, :, cache_seqs_ref[b] :, :] = 0 + import torch.nn.functional as F + + o_ref = F.scaled_dot_product_attention( + q_ref, + k_ref, + v_ref, + attn_mask=None, + dropout_p=0.0, + scale=softmax_scale, + is_causal=False, + ) + s_ref = torch.einsum("bhld,bhsd->bhls", q_ref, k_ref) + s_ref_max, s_ref_max_pos = torch.max(s_ref, dim=-1, keepdim=True) + softmax_scale_log2 = LOG2_E * softmax_scale + s_ref_sum = torch.sum( + torch.exp2((s_ref - s_ref_max) * softmax_scale_log2), dim=-1, keepdim=True + ) + + lse_ref = s_ref_max * softmax_scale_log2 + torch.log2(s_ref_sum) + lse_ref = lse_ref.squeeze(3).permute(2, 1, 0) + o_ref = o_ref * output_scale + o_ref = o_ref.permute(2, 3, 1, 0) + + return o_ref, lse_ref + + if skip_correction_threshold > 0.0: + print( + "Skipping correction verification since skip_correction_threshold is greater than 0.0..." + ) + skip_ref_check = True + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + torch.cuda.synchronize() + + print("Verifying results...") + if in_dtype == cutlass.Float8E4M3FN: + tolerance = 0.13 + o_ref, lse_ref = torch_reference_mla( + q_latent_ref, + q_rope_ref, + c_latent_ref, + c_rope_ref, + page_table, + cache_seqs, + softmax_scale, + output_scale, + ) + + if out_dtype in [cutlass.Float8E5M2, cutlass.Float8E4M3FN]: + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + cutlass.Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o, o_fp32) + o = o_fp32_torch.cpu() + ref_fp8, _ = cutlass_torch.cute_tensor_like( + torch.empty( + *o_ref.permute(3, 2, 0, 1).shape, dtype=torch.uint8 + ).permute(2, 3, 1, 0), + out_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + o_ref_gpu = o_ref.cuda() + o_ref_f32 = from_dlpack(o_ref_gpu).mark_layout_dynamic(leading_dim=1) + + # convert ref : f32 -> fp8 -> f32 + cute.testing.convert(o_ref_f32, ref_fp8) + cute.testing.convert(ref_fp8, o_ref_f32) + + o_ref = o_ref_gpu.cpu() + else: + o = o_torch.cpu().to(torch.float32) + lse = lse_torch.cpu() + lse_ref = lse_ref.to(cutlass.torch.dtype(lse_dtype)) + # Assert close results + torch.testing.assert_close(o, o_ref, atol=tolerance, rtol=1e-05) + torch.testing.assert_close(lse, lse_ref, atol=tolerance, rtol=1e-05) + print("Results verified successfully!") + + def generate_tensors(): + _, cache_seqs, _ = create_cache_seqs(batch_size, seq_len_k, is_var_seq) + _, page_table, _ = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size + ) + _split_kv, _, block_split_kvs, _ = create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + + _, q_latent, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, q_rope, _ = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + _, c_latent, _ = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, c_rope, _ = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, o, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, lse, _ = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, _split_kv, acc_dtype + ) + return testing.JitArguments( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + _split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + q_latent_torch.numel() * q_latent_torch.element_size() + + q_rope_torch.numel() * q_rope_torch.element_size() + + c_latent_torch.numel() * c_latent_torch.element_size() + + c_rope_torch.numel() * c_rope_torch.element_size() + + o_torch.numel() * o_torch.element_size() + + lse_torch.numel() * lse_torch.element_size() + + cache_seqs_torch.numel() * cache_seqs_torch.element_size() + ) + one_workspace_bytes += ( + page_table_torch.numel() * page_table_torch.element_size() + ) + if is_var_split_kv: + one_workspace_bytes += ( + block_split_kvs_torch.numel() * block_split_kvs_torch.element_size() + ) + if workspace_torch is not None: + one_workspace_bytes += ( + workspace_torch.numel() * workspace_torch.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + avg_time_us = testing.benchmark( + compiled_mla, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return avg_time_us # Return execution time in microseconds + + +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." + ) + + def parse_mma_tiler(s: str) -> Tuple[int, int, Tuple[int, int]]: + ret = parse_comma_separated_ints(s) + if len(ret) != 2: + raise argparse.ArgumentTypeError( + "Invalid format. Expected 2 comma-separated integers." + ) + return (ret[0], ret[1]) + + parser = argparse.ArgumentParser(description="Example of MLA on Blackwell.") + + parser.add_argument( + "--in_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + help="Input data type", + ) + + parser.add_argument( + "--out_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + help="Output data type", + ) + + parser.add_argument( + "--acc_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="Accumulator data type", + ) + + parser.add_argument( + "--lse_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="LSE data type", + ) + parser.add_argument( + "--mma_qk_tiler_mn", + type=parse_mma_tiler, + default=(128, 128), + help="MMA tile shape (H, K)", + ) + parser.add_argument( + "--mma_pv_tiler_mn", + type=parse_mma_tiler, + default=(128, 256), + help="MMA tile shape (H, D)", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--batch_size", + type=int, + default=1, + help="Batch size", + ) + + parser.add_argument( + "--seq_len_q", + type=int, + default=1, + help="Sequence length of Q", + ) + + parser.add_argument( + "--seq_len_k", + type=int, + default=128, + help="Sequence length of K/V", + ) + + parser.add_argument( + "--num_heads", + type=int, + default=128, + help="Number of heads of Q", + ) + + parser.add_argument( + "--latent_dim", + type=int, + default=512, + help="Latent dimension of Q/C", + ) + + parser.add_argument( + "--rope_dim", + type=int, + default=64, + help="Rope dimension of Q/C", + ) + + parser.add_argument( + "--is_var_seq", + action="store_true", + help="Use variable length of sequence length or not", + ) + + parser.add_argument( + "--is_var_split_kv", + action="store_true", + help="Use variable length of split kv or not", + ) + + parser.add_argument( + "--page_size", + type=int, + default=128, + help="Page size of page table", + ) + + parser.add_argument( + "--split_kv", + type=int, + default=-1, + help="Split KV setting", + ) + + parser.add_argument( + "--softmax_scale", + type=float, + default=0.0416, + help="Scaling factor to scale softmax", + ) + + parser.add_argument( + "--output_scale", + type=float, + default=1.0, + help="Scaling factor to scale output", + ) + + parser.add_argument( + "--skip_correction_threshold", + type=float, + default=0.0, + help="Skip correction threshold", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-02, help="Tolerance for validation" + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + help="Use cold L2 cache", + ) + + args = parser.parse_args() + + run( + args.batch_size, + args.seq_len_q, + args.seq_len_k, + args.num_heads, + args.latent_dim, + args.rope_dim, + args.in_dtype, + args.out_dtype, + args.acc_dtype, + args.lse_dtype, + args.mma_qk_tiler_mn, + args.mma_pv_tiler_mn, + args.split_kv, + args.is_persistent, + args.is_var_seq, + args.is_var_split_kv, + args.page_size, + args.softmax_scale, + args.output_scale, + args.skip_correction_threshold, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py b/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py new file mode 100644 index 00000000..e6383ef8 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py @@ -0,0 +1,4341 @@ +# 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 sys +import argparse +import math +from typing import Type, Tuple, Optional +from types import SimpleNamespace + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode +import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.arch import Arch +from cutlass.cutlass_dsl import BaseDSL + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mla.mla_helpers import ( + ceil_div, + MAX_SPLITS, + LOG2_E, + MLAStaticTileScheduler, + MLAStaticTileSchedulerParams, + create_mla_static_tile_scheduler, + create_mla_static_tile_scheduler_params, +) + +""" +A Multi-Head Latent Attention (MLA) example using fp8 as input/output for the NVIDIA Blackwell SM100 architecture using CUTE DSL + +This example demonstrates an implementation of inference of multi-head latent attention using a TMA + Blackwell +SM100 TensorCore warp-specialized persistent kernel. The implementation integrates the (Qc + Qr)*(Kc + Kr)^T +matrix multiplication, softmax normalization, and softmax((Qc + Qr)*(Kc + Kr)^T)*Vc into a single kernel. +The kernel provides support for page table storage and variable-length KV cache sequences. It implements KV splitting +functionality to minimize latency when processing long KV sequences. + +The kernel implements key optimizations including: +- Warp specialization for different computation phases (load, MMA, softmax, correction, epilogue) +- Pipeline stages between different warps for overlapping computation and memory access +- Support for different precision data types +- Two sub-kernels (split KV kernel and reduction kernel) that enable split KV processing + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mla_fp8.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float8E4M3FN --out_dtype Float8E4M3FN \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent + +The above example runs Multi-Head Latent Attention (MLA) with the following configuration: +- Batch size: 4 +- Sequence length of Q: 1 +- Sequence length of K: 1024 +- Latent dimension: 512 +- RoPE dimension: 64 +- Number of heads: 128 +- Data types: Float8E4M3FN (input), Float8E4M3FN (output), Float32 (accumulation and LSE) + +It utilizes page table storage for the KV cache and enables both variable-length KV cache sequences +and variable split KV processing with persistent scheduling. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mla_fp8.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float8E4M3FN --out_dtype Float8E4M3FN \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent --warmup_iterations 3 \ + --iterations 10 --skip_ref_check + +Constraints for this example: +* Data type requirements: + - Input/output: Float8E4M3FN + - Accumulation and LSE: Float32 +* Fixed architecture parameters: + - Number of attention heads: 128 + - Latent dimension: 512 + - RoPE dimension: 64 +* Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) +* Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) +* Query sequence length must be 1-4 +* Only supports 2-CTA instructions +* Variable sequence length requires page table storage enabled +""" + + +class BlackwellMultiHeadLatentAttentionForwardFP8: + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + max_active_clusters: int, + page_size: int, + skip_correction_threshold: float, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + ): + """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. + + :param acc_dtype: Data type for accumulation S and O + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Data type for output LSE + :type lse_dtype: Type[cutlass.Numeric] + :param mma_s_tiler: The (H, K) tile shape of the MMA instruction for S + :type mma_s_tiler: Tuple[int, int] + :param mma_p_tiler: The (H, D) tile shape of the MMA instruction for P + :type mma_p_tiler: Tuple[int, int] + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: int + :param page_size: The page size + :type page_size: int + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split KV + :type is_var_split_kv: bool + """ + + self.latent_dim = 512 + self.rope_dim = 64 + self.acc_dtype = acc_dtype + self.lse_dtype = lse_dtype + self.mma_qk_tiler_mn = mma_qk_tiler_mn + self.mma_pv_tiler_mn = mma_pv_tiler_mn + self.max_active_clusters = max_active_clusters + self.skip_correction_threshold = skip_correction_threshold + self.is_persistent = is_persistent + self.page_size = page_size + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv + self.cluster_shape_mnk = (2, 1, 1) + self.use_2cta_instrs = True + # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), + # while warps 2-3 handle accumulation for second half [n/2, n) + self.warps_in_n = 2 + self.num_compute_warps = 4 + self.threads_per_warp = 32 + mma_qk_tiler_k = self.rope_dim * 2 + self.mma_qk_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + mma_qk_tiler_k, + ) + self.mma_qk_rope_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + self.rope_dim, + ) + self.mma_pv_tiler = ( + self.mma_pv_tiler_mn[0], + self.mma_pv_tiler_mn[1], + self.mma_qk_tiler[1] * self.mma_qk_tiler[2] // self.mma_pv_tiler_mn[1], + ) + self.iterations_qk_latent = self.latent_dim // self.mma_qk_tiler[2] + self.iterations_qk_rope = 1 + self.iterations_qk = self.iterations_qk_latent + self.iterations_qk_rope + self.iterations_pv_k = self.mma_qk_tiler[1] // self.mma_pv_tiler[2] + self.iterations_pv_n = self.latent_dim // self.mma_pv_tiler[1] + + # Set specialized warp ids + self.compute_warp_ids = (0, 1, 2, 3) + self.correction_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + self.load_tma_k_warp_id = 9 + self.load_tma_v_warp_id = 10 + self.empty_warp_ids = (11,) + self.threads_per_cta = self.threads_per_warp * len( + ( + self.mma_warp_id, + self.load_tma_k_warp_id, + self.load_tma_v_warp_id, + *self.compute_warp_ids, + *self.correction_warp_ids, + *self.empty_warp_ids, + ) + ) + + # register settings + self.softmax_reg_num = 192 + self.correction_reg_num = 256 + self.other_reg_num = 48 + # Named barriers + self.tmem_ptr_sync_bar = pipeline.NamedBarrier( + barrier_id=1, + num_threads=( + self.threads_per_warp + + self.threads_per_warp * self.num_compute_warps * 2 + ), + ) + self.softmax_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=2, num_threads=(self.threads_per_warp * self.num_compute_warps) + ) + self.epilogue_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=3, num_threads=(self.threads_per_warp * self.num_compute_warps) + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the MLA kernel operation. + + This method initializes and configures various attributes required for the + execution of the multi-head latent attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.load_q_stage = 1 + self.load_k_stage = 3 + self.load_v_stage = 2 + self.mma_s_stage = 2 + self.p_mma_stage = 2 + self.p_cor_stage = 2 + self.mma_o_stage = 2 + + self.tmem_o_offset = self.mma_s_stage * self.mma_qk_tiler[1] // self.warps_in_n + self.correction_factor_offset = ( + self.tmem_o_offset + self.latent_dim // self.warps_in_n + ) + + @cute.jit + def __call__( + self, + q_latent: cute.Tensor, + q_rope: cute.Tensor, + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_table: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: Optional[cute.Tensor], + block_split_kvs: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + stream: cuda.CUstream, + ): + """Execute the Multi-Head Latent Attention operation on the provided tensors. + + The method handles: + 1. Initialization of workspace for temporary split KV buffers + 2. Validation of tensor data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch(split KV kernel and reduction kernel) with appropriate parameters + + :param q_latent: The query tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type q_latent: cute.Tensor + :param q_rope: The query RoPE tensor with shape [num_head, rope_dim, seq_len_q, batch_size] + :type q_rope: cute.Tensor + :param c_latent: The key tensor with shape [seq_len_k, latent_dim, batch_size] + :type c_latent: cute.Tensor + :param c_rope: The key RoPE tensor with shape [seq_len_k, rope_dim, batch_size] + :type c_rope: cute.Tensor + :param page_table: The page table tensor with shape [page_count, batch_size] + :type page_table: cute.Tensor + :param o: The output tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type o: cute.Tensor + :param lse: The LSE tensor with shape [num_head, seq_len_q, batch_size] + :type lse: cute.Tensor + :param workspace: The workspace tensor with 1-d shape prepared for acc_o and acc_lse + :type workspace: cute.Tensor + :param split_kv: The scalar factor for split KV + :type split_kv: cutlass.Int32 + :param cache_seqs: The cache sequences tensor with shape [batch_size] + :type cache_seqs: cute.Tensor + :param block_split_kvs: The block split KV tensor with shape [batch_size] + :type block_split_kvs: cute.Tensor + :param softmax_scale: The scale factor for softmax + :type softmax_scale: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param stream: The CUDA stream to execute the kernel on + :type stream: cuda.CUstream + + :raises TypeError: If tensor data types don't match or aren't supported + """ + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q_latent.element_type + self.k_dtype = c_latent.element_type + self.v_dtype = c_latent.element_type + self.o_dtype = o.element_type + + # check type consistency + if cutlass.const_expr( + self.q_dtype != self.k_dtype or self.q_dtype != self.v_dtype + ): + raise TypeError( + f"Type mismatch: {self.q_dtype} != {self.k_dtype} or {self.q_dtype} != {self.v_dtype}" + ) + # check leading dimensions of input/output + if cutlass.const_expr(q_latent.stride[1] != 1 or q_rope.stride[1] != 1): + raise ValueError("q_latent and q_rope must have leading dimension 1") + if cutlass.const_expr(c_latent.stride[1] != 1 or c_rope.stride[1] != 1): + raise ValueError("c_latent and c_rope must have leading dimension 1") + if cutlass.const_expr(o.stride[1] != 1): + raise ValueError("o must have leading dimension 1") + if cutlass.const_expr(lse.stride[0] != 1): + raise ValueError("lse must have leading dimension 0") + + acc_o, acc_lse = self.initialize_workspace( + q_latent.shape[0], + q_latent.shape[1], + q_latent.shape[2], + q_latent.shape[3], + split_kv, + self.acc_dtype, + workspace, + ) + + c_latent_tranpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor( + c_latent.iterator, c_latent_tranpose_layout + ) + + self.q_major_mode = OperandMajorMode.K + self.k_major_mode = OperandMajorMode.K + self.v_major_mode = OperandMajorMode.MN + + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + # the intermediate tensor p is from smem & k-major + p_major_mode = OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.acc_dtype, + cta_group, + self.mma_qk_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.acc_dtype, + cta_group, + self.mma_pv_tiler[:2], + ) + + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + + self.epi_tile = self.mma_pv_tiler[:2] + + q_latent_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_tiler, + self.q_dtype, + (self.iterations_qk_latent * self.load_q_stage), + ) + q_latent_smem_layout_staged = cute.logical_divide( + q_latent_smem_layout_staged, (None, None, None, self.iterations_qk_latent) + ) + q_rope_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.q_dtype, + self.load_q_stage, + ) + + kc_latent_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_tiler, + self.k_dtype, + (self.iterations_qk_latent * self.load_k_stage), + ) + kc_page_tile_size = min( + self.page_size, qk_tiled_mma.op.shape_mnk[0] // qk_tiled_mma.thr_id.shape + ) + kc_latent_smem_layout_staged = cute.logical_divide( + kc_latent_smem_layout_staged, (None, None, None, self.iterations_qk_latent) + ) + + kc_latent_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + (self.mma_qk_tiler[0] // qk_tiled_mma.thr_id.shape, self.mma_qk_tiler[2]), + self.k_dtype, + (self.iterations_qk_latent * self.load_k_stage), + ) + kc_latent_smem_layout_for_tma = cute.tiled_divide( + kc_latent_smem_layout_for_tma, (kc_page_tile_size, self.mma_qk_tiler[2]) + ) + kc_latent_smem_layout_for_tma = cute.logical_divide( + kc_latent_smem_layout_for_tma, (None, None, None, self.iterations_qk_latent) + ) + + kc_rope_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.k_dtype, + self.load_k_stage, + ) + kc_rope_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + ( + self.mma_qk_rope_tiler[0] // qk_tiled_mma.thr_id.shape, + self.mma_qk_rope_tiler[2], + ), + self.k_dtype, + (self.iterations_qk_rope * self.load_k_stage), + ) + kc_rope_smem_layout_for_tma = cute.tiled_divide( + kc_rope_smem_layout_for_tma, (kc_page_tile_size, self.mma_qk_rope_tiler[2]) + ) + + p_smem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.mma_pv_tiler, + self.q_dtype, + (self.iterations_pv_k * self.p_mma_stage), + ) + p_smem_layout_staged = cute.logical_divide( + p_smem_layout_staged, (None, None, None, self.iterations_pv_k) + ) + + vc_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.mma_pv_tiler, + self.v_dtype, + (self.iterations_pv_k * self.iterations_pv_n * self.load_v_stage), + ) + vc_smem_layout_staged = cute.logical_divide( + cute.logical_divide( + vc_smem_layout_staged, + (None, None, None, self.iterations_pv_k * self.iterations_pv_n), + ), + (None, None, None, (self.iterations_pv_n, None)), + ) + vc_page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + vc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.MN, + (self.mma_pv_tiler[1] // pv_tiled_mma.thr_id.shape, self.mma_pv_tiler[2]), + self.v_dtype, + (self.iterations_pv_k * self.iterations_pv_n * self.load_v_stage), + ) + vc_smem_layout_for_tma = cute.tiled_divide( + vc_smem_layout_for_tma, + ( + pv_tiled_mma.op.shape_mnk[1] // pv_tiled_mma.thr_id.shape, + vc_page_tile_size, + ), + ) + vc_smem_layout_for_tma = cute.logical_divide( + cute.logical_divide( + vc_smem_layout_for_tma, + (None, None, None, self.iterations_pv_k * self.iterations_pv_n), + ), + (None, None, None, (self.iterations_pv_n, None)), + ) + # TMA load for Q latent and rope + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + + q_smem_layout = cute.select(q_latent_smem_layout_staged, mode=[0, 1, 2]) + + tma_atom_q_latent, tma_tensor_q_latent = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_latent, + q_smem_layout, + self.mma_qk_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + q_rope_smem_layout = cute.select(q_rope_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q_rope, tma_tensor_q_rope = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_rope, + q_rope_smem_layout, + self.mma_qk_rope_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + # TMA load for c latent and k rope + kc_smem_layout = cute.select(kc_latent_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent, tma_tensor_c_latent = self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + kc_rope_smem_layout = cute.select(kc_rope_smem_layout_for_tma, mode=[0]) + tma_atom_c_rope, tma_tensor_c_rope = self.make_paged_tiled_tma_atom( + tma_load_op, + c_rope, + kc_rope_smem_layout, + (self.mma_qk_rope_tiler[1], self.mma_qk_rope_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + + # TMA load for c latent transpose + vc_smem_layout = cute.select(vc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent_transpose, tma_tensor_c_latent_transpose = ( + self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent_transpose, + vc_smem_layout, + (self.mma_pv_tiler[1], self.mma_pv_tiler[2]), + pv_tiled_mma, + is_k_load=False, + ) + ) + + q_latent_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_smem_layout) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_latent + ) + q_rope_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_rope_smem_layout) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_rope + ) + kc_latent_copy_size = ( + cute.size_in_bytes( + self.k_dtype, + cute.select(kc_latent_smem_layout_staged, mode=[0, 1, 2]), + ) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_latent + ) + kc_rope_copy_size = ( + cute.size_in_bytes( + self.k_dtype, + cute.select(kc_rope_smem_layout_staged, mode=[0, 1, 2]), + ) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_rope + ) + vc_copy_size = ( + cute.size_in_bytes( + self.v_dtype, cute.select(vc_smem_layout_staged, mode=[0, 1, 2]) + ) + * cute.size(pv_tiled_mma.thr_id.shape) + * self.iterations_pv_n + * self.iterations_pv_k + ) + + self.tma_copy_q_bytes = q_latent_copy_size + q_rope_copy_size + self.tma_copy_kc_bytes = kc_latent_copy_size + kc_rope_copy_size + self.tma_copy_vc_bytes = vc_copy_size + + tile_sched_params, grid = self._compute_grid( + o, + split_kv, + self.cluster_shape_mnk, + self.max_active_clusters, + self.is_persistent, + ) + + @cute.struct + class SplitKVKernelSharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_q_stage * 2] + load_k_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_k_stage * 2] + load_v_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_v_stage * 2] + mma_s_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_s_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_mma_stage * 2] + p_cor_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_cor_stage * 2] + mma_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_o_stage * 2] + + # Smem tensors + smem_p: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, cute.cosize(p_smem_layout_staged)], + 1024, + ] + smem_kc_latent: cute.struct.Align[ + cute.struct.MemRange[ + self.k_dtype, cute.cosize(kc_latent_smem_layout_staged) + ], + 1024, + ] + + smem_kc_rope: cute.struct.Align[ + cute.struct.MemRange[ + self.k_dtype, cute.cosize(kc_rope_smem_layout_staged) + ], + 1024, + ] + smem_q_latent: cute.struct.Align[ + cute.struct.MemRange[ + self.q_dtype, cute.cosize(q_latent_smem_layout_staged) + ], + 1024, + ] + smem_q_rope: cute.struct.Align[ + cute.struct.MemRange[ + self.q_dtype, cute.cosize(q_rope_smem_layout_staged) + ], + 1024, + ] + smem_vc: cute.struct.Align[ + cute.struct.MemRange[self.v_dtype, cute.cosize(vc_smem_layout_staged)], + 1024, + ] + softmax_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp + ] + epilogue_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp + ] + + # Tmem dealloc cluster barrier + tmem_dealloc_mbar_ptr: cutlass.Int64 + + # Tmem holding buffer + tmem_holding_buf: cutlass.Int32 + + softmax_scale_log2 = softmax_scale * LOG2_E + + self.split_kv_kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q_latent, + tma_tensor_q_latent, + tma_atom_q_rope, + tma_tensor_q_rope, + tma_atom_c_latent, + tma_tensor_c_latent, + tma_atom_c_rope, + tma_tensor_c_rope, + tma_atom_c_latent_transpose, + tma_tensor_c_latent_transpose, + page_table, + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale_log2, + output_scale, + q_latent_smem_layout_staged, + q_rope_smem_layout_staged, + kc_latent_smem_layout_staged, + kc_rope_smem_layout_staged, + p_smem_layout_staged, + vc_smem_layout_staged, + kc_latent_smem_layout_for_tma, + kc_rope_smem_layout_for_tma, + vc_smem_layout_for_tma, + cta_layout_vmnk, + tile_sched_params, + SplitKVKernelSharedStorage, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + smem=SplitKVKernelSharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + if cutlass.const_expr(acc_o is not None): + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + ).launch( + grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), + block=[self.threads_per_warp * self.num_compute_warps, 1, 1], + smem=MAX_SPLITS * self.acc_dtype.width // 8, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.jit + def make_paged_tiled_tma_atom( + self, + tma_load_op: cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp, + gmem: cute.Tensor, + smem_layout: cute.Layout, + mma_tiler, + tiled_mma: cute.TiledMma, + is_k_load: bool, + ): + ident = cute.make_identity_layout(gmem.shape) + g_tile = cute.composition(ident, mma_tiler) + cta_mn = mma_tiler[0] // tiled_mma.thr_id.shape + cta_v_map = cute.flat_divide(g_tile, (cta_mn,)) + cta_v_map = cute.select(cta_v_map, mode=[0, 2]) + page_tile_size = ( + min(self.page_size, cta_mn) + if is_k_load + else min(self.page_size, mma_tiler[1]) + ) + cta_v_map = cute.zipped_divide( + cta_v_map, + (page_tile_size, mma_tiler[1]) if is_k_load else (cta_mn, page_tile_size), + ) + cta_v_map = cute.select(cta_v_map, mode=[0]) + from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir + + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( + gmem.value, + smem_layout.value, + cta_v_map, + tma_load_op._to_ir(), + num_multicast=1, + ) + return ( + cute.CopyAtom( + tma_load_op, cpasync.CopyBulkTensorTileG2SNonExecTrait(res[0]) + ), + res[1], + ) + + @cute.kernel + def split_kv_kernel( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tma_atom_q_latent: Optional[cute.CopyAtom], + mQL: cute.Tensor, + tma_atom_q_rope: Optional[cute.CopyAtom], + mQR: cute.Tensor, + tma_atom_c_latent: Optional[cute.CopyAtom], + mCL: cute.Tensor, + tma_atom_c_rope: Optional[cute.CopyAtom], + mKR: cute.Tensor, + tma_atom_c_latent_transpose: Optional[cute.CopyAtom], + mCLT: cute.Tensor, + mPT: cute.Tensor, + mO: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + mAccO: Optional[cute.Tensor], + mAccLSE: Optional[cute.Tensor], + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + q_latent_smem_layout_staged: cute.ComposedLayout, + q_rope_smem_layout_staged: cute.ComposedLayout, + kc_latent_smem_layout_staged: cute.ComposedLayout, + kc_rope_smem_layout_staged: cute.ComposedLayout, + p_smem_layout_staged: cute.ComposedLayout, + vc_smem_layout_staged: cute.ComposedLayout, + kc_latent_smem_layout_for_tma: Optional[cute.ComposedLayout], + kc_rope_smem_layout_for_tma: Optional[cute.ComposedLayout], + vc_smem_layout_for_tma: Optional[cute.ComposedLayout], + cta_layout_vmnk: cute.Layout, + tile_sched_params: MLAStaticTileSchedulerParams, + SharedStorage: cutlass.Constexpr, + ): + """The device split_kv kernel implementation of the Multi-Head Latent Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the MLA computation: + 1. Load warp: Loads Q/C latent/rope data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Compute warps: Compute softmax and do rescaling on accumulators, and store the intermediate/final results + to global memory + + The kernel produces either intermediate or final results of the MLA computation based on the split_kv parameter. + When split_kv is 1, the kernel generates the final results directly. Otherwise, it produces intermediate results + that will later be combined by a reduction kernel. + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases. + + :param tiled_mma_qk: Tiled MMA for Q*K^T + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: Tiled MMA for P*V + :type tiled_mma_pv: cute.TiledMma + :param tma_atom_q_latent: TMA copy atom for query latent tensor + :type tma_atom_q_latent: cute.CopyAtom + :param mQL: query latent tensor + :type mQL: cute.Tensor + :param tma_atom_q_rope: TMA copy atom for query rope tensor + :type tma_atom_q_rope: cute.CopyAtom + :param mKR: Compressed rope tensor + :type mKR: cute.Tensor + :param tma_atom_c_latent: TMA copy atom for c latent tensor + :type tma_atom_c_latent: cute.CopyAtom + :param mCL: Compressed latent tensor + :type mCL: cute.Tensor + :param tma_atom_c_rope: TMA copy atom for c rope tensor + :type tma_atom_c_rope: cute.CopyAtom + :param mCLT: Compressed latent transpose tensor + :type mCLT: cute.Tensor + :param mPT: Page table tensor + :type mPT: cute.Tensor + :param mO: Output tensor + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor + :type mLSE: cute.Tensor + :param mAccO: Intermediate accumulator output tensor + :type mAccO: cute.Tensor + :param mAccLSE: Intermediate accumulator log-sum-exp tensor + :type mAccLSE: cute.Tensor + :param split_kv: The split_kv parameter + :type split_kv: cutlass.Int32 + :param cache_seqs: The variable sequence length tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: The per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param softmax_scale_log2: The log2 scale factor for softmax + :type softmax_scale_log2: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param q_latent_smem_layout_staged: Shared memory layout for query tensor + :type q_latent_smem_layout_staged: cute.ComposedLayout + :param q_rope_smem_layout_staged: Shared memory layout for query rope tensor + :type q_rope_smem_layout_staged: cute.ComposedLayout + :param kc_latent_smem_layout_staged: Shared memory layout for key tensor + :type kc_latent_smem_layout_staged: cute.ComposedLayout + :param kc_rope_smem_layout_staged: Shared memory layout for key rope tensor + :type kc_rope_smem_layout_staged: cute.ComposedLayout + :param p_smem_layout_staged: Shared memory layout for probability matrix + :type p_smem_layout_staged: cute.ComposedLayout + :param vc_smem_layout_staged: Shared memory layout for value tensor + :type vc_smem_layout_staged: cute.ComposedLayout + :param cta_layout_vmnk: Layout for compute threads + :type cta_layout_vmnk: cute.Layout + :param tile_sched_params: Scheduling parameters for work distribution + :type tile_sched_params: MLAStaticTileSchedulerParams + :param SharedStorage: Shared storage for the kernel + :type SharedStorage: cutlass.Constexpr + """ + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + + # Prefetch tma descriptor + if warp_idx == self.mma_warp_id: + cpasync.prefetch_descriptor(tma_atom_q_latent) + cpasync.prefetch_descriptor(tma_atom_q_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent) + cpasync.prefetch_descriptor(tma_atom_c_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent_transpose) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_bar, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + load_q_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_q_stage, + self.tma_copy_q_bytes, + ) + load_k_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_k_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_k_stage, + self.tma_copy_kc_bytes, + ) + load_v_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_v_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_v_stage, + self.tma_copy_vc_bytes, + ) + mma_s_pipeline = self.make_and_init_mma_s_pipeline( + storage.mma_s_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + p_mma_pipeline = self.make_and_init_p_mma_pipeline( + storage.p_mma_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + p_cor_pipeline = self.make_and_init_p_cor_pipeline( + storage.p_cor_mbar_ptr.data_ptr() + ) + mma_o_pipeline = self.make_and_init_mma_o_pipeline( + storage.mma_o_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) + + # Generate smem tensor Q/KC/VC/exchange + # (MMA, MMA_H, MMA_R, PIPE) + sQ = storage.smem_q_latent.get_tensor( + q_latent_smem_layout_staged.outer, swizzle=q_latent_smem_layout_staged.inner + ) + sQ_rope = storage.smem_q_rope.get_tensor( + q_rope_smem_layout_staged.outer, swizzle=q_rope_smem_layout_staged.inner + ) + # (MMA, MMA_K, MMA_R, PIPE) + sKC = storage.smem_kc_latent.get_tensor( + kc_latent_smem_layout_staged.outer, + swizzle=kc_latent_smem_layout_staged.inner, + ) + sKC_rope = storage.smem_kc_rope.get_tensor( + kc_rope_smem_layout_staged.outer, swizzle=kc_rope_smem_layout_staged.inner + ) + sKC_for_tma = storage.smem_kc_latent.get_tensor( + kc_latent_smem_layout_for_tma.outer, + swizzle=kc_latent_smem_layout_for_tma.inner, + ) + sKC_rope_for_tma = storage.smem_kc_rope.get_tensor( + kc_rope_smem_layout_for_tma.outer, swizzle=kc_rope_smem_layout_for_tma.inner + ) + # (MMA, MMA_D, MMA_K, PIPE) + sVC = storage.smem_vc.get_tensor( + vc_smem_layout_staged.outer, swizzle=vc_smem_layout_staged.inner + ) + sVC_for_tma = storage.smem_vc.get_tensor( + vc_smem_layout_for_tma.outer, swizzle=vc_smem_layout_for_tma.inner + ) + # (MMA, MMA_H, MMA_K) + sP = storage.smem_p.get_tensor( + p_smem_layout_staged.outer, swizzle=p_smem_layout_staged.inner + ) + # (compute_threads,) + softmax_smem_exchange = storage.softmax_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp) + ) + epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load warps, including page table and data tensors + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.empty_warp_ids[0] and warp_idx <= self.empty_warp_ids[-1]: + cute.arch.setmaxregister_decrease(self.other_reg_num) + + if warp_idx == self.load_tma_k_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_q_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_q_stage + ) + load_k_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_k_stage + ) + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_k_pipeline=load_k_pipeline, + load_v_pipeline=load_v_pipeline, + mPT=mPT, + ) + tma_qk_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + tma_atom_q_latent=tma_atom_q_latent, + tma_atom_q_rope=tma_atom_q_rope, + tma_atom_c_latent=tma_atom_c_latent, + tma_atom_c_rope=tma_atom_c_rope, + mQL=mQL, + mQR=mQR, + mCL=mCL, + mKR=mKR, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC_for_tma, + sKC_rope=sKC_rope_for_tma, + ) + # Load tma + load_q_producer_state, load_k_producer_state = self.load_tma_qk( + tma_common_params, + tma_qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_k_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + load_q_pipeline.producer_tail(load_q_producer_state) + load_k_pipeline.producer_tail(load_k_producer_state) + + if warp_idx == self.load_tma_v_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_v_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_v_stage + ) + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_v_pipeline=load_v_pipeline, + mPT=mPT, + ) + tma_pv_params = SimpleNamespace( + tiled_mma_pv=tiled_mma_pv, + tma_atom_c_latent_transpose=tma_atom_c_latent_transpose, + mCLT=mCLT, + sVC=sVC_for_tma, + ) + # Load tma + load_v_producer_state = self.load_tma_v( + tma_common_params, + tma_pv_params, + k_index, + k_tile_count, + load_v_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_v_pipeline.producer_tail(load_v_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA warp + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + # Alloc tensor memory buffer + tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100")) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + load_q_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_q_stage + ) + load_k_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_k_stage + ) + load_v_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_v_stage + ) + mma_s_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_s_stage + ) + p_mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_mma_stage + ) + mma_o_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_o_stage + ) + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + mma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_k_pipeline=load_k_pipeline, + load_v_pipeline=load_v_pipeline, + tmem_ptr=tmem_ptr, + is_leader_cta=is_leader_cta, + L=mCL.shape[1], + ) + mma_qk_params = SimpleNamespace( + mma_s_pipeline=mma_s_pipeline, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC, + sKC_rope=sKC_rope, + ) + mma_pv_params = SimpleNamespace( + p_mma_pipeline=p_mma_pipeline, + mma_o_pipeline=mma_o_pipeline, + sP=sP, + sVC=sVC, + ) + ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma( + mma_common_params, + mma_qk_params, + mma_pv_params, + k_tile_count, + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + mma_s_pipeline.producer_tail(mma_s_producer_state) + mma_o_pipeline.producer_tail(mma_o_producer_state) + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Compute warp + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx >= self.compute_warp_ids[0] + and warp_idx <= self.compute_warp_ids[-1] + ): + cute.arch.setmaxregister_increase(self.softmax_reg_num) + mma_s_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_s_stage + ) + p_mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_mma_stage + ) + p_cor_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_cor_stage + ) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage + ) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=softmax_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, + p_cor_pipeline=p_cor_pipeline, + ) + compute_softmax_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + sP=sP, + mma_s_pipeline=mma_s_pipeline, + p_mma_pipeline=p_mma_pipeline, + softmax_scale_log2=softmax_scale_log2, + ) + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = ( + self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_cor_pipeline.producer_tail(p_cor_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # Correction warp + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx >= self.correction_warp_ids[0] + and warp_idx <= self.correction_warp_ids[-1] + ): + cute.arch.setmaxregister_increase(self.correction_reg_num) + p_cor_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_cor_stage + ) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage + ) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + 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: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=epilogue_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + H=mQL.shape[0], + tmem_ptr=tmem_ptr, + tidx=tidx, + tiled_mma_pv=tiled_mma_pv, + p_cor_pipeline=p_cor_pipeline, + mma_o_pipeline=mma_o_pipeline, + ) + compute_epilogue_params = SimpleNamespace( + output_scale=output_scale, + softmax_scale_log2=softmax_scale_log2, + mAccLSE=mAccLSE, + mLSE=mLSE, + ) + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + return + + @cute.kernel + def reduction_kernel( + self, + mO: cute.Tensor, + mLSE: cute.Tensor, + mAccO: cute.Tensor, + mAccLSE: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + ): + """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results + from multiple split_kv blocks into final outputs. + + :param mO: Output tensor for storing final results + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor for storing final LSE values + :type mLSE: cute.Tensor + :param mAccO: Accumulated output tensor from split_kv blocks + :type mAccO: cute.Tensor + :param mAccLSE: Accumulated LSE tensor from split_kv blocks + :type mAccLSE: cute.Tensor + :param split_kv: Number of split_kv blocks + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor (for variable split_kv) + :type block_split_kvs: cute.Tensor + """ + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + blk_coord = (bidx, bidy, bidz) + local_split_kv = ( + block_split_kvs[blk_coord[2]] if self.is_var_split_kv else split_kv + ) + k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) + local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) + + # Alloc shared memory + smem = utils.SmemAllocator() + storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) + lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) + smem_lse_scale = cute.make_tensor(lse_scale_ptr, cute.make_layout(MAX_SPLITS)) + + gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) + lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) + + local_lse = cute.make_rmem_tensor( + cute.make_layout(lse_per_thread), self.lse_dtype + ) + lse_max = -self.lse_dtype.inf + # find the max lse + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + local_lse[i] = ( + gLSE[split_kv_idx] + if cute.elem_less(split_kv_idx, local_split_kv) + else -self.lse_dtype.inf + ) + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) + lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 0.0 + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if not sum_lse == self.lse_dtype(0.0) or sum_lse != sum_lse + else self.lse_dtype.inf + ) + if tidx == 0: + mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[split_kv_idx] = cute.math.exp2( + local_lse[i] - global_lse, fastmath=True + ) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( + self.latent_dim, self.threads_per_warp * self.num_compute_warps + ) + gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] + rAccO = cute.make_rmem_tensor( + cute.make_layout(elements_per_thread), self.acc_dtype + ) + rO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), self.o_dtype) + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] + return + + @staticmethod + def get_split_kv( + B: int, S: int, K: int, mma_qk_tiler_mn: tuple, max_active_blocks: int + ) -> int: + """Get the proper split_kv value for the MLA kernel based on parameters. + + :param B: Batch size + :type B: int + :param S: Sequence length + :type S: int + :param K: Sequence length + :type K: int + :param mma_qk_tiler_mn: MLA tiling parameters + :type mma_qk_tiler_mn: tuple + :param max_active_blocks: Maximum number of active blocks + :type max_active_blocks: int + :return: Split_kv value + :rtype: int + """ + max_splits = ceil_div(K, mma_qk_tiler_mn[1]) + blocks_per_batch = max(1, max_active_blocks // B // (S * 2)) + split_heur = min(max_splits, blocks_per_batch) + k_waves = ceil_div(max_splits, split_heur) + split_wave_aware = ceil_div(max_splits, k_waves) + max_split_kv = 32 + return min(split_wave_aware, max_split_kv) + + @cute.jit + def get_k_tile_count( + self, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + blk_coord: cute.Coord, + ) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Get the current k_index, k_tile_count, and local split_kv value for the MLA kernel. + + :param split_kv: Split_kv value + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param blk_coord: Block coordinate + :type blk_coord: cute.Coord + :return: k_index, k_tile_count, split_kv + :rtype: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32] + """ + K = cache_seqs[blk_coord[2]] + if cutlass.const_expr(self.is_var_split_kv): + split_kv = block_split_kvs[blk_coord[2]] + + k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv) + k_index = blk_coord[3] * k_tile_per_cta + k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index) + return k_index, k_tile_count, split_kv + + @cute.jit + def load_tma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState | None = None, + load_k_producer_state: pipeline.PipelineState | None = None, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load wrap to load Q/K tensors. Updates the load qk producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_k_producer_state: The load k producer state + :type load_k_producer_state: pipeline.PipelineState + + :return: The load q producer state and load k producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for QK TMA load + # (bM, bK, rM, rK, rL) + mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) + gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) + mma_qk_tiler_mk_rope = cute.select(self.mma_qk_rope_tiler, mode=[0, 2]) + gQR = cute.flat_divide(qk_params.mQR, mma_qk_tiler_mk_rope) + + thr_mma_qk = qk_params.tiled_mma_qk.get_slice( + common_params.blk_coord[0] % cute.size(qk_params.tiled_mma_qk.thr_id) + ) + tSgQL = thr_mma_qk.partition_A(gQL) + tSgQR = thr_mma_qk.partition_A(gQR) + + cta_m = min( + qk_params.tiled_mma_qk.op.shape_mnk[0] + // qk_params.tiled_mma_qk.thr_id.shape, + self.page_size, + ) + page_tile_size = min(self.page_size, cta_m) + gCL = cute.tiled_divide(qk_params.mCL, (page_tile_size, self.mma_qk_tiler[2])) + tSgCL = ( + gCL[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] + if cta_m < self.page_size + else gCL[None, 0, None, None] + ) + gKR = cute.tiled_divide( + qk_params.mKR, (page_tile_size, self.mma_qk_rope_tiler[2]) + ) + tSgKR = ( + gKR[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] + if cta_m < self.page_size + else gKR[None, 0, None, None] + ) + # tma partition for q, k latent/rope + + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tQsQ, tQLgQL_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_latent, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ, 0, 3), + cute.group_modes(tSgQL, 0, 3), + ) + + tQsQ_rope, tQRgQR_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_rope, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ_rope, 0, 3), + cute.group_modes(tSgQR, 0, 3), + ) + tKCsKC, tCLgCL = cpasync.tma_partition( + qk_params.tma_atom_c_latent, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgCL, + ) + + tKCsKC_rope, tKRgKR = cpasync.tma_partition( + qk_params.tma_atom_c_rope, + 0, + cute.make_layout(1), + qk_params.sKC_rope, + tSgKR, + ) + + tQLgQL = tQLgQL_mkl[ + None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] + ] + tQRgQR = tQRgQR_mkl[ + None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] + ] + + # set extra params + common_params.mPT = mPT + qk_params.tQLgQL = tQLgQL + qk_params.tQRgQR = tQRgQR + qk_params.tCLgCL = tCLgCL + qk_params.tKRgKR = tKRgKR + qk_params.tQsQ = tQsQ + qk_params.tQsQ_rope = tQsQ_rope + qk_params.tKCsKC = tKCsKC + qk_params.tKCsKC_rope = tKCsKC_rope + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + load_q_producer_state, load_k_producer_state = self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_k_producer_state, + load_q=k_tile_count_init == k_tile_count, + ) + k_index += 1 + k_tile_count -= 1 + + return load_q_producer_state, load_k_producer_state + + @cute.jit + def load_tma_v( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_v_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load wrap to load V tensors. Updates the load v producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_v_producer_state: The load v producer state + :type load_v_producer_state: pipeline.PipelineState + + :return: The load v producer state + :rtype: pipeline.PipelineState + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for V TMA load + page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + gCLT = cute.flat_divide(v_params.mCLT, (self.mma_pv_tiler[1], page_tile_size)) + cta_n = self.mma_pv_tiler[1] // v_params.tiled_mma_pv.thr_id.shape + gCLT = cute.logical_divide(gCLT, (cta_n,))[ + (None, common_params.blk_coord[0]), None, None, None, None + ] + tOgCLT = cute.tiled_divide(gCLT, (cta_n, page_tile_size)) + tOgCLT = tOgCLT[None, 0, 0, None, None, None] + # tma partition for vc + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tVCsVC, tCLTgCLT = cpasync.tma_partition( + v_params.tma_atom_c_latent_transpose, + 0, + cute.make_layout(1), + v_params.sVC, + tOgCLT, + ) + + # set extra params + common_params.mPT = mPT + v_params.tCLTgCLT = tCLTgCLT + v_params.tVCsVC = tVCsVC + + while k_tile_count > 0: + load_v_producer_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index, + load_v_producer_state, + ) + k_index += 1 + k_tile_count -= 1 + return load_v_producer_state + + @cute.jit + def load_tma_qk_one_k_tile( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_k_producer_state: pipeline.PipelineState, + load_q: bool, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load one k-tile of Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_k_producer_state: The load kv producer state + :type load_k_producer_state: pipeline.PipelineState + :param load_q: Whether to load q + :type load_q: bool + + :return: The load q producer state and load kv producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = ceil_div( + self.mma_qk_tiler[1] // self.page_size, qk_params.tiled_mma_qk.thr_id.shape + ) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = ( + common_params.mPT[k_index] + if self.mma_qk_tiler[1] // self.page_size == 1 + else common_params.mPT[ + ( + k_index * qk_params.tiled_mma_qk.thr_id.shape + + common_params.blk_coord[0] + ) + * page_per_tile + + i + ] + ) + # load q once at first iteration + load_q_pipeline = common_params.load_q_pipeline + if load_q: + # get the mbar ptr from pipeline. + tma_bar_ptr = load_q_pipeline.producer_get_barrier(load_q_producer_state) + # expect the extra bytes for q. + load_q_pipeline.producer_acquire(load_q_producer_state) + for i in cutlass.range_constexpr(self.iterations_qk_latent): + # load q latent + cute.copy( + qk_params.tma_atom_q_latent, + qk_params.tQLgQL[None, 0, i], + qk_params.tQsQ[None, (i, 0)], + tma_bar_ptr=tma_bar_ptr, + ) + for i in cutlass.range_constexpr(self.iterations_qk_rope): + # load q rope + cute.copy( + qk_params.tma_atom_q_rope, + qk_params.tQRgQR[None, 0, i], + qk_params.tQsQ_rope[None, i], + tma_bar_ptr=tma_bar_ptr, + ) + load_q_producer_state.advance() + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_k_pipeline.producer_get_barrier( + load_k_producer_state + ) + common_params.load_k_pipeline.producer_acquire(load_k_producer_state) + for i in range(self.iterations_qk_latent): + for k in range(page_per_tile): + # load k latent + cute.copy( + qk_params.tma_atom_c_latent, + qk_params.tCLgCL[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, (i, load_k_producer_state.index)], + tma_bar_ptr=tma_bar_ptr, + ) + + for i in cutlass.range_constexpr(self.iterations_qk_rope): + for k in cutlass.range_constexpr(page_per_tile): + # load k rope + cute.copy( + qk_params.tma_atom_c_rope, + qk_params.tKRgKR[None, i, k_idx[k]], + qk_params.tKCsKC_rope[None, k, 0, load_k_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_k_producer_state.advance() + + return load_q_producer_state, load_k_producer_state + + @cute.jit + def load_tma_v_one_k_tile( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + load_v_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load one k-tile of compressed latent transpose tensor(v). Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The load tma v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param load_v_producer_state: The load v producer state + :type load_v_producer_state: pipeline.PipelineState + + :return: The load qkv producer state + :rtype: pipeline.PipelineState + """ + page_per_tile = self.mma_pv_tiler[2] * self.iterations_pv_k // self.page_size + page_per_subtile = ceil_div(page_per_tile, self.iterations_pv_k) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = ( + common_params.mPT[k_index] + if page_per_tile == 1 + else common_params.mPT[k_index * page_per_tile + i] + ) + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_v_pipeline.producer_get_barrier( + load_v_producer_state + ) + common_params.load_v_pipeline.producer_acquire(load_v_producer_state) + for j in cutlass.range_constexpr(self.iterations_pv_n): + for i in cutlass.range_constexpr(self.iterations_pv_k): + if cutlass.const_expr(page_per_tile > 1): + for k in cutlass.range_constexpr(page_per_subtile): + k_idx_i = k_idx[k + i * page_per_subtile] + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[None, j, 0, k_idx_i], + v_params.tVCsVC[ + None, 0, k, ((j, i), load_v_producer_state.index) + ], + tma_bar_ptr=tma_bar_ptr, + ) + else: + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[None, j, i, k_idx[0]], + v_params.tVCsVC[ + None, 0, 0, ((j, i), load_v_producer_state.index) + ], + tma_bar_ptr=tma_bar_ptr, + ) + load_v_producer_state.advance() + return load_v_producer_state + + @cute.jit + def mma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + pv_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_k_consumer_state: pipeline.PipelineState, + load_v_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """MMA warp to compute the result of Q*K^T and P*V. Updates the tiled mma and pipeline states. + + :param common_params: The common parameters for mma qk and pv + :type common_params: SimpleNamespace + :param qk_params: The mma qk parameters + :type qk_params: SimpleNamespace + :param pv_params: The mma pv parameters + :type pv_params: SimpleNamespace + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_k_consumer_state: The load k consumer state + :type load_k_consumer_state: pipeline.PipelineState + :param load_v_consumer_state: The load v consumer state + :type load_v_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + :param p_mma_consumer_state: The p mma consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The mma o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the tiled mma pv, the load q consumer state, the load k consumer state, the load v consumer state, the mma s producer state, the p mma consumer state, and the mma o producer state + :rtype: tuple[cute.TiledMma, cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + tSrQ = tiled_mma_qk.make_fragment_A(qk_params.sQ) + tSrQ_rope = tiled_mma_qk.make_fragment_A(qk_params.sQ_rope) + tSrKC = tiled_mma_qk.make_fragment_B(qk_params.sKC) + tSrKC_rope = tiled_mma_qk.make_fragment_B(qk_params.sKC_rope) + tOrP = tiled_mma_pv.make_fragment_A(pv_params.sP) + tOrVC = tiled_mma_pv.make_fragment_B(pv_params.sVC) + + tStS_shape = tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1]) + ) + tStS_staged_fake = tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage) + ) + # use real tmem ptr for tStS + tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout) + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1]) + ) + # mma O has 1 stage. + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO_staged = cute.make_tensor( + tStS_staged.iterator + self.tmem_o_offset, tOtO_layout + ) + + # set more parameters + qk_params.tSrQ = tSrQ + qk_params.tSrQ_rope = tSrQ_rope + qk_params.tSrKC = tSrKC + qk_params.tSrKC_rope = tSrKC_rope + qk_params.tStS_staged = tStS_staged + pv_params.tOrP = tOrP + pv_params.tOrVC = tOrVC + pv_params.tOtO_staged = tOtO_staged + + # mma O accumulates on K, so the accumlate flag is set to False once before all K blocks. + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + if common_params.is_leader_cta: + load_q_release_state = load_q_consumer_state.clone() + ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + wait_q=True, + ) + k_tile_count -= 1 + + while k_tile_count > 0: + ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + wait_q=False, + ) + ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + k_tile_count -= 1 + # release q consumer states + load_q_pipeline.consumer_release(load_q_release_state) + load_q_release_state.advance() + ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + return ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def mma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + tiled_mma_qk: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_k_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + wait_q: bool, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for Q*K^T. Updates the tiled MMA QK and pipeline states. + + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_k_consumer_state: The load k consumer state + :type load_k_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the load q consumer state, the load k consumer state, and the mma s producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + tStS = qk_params.tStS_staged[None, None, None, mma_s_producer_state.index] + + qk_params.mma_s_pipeline.producer_acquire(mma_s_producer_state) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + load_k_pipeline = common_params.load_k_pipeline + if cutlass.const_expr(wait_q): + load_q_pipeline.consumer_wait(load_q_consumer_state) + load_k_pipeline.consumer_wait(load_k_consumer_state) + for q_stage in range(self.iterations_qk_latent): + kc_stage = load_k_consumer_state.index + for k_block in cutlass.range_constexpr(cute.size(qk_params.tSrQ.shape[2])): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ[None, None, k_block, (q_stage, 0)], + qk_params.tSrKC[None, None, k_block, (q_stage, kc_stage)], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + + for q_stage in range(self.iterations_qk_rope): + kc_stage = load_k_consumer_state.index + for k_block in cutlass.range_constexpr( + self.rope_dim // tiled_mma_qk.shape_mnk[2] + ): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ_rope[None, None, k_block, q_stage], + qk_params.tSrKC_rope[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_k_pipeline.consumer_release(load_k_consumer_state) + load_k_consumer_state.advance() + if cutlass.const_expr(wait_q): + load_q_consumer_state.advance() + + qk_params.mma_s_pipeline.producer_commit(mma_s_producer_state) + mma_s_producer_state.advance() + return ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) + + @cute.jit + def mma_pv( + self, + common_params: SimpleNamespace, + pv_params: SimpleNamespace, + tiled_mma_pv: cute.TiledMma, + load_v_consumer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for P*V. Updates the tiled mma pv and pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param pv_params: The pv parameters + :type pv_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_v_consumer_state: The load v consumer state + :type load_v_consumer_state: pipeline.PipelineState + :param p_mma_consumer_state: The P MMA consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The MMA o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma pv, the load v consumer state, the P MMA consumer state, and the MMA o producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + pv_params.p_mma_pipeline.consumer_wait(p_mma_consumer_state) + load_v_pipeline = common_params.load_v_pipeline + accumulate_flag = tiled_mma_pv.get(tcgen05.Field.ACCUMULATE) + mma_o_pipeline = pv_params.mma_o_pipeline + + load_v_pipeline.consumer_wait(load_v_consumer_state) + vc_stage = load_v_consumer_state.index + for acc_stage in range(self.iterations_pv_n): + mma_o_pipeline.producer_acquire(mma_o_producer_state) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, accumulate_flag) + for p_stage in range(self.iterations_pv_k): + tOtO = pv_params.tOtO_staged[None, None, None, acc_stage] + for k_block in cutlass.range_constexpr(pv_params.tOrP.shape[2]): + cute.gemm( + tiled_mma_pv, + tOtO, + pv_params.tOrP[ + None, + None, + k_block, + (p_stage, p_mma_consumer_state.index), + ], + pv_params.tOrVC[ + None, None, k_block, ((acc_stage, p_stage), vc_stage) + ], + tOtO, + ) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, True) + + mma_o_pipeline.producer_commit(mma_o_producer_state) + mma_o_producer_state.advance() + load_v_pipeline.consumer_release(load_v_consumer_state) + load_v_consumer_state.advance() + pv_params.p_mma_pipeline.consumer_release(p_mma_consumer_state) + p_mma_consumer_state.advance() + + return ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def compute( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + + :return: The MMA s consumer state, the P MMA producer state, and the P correction producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_total = cute.ceil_div(common_params.K, self.mma_qk_tiler[1]) + + row_max = -self.acc_dtype.inf + row_sum = self.acc_dtype(0) + correction_factor = self.acc_dtype(1) + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # no mask applied + while k_tile_count > 1: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + False, + False, + ) + k_index = k_index + 1 + k_tile_count = k_tile_count - 1 + + # mask applied + if cutlass.const_expr(common_params.mAccO is not None): + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + k_index == k_tile_total - 1, + True, + ) + else: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + True, + True, + ) + + return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state + + @cute.jit + def correction( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + p_cor_consumer_state: pipeline.PipelineState, + mma_o_consumer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + :param mma_o_consumer_state: The MMA o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, and the MMA o consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction = ( + self.get_correction_factor(common_params, p_cor_consumer_state) + ) + if k_tile_count_init != k_tile_count: + mma_o_consumer_state = self.rescale( + common_params, + mma_o_consumer_state, + correction_factor, + no_correction, + ) + k_tile_count = k_tile_count - 1 + if k_tile_count == 0: + mma_o_consumer_state = self.epilogue( + common_params, + epilogue_params, + mma_o_consumer_state, + row_sum, + row_max, + ) + return p_cor_consumer_state, mma_o_consumer_state + + @cute.jit + def exchange_p_cor_metadata( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + correction_factor: cutlass.Float32, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + row_max_new: cutlass.Float32, + tAcc: cute.Tensor, + tidx: cutlass.Int32, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, cutlass.Float32]: + """Compute the correction factor for the last k tile.""" + no_correction = 0 + if ( + row_max_new - row_max + ) * softmax_params.softmax_scale_log2 <= self.skip_correction_threshold: + no_correction = 1 + row_max_new = row_max + + # pad for 4x32b + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.mma_s_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, + corr_layout, + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype + ) + corr_tmem_store_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_store_atom, tCor) + corr_tmem_store_thr_copy = corr_tmem_store_tiled_copy.get_slice(tidx) + cCor_for_copy = corr_tmem_store_thr_copy.partition_S(cCor) + tCor_for_copy = corr_tmem_store_thr_copy.partition_D(tCor) + rCor = cute.make_fragment_like( + cCor_for_copy[None, None, None, 0], self.acc_dtype + ) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout + ) + rCor[0] = row_sum + rCor[1] = row_max_new + rCor[2] = correction_factor + rCor_int[3] = no_correction + + cute.copy( + corr_tmem_store_tiled_copy, + rCor, + tCor_for_copy[None, None, None, p_cor_producer_state.index], + ) + # fence between tmem store and correction warp + cute.arch.fence_view_async_tmem_store() + common_params.p_cor_pipeline.producer_commit(p_cor_producer_state) + p_cor_producer_state.advance() + return p_cor_producer_state, row_max_new + + @cute.jit + def softmax( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + row_max: cutlass.Float32, + row_sum: cutlass.Float32, + correction_factor: cutlass.Float32, + is_last_tile: bool, + is_local_last_tile: cutlass.Boolean, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + ]: + """Softmax for one k-tile. Updates the related pipeline states and returns the computed results. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + :param row_max: The row max + :type row_max: cutlass.Float32 + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param is_last_tile: Whether the last tile + :type is_last_tile: bool + :param is_local_last_tile: Whether the last tile is local + :type is_local_last_tile: cutlass.Boolean + + :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32] + """ + + softmax_params.p_mma_pipeline.producer_acquire(p_mma_producer_state) + softmax_params.mma_s_pipeline.consumer_wait(mma_s_consumer_state) + + # load S from tmem + tStS_shape = softmax_params.tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1]) + ) + tStS_staged_fake = softmax_params.tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage) + ) + tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout) + tStS = tStS_staged[None, None, None, mma_s_consumer_state.index] + + tAcc = tStS[(None, None), 0, 0] + cta_qk_tiler = ( + self.mma_qk_tiler[0] // self.cluster_shape_mnk[0], + self.mma_qk_tiler[1], + self.mma_qk_tiler[2], + ) + cS = cute.make_identity_tensor(cute.select(cta_qk_tiler, mode=[0, 1])) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + + tmem_thr_copy = tmem_tiled_copy.get_slice(tidx) + tTR_tAcc = tmem_thr_copy.partition_S(tAcc) + tTR_tS = tmem_thr_copy.partition_D(cS) + + tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) + + row_max_new = row_max + arch = BaseDSL._get_dsl().get_arch_enum() + if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): + cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + if is_last_tile: + tTR_rAcc[i] = ( + tTR_rAcc[i] + if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) + else -self.acc_dtype.inf + ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) + elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + tmem_load_red_atom = cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp( + tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX + ), + self.acc_dtype, + ) + tmem_red_tiled_copy = tcgen05.make_tmem_copy(tmem_load_red_atom, tAcc) + tmem_red_thr_copy = tmem_red_tiled_copy.get_slice(tidx) + tTR_tAcc_red = tmem_red_thr_copy.partition_S(tAcc) + tTR_tS_red = tmem_red_thr_copy.partition_D(cS) + tTR_rAcc_red = cute.make_fragment_like(tTR_tS_red, self.acc_dtype) + tTR_rMax = cute.make_rmem_tensor( + cute.make_layout((1, tTR_tS_red.shape[1], tTR_tS_red.shape[2])), + self.acc_dtype, + ) + cute.copy( + tmem_red_tiled_copy, + tTR_tAcc_red, + (tTR_rAcc_red, tTR_rMax), + ) + tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) + if is_last_tile: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[i] = ( + tTR_rAcc[i] + if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) + else -self.acc_dtype.inf + ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce( + cute.ReductionOp.MAX, row_max_new, 0 + ) + else: + row_max_new = cute.arch.fmax(row_max_new, tTR_rMax[0]) + + # if warps in N is 2, reduce row_max across warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_max_new + self.softmax_exchange_sync_bar.wait() + row_max_new = cute.arch.fmax( + row_max_new, + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp) + ], + ) + + # find correction factor + correction_factor = cute.math.exp2( + (row_max - row_max_new) * softmax_params.softmax_scale_log2, fastmath=True + ) + # split kv case + if cutlass.const_expr(not is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # softmax + fma_b = softmax_params.softmax_scale_log2 + fma_c = (0.0 - row_max_new) * softmax_params.softmax_scale_log2 + + for i in cutlass.range(cute.size(tTR_rAcc), vectorize=True, unroll_full=True): + tTR_rAcc[i] = tTR_rAcc[i] * fma_b + fma_c + tTR_rAcc[i] = cute.math.exp2(tTR_rAcc[i], fastmath=True) + + tTR_rS = cute.make_fragment_like(tTR_tS, self.q_dtype) + + # quantize + tTR_rS.store(tTR_rAcc.load().to(self.q_dtype)) + + # create sP + sP = softmax_params.sP[None, None, None, (None, p_mma_producer_state.index)] + sP_mk_view = cute.make_tensor( + sP.iterator, + cute.make_layout( + ( + (sP.shape[0][0], sP.shape[1]), + (sP.shape[0][1], sP.shape[2], sP.shape[3]), + ), + stride=( + (sP.stride[0][0], sP.stride[1]), + (sP.stride[0][1], sP.stride[2], sP.stride[3]), + ), + ), + ) + # change to PISL + sP_wo_swizzle_iter = cute.recast_ptr(sP.iterator, swizzle_=None) + swizzle_bits = ( + int(math.log2(self.mma_pv_tiler[2] * self.q_dtype.width // 8 // 32)) + 1 + ) + swizzle_base = 3 if self.q_dtype.width == 16 else 4 + sP_swizzle = cute.make_swizzle(swizzle_bits, swizzle_base, 3) + sP_mk_view = cute.make_tensor( + sP_wo_swizzle_iter, + cute.make_composed_layout(sP_swizzle, 0, sP_mk_view.layout), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_copy) + smem_thr_copy = smem_tiled_copy.get_slice(tidx) + rP_copy_view = smem_thr_copy.retile(tTR_rS) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + + # fence between smem store and mma o + cute.arch.fence_view_async_shared() + softmax_params.p_mma_pipeline.producer_commit(p_mma_producer_state) + p_mma_producer_state.advance() + + # row_sum, using `add_packed_f32x2` to reduce the number of instructions + row_sum = row_sum * correction_factor + row_sum_vec = (0.0, 0.0) + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + row_sum_vec = cute.arch.add_packed_f32x2( + row_sum_vec, (tTR_rAcc[i], tTR_rAcc[i + 1]) + ) + row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum + + # split kv case + if cutlass.const_expr(is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # store correction factor/row_sum/row_max to tmem for correction warp + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # fence between tmem load and mma s + cute.arch.fence_view_async_tmem_load() + + softmax_params.mma_s_pipeline.consumer_release(mma_s_consumer_state) + mma_s_consumer_state.advance() + + return ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max_new, + row_sum, + correction_factor, + ) + + @cute.jit + def _tmem_load_partition( + self, common_params: SimpleNamespace, tiled_mma_pv: cute.TiledMma, iter_n: int + ) -> tuple[ + cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma + ]: + """Tensor memory load partition for rescale and epilogue. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param iter_n: The iteration number + :type iter_n: int + + :return: The tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv + :rtype: tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma] + """ + + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1]) + ) + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO = cute.make_tensor( + common_params.tmem_ptr + self.tmem_o_offset, tOtO_layout + ) + tOtO = tOtO[None, None, None, iter_n] + + tAcc = tOtO[(None, None), 0, 0] + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_load_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + tmem_load_thr_copy = tmem_load_tiled_copy.get_slice( + common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + ) + + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + # Flatten divide and partition global tensors for O + cta_pv_tiler_mn = cute.select(cta_pv_tiler, mode=[0, 1]) + + gO = None + if cutlass.const_expr(common_params.mAccO is not None): + gO = cute.local_tile( + common_params.mAccO[None, common_params.blk_coord[3], None, None, None], + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor( + common_params.mAccO[ + None, common_params.blk_coord[3], None, None, None + ].shape + ), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + else: + gO = cute.local_tile( + common_params.mO, + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor(common_params.mO.shape), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + tTR_tAcc = tmem_load_thr_copy.partition_S(tAcc) + tTR_gO = tmem_load_thr_copy.partition_D(gO) + tTR_cO = tmem_load_thr_copy.partition_D(cO) + tTR_rAcc = cute.make_fragment_like(tTR_gO, self.acc_dtype) + return tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc + + def get_correction_factor( + self, + common_params: SimpleNamespace, + p_cor_consumer_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + cutlass.Int32, + ]: + """Get the correction factor from the P correction consumer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, the row_sum, the row_max, and the correction factor + :rtype: tuple[pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Int32] + """ + common_params.p_cor_pipeline.consumer_wait(p_cor_consumer_state) + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + # load correction factor + _, tAcc, _, _, _, _ = self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, 0 + ) + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.p_cor_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, corr_layout + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype + ) + corr_tmem_load_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_load_atom, tCor) + corr_tmem_load_thr_copy = corr_tmem_load_tiled_copy.get_slice(tidx) + tCor_for_copy = corr_tmem_load_thr_copy.partition_S(tCor) + cCor_for_copy = corr_tmem_load_thr_copy.partition_D(cCor) + rCor = cute.make_fragment_like( + cCor_for_copy[None, None, None, 0], self.acc_dtype + ) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout + ) + cute.copy( + corr_tmem_load_tiled_copy, + tCor_for_copy[None, None, None, p_cor_consumer_state.index], + rCor, + ) + row_sum = rCor[0] + row_max = rCor[1] + correction_factor = rCor[2] + no_correction = rCor_int[3] + + common_params.p_cor_pipeline.consumer_release(p_cor_consumer_state) + p_cor_consumer_state.advance() + return p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction + + @cute.jit + def rescale( + self, + common_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + correction_factor: cutlass.Float32, + no_correction: cutlass.Int32, + ) -> pipeline.PipelineState: + """Rescale for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param no_correction: Whether to apply correction factor + :type no_correction: cutlass.Int32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + skip_correction = cute.arch.vote_all_sync(no_correction == 1) + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + if not skip_correction: + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, iter_n + ) + ) + + # tmem store tiled copy + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_store_tiled_copy = tcgen05.make_tmem_copy(tmem_store_atom, tAcc) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + # rescale, using `mul_packed_f32x2` to reduce the number of instructions + for i in cutlass.range( + cute.size(tTR_rAcc), vectorize=True, unroll_full=True + ): + tTR_rAcc[i] = tTR_rAcc[i] * correction_factor + + # store o to tensor memory for next k tile + cute.copy(tmem_store_tiled_copy, tTR_rAcc, tTR_tAcc) + + cute.arch.fence_view_async_tmem_store() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + @cute.jit + def epilogue( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + ) -> pipeline.PipelineState: + """Epilogue for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param row_max: The row max + :type row_max: cutlass.Float32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + + # exchange row_sum between warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_sum + self.epilogue_exchange_sync_bar.wait() + # (64, 2) + row_sum = ( + row_sum + + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp) + ] + ) + # mma_o pipeline consumer wait + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, iter_n + ) + ) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + + # apply output scale and normalize by row_sum + for i in cutlass.range( + cute.size(tTR_rAcc), vectorize=True, unroll_full=True + ): + tTR_rAcc[i] = ( + tTR_rAcc[i] + * epilogue_params.output_scale + * cute.arch.rcp_approx(row_sum) + ) + + # store o to global memory + tR2G_rO_src = None + tR2G_rO_dst = tTR_gO + if cutlass.const_expr(common_params.mAccO is None): + tR2G_rO_src = cute.make_fragment_like(tTR_gO, self.o_dtype) + # using final output dtype for o + tR2G_rO_src.store(tTR_rAcc.load().to(self.o_dtype)) + else: + # using accumulate dtype for o + tR2G_rO_src = tTR_rAcc + + if cute.elem_less(tTR_cO[0][0], common_params.H): + cute.autovec_copy( + tR2G_rO_src, + tR2G_rO_dst, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, + ) + + # store the lse to global memory + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + gLSE = None + cLSE = None + if cutlass.const_expr(epilogue_params.mAccLSE is None): + gLSE = cute.local_tile( + epilogue_params.mLSE, + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor(epilogue_params.mLSE.shape), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + + else: + gLSE = cute.local_tile( + epilogue_params.mAccLSE[ + None, common_params.blk_coord[3], None, None + ], + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor( + epilogue_params.mAccLSE[ + None, common_params.blk_coord[3], None, None + ].shape + ), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + lse = ( + cute.math.log2(row_sum, fastmath=True) + + epilogue_params.softmax_scale_log2 * row_max + ) + if cutlass.const_expr(self.warps_in_n == 2): + if cute.elem_less(cLSE[tidx][0], common_params.H): + gLSE[tidx] = lse + + cute.arch.fence_view_async_tmem_load() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + def make_and_init_load_qkv_pipeline( + self, load_qkv_mbar_ptr, cta_layout_vmnk, load_stages, tx_count + ) -> pipeline.PipelineTmaUmma: + """Create and initialize the tma load qkv pipeline. + + :param load_qkv_mbar_ptr: The load qkv mbar pointer + :type load_qkv_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + :param load_stages: The load stages + :type load_stages: list[int] + :param tx_count: The tx count + :type tx_count: int + + :return: The tma load qkv pipeline + :rtype: pipeline.PipelineTmaUmma + """ + load_qkv_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_tma_k_warp_id]) + ) + load_qkv_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_qkv_mbar_ptr, + num_stages=load_stages, + producer_group=load_qkv_producer_group, + consumer_group=load_qkv_consumer_group, + tx_count=tx_count, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_mma_s_pipeline( + self, mma_s_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma s pipeline. + + :param mma_s_mbar_ptr: The mma s mbar pointer + :type mma_s_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma s pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_s_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + consumer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + mma_s_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_s_mbar_ptr, + num_stages=self.mma_s_stage, + producer_group=mma_s_producer_group, + consumer_group=mma_s_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_mma_pipeline( + self, p_mma_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p mma pipeline. + + :param p_mma_mbar_ptr: The p mma mbar pointer + :type p_mma_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The p mma pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + p_mma_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_mma_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineAsyncUmma.create( + barrier_storage=p_mma_mbar_ptr, + num_stages=self.p_mma_stage, + producer_group=p_mma_producer_group, + consumer_group=p_mma_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_cor_pipeline( + self, p_cor_mbar_ptr + ) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p correction pipeline. + + :param p_cor_mbar_ptr: The p correction mbar pointer + :type p_cor_mbar_ptr: cute.Tensor + + :return: The p correction pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = self.threads_per_warp * len(self.compute_warp_ids) + p_cor_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_cor_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + return pipeline.PipelineAsync.create( + barrier_storage=p_cor_mbar_ptr, + num_stages=self.p_cor_stage, + producer_group=p_cor_producer_group, + consumer_group=p_cor_consumer_group, + defer_sync=True, + ) + + def make_and_init_mma_o_pipeline( + self, mma_o_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma o pipeline. + + :param mma_o_mbar_ptr: The mma o mbar pointer + :type mma_o_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma o pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_o_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + consumer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + mma_o_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_o_mbar_ptr, + num_stages=self.mma_o_stage, + producer_group=mma_o_producer_group, + consumer_group=mma_o_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + @staticmethod + def _compute_grid( + o: cute.Tensor, + split_kv: cutlass.Int32, + cluster_shape_mnk: Tuple[int, int, int], + max_active_clusters: int, + is_persistent: bool, + ) -> Tuple[MLAStaticTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid shape 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] + + :return: Tile scheduler parameters and grid shape. + :rtype: tuple[MLAStaticTileSchedulerParams, tuple[int, int, int]] + """ + o_shape = o.shape + tile_sched_params = create_mla_static_tile_scheduler_params( + is_persistent, + cute.size(o_shape[3]), + cute.size(o_shape[2]), + cluster_shape_mnk, + split_kv, + ) + grid = MLAStaticTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def get_workspace_size( + H: int, + S: int, + D: int, + B: int, + split_kv: int, + acc_dtype: Type[cutlass.Numeric], + ) -> int: + """Get the extra workspace(device memory) size for the MLA kernel when split_kv is not 1. + + :param H: The height of the output tensor C + :type H: int + :param S: The sequence length of the output tensor C + :type S: int + :param D: The depth of the output tensor C + :type D: int + :param B: The batch size of the output tensor C + :type B: int + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + + :return: The workspace size for the MLA kernel + :rtype: int + """ + if split_kv == 1: + return 0 + return B * H * S * split_kv * (D + 1) * acc_dtype.width // 8 + + @cute.jit + def initialize_workspace( + self, + H: cutlass.Int32, + D: cutlass.Int32, + S: cutlass.Int32, + B: cutlass.Int32, + split_kv: cutlass.Int32, + acc_dtype: Type[cutlass.Numeric], + workspace: cute.Tensor, + ) -> tuple[cute.Tensor, cute.Tensor]: + """Initialize the workspace for the MLA kernel. Construct the intermediate tensors + acc_o and acc_lse. + + :param H: The height of the output tensor C + :type H: cutlass.Int32 + :param D: The depth of the output tensor C + :type D: cutlass.Int32 + :param S: The sequence length of the output tensor C + :type S: cutlass.Int32 + :param B: The batch size of the output tensor C + :type B: cutlass.Int32 + :param split_kv: The split key-value of the output tensor C + :type split_kv: cutlass.Int32 + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + :param workspace: The workspace tensor + :type workspace: cute.Tensor + + :return: The output tensor C and the workspace tensor + :rtype: tuple[cute.Tensor, cute.Tensor] + """ + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): + align = 256 // self.q_dtype.width + acc_o_layout = cute.make_layout( + (H, split_kv, D, S, B), + stride=( + cute.assume(split_kv * D, align), + cute.assume(D, align), + 1, + cute.assume(split_kv * H * D, align), + cute.assume(H * split_kv * S * D, align), + ), + ) + acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype) + acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) + acc_lse_layout = cute.make_layout( + (H, split_kv, S, B), + stride=(split_kv, 1, H * split_kv, H * split_kv * S), + ) + acc_lse_iter = cute.recast_ptr( + workspace.iterator + cute.cosize(acc_o_layout) * acc_dtype.width // 8, + dtype=acc_dtype, + ) + acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout) + return acc_o, acc_lse + + @staticmethod + def can_implement( + B: int, + S: int, + K: int, + H: int, + L: int, + R: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + ) -> bool: + """Check if the MLA kernel can be implemented. + + :param B: The batch size of the output tensor C + :type B: int + :param S: The sequence length of the output tensor C + :type S: int + :param K: The width of the output tensor KV + :type K: int + :param H: The number of heads of the output tensor C + :type H: int + :param L: The number of latent dimensions of the tensor KV + :type L: int + :param R: The number of rope dimensions of the tensor C_rope + :type R: int + :param in_dtype: The data type of the input tensor + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: The data type of the output tensor + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: The data type of the log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: The tile shape of the query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: The tile shape of the probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: The page size of the page table + :type page_size: int + + :return: Whether the MLA kernel can be implemented + :rtype: bool + """ + if L != 512 or R != 64: + return False + if in_dtype not in [cutlass.Float8E4M3FN]: + return False + if out_dtype not in [cutlass.Float8E4M3FN]: + return False + if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32: + return False + # page size equals 1 is prohibited by tma specification, not 128B aligned. + if mma_qk_tiler_mn[1] % page_size != 0 or page_size == 1: + return False + if mma_qk_tiler_mn[0] != mma_pv_tiler_mn[0] or mma_qk_tiler_mn[0] != 128: + return False + if is_var_split_kv and not is_var_seq: + return False + if H > 128 or (H < 128 and split_kv != 1): + return False + if S <= 0 or S > 4: + return False + if K <= 0: + return False + return True + + +def run( + batch_size: int, + seq_len_q: int, + seq_len_k: int, + num_heads: int, + latent_dim: int, + rope_dim: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + softmax_scale: float, + output_scale: float, + skip_correction_threshold: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool, + **kwargs, +): + """Execute Multi-Head Latent Attention (MLA) on Blackwell architecture and validate results. + + This function creates random input tensors for query latent/rope, compressed latent/rope, and value, + then performs the complete MLA computation pipeline. It supports configurable data types, tiling parameters, + page table, variable sequence length, and variable split_kv. Results can be validated against a PyTorch reference + implementation or run multiple times for performance measurement. + + :param batch_size: Batch size + :type batch_size: int + :param seq_len_q: Sequence length of Q + :type seq_len_q: int + :param seq_len_k: Sequence length of K + :type seq_len_k: int + :param num_heads: Number of heads + :type num_heads: int + :param latent_dim: dimension of query/compressed latent + :type latent_dim: int + :param rope_dim: dimension of query/compressed rope + :type rope_dim: int + :param in_dtype: Input data type for query/compressed latent/rope tensors + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: Output data type for attention output + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type for query-key matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Accumulator data type for log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: Matrix multiply accumulate tile shape (M, N) for query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: Matrix multiply accumulate tile shape (M, N) for probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: Split key-value + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: Page size of the page table + :type page_size: int + :param softmax_scale: Attention score scaling factor + :type softmax_scale: float + :param output_scale: Output scaling factor + :type output_scale: float + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param tolerance: Maximum acceptable error for validation + :type tolerance: float + :param warmup_iterations: Number of warmup iterations + :type warmup_iterations: int + :param iterations: Number of iterations to run for performance testing + :type iterations: int + :param skip_ref_check: Skip validation against reference implementation + :type skip_ref_check: bool + :param use_cold_l2: Whether to use cold L2 cache + :type use_cold_l2: bool + + :raises ValueError: If input shapes are incompatible or head dimension is unsupported + :raises RuntimeError: If GPU is unavailable for computation + """ + + print("Running Blackwell MLA test with:") + print(f" batch_size: {batch_size}") + print(f" seq_len_q: {seq_len_q}") + print(f" seq_len_k: {seq_len_k}") + print(f" num_heads: {num_heads}") + print(f" latent_dim: {latent_dim}") + print(f" rope_dim: {rope_dim}") + print(f" in_dtype: {in_dtype}") + print(f" out_dtype: {out_dtype}") + print(f" acc_dtype: {acc_dtype}") + print(f" mma_qk_tiler_mn: {mma_qk_tiler_mn}") + print(f" mma_pv_tiler_mn: {mma_pv_tiler_mn}") + print(f" split_kv: {split_kv}") + print(f" is_persistent: {is_persistent}") + print(f" is_var_seq: {is_var_seq}") + print(f" is_var_split_kv: {is_var_split_kv}") + print(f" page_size: {page_size}") + print(f" softmax_scale: {softmax_scale}") + print(f" output_scale: {output_scale}") + print(f" skip_correction_threshold: {skip_correction_threshold}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + + import torch + import cutlass.torch as cutlass_torch + + # Prepare pytorch tensors: Q, K, V (random from 0 to 2) and O (all zero) + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + if not BlackwellMultiHeadLatentAttentionForwardFP8.can_implement( + batch_size, + seq_len_q, + seq_len_k, + num_heads, + latent_dim, + rope_dim, + in_dtype, + out_dtype, + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + split_kv, + is_persistent, + is_var_seq, + is_var_split_kv, + page_size, + ): + raise TypeError( + f"Unsupported testcase {batch_size}, {seq_len_q}, {seq_len_k}, {num_heads}, {latent_dim}, {rope_dim}, {in_dtype}, {out_dtype}, {acc_dtype}, {lse_dtype}, {mma_qk_tiler_mn}, {mma_pv_tiler_mn}, {split_kv}, {is_persistent}, {is_var_seq}, {is_var_split_kv}, {page_size}" + ) + + torch.manual_seed(1111) + + def create_data_tensor( + B, + HK, + D, + dtype, + is_dynamic_layout=True, + page_table=None, + cache_seqs=None, + is_lse=False, + seq_len_q=None, + ): + shape = (B, HK, D) + if page_table is not None: + if cache_seqs is not None: + max_seq_len = torch.max(cache_seqs) + shape = (B * ceil_div(max_seq_len, page_size), page_size, D) + else: + shape = (B * ceil_div(HK, page_size), page_size, D) + + if seq_len_q is not None: + shape = (B, seq_len_q, HK, D) + + permute_order = (1, 2, 0) + stride_order = (2, 0, 1) + leading_dim = 1 + if is_lse: + shape = (B, seq_len_q, HK) + permute_order = (2, 1, 0) + stride_order = (2, 1, 0) + leading_dim = 0 + elif seq_len_q is not None: + permute_order = (2, 3, 1, 0) + stride_order = (3, 2, 0, 1) + leading_dim = 1 + + init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) + + torch_dtype = ( + cutlass_torch.dtype(dtype) if dtype != cutlass.Float8E4M3FN else torch.int8 + ) + + # Create dtype torch tensor (cpu) + torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch_dtype, + permute_order=permute_order, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=init_config, + ) + + # Create dtype torch tensor (gpu) + torch_tensor_gpu = torch_tensor_cpu.cuda() + + # Create f32 torch tensor (cpu) + f32_torch_tensor = torch_tensor_cpu.to(dtype=torch.float32) + + # Create dtype cute tensor (gpu) + cute_tensor = from_dlpack(torch_tensor_gpu, assumed_align=16) + cute_tensor.element_type = dtype + if is_dynamic_layout: + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + if not is_lse: + cute_tensor = cute_tensor.mark_compact_shape_dynamic( + mode=leading_dim, + stride_order=stride_order, + divisibility=(128 // dtype.width), + ) + + cute_tensor = cutlass_torch.convert_cute_tensor( + f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=is_dynamic_layout, + ) + + return f32_torch_tensor, cute_tensor, torch_tensor_gpu + + def create_cache_seqs(batch_size, seq_len_k, is_var_seq): + cache_seqs_ref = torch.ones(batch_size, dtype=torch.int32) * seq_len_k + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack(cache_seqs_gpu, assumed_align=16).mark_layout_dynamic() + if is_var_seq: + max_seq_len = seq_len_k + min_seq_len = int(seq_len_k * 0.8) + cache_seqs_ref = cutlass_torch.create_and_permute_torch_tensor( + (batch_size,), + torch.int32, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig( + min_val=min_seq_len, max_val=max_seq_len + 1 + ), + ) + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack( + cache_seqs_gpu, + assumed_align=16, + ).mark_layout_dynamic() + return cache_seqs_ref, cache_seqs, cache_seqs_gpu + + def create_page_table(batch_size, seq_len_k, is_var_seq, page_size): + max_seq_len = seq_len_k if not is_var_seq else torch.max(cache_seqs_ref) + page_count = ceil_div(max_seq_len, page_size) + page_table_ref = torch.empty([batch_size, page_count], dtype=torch.int32) + # use transposed index for page table to make sure the value is in bound of `batch_size * seq_len_block`. In practice, the value could be any positive values. This setting is only for testing purpose. + for b in range(batch_size): + for j in range(page_count): + page_table_ref[b, j] = b + j * batch_size + page_table_gpu = page_table_ref.permute(1, 0).cuda() + page_table = from_dlpack(page_table_gpu, assumed_align=16).mark_layout_dynamic( + leading_dim=0 + ) + return page_table_ref, page_table, page_table_gpu + + def create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ): + block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu = None, None, None + # check if split_kv is valid otherwise do auto setting of split_kv + if is_var_split_kv: + block_split_kvs_ref = torch.zeros([batch_size], dtype=torch.int32) + for b in range(batch_size): + block_split_kvs_ref[b] = ( + BlackwellMultiHeadLatentAttentionForwardFP8.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[b].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + ) + split_kv = torch.max(block_split_kvs_ref).item() + block_split_kvs_gpu = block_split_kvs_ref.cuda() + block_split_kvs = from_dlpack( + block_split_kvs_gpu, assumed_align=16 + ).mark_layout_dynamic() + elif split_kv <= 0: + split_kv = BlackwellMultiHeadLatentAttentionForwardFP8.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[0].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + return split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu + + def create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, split_kv, acc_dtype + ): + workspace_size = BlackwellMultiHeadLatentAttentionForwardFP8.get_workspace_size( + num_heads, + seq_len_q, + latent_dim, + batch_size, + split_kv, + acc_dtype, + ) + + workspace, workspace_torch = None, None + if workspace_size > 0: + workspace_torch = torch.empty([workspace_size], dtype=torch.int8).cuda() + workspace = from_dlpack(workspace_torch, assumed_align=32) + return workspace, workspace_torch + + cache_seqs_ref, cache_seqs, cache_seqs_torch = create_cache_seqs( + batch_size, seq_len_k, is_var_seq + ) + page_table_ref, page_table, page_table_torch = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size + ) + cluster_shape_mnk = (2, 1, 1) + hardware_info = utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mnk[0] * cluster_shape_mnk[1] + ) + split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_torch = ( + create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + ) + + q_latent_ref, q_latent, q_latent_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + q_rope_ref, q_rope, q_rope_torch = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + c_latent_ref, c_latent, c_latent_torch = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + c_rope_ref, c_rope, c_rope_torch = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + o_ref, o, o_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + lse_ref, lse, lse_torch = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, split_kv, acc_dtype + ) + + mla = BlackwellMultiHeadLatentAttentionForwardFP8( + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + max_active_clusters, + page_size, + skip_correction_threshold, + is_persistent, + is_var_seq, + is_var_split_kv, + ) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + stream = cuda.CUstream(torch_stream.cuda_stream) + + # compile mla kernel + compiled_mla = cute.compile( + mla, + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + options="--opt-level 2", + ) + + def torch_reference_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + softmax_scale=1.0, + output_scale=1.0, + ): + # expand and concat q_latent and q_rope to have the dimension of sequence length for q + q_ref = torch.cat([q_latent, q_rope], dim=1).permute(3, 2, 0, 1) + # expand and concat c_latent and c_rope to have the dimension of num_heads for k and v + page_count = page_table_ref.shape[1] + k_ref_paged = ( + torch.cat([c_latent, c_rope], dim=1) + .permute(2, 0, 1) + .reshape(batch_size * page_count, page_size, latent_dim + rope_dim) + ) + v_ref_paged = c_latent.permute(2, 0, 1).reshape( + batch_size * page_count, page_size, latent_dim + ) + + if is_var_seq: + max_seq_len = torch.max(cache_seqs_ref) + else: + max_seq_len = seq_len_k + + k_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim + rope_dim]) + v_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim]) + k_ref = torch.index_select( + k_ref_paged, 0, torch.flatten(page_table_ref) + ).reshape(batch_size, 1, -1, latent_dim + rope_dim)[:, :, :max_seq_len, :] + v_ref = torch.index_select( + v_ref_paged, 0, torch.flatten(page_table_ref) + ).reshape(batch_size, 1, -1, latent_dim)[:, :, :max_seq_len, :] + for b in range(batch_size): + k_ref[b, :, cache_seqs_ref[b] :, :] = 0 + v_ref[b, :, cache_seqs_ref[b] :, :] = 0 + import torch.nn.functional as F + + o_ref = F.scaled_dot_product_attention( + q_ref, + k_ref, + v_ref, + attn_mask=None, + dropout_p=0.0, + scale=softmax_scale, + is_causal=False, + ) + s_ref = torch.einsum("bhld,bhsd->bhls", q_ref, k_ref) + s_ref_max, s_ref_max_pos = torch.max(s_ref, dim=-1, keepdim=True) + softmax_scale_log2 = LOG2_E * softmax_scale + s_ref_sum = torch.sum( + torch.exp2((s_ref - s_ref_max) * softmax_scale_log2), dim=-1, keepdim=True + ) + + lse_ref = s_ref_max * softmax_scale_log2 + torch.log2(s_ref_sum) + lse_ref = lse_ref.squeeze(3).permute(2, 1, 0) + o_ref = o_ref * output_scale + o_ref = o_ref.permute(2, 3, 1, 0) + + return o_ref, lse_ref + + if skip_correction_threshold > 0.0: + print( + "Skipping correction verification since skip_correction_threshold is greater than 0.0..." + ) + skip_ref_check = True + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + torch.cuda.synchronize() + + print("Verifying results...") + if in_dtype == cutlass.Float8E4M3FN: + tolerance = 0.13 + o_ref, lse_ref = torch_reference_mla( + q_latent_ref, + q_rope_ref, + c_latent_ref, + c_rope_ref, + page_table, + cache_seqs, + softmax_scale, + output_scale, + ) + + if out_dtype in [cutlass.Float8E5M2, cutlass.Float8E4M3FN]: + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + cutlass.Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o, o_fp32) + o = o_fp32_torch.cpu() + ref_fp8, _ = cutlass_torch.cute_tensor_like( + torch.empty( + *o_ref.permute(3, 2, 0, 1).shape, dtype=torch.uint8 + ).permute(2, 3, 1, 0), + out_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + o_ref_gpu = o_ref.cuda() + o_ref_f32 = from_dlpack(o_ref_gpu).mark_layout_dynamic(leading_dim=1) + + # convert ref : f32 -> fp8 -> f32 + cute.testing.convert(o_ref_f32, ref_fp8) + cute.testing.convert(ref_fp8, o_ref_f32) + + o_ref = o_ref_gpu.cpu() + else: + o = o_torch.cpu().to(torch.float32) + lse = lse_torch.cpu() + lse_ref = lse_ref.to(cutlass.torch.dtype(lse_dtype)) + # Assert close results + torch.testing.assert_close(o, o_ref, atol=tolerance, rtol=1e-05) + torch.testing.assert_close(lse, lse_ref, atol=tolerance, rtol=1e-05) + print("Results verified successfully!") + + def generate_tensors(): + _, cache_seqs, _ = create_cache_seqs(batch_size, seq_len_k, is_var_seq) + _, page_table, _ = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size + ) + _split_kv, _, block_split_kvs, _ = create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + + _, q_latent, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, q_rope, _ = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + _, c_latent, _ = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, c_rope, _ = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, o, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, lse, _ = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, _split_kv, acc_dtype + ) + return testing.JitArguments( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + _split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + q_latent_torch.numel() * q_latent_torch.element_size() + + q_rope_torch.numel() * q_rope_torch.element_size() + + c_latent_torch.numel() * c_latent_torch.element_size() + + c_rope_torch.numel() * c_rope_torch.element_size() + + o_torch.numel() * o_torch.element_size() + + lse_torch.numel() * lse_torch.element_size() + + cache_seqs_torch.numel() * cache_seqs_torch.element_size() + ) + one_workspace_bytes += ( + page_table_torch.numel() * page_table_torch.element_size() + ) + if is_var_split_kv: + one_workspace_bytes += ( + block_split_kvs_torch.numel() * block_split_kvs_torch.element_size() + ) + if workspace_torch is not None: + one_workspace_bytes += ( + workspace_torch.numel() * workspace_torch.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + avg_time_us = testing.benchmark( + compiled_mla, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return avg_time_us # Return execution time in microseconds + + +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." + ) + + def parse_mma_tiler(s: str) -> Tuple[int, int, Tuple[int, int]]: + ret = parse_comma_separated_ints(s) + if len(ret) != 2: + raise argparse.ArgumentTypeError( + "Invalid format. Expected 2 comma-separated integers." + ) + return (ret[0], ret[1]) + + parser = argparse.ArgumentParser(description="Example of MLA on Blackwell.") + + parser.add_argument( + "--in_dtype", + type=cutlass.dtype, + default=cutlass.Float8E4M3FN, + help="Input data type", + ) + + parser.add_argument( + "--out_dtype", + type=cutlass.dtype, + default=cutlass.Float8E4M3FN, + help="Output data type", + ) + + parser.add_argument( + "--acc_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="Accumulator data type", + ) + + parser.add_argument( + "--lse_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="LSE data type", + ) + parser.add_argument( + "--mma_qk_tiler_mn", + type=parse_mma_tiler, + default=(128, 128), + help="MMA tile shape (H, K)", + ) + parser.add_argument( + "--mma_pv_tiler_mn", + type=parse_mma_tiler, + default=(128, 256), + help="MMA tile shape (H, D)", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--batch_size", + type=int, + default=1, + help="Batch size", + ) + + parser.add_argument( + "--seq_len_q", + type=int, + default=1, + help="Sequence length of Q", + ) + + parser.add_argument( + "--seq_len_k", + type=int, + default=128, + help="Sequence length of K/V", + ) + + parser.add_argument( + "--num_heads", + type=int, + default=128, + help="Number of heads of Q", + ) + + parser.add_argument( + "--latent_dim", + type=int, + default=512, + help="Latent dimension of Q/C", + ) + + parser.add_argument( + "--rope_dim", + type=int, + default=64, + help="Rope dimension of Q/C", + ) + + parser.add_argument( + "--is_var_seq", + action="store_true", + help="Use variable length of sequence length or not", + ) + + parser.add_argument( + "--is_var_split_kv", + action="store_true", + help="Use variable length of split kv or not", + ) + + parser.add_argument( + "--page_size", + type=int, + default=128, + help="Page size of page table", + ) + + parser.add_argument( + "--split_kv", + type=int, + default=-1, + help="Split KV setting", + ) + + parser.add_argument( + "--softmax_scale", + type=float, + default=0.0416, + help="Scaling factor to scale softmax", + ) + + parser.add_argument( + "--output_scale", + type=float, + default=1.0, + help="Scaling factor to scale output", + ) + parser.add_argument( + "--skip_correction_threshold", + type=float, + default=0.0, + help="Threshold to skip correction", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-02, help="Tolerance for validation" + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + help="Use cold L2 cache", + ) + + args = parser.parse_args() + + run( + args.batch_size, + args.seq_len_q, + args.seq_len_k, + args.num_heads, + args.latent_dim, + args.rope_dim, + args.in_dtype, + args.out_dtype, + args.acc_dtype, + args.lse_dtype, + args.mma_qk_tiler_mn, + args.mma_pv_tiler_mn, + args.split_kv, + args.is_persistent, + args.is_var_seq, + args.is_var_split_kv, + args.page_size, + args.softmax_scale, + args.output_scale, + args.skip_correction_threshold, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mla/mla_helpers.py b/examples/python/CuTeDSL/blackwell/mla/mla_helpers.py new file mode 100644 index 00000000..1790b3c8 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mla/mla_helpers.py @@ -0,0 +1,304 @@ +# 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 cutlass +import cutlass.cute as cute + + +class MLAStaticTileSchedulerParams: + def __init__( + self, + is_persistent: bool, + problem_shape_b: cute.Int32, + problem_shape_s: cute.Int32, + cluster_shape_mnk: cute.Shape, + split_kv: cutlass.Int32, + *, + problem_shape_b_fdd: cute.FastDivmodDivisor = None, + problem_shape_s_fdd: cute.FastDivmodDivisor = None, + split_kv_fdd: cute.FastDivmodDivisor = None, + loc=None, + ip=None, + ): + """The static tile scheduler parameters prepared for MLA static tile scheduler. + + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param problem_shape_b: The shape of the problem + :type problem_shape_b: cute.Int32 + :param problem_shape_s: The shape of the problem in sequence length Q dimension + :type problem_shape_s: cute.Int32 + :param cluster_shape_mnk: The shape of the cluster + :type cluster_shape_mnk: cute.Shape + :param split_kv: The scalar factor for split KV + """ + self.is_persistent = is_persistent + self.problem_shape_b = problem_shape_b + self.problem_shape_s = problem_shape_s + self.problem_shape_b_fdd = problem_shape_b_fdd + self.problem_shape_s_fdd = problem_shape_s_fdd + self.cluster_shape_mnk = cluster_shape_mnk + self.split_kv = split_kv + self.split_kv_fdd = split_kv_fdd + if cutlass.const_expr(problem_shape_b_fdd is None): + self.problem_shape_b_fdd = cute.fast_divmod_create_divisor( + problem_shape_b, loc=loc, ip=ip + ) + if cutlass.const_expr(problem_shape_s_fdd is None): + self.problem_shape_s_fdd = cute.fast_divmod_create_divisor( + problem_shape_s, loc=loc, ip=ip + ) + if cutlass.const_expr(split_kv_fdd is None): + self.split_kv_fdd = cute.fast_divmod_create_divisor( + split_kv, loc=loc, ip=ip + ) + self.loc = loc + self.ip = ip + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.problem_shape_b) + values += cutlass.extract_mlir_values(self.problem_shape_s) + values += cutlass.extract_mlir_values(self.split_kv) + values += cutlass.extract_mlir_values(self.problem_shape_b_fdd) + values += cutlass.extract_mlir_values(self.problem_shape_s_fdd) + values += cutlass.extract_mlir_values(self.split_kv_fdd) + return values + + def __new_from_mlir_values__(self, values): + problem_shape_b = cutlass.new_from_mlir_values( + self.problem_shape_b, (values[0],) + ) + problem_shape_s = cutlass.new_from_mlir_values( + self.problem_shape_s, (values[1],) + ) + split_kv = cutlass.new_from_mlir_values(self.split_kv, (values[2],)) + problem_shape_b_fdd = cutlass.new_from_mlir_values( + self.problem_shape_b_fdd, (values[3],) + ) + problem_shape_s_fdd = cutlass.new_from_mlir_values( + self.problem_shape_s_fdd, (values[4],) + ) + split_kv_fdd = cutlass.new_from_mlir_values(self.split_kv_fdd, (values[5],)) + return MLAStaticTileSchedulerParams( + self.is_persistent, + problem_shape_b, + problem_shape_s, + self.cluster_shape_mnk, + split_kv, + problem_shape_b_fdd=problem_shape_b_fdd, + problem_shape_s_fdd=problem_shape_s_fdd, + split_kv_fdd=split_kv_fdd, + loc=self.loc, + ) + + +def create_mla_static_tile_scheduler_params( + is_persistent: bool, + problem_shape_b: cute.Int32, + problem_shape_s: cute.Int32, + cluster_shape_mnk: cute.Shape, + split_kv: cutlass.Int32, +) -> MLAStaticTileSchedulerParams: + return MLAStaticTileSchedulerParams( + is_persistent, problem_shape_b, problem_shape_s, cluster_shape_mnk, split_kv + ) + + +class WorkTileInfo: + def __init__(self, blk_coord: cute.Coord, is_valid: bool): + self.blk_coord = blk_coord + self.is_valid = cutlass.Boolean(is_valid) + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.blk_coord) + values += cutlass.extract_mlir_values(self.is_valid) + return values + + def __new_from_mlir_values__(self, values): + new_tile_idx = cutlass.new_from_mlir_values(self.blk_coord, values[:-1]) + new_is_valid_tile = cutlass.new_from_mlir_values(self.is_valid, [values[-1]]) + return WorkTileInfo(new_tile_idx, new_is_valid_tile) + + @property + def is_valid_tile(self) -> cutlass.Boolean: + return self.is_valid + + @property + def tile_idx(self) -> cute.Coord: + return self.blk_coord + + +class MLAStaticTileScheduler: + def __init__( + self, + params: MLAStaticTileSchedulerParams, + current_work_linear_idx: cutlass.Int32, + blk_coord: cute.Coord, + grid_shape: cute.Shape, + *, + is_valid: bool = True, + loc=None, + ip=None, + ): + """The static tile scheduler for MLA split kv kernel. + Based on `is_persistent`, it provides 2 modes for use: + - Persistent mode: Launch fixed blocks and reschedule the data blocks. + - Non-persistent mode: Launch dynamic blocks and exit when the current work is done. + + :param params: The static tile scheduler parameters + :type params: MLAStaticTileSchedulerParams + :param current_work_linear_idx: The linear index of the current work + :type current_work_linear_idx: cutlass.Int32 + :param blk_coord: The coordinate of the current work + :type blk_coord: cute.Coord + :param grid_shape: The shape of the grid + :type grid_shape: cute.Shape + :param is_valid: Whether the current work is valid + :type is_valid: bool + """ + self.params = params + self.blk_coord = blk_coord + self.grid_shape = grid_shape + self.current_work_linear_idx = current_work_linear_idx + if params.is_persistent: + self.persistent_blk_layout = cute.make_layout( + ( + params.cluster_shape_mnk[0], + params.problem_shape_s, + params.problem_shape_b, + params.split_kv, + ), + loc=loc, + ip=ip, + ) + self.num_blocks = cute.size(self.persistent_blk_layout, loc=loc, ip=ip) + # Used for persistent scheduling + self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip) + else: + self.is_valid = is_valid + self.loc = loc + self.ip = ip + + @staticmethod + def get_grid_shape( + params: MLAStaticTileSchedulerParams, + max_active_clusters: int, + *, + loc=None, + ip=None, + ) -> cute.Shape: + # called by host + grid_shape = ( + params.cluster_shape_mnk[0], + params.problem_shape_b * params.problem_shape_s, + params.split_kv, + ) + if params.is_persistent: + return ( + cutlass.min( + max_active_clusters * cute.size(params.cluster_shape_mnk), + cute.size(grid_shape, loc=loc, ip=ip), + ), + 1, + 1, + ) + else: + return grid_shape + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + is_valid = ( + self.current_work_linear_idx < self.num_blocks + if self.params.is_persistent + else self.is_valid + ) + + if self.params.is_persistent: + current_work_cluster_batch, cluster_idx = ( + self.current_work_linear_idx // self.params.cluster_shape_mnk[0], + self.current_work_linear_idx % self.params.cluster_shape_mnk[0], + ) + current_work_s_batch, s_idx = divmod( + current_work_cluster_batch, self.params.problem_shape_s_fdd + ) + current_work_b_batch, b_idx = divmod( + current_work_s_batch, self.params.problem_shape_b_fdd + ) + _, split_kv_idx = divmod(current_work_b_batch, self.params.split_kv_fdd) + + blk_coord = (cluster_idx, s_idx, b_idx, split_kv_idx) + else: + s_idx, b_idx = divmod(self.blk_coord[1], self.params.problem_shape_b_fdd) + blk_coord = (self.blk_coord[0], s_idx, b_idx, self.blk_coord[2]) + + return WorkTileInfo(blk_coord, is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + return self.get_current_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None): + if self.params.is_persistent: + self.current_work_linear_idx += advance_count * self.num_persistent_sm + else: + self.is_valid = False + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.params) + values.extend(cutlass.extract_mlir_values(self.current_work_linear_idx)) + values.extend(cutlass.extract_mlir_values(self.blk_coord)) + values.extend(cutlass.extract_mlir_values(self.grid_shape)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 13 + new_params = cutlass.new_from_mlir_values(self.params, values[0:6]) + new_current_work_linear_idx = cutlass.new_from_mlir_values( + self.current_work_linear_idx, [values[6]] + ) + new_blk_coord = cutlass.new_from_mlir_values(self.blk_coord, values[7:10]) + new_grid_shape = cutlass.new_from_mlir_values(self.grid_shape, values[10:]) + return MLAStaticTileScheduler( + new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape + ) + + +def create_mla_static_tile_scheduler( + params: MLAStaticTileSchedulerParams, + blk_coord: cute.Coord, + grid_shape: cute.Shape, +) -> MLAStaticTileScheduler: + return MLAStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape) + + +LOG2_E = 1.4426950408889634074 +# avoid register indexing on array. +MAX_SPLITS = 256 + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b diff --git a/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py new file mode 100644 index 00000000..39f14848 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py @@ -0,0 +1,3039 @@ +# 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 argparse +from typing import Optional, Type, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm103_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import from_dlpack +from dataclasses import dataclass, field + +""" +This example provides an experimental implementation of the SM103 batched 3xFP4 blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. + +A high-performance persistent batched 3xFP4 blockscaled GEMM example for the NVIDIA Blackwell SM103 architecture +using CUTE DSL. + - Matrix A is MxKxL, L is batch dimension, A can only be row-major("K") for MXF4/NVF4 input type + - Matrix B is NxKxL, L is batch dimension, B can only be row-major("K") for MXF4/NVF4 input type + - Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") + - Matrix SFA layout is filled internally according to A shape and sm103_BlockScaledBasicChunk, which has M×ceil_div(K, sf_vec_size)×L elements respectively + - Matrix SFB layout is filled internally according to B shape and sm103_BlockScaledBasicChunk, which has N×ceil_div(K, sf_vec_size)×L elements respectively + +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 with separate TMA warps for A/B and scale factors + - Utilizes circular buffer technique for optimal memory and computation overlap + +This GEMM works as follows: + 1. TMA A/B warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. + 2. TMA SF warp: Load scale factor A/B from global memory (GMEM) to shared memory (SMEM) using TMA operations. + 3. MMA warp: + - Load scale factor A/B from shared memory (SMEM) to tensor memory (TMEM) using tcgen05.cp instruction. + - Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction to deal with circular buffering. + 4. Epilogue warps: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Store C matrix directly 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)) + +SM103 tcgen05.mma.kind.block_scale instructions operate as follows: + - Read matrix A from two SMEM buffers(current buffer and next buffer) + - Read matrix B from two SMEM buffers(current buffer and next buffer) + - Read scalefactor A from TMEM + - Read scalefactor B from TMEM + - 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 shown below: + +.. code-block:: bash + + python examples/blackwell/sm103_dense_blockscaled_gemm_persistent.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,4 \ + --mnkl 4096,4096,6144,1 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/sm103_dense_blockscaled_gemm_persistent.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,4 \ + --mnkl 4096,4096,6144,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + +Constraints: + - Supported input data types: mxf4, nvf4 + - see detailed valid dtype combinations in below Sm103BlockScaledPersistentDenseGemmKernel class documentation + - A/B tensor must have the same data type + - Mma tiler M must be 128 or 256(use_2cta_instrs) + - Mma tiler N must be 128 or 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs) + - The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 32 for MXF4/NVF4. +""" + + +class Sm103BlockScaledPersistentDenseGemmKernel: + """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for FP4 data types + and architectural features specific to Blackwell SM103 GPUs with persistent tile scheduling and warp specialization. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :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] + + + :note: In current version, A and B tensor must have the same data type + - i.e., Float4E2M1FN for A and Float4E2M1FN for B is not supported + + :note: Supported combinations of A/B data types, SF data typs and SF vector size: + - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16 + + :note: Supported accumulator data types: + - Float32 + + :note: Supported C data types: + - Float32 + - Float16/BFloat16 + - Float8E4M3FN/Float8E5M2 + :note: Constraints: + - MMA tiler M must be 128 or 256 (use_2cta_instrs) + - MMA tiler N must be 128/256 + - Cluster shape M must be multiple of 2 if Mma tiler M is 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors + + Example: + >>> gemm = Sm103BlockScaledPersistentDenseGemmKernel( + ... sf_vec_size=16, + ... mma_tiler_mn=(256, 256), + ... cluster_shape_mn=(2, 4) + ... ) + >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, max_active_clusters, stream) + """ + + def __init__( + self, + sf_vec_size: int, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_tma_store: bool, + ): + """Initializes the configuration for a Blackwell SM103 3xFP4 GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator, always set to Float32 + - sf_vec_size: Scalefactor A/B vector size. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + """ + self.acc_dtype = cutlass.Float32 + self.sf_vec_size = sf_vec_size + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + self.use_tma_store = use_tma_store + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.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_ab_warp_id = 5 + self.tma_sf_warp_id = 6 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_ab_warp_id, + self.tma_sf_warp_id, + *self.epilogue_warp_id, + ) + ) + # Set barrier id for 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_103") + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_103") + self.sf_buffers_per_tile_k = 4 if self.sf_vec_size == 16 else 2 + + def _setup_attributes(self): + """Set up kernel attributes that depend on runtime tensor 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/SFA/SFB + - Computing epilogue subtile + - Setting up A/B/SFA/SFB/C stage counts in shared memory + - Computing A/B/SFA/SFB/C shared memory layout + """ + # Compute mma instruction shapes + # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) + self.mma_inst_shape_mn = (self.mma_tiler[0], self.mma_tiler[1]) + + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma = self.sm103_make_blockscaled_trivial_tiled_mma( + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + dummy_tiled_mma_sfb = self.sm103_make_blockscaled_trivial_tiled_mma( + self.sf_dtype, + self.sf_vec_size, + tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + + # Compute mma/cluster/tile shapes + self.mma_tiler = ( + self.mma_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + 768, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_layout_vmnk.shape[0]), + self.mma_tiler[1], + self.mma_tiler[2], + ) + blk_mn = 128 + self.cta_n_sf = cute.round_up(cute.size(self.cta_tile_shape_mnk[1]), blk_mn) + self.mma_sf_tiler = ( + self.cta_tile_shape_mnk[0], + self.cta_n_sf, + self.cta_tile_shape_mnk[2] // self.sf_buffers_per_tile_k, + ) + + self.sf_atom = self.Sm103BlockScaledBasicChunk( + self.sf_vec_size, tiled_mma.op.a_major_mode + ).layout + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (dummy_tiled_mma_sfb.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.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute epilogue subtile + self.epi_tile = sm103_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + self.num_acc_stage, self.num_ab_stage, self.num_sf_stage, self.num_c_stage = ( + self._compute_stages( + tiled_mma, + self.mma_tiler, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + self.use_tma_store, + ) + ) + + # Compute A/B/SFA/SFB/C shared memory layout + # ((CTA_MMA_M,16bytes),1,8,num_ab_stage) + self.a_smem_layout_staged = self.sm103_make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.num_ab_stage, + ) + + # ((CTA_MMA_M,16bytes),1,8,3) + self.a_smem_layout_staged_tma = self.sm103_make_smem_layout_a( + tiled_mma, + self.mma_tiler, + 3, + ) + + # ((CTA_MMA_N,16bytes),1,8,num_ab_stage) + self.b_smem_layout_staged = self.sm103_make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.num_ab_stage, + ) + + # ((CTA_MMA_N,16bytes),1,8,3) + self.b_smem_layout_staged_tma = self.sm103_make_smem_layout_b( + tiled_mma, + self.mma_tiler, + 3, + ) + + # (((8,4,4),(sf_vec_size,4)),1,3,num_sf_stage) + self.sfa_smem_layout_staged = self.sm103_make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_sf_stage, + ) + + # (((32,4,2),(sf_vec_size,4)),1,3,num_sf_stage) + self.sfb_smem_layout_staged = self.sm103_make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_sf_stage, + ) + self.c_smem_layout_staged = None + if self.use_tma_store: + self.c_smem_layout_staged = sm103_utils.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case + self.overlapping_accum = self.num_acc_stage == 1 and not self.use_tma_store + self.epi_tile_n = cute.size(self.epi_tile[1]) + + if self.overlapping_accum: + # Compute SF TMEM column count from a scale factor layout. + # Column count = cosize of Int32-recast layout & 0xFFFF, + # mirroring the computation in find_tmem_tensor_col_offset. + def _sf_tmem_cols(make_tmem_layout_fn, smem_layout_staged): + layout = make_tmem_layout_fn( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(smem_layout_staged, (None, None, None, 0)), + ) + return ( + cute.cosize(cute.recast_layout(32, self.sf_dtype.width, layout)) + & 0xFFFF + ) + + self.num_sfa_tmem_cols = _sf_tmem_cols( + blockscaled_utils.make_tmem_layout_sfa, self.sfa_smem_layout_staged + ) + self.num_sfb_tmem_cols = _sf_tmem_cols( + blockscaled_utils.make_tmem_layout_sfb, self.sfb_smem_layout_staged + ) + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + # Release accumulator buffer early in epilogue when overlapping + self.iter_acc_early_release_in_epilogue = ( + self.num_sf_tmem_cols // self.epi_tile_n + ) + + @cute.jit + def __call__( + self, + a_tensor: cute.Tensor, + b_tensor: cute.Tensor, + sfa_tensor: cute.Tensor, + sfb_tensor: cute.Tensor, + c_tensor: 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_tensor: Input tensor A + :type a_tensor: cute.Tensor + :param b_tensor: Input tensor B + :type b_tensor: cute.Tensor + :param sfa_tensor: Scale factor tensor A + :type sfa_tensor: cute.Tensor + :param sfb_tensor: Scale factor tensor B + :type sfb_tensor: cute.Tensor + :param c_tensor: Output tensor C + :type c_tensor: 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. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type + self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type + self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type + self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) + # 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() + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + sfa_layout = cute.tile_to_shape(self.sf_atom, a_tensor.shape, (2, 1, 3)) + sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) + + sfb_layout = cute.tile_to_shape(self.sf_atom, b_tensor.shape, (2, 1, 3)) + sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) + + tiled_mma = self.sm103_make_blockscaled_trivial_tiled_mma( + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + dummy_tiled_mma_sfb = self.sm103_make_blockscaled_trivial_tiled_mma( + self.sf_dtype, + self.sf_vec_size, + tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm103_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + # casting layout as uint8 for multicast + a_smem_layout_tma_ready = self.adapt_layout_for_tma_ab( + self.a_smem_layout_staged_tma + ) + a_tensor_uint8 = cute.recast_tensor(a_tensor, cutlass.Uint8) + tma_atom_a, tma_tensor_a = cute.nvgpu.cpasync.make_tiled_tma_atom( + a_op, + a_tensor_uint8, + a_smem_layout_tma_ready, + # 384 corresponds to the number of uint8 elements along the K dimension processed in a single MMA mainloop iteration. + (cute.size(tiled_mma.tv_layout_A[1][0]), 384), + self.cluster_shape_mn[1], + internal_type=cutlass.Uint8, + ) + + # Setup TMA load for B + b_op = sm103_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + # casting layout as uint8 for multicast + b_smem_layout_tma_ready = self.adapt_layout_for_tma_ab( + self.b_smem_layout_staged_tma + ) + b_tensor_uint8 = cute.recast_tensor(b_tensor, cutlass.Uint8) + tma_atom_b, tma_tensor_b = cute.nvgpu.cpasync.make_tiled_tma_atom( + b_op, + b_tensor_uint8, + b_smem_layout_tma_ready, + (cute.size(tiled_mma.tv_layout_B[1][0]), 384), + self.cluster_shape_mn[0] // cute.size(tiled_mma.thr_id.shape), + internal_type=cutlass.Uint8, + ) + + # Setup TMA load for SFA + sfa_op = sm103_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + sfa_smem_layout_tma_ready = self.adapt_layout_for_tma_sf(sfa_smem_layout) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.cpasync.make_tiled_tma_atom( + sfa_op, + sfa_tensor, + sfa_smem_layout_tma_ready, + (self.mma_sf_tiler[0], self.mma_sf_tiler[2]), + self.cluster_shape_mn[1], + internal_type=cutlass.Uint8, + ) + + # Setup TMA load for SFB + sfb_op = sm103_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + sfb_smem_layout_tma_ready = self.adapt_layout_for_tma_sf(sfb_smem_layout) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.cpasync.make_tiled_tma_atom( + sfb_op, + sfb_tensor, + sfb_smem_layout_tma_ready, + (self.mma_sf_tiler[1], self.mma_sf_tiler[2]), + self.cluster_shape_mn[0] // cute.size(dummy_tiled_mma_sfb.thr_id), + internal_type=cutlass.Uint8, + ) + + # 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_tensor, + epi_smem_layout, + self.epi_tile, + ) + + a_copy_size = cute.size_in_bytes( + cutlass.Uint8, + cute.slice_(self.a_smem_layout_staged_tma, (None, None, None, 0)), + ) + b_copy_size = cute.size_in_bytes( + cutlass.Uint8, + cute.slice_(self.b_smem_layout_staged_tma, (None, None, None, 0)), + ) + sfa_copy_size = cute.size_in_bytes( + cutlass.Uint8, + cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)), + ) + sfb_copy_size = cute.size_in_bytes( + cutlass.Uint8, + cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)), + ) + self.num_tma_load_bytes_ab = (a_copy_size + b_copy_size) * atom_thr_size + self.num_tma_load_bytes_sf = (sfa_copy_size + sfb_copy_size) * atom_thr_size + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c_tensor, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + self.buffer_align_bytes = 1024 + + # 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] + sf_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_sf_stage] + sf_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_sf_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 + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Uint8, 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[ + cutlass.Uint8, cute.cosize(self.b_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sSFA: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Uint8, cute.cosize(self.sfa_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sSFB: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Uint8, cute.cosize(self.sfb_smem_layout_staged) + ], + 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_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c_tensor, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_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, + min_blocks_per_mp=1, + ) + 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_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + 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_ab_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) + if warp_idx == self.tma_sf_warp_id: + cpasync.prefetch_descriptor(tma_atom_sfa) + cpasync.prefetch_descriptor(tma_atom_sfb) + + 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 + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, sfa+sfb full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize mainloop ab_producer and ab_consumer + ab_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_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_producer_group, + consumer_group=ab_consumer_group, + tx_count=self.num_tma_load_bytes_ab, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + # Initialize mainloop sf_producer and sf_consumer + sf_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_sf_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + sf_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_sf_tma_producer + ) + sf_producer, sf_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.sf_full_mbar_ptr.data_ptr(), + num_stages=self.num_sf_stage, + producer_group=sf_producer_group, + consumer_group=sf_consumer_group, + tx_count=self.num_tma_load_bytes_sf, + 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=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/SFA/SFB/C + # + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_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 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (BLK_M, BLK_K, m, k, l) + gA_mkl = cute.local_tile( + mA_mkl, + cute.slice_((self.mma_tiler[0], self.mma_tiler[1], 384), (None, 0, None)), + (None, None, None), + ) + # (BLK_N, BLK_K, n, k, l) + gB_nkl = cute.local_tile( + mB_nkl, + cute.slice_((self.mma_tiler[0], self.mma_tiler[1], 384), (0, None, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + mSFA_mkl, + cute.slice_(self.mma_sf_tiler, (None, 0, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + mSFB_nkl, + cute.slice_(self.mma_sf_tiler, (0, None, None)), + (None, None, None), + ) + 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) + + # create tCgA_tmp + tCgA_mkl_tmp = thr_mma.partition_A(gA_mkl) + tCgA_layout = self.append_coalesce_layout(tCgA_mkl_tmp.layout) + cta_tCgA = cute.make_tensor(tCgA_mkl_tmp.iterator, tCgA_layout) + # ((CTA_MMA_M,256),Rest_MMA_M,Rest_MMA_K, m, k, l) + tCgA = cute.make_tensor( + cta_tCgA.iterator, + cute.tiled_divide( + cta_tCgA.layout, (cute.size(tiled_mma.tv_layout_A[1][0]), 128) + ), + ) + + tCgB_nkl_tmp = thr_mma.partition_B(gB_nkl) + tCgB_layout = self.append_coalesce_layout(tCgB_nkl_tmp.layout) + cta_tCgB = cute.make_tensor(tCgB_nkl_tmp.iterator, tCgB_layout) + # ((CTA_MMA_N,256),Rest_MMA_N, Rest_MMA_K, n, k, l) + tCgB = cute.make_tensor( + cta_tCgB.iterator, + cute.tiled_divide( + cta_tCgB.layout, (cute.size(tiled_mma.tv_layout_B[1][0]), 128) + ), + ) + + tCgSFA = cute.make_tensor( + gSFA_mkl.iterator, + cute.tiled_divide( + gSFA_mkl.layout, (self.mma_sf_tiler[0], self.mma_sf_tiler[2]) + ), + ) + + tCgSFB = cute.make_tensor( + gSFB_nkl.iterator, + cute.tiled_divide( + gSFB_nkl.layout, (self.mma_sf_tiler[1], self.mma_sf_tiler[2]) + ), + ) + tCgC = thr_mma.partition_C(gC_mnl) + + # Create identity tensor for C to use in epilogue predication + idC = cute.make_identity_tensor(mC_mnl.shape) + cC_mnl = cute.local_tile( + idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCcC = thr_mma.partition_C(cC_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 + ) + + 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, 1), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + 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, 1), + ) + + # TMA partition for scale factor A + sfa_cta_layout = a_cta_layout + tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA_compact = cute.filter_zeros(tAsSFA) + + # TMA partition for scale factor B + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB_compact = cute.filter_zeros(tBsSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + if cutlass.const_expr(self.overlapping_accum): + num_acc_stage_overlapped = 2 + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, num_acc_stage_overlapped) + ) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = cute.make_tensor( + tCtAcc_fake.iterator, + cute.make_layout( + tCtAcc_fake.shape, + stride=( + tCtAcc_fake.stride[0], + tCtAcc_fake.stride[1], + tCtAcc_fake.stride[2], + (self.cta_tile_shape_mnk[1] - self.num_sf_tmem_cols) + * tCtAcc_fake.stride[0][1], + ), + ), + ) + else: + # (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=self.cluster_shape_mn) + + # + # Construct the scheduler + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp for A/B tensors + # + if warp_idx == self.tma_ab_warp_id: + # + # Persistent tile scheduling loop for AB loads + # + buffers_per_k_tile = 3 + + 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 + # + tAgA_slice = tAgA[ + ( + None, + None, + None, + mma_tile_coord_mnl[0], + None, + mma_tile_coord_mnl[2], + ) + ] + tBgB_slice = tBgB[ + ( + None, + None, + 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 = cutlass.Boolean(1) + peek_ab_empty_status = ab_producer.try_acquire() + + # + # TMA load loop for A/B tensors + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Load buffers_per_k_tile buffers + for buffer in cutlass.range(buffers_per_k_tile, unroll_full=True): + # Acquire next empty AB buffer + ab_empty = ab_producer.acquire_and_advance(peek_ab_empty_status) + + # TMA load A/B + cute.copy( + tma_atom_a, + cute.group_modes( + tAgA_slice[(None, None, buffer, k_tile)], 0, 2 + ), + tAsA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + cute.group_modes( + tBgB_slice[(None, None, buffer, k_tile)], 0, 2 + ), + tBsB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for next buffer + peek_ab_empty_status = cutlass.Boolean(1) + # Check if we're not at the last buffer of the last k_tile + if not ( + (k_tile == k_tile_cnt - 1) + and (buffer == buffers_per_k_tile - 1) + ): + 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() + + # Signal end of AB loads + ab_producer.tail() + + # + # Specialized TMA load warp for scale factor tensors + # + if warp_idx == self.tma_sf_warp_id: + # + # Persistent tile scheduling loop for SF loads + # + 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], + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + tAgSFA_slice = tAgSFA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + tBgSFB_slice = tBgSFB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) SF buffer empty + sf_producer.reset() + peek_sf_empty_status = cutlass.Boolean(1) + peek_sf_empty_status = sf_producer.try_acquire() + + # + # TMA load loop for scale factors + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Load SF stages based on sf_buffers_per_tile_k + for sf_stage in cutlass.range( + self.sf_buffers_per_tile_k, unroll_full=True + ): + # Acquire next empty SF buffer + sf_empty = sf_producer.acquire_and_advance(peek_sf_empty_status) + + tAgSFA_compact = cute.filter_zeros( + tAgSFA_slice[ + (None, k_tile * self.sf_buffers_per_tile_k + sf_stage) + ] + ) + tBgSFB_compact = cute.filter_zeros( + tBgSFB_slice[ + (None, k_tile * self.sf_buffers_per_tile_k + sf_stage) + ] + ) + + # TMA load SFA/SFB for this SF stage + cute.copy( + tma_atom_sfa, + tAgSFA_compact, + tAsSFA_compact[(None, sf_empty.index)], + tma_bar_ptr=sf_empty.barrier, + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_compact, + tBsSFB_compact[(None, sf_empty.index)], + tma_bar_ptr=sf_empty.barrier, + mcast_mask=sfb_full_mcast_mask, + ) + + # Peek (try_wait) SF buffer empty for next stage + peek_sf_empty_status = cutlass.Boolean(1) + # Check if we're not at the last stage of the last k_tile + if not ( + k_tile == k_tile_cnt - 1 + and sf_stage == self.sf_buffers_per_tile_k - 1 + ): + peek_sf_empty_status = sf_producer.try_acquire() + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Signal end of SF loads + sf_producer.tail() + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # Make accumulator tmem tensor + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Make SFA tmem tensor + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + dtype=self.sf_dtype, + ) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + + MMA_M = self.cta_tile_shape_mnk[0] + MMA_N_SF = self.cta_n_sf + MMA_K_SF = self.cta_tile_shape_mnk[2] // 2 + mnBasicBlockShape = (32, 4) + kBasicBlockShape_single = (self.sf_vec_size, 1) + mma_iter_SFA_shape = ( + (mnBasicBlockShape, MMA_M // 128), + kBasicBlockShape_single, + ) + sSFA_iter_shape = (mma_iter_SFA_shape, 1, MMA_K_SF // self.sf_vec_size) + sSFA_iter_layout = cute.make_layout(sSFA_iter_shape) + mma_iter_SFB_shape = ( + (mnBasicBlockShape, MMA_N_SF // 128), + kBasicBlockShape_single, + ) + sSFB_iter_shape = (mma_iter_SFB_shape, 1, MMA_K_SF // self.sf_vec_size) + sSFB_iter_layout = cute.make_layout(sSFB_iter_shape) + + tCtSFA_layout_mma = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, self.mma_tiler, self.sf_vec_size, sSFA_iter_layout + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + tCtSFA_mma = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout_mma) + + # Make SFB tmem tensor + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=self.sf_dtype, + ) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB_layout_mma = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, self.mma_tiler, self.sf_vec_size, sSFB_iter_layout + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + tCtSFB_mma = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout_mma) + + # + # Partition for S2T copy of SFA/SFB + # + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + # + # Persistent tile scheduling loop + # + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + MmasPerSfBuffer = 8 // self.sf_buffers_per_tile_k + sf_stride = 6 if self.sf_vec_size == 16 else 3 + + 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], + ) + + # Get accumulator stage index + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_producer_state.phase ^ 1 + else: + acc_stage_index = acc_producer_state.index + + # Set tensor memory buffer for current tile + tCtAcc = tCtAcc_base[(None, 0, 0, acc_stage_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() + + # Peek (try_wait) SF buffer full + sf_consumer.reset() + peek_sf_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_sf_full_status = sf_consumer.try_wait() + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + is_first_iteration = True + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + if is_leader_cta: + # Conditionally load SFA/SFB for MMA0/MMA1 depending on sf_vec_size + if 0 % MmasPerSfBuffer == 0: + sf_full = sf_consumer.wait_and_advance(peek_sf_full_status) + s2t_stage_coord = ( + None, + None, + None, + None, + sf_full.index, + ) + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t[s2t_stage_coord], + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t[s2t_stage_coord], + tCtSFB_compact_s2t, + ) + sf_full.release() + peek_sf_full_status = cutlass.Boolean(1) + peek_sf_full_status = sf_consumer.try_wait() + + # Wait for A/B data to be ready(MMA0, MMA1, part of MMA2) + ab_full0 = ab_consumer.wait_and_advance(peek_ab_full_status) + + # peek for next stage (MMA2, MMA3, MMA4, part of MMA5) + peek_ab_full_status = cutlass.Boolean(1) + peek_ab_full_status = ab_consumer.try_wait() + + # delay the acc acquire to ublock tmem + if is_first_iteration: + acc_pipeline.producer_acquire(acc_producer_state) + is_first_iteration = False + + # MMA0 + k_block_coord_cur = (None, 0, 0, ab_full0.index) + k_block_coord_next = (None, 0, 0, ab_full0.index) + sf_kblock_coord = (None, None, 0 % MmasPerSfBuffer * sf_stride) + tiled_mma.set( + tcgen05.Field.SFA, tCtSFA_mma[sf_kblock_coord].iterator + ) + tiled_mma.set( + tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator + ) + self.make_desc_and_call_mma( + tiled_mma, + tCtAcc, + sA[k_block_coord_cur], + sA[k_block_coord_next], + sB[k_block_coord_cur], + sB[k_block_coord_next], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # MMA1 + k_block_coord_cur = (None, 0, 3, ab_full0.index) + k_block_coord_next = (None, 0, 0, ab_full0.index) + sf_kblock_coord = (None, None, 1 % MmasPerSfBuffer * sf_stride) + tiled_mma.set( + tcgen05.Field.SFA, tCtSFA_mma[sf_kblock_coord].iterator + ) + tiled_mma.set( + tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator + ) + self.make_desc_and_call_mma( + tiled_mma, + tCtAcc, + sA[k_block_coord_cur], + sA[k_block_coord_next], + sB[k_block_coord_cur], + sB[k_block_coord_next], + tCtAcc, + ) + + # Conditionally load SFA/SFB for MMA2/MMA3 + if 2 % MmasPerSfBuffer == 0: + sf_full = sf_consumer.wait_and_advance(peek_sf_full_status) + s2t_stage_coord = ( + None, + None, + None, + None, + sf_full.index, + ) + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t[s2t_stage_coord], + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t[s2t_stage_coord], + tCtSFB_compact_s2t, + ) + sf_full.release() + peek_sf_full_status = cutlass.Boolean(1) + peek_sf_full_status = sf_consumer.try_wait() + + # Wait for A/B data to be ready(MMA2, MMA3, MMA4, part of MMA5) + ab_full1 = ab_consumer.wait_and_advance(peek_ab_full_status) + + # peek for next stage (part of MMA5, MMA6, MMA7) + peek_ab_full_status = cutlass.Boolean(1) + peek_ab_full_status = ab_consumer.try_wait() + + # MMA2 + k_block_coord_cur = (None, 0, 6, ab_full0.index) + k_block_coord_next = (None, 0, 0, ab_full1.index) + sf_kblock_coord = (None, None, 2 % MmasPerSfBuffer * sf_stride) + tiled_mma.set( + tcgen05.Field.SFA, tCtSFA_mma[sf_kblock_coord].iterator + ) + tiled_mma.set( + tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator + ) + self.make_desc_and_call_mma( + tiled_mma, + tCtAcc, + sA[k_block_coord_cur], + sA[k_block_coord_next], + sB[k_block_coord_cur], + sB[k_block_coord_next], + tCtAcc, + ) + + # Release stage_ab_0 as it is no longer needed + ab_full0.release() + + # MMA3 + k_block_coord_cur = (None, 0, 1, ab_full1.index) + k_block_coord_next = (None, 0, 0, ab_full1.index) + sf_kblock_coord = (None, None, 3 % MmasPerSfBuffer * sf_stride) + tiled_mma.set( + tcgen05.Field.SFA, tCtSFA_mma[sf_kblock_coord].iterator + ) + tiled_mma.set( + tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator + ) + self.make_desc_and_call_mma( + tiled_mma, + tCtAcc, + sA[k_block_coord_cur], + sA[k_block_coord_next], + sB[k_block_coord_cur], + sB[k_block_coord_next], + tCtAcc, + ) + + # Conditionally load SFA/SFB for MMA4/MMA5 + if 4 % MmasPerSfBuffer == 0: + sf_full = sf_consumer.wait_and_advance(peek_sf_full_status) + s2t_stage_coord = ( + None, + None, + None, + None, + sf_full.index, + ) + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t[s2t_stage_coord], + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t[s2t_stage_coord], + tCtSFB_compact_s2t, + ) + sf_full.release() + peek_sf_full_status = cutlass.Boolean(1) + peek_sf_full_status = sf_consumer.try_wait() + + # MMA4 + k_block_coord_cur = (None, 0, 4, ab_full1.index) + k_block_coord_next = (None, 0, 0, ab_full1.index) + sf_kblock_coord = (None, None, 4 % MmasPerSfBuffer * sf_stride) + tiled_mma.set( + tcgen05.Field.SFA, tCtSFA_mma[sf_kblock_coord].iterator + ) + tiled_mma.set( + tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator + ) + self.make_desc_and_call_mma( + tiled_mma, + tCtAcc, + sA[k_block_coord_cur], + sA[k_block_coord_next], + sB[k_block_coord_cur], + sB[k_block_coord_next], + tCtAcc, + ) + + # Wait for A/B data to be ready(part of MMA5, MMA6, MMA7) + ab_full2 = ab_consumer.wait_and_advance(peek_ab_full_status) + + # peek for next loop's first stage (MMA0, MMA1, part of MMA2) + peek_ab_full_status = cutlass.Boolean(1) + if k_tile + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + # MMA5 + k_block_coord_cur = (None, 0, 7, ab_full1.index) + k_block_coord_next = (None, 0, 0, ab_full2.index) + sf_kblock_coord = (None, None, 5 % MmasPerSfBuffer * sf_stride) + tiled_mma.set( + tcgen05.Field.SFA, tCtSFA_mma[sf_kblock_coord].iterator + ) + tiled_mma.set( + tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator + ) + self.make_desc_and_call_mma( + tiled_mma, + tCtAcc, + sA[k_block_coord_cur], + sA[k_block_coord_next], + sB[k_block_coord_cur], + sB[k_block_coord_next], + tCtAcc, + ) + + # Conditionally load SFA/SFB for MMA6/MMA7 + if 6 % MmasPerSfBuffer == 0: + sf_full = sf_consumer.wait_and_advance(peek_sf_full_status) + s2t_stage_coord = ( + None, + None, + None, + None, + sf_full.index, + ) + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t[s2t_stage_coord], + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t[s2t_stage_coord], + tCtSFB_compact_s2t, + ) + sf_full.release() + peek_sf_full_status = cutlass.Boolean(1) + if k_tile + 1 < k_tile_cnt: + peek_sf_full_status = sf_consumer.try_wait() + + ab_full1.release() + + # MMA6 + k_block_coord_cur = (None, 0, 2, ab_full2.index) + k_block_coord_next = (None, 0, 0, ab_full2.index) + sf_kblock_coord = (None, None, 6 % MmasPerSfBuffer * sf_stride) + tiled_mma.set( + tcgen05.Field.SFA, tCtSFA_mma[sf_kblock_coord].iterator + ) + tiled_mma.set( + tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator + ) + self.make_desc_and_call_mma( + tiled_mma, + tCtAcc, + sA[k_block_coord_cur], + sA[k_block_coord_next], + sB[k_block_coord_cur], + sB[k_block_coord_next], + tCtAcc, + ) + + # MMA7 + k_block_coord_cur = (None, 0, 5, ab_full2.index) + k_block_coord_next = (None, 0, 0, ab_full2.index) + sf_kblock_coord = (None, None, 7 % MmasPerSfBuffer * sf_stride) + tiled_mma.set( + tcgen05.Field.SFA, tCtSFA_mma[sf_kblock_coord].iterator + ) + tiled_mma.set( + tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator + ) + self.make_desc_and_call_mma( + tiled_mma, + tCtAcc, + sA[k_block_coord_cur], + sA[k_block_coord_next], + sB[k_block_coord_cur], + sB[k_block_coord_next], + tCtAcc, + ) + + ab_full2.release() + + 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) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + if cutlass.const_expr(self.use_tma_store): + assert tma_atom_c is not None and sC is not None + 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 + ) + 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], + ) + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + tCcC_base=tCcC, + mC_mnl=mC_mnl, + overlapping_accum=self.overlapping_accum, + ) + + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete + c_pipeline.producer_tail() + else: + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(acc_tmem_ptr) + + @staticmethod + def make_desc_and_call_mma( + tiled_mma: cute.TiledMma, + d: cute.Tensor, + sA_cur: cute.Tensor, + sA_next: cute.Tensor, + sB_cur: cute.Tensor, + sB_next: cute.Tensor, + c: cute.Tensor, + ) -> None: + """Specialized GEMM for circular-buffered A/B from SMEM. + + Performs D <- A * B + C where A and B are described by circular SMEM + descriptors constructed from the (current, next) buffers. C and D may alias. + + Some tcgen05 MMAs require explicitly toggling an accumulate field outside of + this routine; the caller is responsible for that. + + All tensors must already be partitioned for the provided tiled MMA. + + For MMA Atoms that require single-threaded execution, the gemm op automatically handles thread + election internally. Manual thread selection is not required in such cases. + + :param atom: MMA atom + :type atom: cute.MmaAtom + :param d: Destination tensor + :type d: cute.Tensor + :param sA_cur: Current shared memory tensor for operand A + :type sA_cur: cute.Tensor + :param sA_next: Next shared memory tensor for operand A, used for circular buffering + :type sA_next: cute.Tensor + :param sB_cur: Current shared memory tensor for operand B + :type sB_cur: cute.Tensor + :param sB_next: Next shared memory tensor for operand B, used for circular buffering + :type sB_next: cute.Tensor + :param c: Third source tensor + :type c: cute.Tensor + :return: None + :rtype: None + """ + a_desc = tcgen05.make_umma_smem_desc( + sA_cur.iterator, + sA_cur.layout, + "k" if tiled_mma.op.a_major_mode.name == "K" else "mn", + next_src=sA_next.iterator, + ) + b_desc = tcgen05.make_umma_smem_desc( + sB_cur.iterator, + sB_cur.layout, + "k" if tiled_mma.op.b_major_mode.name == "K" else "mn", + next_src=sB_next.iterator, + ) + + view_layout = cute.make_layout(1, stride=0) + a_tensor = cute.make_tensor(a_desc, view_layout) + b_tensor = cute.make_tensor(b_desc, view_layout) + return cute.mma_atom_call(tiled_mma, d, a_tensor, b_tensor, c) + + @staticmethod + def sm103_make_blockscaled_trivial_tiled_mma( + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + cta_group: tcgen05.CtaGroup, + mma_tiler_mn: Tuple[int, int], + a_source: tcgen05.OperandSource = tcgen05.OperandSource.SMEM, + ) -> cute.TiledMma: + """Create a blockscaled trivial tiled MMA for SM103 (3xFP4), K fixed to 96. + + Returns a tcgen05 MMA configured for the given (M, N) tiler and CTA group. + + :param sf_dtype: Data type of the scale factor (typically 8-bit) + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param cta_group: The CTA group configuration + :type cta_group: tcgen05.CtaGroup + :param mma_tiler_mn: The MMA tiler dimensions (M, N) + :type mma_tiler_mn: Tuple[int, int] + :param a_source: Source location for operand A (SMEM by default) + :type a_source: tcgen05.OperandSource + + :return: A tiled MMA atom configured for SM103 blockscaled operations + :rtype: cute.TiledMma + + :raises TypeError: If the data type is not supported. + :raises ValueError: If the sf_vec_size is not supported. + """ + if sf_vec_size == 32: + mma_op = tcgen05.SM103MmaMXF4Op( + (*mma_tiler_mn, 96), + cta_group, + a_source, + ) + elif sf_vec_size == 16: + mma_op = tcgen05.SM103MmaMXF4NVF4Op( + sf_dtype, + (*mma_tiler_mn, 96), + cta_group, + a_source, + ) + else: + raise ValueError( + f"Unsupported sf_vec_size: {sf_vec_size}. Expected 16 or 32." + ) + return cute.make_tiled_mma(cute.make_mma_atom(mma_op)) + + # Utils + @staticmethod + def sm103_make_smem_layout_a( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: cute.Tile, + num_stages: int, + ) -> Union[cute.Layout, cute.ComposedLayout]: + """ + Create the SMEM layout for operand A using K_SW128 and Uint8. + + This function creates a SMEM layout for operand A using the make_smem_layout_atom function with K_SW128 kind and Uint8 element type. + + :param tiled_mma: The tiled MMA atom + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The mma tiler shape (M, N, K) + :type mma_tiler_mnk: cute.Tile + :param num_stages: The number of stages + :type num_stages: int + + :return: SMEM layout for operand A + :rtype: cute.Layout + """ + is_k_major = tiled_mma.op.a_major_mode == tcgen05.OperandMajorMode.K + a_smem_layout_staged = tcgen05.tile_to_mma_shape( + tcgen05.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.Uint8 + ), + cute.append( + ( + ( + mma_tiler_mnk[0] + // cute.size(tiled_mma.thr_layout_vmnk.shape[0]), + 16, + ), + 1, + 8, + ), + num_stages, + ), + order=((1, 0, 2) if not is_k_major else (0, 1, 2)), + ) + + return a_smem_layout_staged + + @staticmethod + def sm103_make_smem_layout_b( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: cute.Tile, + num_stages: int, + ) -> Union[cute.Layout, cute.ComposedLayout]: + """ + Create the SMEM layout for operand B using K_SW128 and Uint8. + + This function creates a SMEM layout for operand B using the make_smem_layout_atom function with K_SW128 kind and Uint8 element type. + + :param tiled_mma: The tiled MMA atom + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The mma tiler shape (M, N, K) + :type mma_tiler_mnk: cute.Tile + :param num_stages: The number of stages + :type num_stages: int + + :return: SMEM layout for operand B + :rtype: cute.Layout + """ + is_k_major = tiled_mma.op.b_major_mode == tcgen05.OperandMajorMode.K + b_smem_layout_staged = tcgen05.tile_to_mma_shape( + tcgen05.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.Uint8 + ), + cute.append( + ((mma_tiler_mnk[1] // cute.size(tiled_mma.thr_id.shape), 16), 1, 8), + num_stages, + ), + order=((1, 0, 2) if not is_k_major else (0, 1, 2)), + ) + return b_smem_layout_staged + + @dataclass(frozen=True) + class Sm103BlockScaledBasicChunk: + """ + Basic scale-factor atom layout decided by tcgen05 BlockScaled MMA Ops on SM103. + + Represents the fixed layout pattern for scale factors used by tcgen05 + BlockScaled MMA Ops on SM103. The layout is determined by the instruction + specification and is not configurable. + """ + + sf_vec_size: int + major_mode: tcgen05.OperandMajorMode = tcgen05.OperandMajorMode.K + _layout: cute.Layout = field(init=False, repr=False) + + def __post_init__(self) -> None: + if self.major_mode == tcgen05.OperandMajorMode.K: + atom_shape = ((8, 4, 4), (self.sf_vec_size, 4)) + atom_stride = ((16, 128, 4), (0, 1)) + else: + atom_shape = ((self.sf_vec_size, 4), (8, 4, 4)) + atom_stride = ((0, 1), (16, 128, 4)) + + object.__setattr__( + self, "_layout", cute.make_layout(shape=atom_shape, stride=atom_stride) + ) + + @property + def layout(self) -> cute.Layout: + return self._layout + + @staticmethod + def sm103_make_smem_layout_sfa( + tiled_mma: cute.TiledMma, + mma_tiler: cute.Tile, + sf_vec_size: int, + num_stages: int, + ) -> cute.Layout: + """ + Make SMEM layout for SFA based on: + 1) Sm103BlockScaledBasicChunk, 2) MMA tiler, 3) sf_vec_size, 4) stages. + + :param tiled_mma: The tiled MMA + :type tiled_mma: cute.TiledMma + :param mma_tiler: The mma tiler shape + :type mma_tiler: cute.Tile + :param sf_vec_size: The scale factor vector size + :type sf_vec_size: int + :param num_stages: The number of stages + :type num_stages: int + + :return: Smem layout for SFA + :rtype: cute.Layout + """ + mma_shape_mk = tiled_mma.partition_shape_A((mma_tiler[0], mma_tiler[2])) + sf_atom = Sm103BlockScaledPersistentDenseGemmKernel.Sm103BlockScaledBasicChunk( + sf_vec_size, tiled_mma.op.a_major_mode + ).layout + k_divisor = 4 if sf_vec_size == 16 else 2 + mma_sfa_tiler = ( + mma_shape_mk[0][0] * mma_shape_mk[1], + mma_shape_mk[0][1] * mma_shape_mk[2] // k_divisor, + ) + sfa_smem_atom_layout = cute.tiled_product( + sf_atom, + cute.make_layout( + cute.shape_div(mma_sfa_tiler, cute.product_each(sf_atom.shape)) + ), + ) + sfa_smem_layout_staged = cute.make_layout( + shape=cute.append(sfa_smem_atom_layout.shape, num_stages), + stride=cute.append( + sfa_smem_atom_layout.stride, + cute.size(cute.filter_zeros(sfa_smem_atom_layout)), + ), + ) + return sfa_smem_layout_staged + + @staticmethod + def sm103_make_smem_layout_sfb( + tiled_mma: cute.TiledMma, + mma_tiler: cute.Tile, + sf_vec_size: int, + num_stages: int, + ) -> cute.Layout: + """ + Make SMEM layout for SFB based on the basic chunk, MMA tiler, sf_vec_size, stages. + + :param tiled_mma: The tiled MMA + :type tiled_mma: cute.TiledMma + :param mma_tiler: The mma tiler shape + :type mma_tiler: cute.Tile + :param sf_vec_size: The scale factor vector size + :type sf_vec_size: int + :param num_stages: The number of stages + :type num_stages: int + + :return: Smem layout for SFB + :rtype: cute.Layout + """ + sf_atom = Sm103BlockScaledPersistentDenseGemmKernel.Sm103BlockScaledBasicChunk( + sf_vec_size, tiled_mma.op.a_major_mode + ).layout + k_divisor = 4 if sf_vec_size == 16 else 2 + mma_sfb_tiler = (mma_tiler[1], mma_tiler[2] // k_divisor) + if mma_sfb_tiler[0] == 128: + sfb_smem_atom_layout = cute.tiled_product( + sf_atom, + cute.make_layout( + cute.shape_div(mma_sfb_tiler, cute.product_each(sf_atom.shape)) + ), + ) + else: + sf_k_major_atom256 = cute.make_layout( + shape=( + (32, 4, 2), + (sf_vec_size, 4), + ), + stride=( + (16, 4, mma_sfb_tiler[1] // sf_vec_size // 4 * 512), + (0, 1), + ), + ) + sfb_smem_atom_layout = cute.tiled_product( + sf_k_major_atom256, + cute.make_layout( + cute.shape_div( + mma_sfb_tiler, cute.product_each(sf_k_major_atom256.shape) + ) + ), + ) + + sfb_smem_layout_staged = cute.make_layout( + shape=cute.append(sfb_smem_atom_layout.shape, num_stages), + stride=cute.append( + sfb_smem_atom_layout.stride, + cute.size(cute.filter_zeros(sfb_smem_atom_layout)), + ), + ) + return sfb_smem_layout_staged + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). + + :param sSF: The scale factor tensor in smem + :type sSF: cute.Tensor + :param tSF: The scale factor tensor in tmem + :type tSF: cute.Tensor + + :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: + - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) + - tCsSF_compact_s2t: The partitioned scale factor tensor in smem + - tSF_compact_s2t: The partitioned scale factor tensor in tmem + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSF_compact = cute.filter_zeros(sSF) + # (MMA, MMA_MN, MMA_K) + tCtSF_compact = cute.filter_zeros(tSF) + tCtSF_compact_copy = cute.make_tensor( + tCtSF_compact.iterator, + cute.append( + cute.append(tCtSF_compact[(None, 0, 0)].layout, cute.make_layout((1))), + cute.make_layout(1), + ), + ) + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact_copy) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler: Tuple[int, int, int], + epi_tile: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + ) -> Tuple[int, int, int]: + """Computes the number of stages for A/B and SF operands based on heuristics. + + SM103 requires separate stage counts for AB and SF pipelines. + + :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 tiler. + :type mma_tiler: tuple[int, int, int] + :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 sf_dtype: Data type of Scale factor. + :type sf_dtype: type[cutlass.Numeric] + :param sf_vec_size: Scale factor vector size. + :type sf_vec_size: int + :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, SF stages) + :rtype: tuple[int, int, int] + """ + # ACC stages - same as SM100 dense blockscaled gemm + num_acc_stage = 1 if mma_tiler[1] == 256 else 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, SFA, SFB + a_smem_layout_stage_one = ( + Sm103BlockScaledPersistentDenseGemmKernel.sm103_make_smem_layout_a( + tiled_mma, + mma_tiler, + 1, + ) + ) + b_smem_layout_staged_one = ( + Sm103BlockScaledPersistentDenseGemmKernel.sm103_make_smem_layout_b( + tiled_mma, + mma_tiler, + 1, + ) + ) + sfa_smem_layout_staged_one = ( + Sm103BlockScaledPersistentDenseGemmKernel.sm103_make_smem_layout_sfa( + tiled_mma, + mma_tiler, + sf_vec_size, + 1, + ) + ) + sfb_smem_layout_staged_one = ( + Sm103BlockScaledPersistentDenseGemmKernel.sm103_make_smem_layout_sfb( + tiled_mma, + mma_tiler, + sf_vec_size, + 1, + ) + ) + + c_smem_layout_staged_one = sm103_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + + ab_bytes_per_stage = cute.size_in_bytes( + cutlass.Uint8, a_smem_layout_stage_one + ) + cute.size_in_bytes(cutlass.Uint8, b_smem_layout_staged_one) + sf_bytes_per_stage = cute.size_in_bytes( + sf_dtype, sfa_smem_layout_staged_one + ) + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + + mbar_helpers_bytes = 1024 + + num_ab_stage = ( + smem_capacity // occupancy + - (mbar_helpers_bytes + sf_bytes_per_stage + c_bytes) + ) // ab_bytes_per_stage + + num_sf_stage = ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * mbar_helpers_bytes + - occupancy * c_bytes + ) // (occupancy * sf_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: + # xinyu TODO: not sure if aligned with c++ + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * sf_bytes_per_stage * num_sf_stage + - occupancy * mbar_helpers_bytes + - occupancy * c_bytes + ) // (occupancy * c_bytes_per_stage) + + return num_acc_stage, num_ab_stage, num_sf_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 is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + ) -> bool: + """ + Check if the dtypes and sf_vec_size are valid combinations + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes and sf_vec_size are valid, False otherwise + :rtype: bool + """ + is_valid = True + + # Check valid ab_dtype + if ab_dtype != cutlass.Float4E2M1FN: + is_valid = False + + # Check valid sf_vec_size + if sf_vec_size not in {16, 32}: + is_valid = False + + # Check valid sf_dtype + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + is_valid = False + + # Check valid sf_dtype and sf_vec_size combinations + if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32: + is_valid = False + + # Check valid c_dtype + if c_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + return is_valid + + @staticmethod + def is_valid_layouts( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if layouts and dtypes are valid combinations + + :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 dimension of the A tensor + :type a_major: str + :param b_major: The major dimension of the B tensor + :type b_major: str + :param c_major: The major dimension of the C tensor + :type c_major: str + + :return: True if the layouts are valid, False otherwise + :rtype: bool + """ + is_valid = True + + if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"): + is_valid = False + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ) -> bool: + """ + Check if the mma tiler and cluster shape are valid + + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + + :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 mma_tiler_mn[0] in [128, 256]: + is_valid = False + if not mma_tiler_mn[1] in [128, 256]: + is_valid = False + # Skip illegal cluster shape + if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + # Special cluster shape check for scale factor multicasts. + # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + 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_alignment( + dtype, is_mode0_major, tensor_shape, alignment_bytes + ): + """Check if tensor satisfies the required byte alignment. + + :param dtype: Data type of the tensor + :param is_mode0_major: Whether mode 0 is the major (contiguous) mode + :param tensor_shape: Shape of the tensor (mode0, mode1, batch) + :param alignment_bytes: Required alignment in bytes (e.g., 16 or 32) + :return: True if alignment is satisfied + """ + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + # Calculate number of contiguous elements needed for alignment + # alignment_bytes * 8 (bits per byte) / dtype.width (bits per element) + num_contiguous_elements = alignment_bytes * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + # Check A/B tensors for 16B alignment + # Check C tensor for 32B alignment + if ( + not check_contigous_alignment(ab_dtype, a_major == "m", (m, k, l), 16) + or not check_contigous_alignment(ab_dtype, b_major == "n", (n, k, l), 16) + or not check_contigous_alignment(c_dtype, c_major == "m", (m, n, l), 32) + ): + is_valid = False + return is_valid + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + c_major: str, + use_tma_store: bool, + ) -> bool: + """ + Check if the gemm can be implemented + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + :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 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 gemm can be implemented, False otherwise + :rtype: bool + """ + can_implement = True + # Skip unsupported types + if not Sm103BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype, sf_dtype, sf_vec_size, c_dtype + ): + can_implement = False + # Skip unsupported layouts + if not Sm103BlockScaledPersistentDenseGemmKernel.is_valid_layouts( + ab_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not Sm103BlockScaledPersistentDenseGemmKernel.is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn, cluster_shape_mn + ): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not Sm103BlockScaledPersistentDenseGemmKernel.is_valid_tensor_alignment( + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + return can_implement + + # Helper function for append and coalesce layout + @staticmethod + def append_coalesce_layout(layout): + # coalesce is like: cutlass/python/pycute/layout.py:coalesce + part1 = cute.coalesce(cute.append(layout[0][0], layout[1])) + part2 = cute.coalesce(cute.append(layout[0][1], layout[2])) + result = cute.append(part1, part2) + result = cute.append(result, layout[3]) + result = cute.append(result, layout[4]) + result = cute.append(result, layout[5]) + return result + + @staticmethod + def adapt_layout_for_tma_ab(composed_layout): + # input: S<3,4,3> o 0 o ((128,16),1,8,3):((128,1),0,16,16384) + # output: S<3,4,3> o 0 o (128,(128,3)):(128,(1,16384)) + # for ctaValueMap: (128,384):(1@0,1@1) + layout = composed_layout.outer + part1 = cute.coalesce(cute.append(layout[0][0], layout[1])) + part2 = cute.coalesce(cute.append(layout[0][1], layout[2])) + part3 = cute.append(part2, layout[3]) + result = cute.append(part1, part3) + return cute.make_composed_layout( + composed_layout.inner, composed_layout.offset, result + ) + + @staticmethod + def adapt_layout_for_tma_sf(layout): + # TODO: need ethan check this + # input: (((8,4,4),(16,4)),1,3):(((16,128,4),(0,1)),0,512) + # output: ((32,4),(16,4,3)):((16,4),(0,1,512)) + # for ctaValueMap: ((8,4,4),(16,4,3)):((1@0@0@0,1@1@0@0,1@2@0@0),(1@0@0@1,1@1@0@1,1@1@1)) + part1 = cute.coalesce(cute.append(layout[0][0], layout[1])) + part2 = cute.coalesce(cute.append(layout[0][1], layout[2])) + result = cute.append(cute.group_modes(part1, 0, cute.rank(part1)), part2) + return result + + +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_tensor: cute.Tensor, + sf_mma_tensor: cute.Tensor, +): + """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout""" + # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) + # group to ((32, 4, rest_m), (4, rest_k), l) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + 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 persistent batched dense blockscaled 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 sf_dtype: Data type for scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: Vector size for scale factor tensor + :type sf_vec_size: int + :param c_dtype: Data type for output tensor C + :type c_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. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :param use_2cta_instrs: Whether to use 2CTA instructions. + :type use_2cta_instrs: bool, optional + :param use_tma_store: Whether to use TMA store. + :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(f"Running Sm103 Persistent 3xfp4 Dense BlockScaled GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}") + print(f"C dtype: {c_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"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'}") + + import torch + import cutlass.torch as cutlass_torch + + # Unpack parameters + m, n, k, l = mnkl + + # Skip unsupported testcase + if not Sm103BlockScaledPersistentDenseGemmKernel.can_implement( + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + c_major, + use_tma_store, + ): + raise TypeError( + f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}, " + f"use_tma_store: {use_tma_store}" + ) + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Create tensor A/B/C + a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) + b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) + c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) + + a_tensor, a_torch = cutlass_torch.cute_tensor_like( + a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, b_torch = cutlass_torch.cute_tensor_like( + b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=32 + ) + + # Mark tensor with byte alignment divisibility + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=64 if ab_dtype == cutlass.Float4E2M1FN else 32, + ) + + # Create scale factor tensor SFA/SFB + def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + ref_shape = (l, mn, sf_k) + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + ref_permute_order = (1, 2, 0) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + # Create f32 ref torch tensor (cpu) + ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + ref_shape, + torch.float32, + permute_order=ref_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=1, + max_val=3, + ), + ) + + # Create f32 cute torch tensor (cpu) + cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.SCALAR, + init_config=cutlass_torch.ScalarInitConfig(value=1.0), + ) + + # convert ref f32 tensor to cute f32 tensor + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + from_dlpack(ref_f32_torch_tensor_cpu), + from_dlpack(cute_f32_torch_tensor_cpu), + ) + cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() + + # reshape makes memory contiguous + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(l, mn, sf_k, sf_vec_size) + .reshape(l, mn, sf_k * sf_vec_size) + .permute(*ref_permute_order) + ) + # prune to mkl for reference check. + ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] + + # Create dtype cute torch tensor (cpu) + cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( + cute_f32_torch_tensor_cpu, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + # Convert f32 cute tensor to dtype cute tensor + cute_tensor = cutlass_torch.convert_cute_tensor( + cute_f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=True, + ) + return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor + + sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype + ) + sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype + ) + + # Configure gemm kernel + gemm = Sm103BlockScaledPersistentDenseGemmKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + ) + + # Compute max active clusters on current device + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + # Compile gemm kernel + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + max_active_clusters, + current_stream, + ) + # Compute reference result + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_gemm( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, current_stream + ) + print("Verifying results...") + res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) + res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + + # Convert c back to f32 for comparison. + c_ref_device = c_ref.cuda() + cute.testing.convert( + c_tensor, + from_dlpack(c_ref_device, assumed_align=32).mark_layout_dynamic( + leading_dim=(1 if c_major == "n" else 0) + ), + ) + c_ref = c_ref_device.cpu() + + if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Convert ref : f32 -> f8 -> f32 + ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute( + 1, 2, 0 + ) + ref_f8 = from_dlpack(ref_f8_, assumed_align=32).mark_layout_dynamic( + leading_dim=1 + ) + ref_f8.element_type = c_dtype + ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() + ref_tensor = from_dlpack(ref_device, assumed_align=32).mark_layout_dynamic( + leading_dim=1 + ) + cute.testing.convert(ref_tensor, ref_f8) + cute.testing.convert(ref_f8, ref_tensor) + ref = ref_device.cpu() + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + def generate_tensors(): + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, _ = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=32 + ) + + # Mark tensor to be byte aligned + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=64 if ab_dtype == cutlass.Float4E2M1FN else 32, + ) + + _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype) + _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype) + return cute.testing.JitArguments( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, current_stream + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch.numel() * a_torch.element_size() + + b_torch.numel() * b_torch.element_size() + + sfa_torch.numel() * sfa_torch.element_size() + + sfb_torch.numel() * sfb_torch.element_size() + + c_torch.numel() * c_torch.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +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 Sm103 3xfp4 Dense Persistent BlockScaled GEMM." + ) + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(4096, 4096, 6144, 2), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=parse_comma_separated_ints, + default=(256, 256), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=parse_comma_separated_ints, + default=(2, 4), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) + parser.add_argument("--sf_vec_size", type=int, default=16) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) + parser.add_argument("--a_major", choices=["k"], type=str, default="k") + parser.add_argument("--b_major", choices=["k"], type=str, default="k") + parser.add_argument("--c_major", choices=["n"], 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") + + run( + args.mnkl, + args.ab_dtype, + args.sf_dtype, + args.sf_vec_size, + args.c_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_tma_store, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py index 341b47d2..6f05b56e 100644 --- a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual @@ -9,13 +9,11 @@ # its affiliates is strictly prohibited. import argparse -import torch from typing import Tuple import cutlass import cutlass.cute as cute import cutlass.utils as utils -import cutlass.torch as cutlass_torch import cutlass.pipeline as pipeline from cutlass.cute.nvgpu import cpasync, tcgen05 import cutlass.utils.blackwell_helpers as sm100_utils @@ -32,9 +30,8 @@ with optimizations for challenges that may arise with other problem sizes. To run this example: .. code-block:: bash - python examples/blackwell/tutorial_fp16_gemm_0.py \ - --mnk 8192,8192,8192 \ - --tolerance 1e-01 + python examples/blackwell/tutorial_gemm/fp16_gemm_0.py \ + --mnk 8192,8192,8192 Constraints for this example: * The problem size of m and n must be divisible by the tile size m & n (128, 256) @@ -128,7 +125,8 @@ def kernel( num_stages=acc_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, threads_per_cta + pipeline.Agent.Thread, + threads_per_cta, ), barrier_storage=storage.acc_mbar_ptr.data_ptr(), ).make_participants() @@ -141,15 +139,15 @@ def kernel( # (bM, bN) gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None)) thr_mma = tiled_mma.get_slice(0) - # (MMA, MMA_M, MMA_K, RestK) + # (MMA, MMA_M, MMA_K) tCgA = thr_mma.partition_A(gA) - # (MMA, MMA_N, MMA_K, RestK) + # (MMA, MMA_N, MMA_K) tCgB = thr_mma.partition_B(gB) # (MMA, MMA_M, MMA_N) tCgC = thr_mma.partition_C(gC) - # (MMA, MMA_M, MMA_K, STAGE) + # (MMA, MMA_M, MMA_K) tCrA = tiled_mma.make_fragment_A(sA) - # (MMA, MMA_N, MMA_K, STAGE) + # (MMA, MMA_N, MMA_K) tCrB = tiled_mma.make_fragment_B(sB) # (MMA, MMA_M, MMA_N) acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) @@ -188,7 +186,7 @@ def kernel( # (EpiTile, NumTiles) gC_epi = cute.zipped_divide(tCgC, epi_tiler) - # Every thread loads 32x128 bits + # Every thread loads 64 x fp32 tmem_atom = cute.make_copy_atom( tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32, @@ -273,11 +271,7 @@ def kernel( @cute.jit -def host_function( - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, -): +def host_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor): # Construct tiled MMA op = tcgen05.MmaF16BF16Op( io_dtype, @@ -354,6 +348,10 @@ def run_dense_gemm( mnk: Tuple[int, int, int], tolerance: float, ): + global torch, cutlass_torch + import torch + import cutlass.torch as cutlass_torch + print("===================================================================") print("Running Blackwell fp16 GEMM example 0 with:") print(f" mnk: {mnk}") @@ -393,12 +391,7 @@ def run_dense_gemm( ) # Entry point to the host JIT function - host_function( - a_tensor, - b_tensor, - c_tensor, - no_cache=True, - ) + host_function(a_tensor, b_tensor, c_tensor, no_cache=True) # Compute reference result and verify ref = (torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))).cpu() @@ -418,7 +411,11 @@ if __name__ == "__main__": "Invalid format. Expected comma-separated integers." ) - if not torch.cuda.is_available(): + from cuda.bindings import driver as cu_driver + + cu_driver.cuInit(0) + err, device_count = cu_driver.cuDeviceGetCount() + if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1: raise RuntimeError("A GPU is required to run this example") parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 0") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py index 1d1d98ae..12e74fa0 100644 --- a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual @@ -13,15 +13,14 @@ import argparse -import torch from typing import Tuple import cutlass import cutlass.cute as cute import cutlass.utils as utils -import cutlass.torch as cutlass_torch import cutlass.pipeline as pipeline from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack """ @@ -78,7 +77,6 @@ acc_stage = 1 @cute.struct class SharedStorage: - # each stage has 2 kinds of barrier, i.e. empty & full ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2] acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stage * 2] tmem_dealloc_mbar_ptr: cutlass.Int64 @@ -174,15 +172,15 @@ def kernel( # (bM, bN) gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None)) thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0]) - # (MMA, MMA_M, MMA_K, RestK) + # (MMA, MMA_M, MMA_K) tCgA = thr_mma.partition_A(gA) - # (MMA, MMA_N, MMA_K, RestK) + # (MMA, MMA_N, MMA_K) tCgB = thr_mma.partition_B(gB) # (MMA, MMA_M, MMA_N) tCgC = thr_mma.partition_C(gC) - # (MMA, MMA_M, MMA_K, STAGE) + # (MMA, MMA_M, MMA_K) tCrA = tiled_mma.make_fragment_A(sA) - # (MMA, MMA_N, MMA_K, STAGE) + # (MMA, MMA_N, MMA_K) tCrB = tiled_mma.make_fragment_B(sB) # (MMA, MMA_M, MMA_N) acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) @@ -256,9 +254,9 @@ def kernel( tDgC = tmem_thr_copy.partition_D(gC_epi) # (TmemCpy,NumTmemCpy) - tCrAcc = cute.make_rmem_tensor_like(tDgC[None, None, 0], acc_dtype) + tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype) # (TmemCpy,NumTmemCpy) - tCrC = cute.make_rmem_tensor_like(tDgC[None, None, 0], io_dtype) + tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype) # # 2. Main loop @@ -356,13 +354,13 @@ def host_function( tiled_mma = cute.make_tiled_mma(op) # Construct SMEM layouts for A and B - a_smem_layout = utils.sm100.make_smem_layout_a( + a_smem_layout = sm100_utils.make_smem_layout_a( tiled_mma, mma_tiler_mnk, a.element_type, ab_stages, ) - b_smem_layout = utils.sm100.make_smem_layout_b( + b_smem_layout = sm100_utils.make_smem_layout_b( tiled_mma, mma_tiler_mnk, b.element_type, @@ -383,7 +381,7 @@ def host_function( a_smem_layout_one_stage, mma_tiler_mnk, tiled_mma, - cta_layout_vmnk.shape, + cta_layout_vmnk.shape, # take the layout and extract the shape internally ) b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B( op, @@ -438,6 +436,10 @@ def run_dense_gemm( mnk: Tuple[int, int, int], tolerance: float, ): + global torch, cutlass_torch + import torch + import cutlass.torch as cutlass_torch + print("===================================================================") print("Running Blackwell fp16 GEMM example 1 with:") print(f" mnk: {mnk}") @@ -501,7 +503,11 @@ if __name__ == "__main__": "Invalid format. Expected comma-separated integers." ) - if not torch.cuda.is_available(): + from cuda.bindings import driver as cu_driver + + cu_driver.cuInit(0) + err, device_count = cu_driver.cuDeviceGetCount() + if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1: raise RuntimeError("A GPU is required to run this example") parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 1") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py new file mode 100644 index 00000000..f0dfeea1 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py @@ -0,0 +1,778 @@ +# 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 argparse +import os +import sys +from typing import Type, Tuple +import cuda.bindings.driver as cuda + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import make_ptr + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + examples_dir = os.path.join(current_dir, "..", "..") + if examples_dir not in sys.path: + sys.path.insert(0, examples_dir) + +from blackwell.tutorial_gemm.utils import create_parser, run + +mma_tiler_mn = (128, 256) +mma_inst_shape_k = 64 +ab_dtype = cutlass.Float4E2M1FN +sf_dtype = cutlass.Float8E4M3FN +c_dtype = cutlass.Float16 +sf_vec_size = 16 + +""" +The first tutorial NVFP4 block-scaled batched GEMM demonstrating a simple kernel implementation in CuTeDSL + +This example demonstrates the kernel implementation of block-scaled batched GEMM with NVFP4 data type. +With large tile sizes (128x256x256), it can achieve very high performance on 8k×8k×8k problem sizes. +It can serve as a starting point to help users quickly experiment with optimizations for +challenges that may arise with other problem sizes. + +To run this example: +.. code-block:: bash + + python examples/blackwell/tutorial_gemm/nvfp4_gemm_0.py \ + --mnkl 8192,8192,8192,1 --do_benchmark + +Constraints for this example: +* The problem size of m, n and k must be divisible by the tile size m&n&k (128,256,256) +* The scaling factor vector size is 16. +* The A/B matrices have data contiguous on the k dimension. +* The C matrix has data contiguous on the n dimension. +* The A/B matrix data type is Float4E2M1FN. +* The SFA/SFB matrix data type is Float8E4M3FN. +""" + + +class Sm100BlockScaledDenseGemmKernel: + def __init__(self): + self.threads_per_cta = 128 + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = 512 + + # set stages for ab_pipeline and acc_pipeline + self.num_acc_stage = 1 + self.num_ab_stage = 4 + + @cute.jit + def __call__( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + problem_size: tuple, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # setup static attributes before smem/grid/tma computation + self.c_layout = utils.LayoutEnum.ROW_MAJOR + m, n, k, l = problem_size + + # Setup attributes that depend on gemm inputs + mma_inst_tile_k = 4 + self.mma_tiler = ( + mma_tiler_mn[0], + mma_tiler_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0], + self.mma_tiler[1], + self.mma_tiler[2], + ) + + a_tensor = cute.make_tensor( + a_ptr, + cute.make_layout( + (m, cute.assume(k, 32), l), + stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)), + ), + ) + b_tensor = cute.make_tensor( + b_ptr, + cute.make_layout( + (n, cute.assume(k, 32), l), + stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)), + ), + ) + # make address offset of c_tensor 256bit aligned, + # so that epilogue could use vectorized store with larger vector size. + c_tensor = cute.make_tensor( + c_ptr, + cute.make_layout( + (cute.assume(m, 32), cute.assume(n, 16), l), + stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)), + ), + ) + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) + + mma_op = tcgen05.MmaMXF4NVF4Op( + sf_dtype, + (*mma_tiler_mn, mma_inst_shape_k), + tcgen05.CtaGroup.ONE, + tcgen05.OperandSource.SMEM, + ) + tiled_mma = cute.make_tiled_mma(mma_op) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute A/B/SFA/SFB/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + ab_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + ab_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + sf_vec_size, + self.num_ab_stage, + ) + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # TMA load for A + 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( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + a_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for B + 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( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + b_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load for SFA + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # TMA load for SFB + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + sfb_tensor, + sfb_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Compute TMA load bytes + a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Compute grid size + grid = ( + cute.ceil_div(c_tensor.shape[0], self.cta_tile_shape_mnk[0]), + cute.ceil_div(c_tensor.shape[1], self.cta_tile_shape_mnk[1]), + c_tensor.shape[2], + ) + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + c_tensor, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(1, 1, 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_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + mC_mnl: cute.Tensor, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + tidx, _, _ = cute.arch.thread_idx() + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + + # Coords outside cluster + cta_coord = (bidx, bidy, bidz) + mma_tile_coord_mnl = ( + cta_coord[0] // cute.size(tiled_mma.thr_id.shape), + cta_coord[1], + cta_coord[2], + ) + + # + # Define shared storage for kernel + # + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=ab_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=ab_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = smem.allocate_tensor( + element_type=sf_dtype, + layout=sfa_smem_layout_staged, + byte_alignment=128, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = smem.allocate_tensor( + element_type=sf_dtype, + layout=sfb_smem_layout_staged, + byte_alignment=128, + ) + + # + # Initialize mainloop ab_pipeline, acc_pipeline and their states + # + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_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, + ).make_participants() + acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_cta, + ), + ).make_participants() + + # + # 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) + ) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gSFB_nkl = cute.local_tile( + mSFB_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/SFA/SFB/C + # + # (MMA, MMA_M, MMA_K, RestK) + thr_mma = tiled_mma.get_slice(0) + # (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_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma.partition_B(gSFB_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/SFA/SFB + # + # TMA load A partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load partition for SFA tensor + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + 0, + cute.make_layout(1), + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load partition for SFB tensor + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + 0, + cute.make_layout(1), + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # 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) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + + # + # Alloc tensor memory buffer + # + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + ) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # + # Make SFA/SFB tmem tensor + # + # Get SFA tmem ptr + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc), + dtype=sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + # Get SFB tmem ptr + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + # + # Partition for S2T copy of SFA/SFB + # + # Make S2T CopyAtom + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), + sf_dtype, + ) + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSFA_compact = cute.filter_zeros(sSFA) + # (MMA, MMA_MN, MMA_K) + tCtSFA_compact = cute.filter_zeros(tCtSFA) + tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact) + thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t_sfa, tCsSFA_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact) + + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSFB_compact = cute.filter_zeros(sSFB) + # (MMA, MMA_MN, MMA_K) + tCtSFB_compact = cute.filter_zeros(tCtSFB) + tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact) + thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t_sfb, tCsSFB_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # + # Execute Data copy and Math computation in the k_tile loop + # + if warp_idx == 0: + # Wait for accumulator buffer empty + acc_empty = acc_producer.acquire_and_advance() + # Set ACCUMULATE field to False for the first k_tile iteration + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Execute k_tile loop + for k_tile in cutlass.range( + k_tile_cnt, prefetch_stages=self.num_ab_stage - 2 + ): + # Wait for AB buffer empty + ab_empty = ab_producer.acquire_and_advance() + + # TMA load for A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA[(None, ab_empty.count)], + tAsA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + ) + cute.copy( + tma_atom_b, + tBgB[(None, ab_empty.count)], + tBsB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + ) + cute.copy( + tma_atom_sfa, + tAgSFA[(None, ab_empty.count)], + tAsSFA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + ) + cute.copy( + tma_atom_sfb, + tBgSFB[(None, ab_empty.count)], + tBsSFB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + ) + + # Wait for AB buffer full + ab_full = ab_consumer.wait_and_advance() + + # Copy SFA/SFB to tmem + s2t_stage_coord = (None, None, None, None, ab_full.index) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + 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_full.index, + ) + + # Set SFA/SFB tensor to tiled_mma + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set( + tcgen05.Field.SFA, + tCtSFA[sf_kblock_coord].iterator, + ) + tiled_mma.set( + tcgen05.Field.SFB, + tCtSFB[sf_kblock_coord].iterator, + ) + + 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_full.release() + acc_empty.commit() + + # + # Epilogue + # Partition for epilogue + # + op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE) + copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc) + # (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(tCgC) + # (T2R_M, T2R_N, EPI_M, EPI_N） + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32 + ) + # (T2R_M, T2R_N, EPI_M, EPI_N） + tTR_rC = cute.make_rmem_tensor( + tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype + ) + # STG Atom + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype) + tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)] + + # Release TMEM allocation lock + tmem.relinquish_alloc_permit() + + # Wait for accumulator buffer full + acc_full = acc_consumer.wait_and_advance() + + # Copy accumulator to register + cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc) + acc_vec = epilogue_op(tTR_rAcc.load().to(c_dtype)) + tTR_rC.store(acc_vec) + # Store C to global memory + cute.copy(simt_atom, tTR_rC, tTR_gC) + + acc_full.release() + + # Deallocate TMEM + cute.arch.barrier() + tmem.free(acc_tmem_ptr) + + return + + +def run_nvfp4_gemm( + mnkl: Tuple[int, int, int, int], + tolerance: float, + do_benchmark: bool = False, + warmup_iterations: int = 10, + iterations: int = 100, + use_cold_l2: bool = True, +): + run( + gemm_class=Sm100BlockScaledDenseGemmKernel, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mnk=(1, 1, 1), + mnkl=mnkl, + tolerance=tolerance, + do_benchmark=do_benchmark, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cold_l2=use_cold_l2, + ) + + +if __name__ == "__main__": + parser = create_parser() + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + m, n, k, _ = args.mnkl + if m % mma_tiler_mn[0] != 0: + parser.error("m must be multiples of mma_tiler_mn[0] (got m={})".format(m)) + if n % mma_tiler_mn[1] != 0: + parser.error("n must be multiples of mma_tiler_mn[1] (got n={})".format(n)) + if k % 256 != 0: + parser.error("k must be a multiple of 256 (got k={})".format(k)) + + run_nvfp4_gemm( + args.mnkl, + args.tolerance, + args.do_benchmark, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py new file mode 100644 index 00000000..715de1b7 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py @@ -0,0 +1,934 @@ +# 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. + +# This is the second tutorial nvfp4 GEMM. It builds on the first tutorial by adding 2CTA MMA +# instructions with a 2x1 cluster. + +import argparse +import os +import sys +from typing import Type, Tuple +import cuda.bindings.driver as cuda + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import from_dlpack, make_ptr + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + examples_dir = os.path.join(current_dir, "..", "..") + if examples_dir not in sys.path: + sys.path.insert(0, examples_dir) + +from blackwell.tutorial_gemm.utils import create_parser, run + +mma_tiler_mn = (256, 256) +mma_inst_shape_k = 64 +ab_dtype = cutlass.Float4E2M1FN +sf_dtype = cutlass.Float8E4M3FN +c_dtype = cutlass.Float16 +sf_vec_size = 16 +cluster_shape_mnk = (2, 1, 1) + +""" +The second tutorial further improves the performance of NVFP4 block-scaled batched GEMM +by adding 2CTA instructions and TMA multicast optimizations. + +(1) The 2 CTA instructions could reduce the smem size requirement for B tensor, +increased num_ab_stage and improves the latency hiding capability. + +For both 1CTA and 2CTA, the shared memory (smem) size per stage for the A, sfA, and sfB tensors is the same: +- For the A tensor, each stage requires 128 x 256 x sizeof(float4) = 16KB. +- For the sfA tensor, each stage requires 128 x (256 / 16) x sizeof(float8) = 2KB. +- For the sfB tensor, each stage requires 256 x (256 / 16) x sizeof(float8) = 4KB. + +The situation is different for the B tensor: +- In the 1CTA case, each stage for the B tensor requires 256 x 256 x sizeof(float4) = 32KB. +- In the 2CTA case, only half this size is needed, i.e., 128 x 256 x sizeof(float4) = 16KB. + +Therefore, the maximum number of AB stages is: +- For 1CTA: 227 // (16 + 32 + 2 + 4) = 4 +- For 2CTA: 227 // (16 + 16 + 2 + 4) = 5 + +The latency hiding capability is: +- 1CTA: 512 * (4 - 1) = 1.5K cycles +- 2CTA: 512 * (5 - 1) = 2K cycles + +(2) TMA multicast can help reduce L2 cache traffic. + +Without TMA multicast, the L2 traffic per tile is typically 16KB + 32KB = 48KB (possibly less in practice, depending on hardware optimizations). +With TMA multicast in a cluster of shape (m, n), the L2 traffic per tile is reduced to 16KB / n + 32KB / m. +For example: +- In a 2x1 cluster: 16KB / 1 + 32KB / 2 = 24KB per tile +- In a 4x4 cluster: 16KB / 4 + 32KB / 4 = 12KB per tile + +The first approach offers substantial capacity for hiding latency, whereas the second reduces the time required for data to become ready. +Both could be tried when the workload is latency-bound or limited by memory throughput. + +To run this example: +.. code-block:: bash + + python examples/blackwell/tutorial_gemm/nvfp4_gemm_1.py \ + --mnkl 8192,8192,8192,1 --do_benchmark + +Constraints for this example: +* The problem size of m, n and k must be divisible by the tile size m&n&k (256, 256, 256) +* The scaling factor vector size is 16. +* The A/B matrices have data contiguous on the k dimension. +* The C matrix has data contiguous on the n dimension. +* The A/B matrix data type is Float4E2M1FN. +* The SFA/SFB matrix data type is Float8E4M3FN. +""" + + +class Sm100BlockScaledDenseGemmKernel: + def __init__(self): + self.threads_per_cta = 128 + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = 512 + + # set stages for ab_pipeline and acc_pipeline + self.num_acc_stage = 1 + self.num_ab_stage = 5 + + @cute.jit + def __call__( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + problem_size: tuple, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # setup static attributes before smem/grid/tma computation + self.c_layout = utils.LayoutEnum.ROW_MAJOR + m, n, k, l = problem_size + + self.use_2cta_instrs = False if mma_tiler_mn[0] == 128 else True + + # Setup attributes that depend on gemm inputs + mma_inst_tile_k = 4 + self.mma_tiler = ( + mma_tiler_mn[0], + mma_tiler_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.mma_inst_shape_sfb = ( + mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), + mma_tiler_mn[1], + mma_inst_shape_k, + ) + self.mma_tiler_sfb = ( + self.mma_inst_shape_sfb[0], + self.mma_inst_shape_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + a_tensor = cute.make_tensor( + a_ptr, + cute.make_layout( + (m, cute.assume(k, 32), l), + stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)), + ), + ) + b_tensor = cute.make_tensor( + b_ptr, + cute.make_layout( + (n, cute.assume(k, 32), l), + stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)), + ), + ) + # 256bit aligned. row_major + c_tensor = cute.make_tensor( + c_ptr, + cute.make_layout( + (cute.assume(m, 32), cute.assume(n, 16), l), + stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)), + ), + ) + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) + + mma_op = tcgen05.MmaMXF4NVF4Op( + sf_dtype, + (*mma_tiler_mn, mma_inst_shape_k), + tcgen05.CtaGroup.ONE if not self.use_2cta_instrs else tcgen05.CtaGroup.TWO, + tcgen05.OperandSource.SMEM, + ) + tiled_mma = cute.make_tiled_mma(mma_op) + + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + # Note sfB don't support share among 2ctas + sfb_mma_op = tcgen05.MmaMXF4NVF4Op( + sf_dtype, + self.mma_inst_shape_sfb, + tcgen05.CtaGroup.ONE, + tcgen05.OperandSource.SMEM, + ) + tiled_mma_sfb = cute.make_tiled_mma(sfb_mma_op) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(cluster_shape_mnk), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout(cluster_shape_mnk), + (tiled_mma_sfb.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.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute A/B/SFA/SFB/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + ab_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + ab_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + sf_vec_size, + self.num_ab_stage, + ) + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + cluster_shape_mnk[:2], tiled_mma.thr_id + ) + # TMA load for A + 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_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + cluster_shape_mnk[:2], 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_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load for SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + cluster_shape_mnk[:2], tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + cluster_shape_mnk[:2], tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb_tensor, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Compute TMA load bytes + a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Compute grid size + grid = cute.round_up( + cute.ceil_div( + (c_tensor.layout.shape), + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], 1), + ), + cluster_shape_mnk, + ) + + # Launch the kernel + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + c_tensor, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=cluster_shape_mnk, + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + mC_mnl: cute.Tensor, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + cta_layout_vmnk: cute.Layout, + cta_layout_sfb_vmnk: cute.Layout, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + tidx, _, _ = cute.arch.thread_idx() + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + cta_rank_in_cluster = cute.arch.block_idx_in_cluster() + cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + cta_in_cluster_coord_sfb_vmnk = cta_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + + # Coords outside cluster + mma_tile_coord_vmnk = ( + bidx % cute.size(cta_layout_vmnk, mode=[0]), + bidx // cute.size(cta_layout_vmnk, mode=[0]), + bidy, + bidz, + ) + mma_tile_coord_mnl = mma_tile_coord_vmnk[1:] + is_leader_cta = mma_tile_coord_vmnk[0] == 0 + + # + # Define shared storage for kernel + # + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_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) + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=ab_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=ab_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = smem.allocate_tensor( + element_type=sf_dtype, + layout=sfa_smem_layout_staged, + byte_alignment=128, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = smem.allocate_tensor( + element_type=sf_dtype, + layout=sfb_smem_layout_staged, + byte_alignment=128, + ) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr( + self.is_a_mcast or self.is_b_mcast or self.use_2cta_instrs + ): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cta_layout_sfb_vmnk, cta_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + + # + # Initialize mainloop ab_pipeline, acc_pipeline and their states + # + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_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=cta_layout_vmnk, + ).make_participants() + acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_cta * (2 if self.use_2cta_instrs else 1), + ), + cta_layout_vmnk=cta_layout_vmnk, + ).make_participants() + + # + # 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) + ) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gSFB_nkl = cute.local_tile( + mSFB_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/SFA/SFB/C + # + # (MMA, MMA_M, MMA_K, RestK) + thr_mma = tiled_mma.get_slice(mma_tile_coord_vmnk[0]) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_vmnk[0]) + # (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_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + # tCgSFB = thr_mma.partition_B(gSFB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_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/SFA/SFB + # + # TMA load A partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + # 0, + # cute.make_layout(1), + cta_in_cluster_coord_vmnk[2], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + # 0, + # cute.make_layout(1), + cta_in_cluster_coord_vmnk[1], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load SFA partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + # 0, + # cute.make_layout(1), + cta_in_cluster_coord_vmnk[2], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])), + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load SFB partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + sfb_cta_layout = cute.make_layout( + cute.slice_(cta_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + cta_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # 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) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + + # + # Alloc tensor memory buffer + # + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # + # Make SFA/SFB tmem tensor + # + # Get SFA tmem ptr + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc), + dtype=sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + # Get SFB tmem ptr + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + # + # Partition for S2T copy of SFA/SFB + # + # Make S2T CopyAtom + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp( + tcgen05.CtaGroup.ONE + if not self.use_2cta_instrs + else tcgen05.CtaGroup.TWO + ), + sf_dtype, + ) + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSFA_compact = cute.filter_zeros(sSFA) + # (MMA, MMA_MN, MMA_K) + tCtSFA_compact = cute.filter_zeros(tCtSFA) + tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact) + thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t_sfa, tCsSFA_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact) + + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSFB_compact = cute.filter_zeros(sSFB) + # (MMA, MMA_MN, MMA_K) + tCtSFB_compact = cute.filter_zeros(tCtSFB) + tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact) + thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t_sfb, tCsSFB_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # + # Execute Data copy and Math computation in the k_tile loop + # + if warp_idx == 0: + # Wait for accumulator buffer empty + if is_leader_cta: + acc_producer.acquire_and_advance() + + # Set ACCUMULATE field to False for the first k_tile iteration + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Execute k_tile loop + for k_tile in cutlass.range( + k_tile_cnt, prefetch_stages=self.num_ab_stage - 2 + ): + # Wait for AB buffer empty + ab_empty = ab_producer.acquire_and_advance() + + # TMA load A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA[(None, ab_empty.count)], + tAsA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB[(None, ab_empty.count)], + tBsB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_sfa, + tAgSFA[(None, ab_empty.count)], + tAsSFA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB[(None, ab_empty.count)], + tBsSFB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=sfb_full_mcast_mask, + ) + + if is_leader_cta: + # Wait for AB buffer full + ab_full = ab_consumer.wait_and_advance() + + # Copy SFA/SFB to tmem + s2t_stage_coord = (None, None, None, None, ab_full.index) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + 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_full.index, + ) + + # Set SFA/SFB tensor to tiled_mma + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set( + tcgen05.Field.SFA, + tCtSFA[sf_kblock_coord].iterator, + ) + tiled_mma.set( + tcgen05.Field.SFB, + tCtSFB[sf_kblock_coord].iterator, + ) + + 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_full.release() + if is_leader_cta: + acc_producer.commit() + + # + # Epilogue + # Partition for epilogue + # + # x32 or x128 all is ok. + op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE) + copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc) + # (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(tCgC) + # (T2R_M, T2R_N, EPI_M, EPI_N） + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32 + ) + # (T2R_M, T2R_N, EPI_M, EPI_N） + tTR_rC = cute.make_rmem_tensor( + tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype + ) + # STG Atom + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype) + tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)] + + # Wait for accumulator buffer full + acc_full = acc_consumer.wait_and_advance() + + # Copy accumulator to register + cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc) + acc_vec = epilogue_op(tTR_rAcc.load().to(c_dtype)) + tTR_rC.store(acc_vec) + # Store C to global memory + cute.copy(simt_atom, tTR_rC, tTR_gC) + + acc_full.release() + + # Ensure used buffers are properly synchronized before producer exit. + # This could avoid the invalid dsmem access due to early leading CTA exit. + if warp_idx == 0: + ab_producer.tail() + if is_leader_cta: + acc_producer.tail() + + # Deallocate TMEM + cute.arch.barrier() + tmem.free(acc_tmem_ptr) + + return + + +def run_nvfp4_gemm( + mnkl: Tuple[int, int, int, int], + tolerance: float, + warmup_iterations: int = 10, + iterations: int = 100, + use_cold_l2: bool = True, + do_benchmark: bool = False, +): + run( + gemm_class=Sm100BlockScaledDenseGemmKernel, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mnk=cluster_shape_mnk, + mnkl=mnkl, + tolerance=tolerance, + do_benchmark=do_benchmark, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cold_l2=use_cold_l2, + ) + + +if __name__ == "__main__": + parser = create_parser() + args = parser.parse_args() + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + m, n, k, _ = args.mnkl + if m % mma_tiler_mn[0] != 0: + parser.error("M must be multiples of mma_tiler_mn[0] (got m={})".format(m)) + if n % mma_tiler_mn[1] != 0: + parser.error("N must be multiples of mma_tiler_mn[1] (got n={})".format(n)) + if k % 256 != 0: + parser.error("k must be a multiple of 256 (got k={})".format(k)) + + run_nvfp4_gemm( + args.mnkl, + args.tolerance, + do_benchmark=args.do_benchmark, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/utils.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/utils.py new file mode 100644 index 00000000..f9c96b20 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/utils.py @@ -0,0 +1,366 @@ +# 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 argparse +from typing import Tuple + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +from cutlass.cute.runtime import make_ptr + + +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." + ) + + +def create_parser(): + parser = argparse.ArgumentParser( + description="Example of Sm100 Dense BlockScaled GEMM." + ) + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(8192, 8192, 8192, 8), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--do_benchmark", action="store_true", default=False, help="Do benchmark test" + ) + return parser + + +def ceil_div(a, b): + return (a + b - 1) // b + + +# Helper function to create scale factor tensor SFA/SFB +# for 1x16 block scaled wise use case and follow the layout requirement +# defined in https://docs.nvidia.com/cuda/cublas/index.html?highlight=fp4#d-block-scaling-factors-layout +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_ptr: cute.Pointer, + sf_mma_ptr: cute.Pointer, + mn: int, + sf_k: int, + l: int, + mma_shape: tuple, +): + mma_permute_order = (3, 4, 1, 5, 2, 0) + permuted_shape = tuple(mma_shape[i] for i in mma_permute_order) + cute_layout = cute.make_ordered_layout(permuted_shape, order=(2, 1, 4, 0, 3, 5)) + + sf_ref_tensor = cute.make_tensor( + sf_ref_ptr, cute.make_layout((mn, sf_k, l), stride=(sf_k, 1, mn * sf_k)) + ) + sf_mma_tensor = cute.make_tensor(sf_mma_ptr, cute_layout) + + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + pass + + +def to_blocked(input_matrix): + rows, cols = input_matrix.shape + + # Please ensure rows and cols are multiples of 128 and 4 respectively + n_row_blocks = ceil_div(rows, 128) + n_col_blocks = ceil_div(cols, 4) + + padded = input_matrix + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) + + return rearranged.flatten() + + +def run( + gemm_class, + ab_dtype, + sf_dtype, + c_dtype, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mnk, + mnkl: Tuple[int, int, int, int], + tolerance: float, + warmup_iterations: int = 10, + iterations: int = 100, + use_cold_l2: bool = True, + do_benchmark: bool = False, +): + """ + Prepare A/B/SFA/SFB/C tensors, launch GPU kernel, and reference checking. + """ + print("=" * 60) + print("Launching Blackwell Dense BlockScaled GEMM Test") + print("-" * 60) + print(f"Input dimensions (m, n, k, l): {mnkl}") + print(f" m (rows): {mnkl[0]}") + print(f" n (cols): {mnkl[1]}") + print(f" k (inner): {mnkl[2]}") + print(f" l (batch): {mnkl[3]}") + print(f"Data Types & Precision:") + print(f" Input matrices (A, B): {ab_dtype}") + print(f" Scale factors (SFA, SFB): {sf_dtype}") + print(f" Output matrix (C): {c_dtype}") + print(f" Scale factor vector size: {sf_vec_size}") + print("Tile and cluster configuration:") + print(f" MMA tiler (M, N, K): {mma_tiler_mn}") + print(f" Cluster shape (M, N, K): {cluster_shape_mnk}") + print(f"Validation tolerance: {tolerance}") + print(f"Do benchmark: {do_benchmark}") + print("=" * 60) + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Create tensor A/B/C + a_ref = torch.randint( + 0, 2, (l, m, k // 2), dtype=torch.uint8, device="cuda" + ).permute(1, 2, 0) + b_ref = torch.randint( + 0, 2, (l, n, k // 2), dtype=torch.uint8, device="cuda" + ).permute(1, 2, 0) + # a_ref = torch.ones((l, m, k // 2), dtype=torch.uint8, device="cuda").permute(1, 2, 0) + # b_ref = torch.ones((l, n, k // 2), dtype=torch.uint8, device="cuda").permute(1, 2, 0) + a_ref_f4 = a_ref.view(torch.float4_e2m1fn_x2) + b_ref_f4 = b_ref.view(torch.float4_e2m1fn_x2) + + c_tensor = torch.randn((l, m, n), dtype=torch.float16, device="cuda").permute( + 1, 2, 0 + ) + + # Create a torch tensor for scale factor tensor of A and B + def create_ref_scale_factor_tensor(l, mn, sf_k): + """ + Create the reference scale factor tensor on CPU. + Returns the reshaped/pruned tensor ready for ref computation and its original permuted form. + """ + ref_shape = (l, mn, sf_k) + ref_permute_order = (1, 2, 0) + ref_f8_random_int = torch.randint(1, 3, ref_shape, dtype=torch.int8) + ref_f8_torch_tensor_cpu = ref_f8_random_int.to(dtype=torch.float8_e4m3fn) + # permute to match ref_permute_order + ref_f8_torch_tensor_cpu_permuted = ref_f8_torch_tensor_cpu.permute( + *ref_permute_order + ) + return ref_f8_torch_tensor_cpu_permuted + + # Copy the reference scale factor tensor to the CUTE-format scale factor tensor + def create_cute_scale_factor_tensor(l, mn, sf_k, ref_f8_torch_tensor_cpu_permuted): + """ + Create the CUTE-format scale factor tensor on CUDA based on the reference tensor. + """ + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, # batch size + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + # Generate a random int8 tensor, then convert to float8_e4m3fn + rand_int_tensor = torch.randint(0, 2, mma_shape, dtype=torch.int8) + cute_f8_torch_tensor_cpu = rand_int_tensor.to(dtype=torch.float8_e4m3fn) + # Permute according to mma_permute_order + cute_f8_torch_tensor_cpu = cute_f8_torch_tensor_cpu.permute(*mma_permute_order) + + # Call the helper function to do layout conversion + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + make_ptr( + cutlass.Float8E4M3FN, + ref_f8_torch_tensor_cpu_permuted.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr( + cutlass.Float8E4M3FN, + cute_f8_torch_tensor_cpu.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + mn, + sf_k, + l, + mma_shape, + ) + return cute_f8_torch_tensor_cpu.cuda() + + sf_k = ceil_div(k, sf_vec_size) + sfa_ref = create_ref_scale_factor_tensor(l, m, sf_k) + sfb_ref = create_ref_scale_factor_tensor(l, n, sf_k) + # sfa_ref.fill_(1) + # sfb_ref.fill_(1) + sfa_tensor = create_cute_scale_factor_tensor(l, m, sf_k, sfa_ref) + sfb_tensor = create_cute_scale_factor_tensor(l, n, sf_k, sfb_ref) + + # Configure gemm kernel + gemm = gemm_class() + # Initialize Stream + current_stream = cutlass_torch.default_stream() + a_ptr = make_ptr( + ab_dtype, a_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + b_ptr = make_ptr( + ab_dtype, b_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + c_ptr = make_ptr( + c_dtype, c_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfa_ptr = make_ptr( + sf_dtype, sfa_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfb_ptr = make_ptr( + sf_dtype, sfb_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + # Compile gemm kernel + compiled_gemm = cute.compile( + gemm, + a_ptr, + b_ptr, + sfa_ptr, + sfb_ptr, + c_ptr, + (m, n, k, l), + current_stream, + ) + # Launch GPU kernel + compiled_gemm(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream) + # For batch l, do (m, k, l) @ (n, k, l).T along k for each batch. + # Result: (m, n, l) + # Allocate ref as (l, m, n) with n-contiguous layout, then permute to (m, n, l) + ref = torch.empty( + (l, m, n), + dtype=torch.float16, + device="cuda", + ).permute(1, 2, 0) + for l_idx in range(l): + # Convert the scale factor tensor to blocked format + scale_a = to_blocked(sfa_ref[:, :, l_idx]) + scale_b = to_blocked(sfb_ref[:, :, l_idx]) + # (m, k) @ (n, k).T -> (m, n) + res = torch._scaled_mm( + a_ref_f4[:, :, l_idx], + b_ref_f4[:, :, l_idx].transpose(0, 1), + scale_a.cuda(), + scale_b.cuda(), + bias=None, + out_dtype=torch.float16, + ) + ref[:, :, l_idx] = res + torch.testing.assert_close(c_tensor, ref, atol=tolerance, rtol=1e-02) + + if do_benchmark: + + def generate_tensors(): + a_ptr = make_ptr( + ab_dtype, a_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + b_ptr = make_ptr( + ab_dtype, b_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + c_ptr = make_ptr( + c_dtype, c_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfa_ptr = make_ptr( + sf_dtype, + sfa_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + sfb_ptr = make_ptr( + sf_dtype, + sfb_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + args = cute.testing.JitArguments( + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream + ) + args.add_to_scope([a_ref_f4, b_ref_f4, sfa_tensor, sfb_tensor, c_tensor]) + return args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_ref_f4.numel() * a_ref_f4.element_size() + + b_ref_f4.numel() * b_ref_f4.element_size() + + sfa_tensor.numel() * sfa_tensor.element_size() + + sfb_tensor.numel() * sfb_tensor.element_size() + + c_tensor.numel() * c_tensor.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + # Return execution time in microseconds + time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print(f"Execution time: {time} us") + peta_flops = (4 * m * n * k * l) / (time * 1e-6) / 1e9 / 1000000 + print(f"FLOPS: {peta_flops} PFLOPS") + bytes_transfer = ( + 2 * m * k / 2 * l * a_ref_f4.element_size() + + 2 * n * k / 2 * l * b_ref_f4.element_size() + + 2 * m * n * l * c_tensor.element_size() + + 2 * m * sf_k * l * sfa_tensor.element_size() + + 2 * n * sf_k * l * sfb_tensor.element_size() + ) + print(f"Bytes: {bytes_transfer} Bytes") + bandwidth = bytes_transfer / time * 1e-3 + print(f"BW: {bandwidth} GB/s") diff --git a/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py b/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py index 74d2afa3..7fc7b421 100644 --- a/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py +++ b/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py @@ -29,7 +29,6 @@ import argparse from typing import Tuple, Type -import torch import cuda.bindings.driver as cuda import cutlass @@ -37,7 +36,6 @@ import cutlass.cute as cute import cutlass.cute.testing as testing import cutlass.utils as utils import cutlass.pipeline as pipeline -import cutlass.torch as cutlass_torch import cutlass.utils.hopper_helpers as sm90_utils """ @@ -49,15 +47,15 @@ using CUTE DSL. This GEMM kernel supports the following features: - Utilizes Tensor Memory Access (TMA) for efficient memory operations - - Utilizes non-Tensor Core MMA for matrix multiply-accumulate (MMA) operations + - Utilizes Blackwell MMA for matrix multiply-accumulate (MMA) operations - Supports multi-stage pipeline to overlap computation and memory access This GEMM works as follows: 1. Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. -2. Perform matrix multiply-accumulate (MMA) operations using non-Tensor Core MMA instruction. +2. Perform matrix multiply-accumulate (MMA) operations using Blackwell MMA instruction. 3. Store results from registers (RMEM) to shared memory (SMEM), then to global memory (GMEM) with TMA operations. -Non-Tensor Core MMA instructions operate as follows: +Blackwell MMA instructions operate as follows: - Read matrix A from registers - Read matrix B from registers - Perform MMA operation and store the result in Accumulator(register) @@ -114,9 +112,7 @@ def parse_comma_separated_ints(s: str): def parse_arguments() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Example of MxNxKxL GEMM on Blackwell Geforce." - ) + parser = argparse.ArgumentParser(description="Example of MxNxKxL GEMM on Blackwell Geforce.") parser.add_argument( "--mnkl", @@ -873,7 +869,10 @@ class Sm120GemmKernel: tRS_sD[(None, None, None, epi_buffer)], ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # barrier for sync self.epilog_sync_barrier.arrive_and_wait() @@ -1176,6 +1175,9 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + print("Running Blackwell Geforce Dense GEMM with:") print(f"mnkl: {mnkl}") print( diff --git a/examples/python/CuTeDSL/cute/ffi/jit_argument.py b/examples/python/CuTeDSL/cute/ffi/jit_argument.py index c48cf8a1..58368551 100644 --- a/examples/python/CuTeDSL/cute/ffi/jit_argument.py +++ b/examples/python/CuTeDSL/cute/ffi/jit_argument.py @@ -240,10 +240,11 @@ import os import subprocess import shutil import tempfile -import torch def run_test(tmpdir=None, cmake_args="", cleanup=True): + import torch + try: current_dir = os.path.dirname(os.path.abspath(__file__)) diff --git a/examples/python/CuTeDSL/cute/torch_fake_tensor.py b/examples/python/CuTeDSL/cute/torch_fake_tensor.py index 37064441..60e936bd 100644 --- a/examples/python/CuTeDSL/cute/torch_fake_tensor.py +++ b/examples/python/CuTeDSL/cute/torch_fake_tensor.py @@ -27,8 +27,6 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -import torch - import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack @@ -66,6 +64,7 @@ def print_tensor(t: cute.Tensor): def run(): + import torch from torch._subclasses.fake_tensor import FakeTensorMode shape = (3, 4) diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp index e1f0645b..37045630 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp @@ -1,6 +1,6 @@ // clang-format off /* - * SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: LicenseRef-NvidiaProprietary * * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh index 3f110a4a..9c409724 100755 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh @@ -28,7 +28,7 @@ #!/bin/bash # Set up library paths for runtime -export LD_LIBRARY_PATH=$(python3 -m cutlass.cute.export.aot_config --libdir):$(tvm-ffi-config --libdir) +export LD_LIBRARY_PATH=$(python3 -m cutlass.cute.export.aot_config --libdir):$(tvm-ffi-config --libdir):$LD_LIBRARY_PATH CUDA_HOME=/usr/local/cuda SOURCE_FILE="$(dirname "$0")/aot_use_in_cpp_bundle.cpp" diff --git a/examples/python/CuTeDSL/distributed/all_reduce_tma.py b/examples/python/CuTeDSL/distributed/all_reduce_tma.py new file mode 100644 index 00000000..b3496c2b --- /dev/null +++ b/examples/python/CuTeDSL/distributed/all_reduce_tma.py @@ -0,0 +1,691 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +A Distributed All-Reduce Example using TMA (Tensor Memory Accelerator). + +This example demonstrates distributed all-reduce across multiple GPUs using TMA +for data movement. It serves as a tutorial for TMA-based distributed operations, +not as a performance-optimized implementation. + +Tensor Semantics: + - Input: Logical shape (world_size, S), where S is the per-rank tensor size + - Output: Logical shape (world_size, S), each rank gets the sum of all inputs + +Kernel Parameters: + - input: List of world_size tensors, each with shape S (accessible via NVSHMEM) + - output: Single tensor with shape S, using multicast address for broadcast + +Algorithm (Two-Shot): + 1. Each CTA loads data from all ranks at its assigned tile position (TMA Load) + 2. Accumulates the data locally in registers + 3. Stores the result via TMA multicast (broadcasts to all ranks) + 4. Cross-GPU barrier ensures completion before kernel exit + +Tile Assignment: + - Total tiles = ceil(S / elems_per_cta) + - Each rank processes ceil(total_tiles / world_size) CTAs + - CTA i on rank r processes global_tile_id = r * ctas_per_rank + i + +TMA Usage Notes (for tutorial purposes, not perf-optimal): + - Uses 1D TMA load to load from remote GPU memory via NVSHMEM addresses + - Uses 1D TMA load to store to multicast address for broadcasting to all ranks + - Supports any input shape by flattening to 1D and tiling linearly + - Pipeline with 2 stages overlaps TMA loads across ranks + +To run this example: + +.. code-block:: bash + + torchrun --nproc-per-node 8 examples/distributed/all_reduce_tma.py --shape 1024,1024 + torchrun --nproc-per-node 8 examples/distributed/all_reduce_tma.py --shape 4,6,8,10,12 +""" + +import cutlass +import cutlass.utils as utils +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync + + +class AllReduceTmaKernel: + """ + TMA-based distributed All-Reduce kernel. + + This kernel performs an all-reduce operation across multiple GPUs using TMA + (Tensor Memory Accelerator) for efficient data movement. + + Algorithm (Two-Shot): + 1. Each CTA loads data from all ranks at its assigned tile position + 2. Accumulates the data locally in registers + 3. Stores the result via TMA multicast (broadcasts to all ranks) + 4. Cross-GPU barrier ensures completion before kernel exit + + The input/output tensors can be of any rank, as long as: + - All input tensors and output tensor share the same layout + - The layout is compact (no holes in memory) + + We traverse the tensors linearly in codomain (physical offset) order, + which guarantees consistent logical coordinate access across all tensors. + """ + + _elems_per_cta: int = 128 * 128 # Elements processed per CTA + _tma_threads: int = 32 + _consumer_threads: int = 128 + _threads_per_cta: int = _tma_threads + _consumer_threads + _num_stages: int = 2 + + def __init__(self, dtype): + self.dtype = dtype + + # SMEM layout shape (will be converted to Layout in JIT context) + self.smem_layout_shape = (self._elems_per_cta,) + self.tiler = (self._elems_per_cta,) + + # TMA transaction bytes (computed from dtype size) + # dtype.width is in bits, divide by 8 to get bytes + self.tma_bytes = (dtype.width // 8) * self._elems_per_cta + + # Dynamically create SharedStorage type based on dtype + elems = self._elems_per_cta + stages = self._num_stages + + @cute.struct + class SharedStorage: + mbar_array: cute.struct.MemRange[cutlass.Int64, stages * 2] + smem_buffer: cute.struct.Align[ + cute.struct.MemRange[dtype, elems * stages], # stages 个 tile + 128, + ] + + self._SharedStorage = SharedStorage + + @cute.jit + def __call__( + self, + input_tensors: list[cute.Tensor], + output_tensor_mc: cute.Tensor, + flag: cute.Tensor, + flag_mc: cute.Tensor, + local_rank: cutlass.Constexpr, + world_size: cutlass.Constexpr, + ): + """ + Host-side JIT function: creates TMA descriptors and launches kernel. + + Args: + input_tensors: List of input tensors from each rank (world_size tensors) + output_tensor_mc: Output tensor with multicast address + flag: Synchronization flag (local view) + flag_mc: Synchronization flag (multicast view) + local_rank: This rank's ID + world_size: Total number of ranks + """ + # ====================================================================== + # Layout validation + # ====================================================================== + ref_layout = input_tensors[0].layout + ref_size = cute.size(ref_layout) + ref_cosize = cute.cosize(ref_layout) + + # Check compact: size == cosize (no holes in memory) + assert ref_size == ref_cosize, ( + f"Input tensor must be compact: size={ref_size}, cosize={ref_cosize}" + ) + assert self.tma_bytes % 16 == 0, f"Not aligned to 16B, TMA should not be used." + + # Check all input tensors have the same layout + for i in cutlass.range_constexpr(world_size): + assert input_tensors[i].layout == ref_layout, ( + f"All input tensors must have the same layout. " + f"input_tensors[0].layout={ref_layout}, " + f"input_tensors[{i}].layout={input_tensors[i].layout}" + ) + + # Check output tensor has the same layout + assert output_tensor_mc.layout == ref_layout, ( + f"Output tensor must have the same layout as input tensors. " + f"input layout={ref_layout}, output layout={output_tensor_mc.layout}" + ) + + # ====================================================================== + # Extract tensor info + # ====================================================================== + # Verify dtype matches + assert input_tensors[0].element_type == self.dtype, ( + f"Input tensor dtype mismatch: expected {self.dtype}, " + f"got {input_tensors[0].element_type}" + ) + + total_elems = ref_size + + # Flatten layout: treat tensor as 1D in codomain order + flat_layout = cute.make_layout((total_elems,)) + + # SMEM layout (created in JIT context) + smem_layout = cute.make_layout(self.smem_layout_shape) + + # Create TMA load descriptors (one per rank) + tma_load_op = cpasync.CopyBulkTensorTileG2SOp() + tma_load_atoms = [] + tma_load_tensors = [] + + for i in cutlass.range_constexpr(world_size): + flat_input = cute.make_tensor(input_tensors[i].iterator, flat_layout) + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + tma_load_op, + flat_input, + smem_layout, + self.tiler, + ) + tma_load_atoms.append(tma_atom) + tma_load_tensors.append(tma_tensor) + + # Create TMA store descriptor + tma_store_op = cpasync.CopyBulkTensorTileS2GOp() + flat_output = cute.make_tensor(output_tensor_mc.iterator, flat_layout) + tma_store_atom, tma_store_tensor = cpasync.make_tiled_tma_atom( + tma_store_op, + flat_output, + smem_layout, + self.tiler, + ) + + # Grid calculation + num_tiles_total = cute.ceil_div(total_elems, self._elems_per_cta) + ctas_per_rank = cute.ceil_div(num_tiles_total, world_size) + + # SMEM size from SharedStorage + smem_bytes = self._SharedStorage.size_in_bytes() + + # Launch kernel + self.kernel( + tma_load_atoms, + tma_load_tensors, + tma_store_atom, + tma_store_tensor, + flag, + flag_mc, + local_rank, + world_size, + num_tiles_total, + ctas_per_rank, + ).launch( + grid=[ctas_per_rank, 1, 1], + block=[self._threads_per_cta, 1, 1], + smem=smem_bytes, + ) + + @cute.kernel + def kernel( + self, + # TMA atoms and tensors for loading from each rank + tma_load_atoms: list[cute.CopyAtom], + tma_load_tensors: list[cute.Tensor], + # TMA atom and tensor for storing to multicast address + tma_store_atom: cute.CopyAtom, + tma_store_tensor: cute.Tensor, + # Synchronization flags + flag: cute.Tensor, + flag_mc: cute.Tensor, + # Rank info + local_rank: cutlass.Constexpr, + world_size: cutlass.Constexpr, + # Grid info for tile calculation + num_tiles_total: cutlass.Constexpr, + ctas_per_rank: cutlass.Constexpr, + ): + # ====================================================================== + # Thread/Block indexing + # ====================================================================== + tidx = cute.arch.thread_idx()[0] + bidx = cute.arch.block_idx()[0] + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # ====================================================================== + # SMEM allocation + # ====================================================================== + staged_smem_layout = cute.make_layout((self._elems_per_cta, self._num_stages)) + + smem = utils.SmemAllocator() + storage = smem.allocate(self._SharedStorage) + mbar_ptr = storage.mbar_array.data_ptr() + staged_smem_tensor = storage.smem_buffer.get_tensor(staged_smem_layout) + + # ====================================================================== + # TMA Pipeline setup + # ====================================================================== + producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self._consumer_threads + ) + + tma_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=mbar_ptr, + num_stages=self._num_stages, + producer_group=producer_group, + consumer_group=consumer_group, + tx_count=self.tma_bytes, + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + ) + + global_tile_id = local_rank * ctas_per_rank + bidx + + if global_tile_id < num_tiles_total: + # ====================================================================== + # Warp 0: Producer - TMA Load from all ranks + # ====================================================================== + if warp_idx == 0: + producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._num_stages + ) + + for rank_i in cutlass.range_constexpr(world_size): + tma_pipeline.producer_acquire(producer_state) + + stage_idx = producer_state.index + smem_tile = cute.slice_(staged_smem_tensor, (None, stage_idx)) + + g_tensor_tiled = cute.zipped_divide( + tma_load_tensors[rank_i], self.tiler + ) + g_tile = g_tensor_tiled[(None,), global_tile_id] + + g_tile_flat = cute.group_modes(g_tile, 0, cute.rank(g_tile)) + s_tile_flat = cute.group_modes(smem_tile, 0, cute.rank(smem_tile)) + + s_part, g_part = cute.nvgpu.cpasync.tma_partition( + tma_load_atoms[rank_i], + 0, + cute.make_layout(1), + s_tile_flat, + g_tile_flat, + ) + + cute.copy( + tma_load_atoms[rank_i], + g_part, + s_part, + tma_bar_ptr=tma_pipeline.producer_get_barrier(producer_state), + ) + + tma_pipeline.producer_commit(producer_state) + producer_state.advance() + + # ====================================================================== + # Warp 1-4: Consumer - Load from smem, ADD, Store to smem + # ====================================================================== + else: + consumer_tid = tidx - self._tma_threads + + vec_size = 4 + chunk_size = vec_size * self._consumer_threads + + # ------------------------------------------------------------------ + # Initialize accumulator using stage 0's layout + # ------------------------------------------------------------------ + # (elems, stages) -> (elems,) + smem_tensor_wo_stage = cute.slice_(staged_smem_tensor, (None, 0)) + # (elems,) -> ((thr_vec,), (num_chunks,)) + smem_tensor_tiled_by_thr_vec = cute.zipped_divide( + smem_tensor_wo_stage, (chunk_size,) + ) + # ((thr_vec,), (num_chunks,)) -> (((vec, threads),), (num_chunks,)) + smem_tensor_tiled_by_thr_vec_tiled_by_vec = cute.logical_divide( + smem_tensor_tiled_by_thr_vec, (vec_size,) + ) + # (((vec, threads),), (num_chunks,)) -> ((vec,), (num_chunks,)) + per_thread_smem_tensor = cute.slice_( + smem_tensor_tiled_by_thr_vec_tiled_by_vec, + ((None, consumer_tid), None), + ) + + accum = cute.make_rmem_tensor(per_thread_smem_tensor.layout, self.dtype) + accum.fill(self.dtype(0.0)) + + # ------------------------------------------------------------------ + # Main loop: load from SMEM and accumulate + # ------------------------------------------------------------------ + consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._num_stages + ) + + for rank_i in cutlass.range_constexpr(world_size): + tma_pipeline.consumer_wait(consumer_state) + + stage_idx = consumer_state.index + smem_tile = cute.slice_(staged_smem_tensor, (None, stage_idx)) + + # (elems,) -> ((thr_vec,), (num_chunks,)) + smem_tiled_by_thr_vec = cute.zipped_divide(smem_tile, (chunk_size,)) + # ((thr_vec,), (num_chunks,)) -> (((vec, threads),), (num_chunks,)) + smem_tiled_by_thr_vec_tiled_by_vec = cute.logical_divide( + smem_tiled_by_thr_vec, (vec_size,) + ) + # (((vec, threads),), (num_chunks,)) -> ((vec,), (num_chunks,)) + per_thread_smem_view = cute.slice_( + smem_tiled_by_thr_vec_tiled_by_vec, + ((None, consumer_tid), None), + ) + + fragment = per_thread_smem_view.load() + accum.store(accum.load() + fragment) + + tma_pipeline.sync_object_empty.arrive( + consumer_state.index, tma_pipeline.consumer_mask + ) + consumer_state.advance() + + # Store accumulated result back to SMEM (stage 0) + per_thread_smem_tensor.store(accum.load()) + + # ====================================================================== + # Sync point: all warps meet here + # ====================================================================== + cute.arch.sync_threads() + + # ====================================================================== + # Warp 0: TMA Store to multicast output + # ====================================================================== + if warp_idx == 0: + # Fence to ensure SMEM writes are visible + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + + smem_tile_out = cute.slice_(staged_smem_tensor, (None, 0)) + + g_output_tiled = cute.zipped_divide(tma_store_tensor, self.tiler) + g_output_tile = g_output_tiled[(None,), global_tile_id] + + g_out_flat = cute.group_modes( + g_output_tile, 0, cute.rank(g_output_tile) + ) + s_out_flat = cute.group_modes( + smem_tile_out, 0, cute.rank(smem_tile_out) + ) + + s_part, g_part = cute.nvgpu.cpasync.tma_partition( + tma_store_atom, + 0, + cute.make_layout(1), + s_out_flat, + g_out_flat, + ) + + cute.copy(tma_store_atom, s_part, g_part) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + + # ================================================================== + # Cross-GPU barrier synchronization (thread 0 only) + # ================================================================== + if tidx == 0: + sm_id_linear = ( + cute.arch.block_idx()[0] + + cute.arch.block_idx()[1] * cute.arch.grid_dim()[0] + + cute.arch.block_idx()[2] + * cute.arch.grid_dim()[0] + * cute.arch.grid_dim()[1] + ) + + # Signal completion to all ranks + utils.distributed.multimem_red_add1( + flag_mc.iterator + sm_id_linear, + scope="sys", + order="release", + ) + + # The same idx ctas wait until all peer ranks' ctas complete + utils.distributed.spin_lock_atom_cas_relaxed_wait( + flag.iterator + sm_id_linear, + expected_val=world_size, + reset_val=0, + scope="sys", + ) + + +# ============================================================================= +# HOST-SIDE DRIVER CODE +# ============================================================================= + +import os +import argparse +import math + +import numpy as np +import torch +import torch.distributed as dist +from cuda.core.experimental import Device +from cuda.pathfinder import load_nvidia_dynamic_lib + +from cutlass.cute.runtime import from_dlpack + +try: + import nvshmem.core +except ImportError as exc: + raise ImportError( + "nvshmem4py is required but not installed. Please install it using:\n" + " For CUDA 12: pip install nvshmem4py-cu12\n" + " For CUDA 13: pip install nvshmem4py-cu13\n" + "Note: nvshmem4py version >= 0.1.3 is recommended." + ) from None + +try: + load_nvidia_dynamic_lib("nvshmem_host") +except RuntimeError as exc: + raise ImportError( + "nvshmem lib is required but not installed. Please install it using:\n" + " For CUDA 12: pip install nvidia-nvshmem-cu12\n" + " For CUDA 13: pip install nvidia-nvshmem-cu13\n" + ) from None + + +def torchrun_uid_init_bcast(): + """Initialize NVSHMEM using UniqueID with torchrun as launcher.""" + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + + dev = Device(local_rank) + dev.set_current() + global stream + stream = dev.create_stream() + + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + num_ranks = dist.get_world_size() + + uid = nvshmem.core.get_unique_id(empty=(local_rank != 0)) + uid_bytes = uid._data.view(np.uint8).copy() + uid_tensor = torch.from_numpy(uid_bytes).cuda() + dist.broadcast(uid_tensor, src=0) + dist.barrier() + uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype) + + nvshmem.core.init( + device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid" + ) + + +def torchrun_finalize(): + """Finalize NVSHMEM and destroy process group.""" + nvshmem.core.finalize() + dist.destroy_process_group() + + +def run_all_reduce_tma( + shape: tuple, + skip_ref_check: bool = False, +): + """ + Run the TMA-based All-Reduce kernel. + + Args: + shape: Tensor shape tuple, e.g., (4, 6, 8, 10) + skip_ref_check: If True, skip reference result verification + """ + local_rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + + # Calculate total elements + total_elems = math.prod(shape) + + if local_rank == 0: + print("\nRunning TMA All-Reduce test with:") + print(f" Tensor shape: {shape}") + print(f" Total elements: {total_elems}") + print(f" GPU count: {world_size}") + + # Allocate input tensor (symmetric memory, accessible from all ranks) + local_input_tensor = nvshmem.core.tensor(shape, dtype=torch.float32) + local_input_tensor.random_(0, 100) + + # Get peer tensors (views into each rank's input) + peer_input_tensors = [ + nvshmem.core.get_peer_tensor(local_input_tensor, r) for r in range(world_size) + ] + + if local_rank == 0: + print(f" Input tensor ptr: {local_input_tensor.data_ptr():#x}") + + # Allocate output tensor with multicast address + local_output_tensor = nvshmem.core.tensor(shape, dtype=torch.float32) + local_output_tensor.fill_(0) + output_tensor_mc = nvshmem.core.get_multicast_tensor( + nvshmem.core.Teams.TEAM_NODE, local_output_tensor + ) + + # Allocate synchronization flags + # Flag size = ctas_per_rank (matches kernel's bidx indexing) + elems_per_cta = AllReduceTmaKernel._elems_per_cta + num_tiles = (total_elems + elems_per_cta - 1) // elems_per_cta + ctas_per_rank = (num_tiles + world_size - 1) // world_size + local_flag = nvshmem.core.tensor((ctas_per_rank,), dtype=torch.int32) + local_flag.fill_(0) + flag_mc = nvshmem.core.get_multicast_tensor( + nvshmem.core.Teams.TEAM_NODE, local_flag + ) + + if local_rank == 0: + print(f" Number of tiles: {num_tiles}") + print(f" CTAs per rank: {ctas_per_rank}") + print("Compiling kernel...") + + # Create kernel instance and compile + kernel = AllReduceTmaKernel(cutlass.Float32) + + compiled_func = cute.compile( + kernel, + [from_dlpack(t) for t in peer_input_tensors], + from_dlpack(output_tensor_mc), + from_dlpack(local_flag), + from_dlpack(flag_mc), + local_rank, + world_size, + ) + + if local_rank == 0: + print("Compilation successful!") + + if not skip_ref_check: + if local_rank == 0: + print("Executing kernel...") + + dist.barrier(device_ids=[local_rank]) + compiled_func( + [from_dlpack(t) for t in peer_input_tensors], + from_dlpack(output_tensor_mc), + from_dlpack(local_flag), + from_dlpack(flag_mc), + ) + dist.barrier(device_ids=[local_rank]) + + if local_rank == 0: + print("Verifying results...") + + # Compute expected result: sum of all inputs + expected = sum([t.cpu() for t in peer_input_tensors]) + + # Compare with actual output + torch.testing.assert_close(expected, local_output_tensor.cpu()) + + if local_rank == 0: + print("Results verified successfully!") + + # Cleanup + for i in range(world_size): + if i != local_rank: + nvshmem.core.free_tensor(peer_input_tensors[i]) + + nvshmem.core.free_tensor(output_tensor_mc) + nvshmem.core.free_tensor(flag_mc) + nvshmem.core.free_tensor(local_input_tensor) + nvshmem.core.free_tensor(local_output_tensor) + nvshmem.core.free_tensor(local_flag) + + +def parse_shape(shape_str: str) -> tuple: + """ + Parse shape string into tuple. + Examples: + "1024,1024" -> (1024, 1024) + "2,3,4,5,6,7,8" -> (2, 3, 4, 5, 6, 7, 8) + """ + return tuple(int(x.strip()) for x in shape_str.split(",")) + + +def main(): + parser = argparse.ArgumentParser( + description="TMA-based distributed all-reduce example" + ) + parser.add_argument( + "--shape", + default="1024,1024", + type=str, + help="Tensor shape as comma-separated values, e.g., '1024,1024' or 4,6,8,10,12'", + ) + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference result verification", + ) + + args = parser.parse_args() + shape = parse_shape(args.shape) + + torchrun_uid_init_bcast() + run_all_reduce_tma( + shape=shape, + skip_ref_check=args.skip_ref_check, + ) + torchrun_finalize() + + +if __name__ == "__main__": + main() diff --git a/examples/python/CuTeDSL/experimental/ampere/memcpy_simt_universal_copy.py b/examples/python/CuTeDSL/experimental/ampere/memcpy_simt_universal_copy.py new file mode 100644 index 00000000..13fc1e8a --- /dev/null +++ b/examples/python/CuTeDSL/experimental/ampere/memcpy_simt_universal_copy.py @@ -0,0 +1,155 @@ +# 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 argparse +import torch +import pytest + +from cutlass import cute +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack +import cutlass.utils as utils + + +@cute.experimental.kernel +def memcpy_simt_universal_copy_kernel( + mA: cute.Tensor, mD: cute.Tensor, addend: cute.Float16 +): + tile_mn = cute.core._pack_shape((128, 64)) + gA = cute.zipped_divide(mA, tile_mn) + gD = cute.zipped_divide(mD, tile_mn) + + cta_m, cta_n, cta_l = cute.arch.block_idx() + tid_x, _, _ = cute.arch.thread_idx() + + gA_tile = gA[(None, None), (cta_m, cta_n, cta_l)] + gD_tile = gD[(None, None), (cta_m, cta_n, cta_l)] + + buffer = cute_ext.allocate( + cute.Float16, + cute.AddressSpace.rmem, + cute.make_layout(((8, 1), (1, 8)), stride=((1, 8), (1, 8))), + alignment=16, + ) + + tCgA = cute_ext.partition( + gA_tile, + tid_x, + layout_tv=cute.make_layout(((16, 8), (8, 1)), stride=((8, 128), (1, 1024))), + tiler=cute.core._pack_tile((128, 8)), + ) + + tCgD = cute_ext.partition( + gD_tile, + tid_x, + layout_tv=cute.make_layout(((16, 8), (8, 1)), stride=((8, 128), (1, 1024))), + tiler=cute.core._pack_tile((128, 8)), + ) + + # cute_ext.copy() automatically computes predicates based on the shape of + # the tensor passed to the @cute.experimental.kernel argument + cute_ext.copy( + tCgA, + buffer, + copy_atom=cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + tCgD.element_type, + num_bits_per_copy=128, + ), + ) + + # Update the RMEM tensor in place using elementwise addition. + buffer.store(buffer.load() + addend) + + # cute_ext.copy() automatically computes predicates based on the shape of + # the tensor passed to the @cute.experimental.kernel argument + cute_ext.copy( + buffer, + tCgD, + copy_atom=cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + tCgD.element_type, + num_bits_per_copy=128, + ), + ) + + +@cute.experimental.jit +def memcpy_simt_universal_copy( + src: cute.Tensor, dst: cute.Tensor, addend: cute.Float16 +): + tile_mn = cute.core._pack_shape((128, 64)) + div = cute.tiled_divide(src, tile_mn) + grid = (div.shape[1], div.shape[2], div.shape[3]) + memcpy_simt_universal_copy_kernel(src, dst, addend).launch( + grid=grid, + block=(128, 1, 1), + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_80")), + ) + + +def run_simt_universal_memcpy(M, N, L): + src = torch.randn(L, N, M).permute(2, 1, 0).to(torch.float16).cuda() + dst = torch.randn(L, N, M).permute(2, 1, 0).to(torch.float16).cuda() + + mA = ( + from_dlpack(src, assumed_align=16) + .mark_layout_dynamic(leading_dim=0) + .mark_compact_shape_dynamic( + mode=0, stride_order=src.dim_order(), divisibility=8 + ) + ) + mD = ( + from_dlpack(dst, assumed_align=16) + .mark_layout_dynamic(leading_dim=0) + .mark_compact_shape_dynamic( + mode=0, stride_order=dst.dim_order(), divisibility=8 + ) + ) + addend = 5.0 + + memcpy_simt_universal_copy( + mA, + mD, + cute.Float16(addend), + no_cache=True, + ) + + torch.testing.assert_close(src.cpu() + addend, dst.cpu()) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Example memory copy example using CuTe auto predication features." + ) + parser.add_argument("--mnl", default=[136, 7, 9], nargs="+", type=int) + args = parser.parse_args() + + M, N, L = tuple(args.mnl) + run_simt_universal_memcpy(M, N, L) + print("PASS") diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_block_scaled_gemm.py b/examples/python/CuTeDSL/experimental/blackwell/dense_block_scaled_gemm.py new file mode 100644 index 00000000..f5c2a542 --- /dev/null +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_block_scaled_gemm.py @@ -0,0 +1,1021 @@ +# Copyright (c) 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 argparse +from typing import Type, Tuple +from dataclasses import dataclass +import torch +import cutlass +from cutlass import ( + cute as cute, + utils as utils, +) +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils + +""" + +This is an implementation of dense block scaled GEMM. + +""" +class BlockScaledDenseGemmKernel: + def __init__( + self, + mma_inst_mn: tuple[int, int], + mma_dtype: tuple[Type[cutlass.Numeric], Type[cutlass.Numeric]], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + epilogue_op=lambda x: x, + ): + self.ab_dtype, self.acc_dtype = mma_dtype + self.sf_dtype = sf_dtype + self.sf_vec_size = sf_vec_size + self.mma_inst_shape_mn = mma_inst_mn + self.use_2cta_instrs = False + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.epilogue_op = epilogue_op + # TODO: instead of using max shared memory, we should define a SharedStorage and then + # query its size. + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + + # Stages + self.num_acc_stages = 1 if self.mma_inst_shape_mn[1] == 256 else 2 + + # TODO: provide a computation for this so that it is not fixed; + # fitting as many stages as there is available shared memory + self.num_main_stages = 4 + + self.tma_store_stages = 4 + + @cute.experimental.jit + def __call__( + self, + mA: cute.Tensor, + mSFA: cute.Tensor, + mB: cute.Tensor, + mSFB: cute.Tensor, + mC: cute.Tensor, + ): + tile_mn = (*self.mma_inst_shape_mn, 1) + div = cute.tiled_divide(mC, tile_mn) + grid = (div.shape[1], div.shape[2], div.shape[3]) + self.kernel(mA, mSFA, mB, mSFB, mC).launch( + grid=grid, + # Using a total of 6 warps (1x load + 1x mma + 4x epilogue) + block=(192, 1, 1), + cluster=(1, 1, 1), + smem=self.smem_capacity, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, + mSFA: cute.Tensor, + mB: cute.Tensor, + mSFB: cute.Tensor, + mC: cute.Tensor, + ): + # Prologue + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + cta_m, cta_n, cta_l = cute.arch.block_idx() + + a_dtype: Type[cutlass.Numeric] = mA.element_type + sf_dtype: Type[cutlass.Numeric] = mSFA.element_type + c_dtype: Type[cutlass.Numeric] = mC.element_type + a_major_mode = utils.LayoutEnum.from_tensor(mA).mma_major_mode() + b_major_mode = utils.LayoutEnum.from_tensor(mB).mma_major_mode() + d_layout = utils.LayoutEnum.from_tensor(mC) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + a_dtype, + a_major_mode, + b_major_mode, + sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + mma_tiler_mnk = ( + self.mma_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + tiler_mk = (mma_tiler_mnk[0], mma_tiler_mnk[2]) + tiler_nk = (mma_tiler_mnk[1], mma_tiler_mnk[2]) + tiler_mn = (mma_tiler_mnk[0], mma_tiler_mnk[1]) + + # ((Atom_M, Rest_M),(Atom_K, Rest_K), RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(mA.shape, self.sf_vec_size) + sfa_tensor = cute.make_tensor(mSFA.iterator, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K), RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(mB.shape, self.sf_vec_size) + sfb_tensor = cute.make_tensor(mSFB.iterator, sfb_layout) + + gA = cute.zipped_divide(mA, tiler_mk) + gB = cute.zipped_divide(mB, tiler_nk) + gSFA = cute.zipped_divide(sfa_tensor, tiler_mk) + gSFB = cute.zipped_divide(sfb_tensor, tiler_nk) + gC = cute.zipped_divide(mC, tiler_mn) + + gA_tile = gA[(None, None), (cta_m, None, cta_l)] + gB_tile = gB[(None, None), (cta_n, None, cta_l)] + gSFA_tile = gSFA[(None, None), (cta_m, None, cta_l)] + gSFB_tile = gSFB[(None, None), (cta_n, None, cta_l)] + gC_tile = gC[(None, None), (cta_m, cta_n, cta_l)] + + # Shared memory layouts for A/B/SFA/SFB/D + # (MMA, MMA_M, MMA_K, PIPE) + a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + self.ab_dtype, + self.num_main_stages, + ) + + # (MMA, MMA_N, MMA_K, PIPE) + b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + self.ab_dtype, + self.num_main_stages, + ) + + # (MMA, MMA_M, MMA_K, PIPE) + sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + self.sf_vec_size, + self.num_main_stages, + ) + + # (MMA, MMA_N, MMA_K, PIPE) + sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + self.sf_vec_size, + self.num_main_stages, + ) + + cta_tile_shape_mnk = cute.shape_div( + mma_tiler_mnk, (cute.size(tiled_mma.thr_id.shape), 1, 1) + ) + epi_tile = sm100_utils.compute_epilogue_tile_shape( + cta_tile_shape_mnk, + self.use_2cta_instrs, + d_layout, + c_dtype, + ) + smem_epi_staged_layout = sm100_utils.make_smem_layout_epi( + c_dtype, + d_layout, + epi_tile, + self.tma_store_stages, + ) + + # UMMA ACC TMEM Layout + # ((MMA_M, MMA_N), REST_MMA_M, REST_MMA_N, ACC_STAGES) + tmem_accs_layout = cute_ext.make_tmem_layout_acc( + tiled_mma, mma_tiler_mnk, self.num_acc_stages + ) + + sfa_tmem_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + + sfb_tmem_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + + # Allocate UMMA Buffers + buffer_smem_a = cute_ext.allocate( + self.ab_dtype, + cute.AddressSpace.smem, + a_smem_layout_staged, + alignment=1024, + ) + + buffer_smem_b = cute_ext.allocate( + self.ab_dtype, + cute.AddressSpace.smem, + b_smem_layout_staged, + alignment=1024, + ) + + buffer_smem_sfa = cute_ext.allocate( + self.sf_dtype, + cute.AddressSpace.smem, + sfa_smem_layout_staged, + alignment=1024, + ) + + buffer_smem_sfb = cute_ext.allocate( + self.sf_dtype, + cute.AddressSpace.smem, + sfb_smem_layout_staged, + alignment=1024, + ) + + buffer_tmem_accs = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.tmem, + tmem_accs_layout, + alignment=16, + ) + + buffer_tmem_sfa = cute_ext.allocate( + self.sf_dtype, + cute.AddressSpace.tmem, + sfa_tmem_layout, + alignment=16, + ) + + buffer_tmem_sfb = cute_ext.allocate( + self.sf_dtype, + cute.AddressSpace.tmem, + sfb_tmem_layout, + alignment=16, + ) + + buffer_tmem_sfa_compact = cute.filter_zeros(buffer_tmem_sfa) + buffer_tmem_sfb_compact = cute.filter_zeros(buffer_tmem_sfb) + + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + + tiled_copy_s2t_sfa = cute.nvgpu.tcgen05.make_s2t_copy( + copy_atom_s2t, buffer_tmem_sfa_compact + ) + tiled_copy_s2t_sfb = cute.nvgpu.tcgen05.make_s2t_copy( + copy_atom_s2t, buffer_tmem_sfb_compact + ) + + # Allocate SMEM buffer for C + buffer_smem_d = cute_ext.allocate( + c_dtype, + cute.AddressSpace.smem, + smem_epi_staged_layout, + alignment=1024, + ) + + # Create the TMEM load atom + copy_atom_t2r = sm100_utils.get_tmem_load_op( + cta_tile_shape_mnk, + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + self.use_2cta_instrs, + ) + + # Derive tiled_copy_t2r from the allocated TMEM buffer + accumulators = cute.zipped_divide(buffer_tmem_accs, ((epi_tile), 1)) + acc_epi_div = accumulators[((None, None), 0), 0] + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + + # Derive per-thread RMEM layout for the T2R epilogue copy + gC_tile_epi = cute.flat_divide(gC_tile, epi_tile) + acc_epi_rmem_layout = cute_ext.make_t2r_rmem_layout( + tiled_copy_t2r, gC_tile_epi, tidx + ) + + # Allocate RMEM buffers + buffer_rmem_t2r = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + acc_epi_rmem_layout, + alignment=32, + ) + buffer_rmem_r2s = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + acc_epi_rmem_layout, + alignment=32, + ) + + # TMA -> UMMA + mainloop_pipe = cute_ext.TMAToUMMAPipeline.create( + num_stages=self.num_main_stages, + mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS, + ) + + # UMMA -> TMEM + acc_pipe = cute_ext.UMMAtoAsyncPipeline.create( + num_stages=self.num_acc_stages, + mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS, + consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R, + consumer_arv_count=128, + ) + + # warp assignment: [0]-tma_store, [0-3]-epi, [4]-mma, [5]-tma_load + tma_store_warp_id = 0 + mma_warp_id = 4 + tma_load_warp_id = 5 + is_tma_load_warp = warp_idx == tma_load_warp_id + is_mma_warp = warp_idx == mma_warp_id + is_epi_warp = warp_idx < 4 + + # SMEM -> GMEM + tma_store_pipe = cute_ext.TMAStorePipeline( + stages=self.tma_store_stages, + arv_count=128, + barrier_id=1, + tma_warp_id=tma_store_warp_id, + ) + + k_tile_size = cute.size(gA, mode=[1, 1]) + + if is_tma_load_warp: + for k_tile_idx in cutlass.range(0, k_tile_size, 1, unroll=1): + gA_k = gA_tile[None, None, k_tile_idx] + gB_k = gB_tile[None, None, k_tile_idx] + gSFA_k = gSFA_tile[None, None, k_tile_idx] + gSFB_k = gSFB_tile[None, None, k_tile_idx] + + # Scoped state management - pipeline object manages state internally + ( + producer_stage_token, + stage_idx, + ) = mainloop_pipe.producer_acquire_and_get_stage() + mbar = cute_ext.get_mbarrier(producer_stage_token) + ## producer_body begin ## + buffer_smem_a_sliced = buffer_smem_a[None, None, None, stage_idx] + buffer_smem_b_sliced = buffer_smem_b[None, None, None, stage_idx] + buffer_smem_sfa_sliced = buffer_smem_sfa[None, None, None, stage_idx] + buffer_smem_sfb_sliced = buffer_smem_sfb[None, None, None, stage_idx] + + a_cta_v_map = cute_ext.get_cta_v_map_ab( + mA, mma_tiler_mnk, tiled_mma, "A" + ) + b_cta_v_map = cute_ext.get_cta_v_map_ab( + mB, mma_tiler_mnk, tiled_mma, "B" + ) + sfa_cta_v_map = cute_ext.get_cta_v_map_ab( + sfa_tensor, mma_tiler_mnk, tiled_mma, "SFA" + ) + sfb_cta_v_map = cute_ext.get_cta_v_map_ab( + sfb_tensor, mma_tiler_mnk, tiled_mma, "SFB" + ) + + cute_ext.tma_load( + gA_k, + buffer_smem_a_sliced, + mbar, + cta_v_map=a_cta_v_map, + ) + cute_ext.tma_load( + gB_k, + buffer_smem_b_sliced, + mbar, + cta_v_map=b_cta_v_map, + ) + cute_ext.tma_load( + gSFA_k, + buffer_smem_sfa_sliced, + mbar, + cta_v_map=sfa_cta_v_map, + ) + cute_ext.tma_load( + gSFB_k, + buffer_smem_sfb_sliced, + mbar, + cta_v_map=sfb_cta_v_map, + ) + ## producer_body end ## + mainloop_pipe.producer_commit_and_advance() + + if is_mma_warp: + producer_stage_token, acc_stage_idx = ( + acc_pipe.producer_acquire_and_get_stage() + ) + ## acc_producer_body begin ## + accumulators_sliced = buffer_tmem_accs[None, None, None, acc_stage_idx] + + mma_atom = cute.make_mma_atom(tiled_mma.op) + mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, False) + + filtered_buffer_smem_sfa = cute.filter_zeros(buffer_smem_sfa) + filtered_buffer_smem_sfb = cute.filter_zeros(buffer_smem_sfb) + + for k_tile_idx in cutlass.range(0, k_tile_size, 1, unroll=1): + # Scoped state management - pipeline object manages consumer state internally + ( + _, + mainloop_idx, + ) = mainloop_pipe.consumer_wait_and_get_stage() + ## tma_consumer_body begin ## + buffer_smem_a_sliced_stage = buffer_smem_a[ + (None, None, None, mainloop_idx) + ] + buffer_smem_b_sliced_stage = buffer_smem_b[ + (None, None, None, mainloop_idx) + ] + filtered_buffer_smem_sfa_sliced_stage = filtered_buffer_smem_sfa[ + (None, None, None, mainloop_idx) + ] + filtered_buffer_smem_sfb_sliced_stage = filtered_buffer_smem_sfb[ + (None, None, None, mainloop_idx) + ] + + # Copy SFA/SFB from SMEM to TMEM (UTCCP) + src_partitioned_SFA = cute_ext.partition( + filtered_buffer_smem_sfa_sliced_stage, + cute.Int32(0), + layout_tv=tiled_copy_s2t_sfa.layout_src_tv_tiled, + tiler=cute.core._pack_tile(tiled_copy_s2t_sfa.tiler_mn), + ) + dst_partitioned_SFA = cute_ext.partition( + buffer_tmem_sfa_compact, + cute.Int32(0), + layout_tv=tiled_copy_s2t_sfa.layout_dst_tv_tiled, + tiler=cute.core._pack_tile(tiled_copy_s2t_sfa.tiler_mn), + ) + + cute_ext.copy( + src_partitioned_SFA, dst_partitioned_SFA, copy_atom=copy_atom_s2t + ) + + src_partitioned_SFB = cute_ext.partition( + filtered_buffer_smem_sfb_sliced_stage, + cute.Int32(0), + layout_tv=tiled_copy_s2t_sfb.layout_src_tv_tiled, + tiler=cute.core._pack_tile(tiled_copy_s2t_sfb.tiler_mn), + ) + dst_partitioned_SFB = cute_ext.partition( + buffer_tmem_sfb_compact, + cute.Int32(0), + layout_tv=tiled_copy_s2t_sfb.layout_dst_tv_tiled, + tiler=cute.core._pack_tile(tiled_copy_s2t_sfb.tiler_mn), + ) + + cute_ext.copy( + src_partitioned_SFB, dst_partitioned_SFB, copy_atom=copy_atom_s2t + ) + + for k_block_idx in cutlass.range(mma_inst_tile_k, unroll_full=True): + buffer_smem_a_sliced = buffer_smem_a_sliced_stage[ + None, None, k_block_idx + ] + buffer_smem_b_sliced = buffer_smem_b_sliced_stage[ + None, None, k_block_idx + ] + + cute_ext.dot_block_scaled( + mma_atom, + cute.append_ones(buffer_smem_a_sliced, up_to_rank=3), + buffer_tmem_sfa[None, None, k_block_idx], + cute.append_ones(buffer_smem_b_sliced, up_to_rank=3), + buffer_tmem_sfb[None, None, k_block_idx], + accumulators_sliced, + ) + mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True) + + ## tma_consumer_body end ## + mainloop_pipe.consumer_release_and_advance() + + ## acc_producer_body end ## + acc_pipe.producer_commit_and_advance() + + if is_epi_warp: + _, acc_stage_idx = acc_pipe.consumer_wait_and_get_stage() + ## acc_consume_body begin ## + tmem_acc_stage = buffer_tmem_accs[ + (None, None), 0, 0, acc_stage_idx + ] # (MMA_M, MMA_N) + # (EPI_TILE_M, EPI_TILE_N, EPI_REST_M, EPI_REST_N) + # we have an implicit assumption that EPI_REST_M == 1 + tmem_acc_epi_stage = cute.flat_divide(tmem_acc_stage, epi_tile) + + subtile_cnt = cute.size(tmem_acc_epi_stage.shape, mode=[3]) # EPI_REST_N + for subtile_idx in range(subtile_cnt): + # TMEM -> RMEM + cute_ext.partition_and_copy( + thr_copy_t2r, + tmem_acc_epi_stage[(None, None, 0, subtile_idx)], + buffer_rmem_t2r, + ) + + # RMEM -> RMEM + buffer_rmem_r2s.store( + self.epilogue_op(buffer_rmem_t2r.load().to(c_dtype)) + ) + + # Acquire pipeline stage and synchronize before RMEM->SMEM copy + tma_store_pipe.acquire_sync() + tma_store_idx = tma_store_pipe.get_index() + + # RMEM -> SMEM + tiled_copy_r2s = cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype), + tiled_copy_t2r, + ) + cute_ext.partition_and_copy( + tiled_copy_r2s.get_slice(tidx), + buffer_rmem_r2s, + buffer_smem_d[None, None, tma_store_idx], + ) + + # Fence SMEM writes and synchronize before TMA store + tma_store_pipe.commit_sync() + + # SMEM -> GMEM (only designated TMA store warp performs TMA store) + if warp_idx == tma_store_warp_id: + c_cta_v_map = cute_ext.get_cta_v_map_c(mC, epi_tile) + cute_ext.tma_store( + buffer_smem_d[None, None, tma_store_idx], + gC_tile_epi[(None, None, 0, subtile_idx)], + cta_v_map=c_cta_v_map, + ) + + # Release pipeline stage and advance + tma_store_pipe.release_advance() + + tma_store_pipe.tail() + acc_pipe.consumer_release_and_advance() + + +@cute.experimental.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_tensor: cute.Tensor, + sf_mma_tensor: cute.Tensor, +): + """ + Convert scale factor tensor from MKL layout to mma specification + M(32x4xrest_m)xK(4xrest_k)xL layout + """ + # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) + # group to ((32, 4, rest_m), (4, rest_k), l) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + + +# TODO: add residual support (C) +@dataclass +class BlockScaledGemmTestbed: + """ + Testbed for block-scaled GEMM operations on Blackwell (SM100) architecture. + + This class manages test data and tensors for block-scaled matrix multiplication: + D = (A * scale_factor_A) @ (B * scale_factor_B) + + The testbed maintains three representations of each tensor: + 1. Reference tensors (f32 on CPU) - used for reference computation and validation + 2. CUTE tensors - device tensors passed directly to CUDA kernels + 3. PyTorch tensors - mirrors of CUTE tensors for host-side operations + + Attributes: + a_ref, b_ref: Reference input matrices (f32 format) + sfa_ref, sfb_ref: Reference scale factors for A and B matrices (f32 format) + d_ref: Reference output matrice (f32 format) + + a_tensor, b_tensor: CUTE tensors for input matrices (device) + sfa_tensor, sfb_tensor: CUTE tensors for scale factors (device) + d_tensor: CUTE tensors for output (device) + + a_torch, b_torch: PyTorch mirrors of A and B CUTE tensors + sfa_torch, sfb_torch: PyTorch mirrors of scale factor CUTE tensors + d_torch: PyTorch mirrors of D CUTE tensors + + The class provides: + - Automatic tensor creation with proper layouts and alignment + - Scale factor tensor generation with block-scaled MMA layout + - Reference checking via einsum-based computation + + Example: + testbed = BlockScaledGemmTestbed( + MNKL=(128, 128, 64, 1), + mma_dtypes=(cutlass.Float16, cutlass.Float16, cutlass.Float32), + c_dtypes=(cutlass.Float16), + sf_dtype=cutlass.Float16, + sf_vec_size=32, + a_major='m', b_major='n', d_major='m' + ) + # ... run kernel with testbed.a_tensor, testbed.b_tensor, etc. + testbed.reference_check() # Validate results + """ + + import torch + + # Reference tensors (all are in f32 format for simplicity of + # reference checks) + a_ref: torch.Tensor + b_ref: torch.Tensor + sfa_ref: torch.Tensor + sfb_ref: torch.Tensor + + # CUTE tensors (to be passed to the device kernel) + a_tensor: cute.Tensor + b_tensor: cute.Tensor + sfa_tensor: cute.Tensor + sfb_tensor: cute.Tensor + d_tensor: cute.Tensor + + # PyTorch tensors (mirrors the CUTE tensors above); these tensors + # can be used on the host, for example if certain trivial epilogue + # needs to be performed. + a_torch: torch.Tensor + b_torch: torch.Tensor + sfa_torch: torch.Tensor + sfb_torch: torch.Tensor + d_torch: torch.Tensor + + def __init__( + self, + MNKL: Tuple[int, int, int, int], + mma_dtypes: tuple[ + Type[cutlass.Numeric], Type[cutlass.Numeric], Type[cutlass.Numeric] + ], + c_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + a_major: str, + b_major: str, + d_major: str, + ): + import cutlass.torch as cutlass_torch + + self.d_major = d_major + + # Problem size + (M, N, K, L) = MNKL + + a_dtype, b_dtype, _ = mma_dtypes + + assert a_major in ("m", "k"), f"a_major must be 'm' or 'k', got {a_major}" + assert b_major in ("n", "k"), f"b_major must be 'n' or 'k', got {b_major}" + assert d_major in ("m", "n"), f"d_major must be 'm' or 'n', got {d_major}" + + self.a_ref = cutlass_torch.matrix(L, M, K, a_major == "m", cutlass.Float32) + self.b_ref = cutlass_torch.matrix(L, N, K, b_major == "n", cutlass.Float32) + self.d_temp = cutlass_torch.matrix(L, M, N, d_major == "m", cutlass.Float32) + + self.a_tensor, self.a_torch = cutlass_torch.cute_tensor_like( + self.a_ref, a_dtype, is_dynamic_layout=True, assumed_align=16 + ) + self.b_tensor, self.b_torch = cutlass_torch.cute_tensor_like( + self.b_ref, b_dtype, is_dynamic_layout=True, assumed_align=16 + ) + self.d_tensor, self.d_torch = cutlass_torch.cute_tensor_like( + self.d_temp, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Mark tensor with element divisibility for 16B alignment + self.a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=32 if a_dtype == cutlass.Float4E2M1FN else 16, + ) + self.b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=32 if b_dtype == cutlass.Float4E2M1FN else 16, + ) + self.d_tensor.mark_compact_shape_dynamic( + mode=1 if d_major == "k" else 0, + stride_order=(2, 0, 1) if d_major == "n" else (2, 1, 0), + divisibility=32 if c_dtype == cutlass.Float4E2M1FN else 16, + ) + + self.sfa_ref, self.sfa_tensor, self.sfa_torch = self.create_scale_factor_tensor( + L, M, K, sf_vec_size, sf_dtype + ) + self.sfb_ref, self.sfb_tensor, self.sfb_torch = self.create_scale_factor_tensor( + L, N, K, sf_vec_size, sf_dtype + ) + + # Create scale factor tensor + @staticmethod + def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): + import torch + import cutlass.torch as cutlass_torch + + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + ref_shape = (l, mn, sf_k) + ref_permute_order = (1, 2, 0) # MKL + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + mma_permute_order = (3, 4, 1, 5, 2, 0) # M(32x4xrest_m)xK(4xrest_k)xL + + # Create f32 ref torch tensor (cpu) + ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + ref_shape, + torch.float32, + permute_order=ref_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=1, + max_val=3, + ), + ) + + # Create f32 cute torch tensor (cpu) + cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=0, + max_val=1, + ), + ) + + # convert ref f32 tensor to cute f32 tensor + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + from_dlpack(ref_f32_torch_tensor_cpu), + from_dlpack(cute_f32_torch_tensor_cpu), + ) + cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() + + # reshape makes memory contiguous + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(l, mn, sf_k, sf_vec_size) + .reshape(l, mn, sf_k * sf_vec_size) + .permute(*ref_permute_order) + ) + # prune to mkl for reference check. + ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] + + # Create dtype cute torch tensor (cpu) + cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( + cute_f32_torch_tensor_cpu, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + # Convert f32 cute tensor to dtype cute tensor + cute_tensor = cutlass_torch.convert_cute_tensor( + cute_f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=True, + ) + return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor + + # Transfers results back to CPU and uses PyTorch's methods to do + # reference checks + def reference_check(self): + import torch + + # Compute reference result, simulate block-scaled GEMV via 2 FFMA + # based elementwise multiplication and 1 FFMA based matmul computations + res_a = torch.einsum("mkl,mkl->mkl", self.a_ref, self.sfa_ref) + res_b = torch.einsum("nkl,nkl->nkl", self.b_ref, self.sfb_ref) + ref_output = torch.einsum("mkl,nkl->mnl", res_a, res_b) + + # Convert d back to f32 for comparison. + d_epi_device = self.d_temp.cuda() + cute.testing.convert( + self.d_tensor, + from_dlpack(d_epi_device, assumed_align=16).mark_layout_dynamic( + leading_dim=(1 if self.d_major == "n" else 0) + ), + ) + + # abs(actual - expected) <= atol + rtol * abs(expected) + torch.testing.assert_close( + d_epi_device.cpu(), ref_output, atol=1e-01, rtol=1e-02 + ) + print("Reference check finished.") + + +def run( + mnkl: Tuple[int, int, int, int], + mma_inst_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + d_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + d_major: str, +): + """Execute a batched block scaled dense GEMM operation on Blackwell architecture. + + This function prepares input tensors, configures and launches the GEMM kernel, + and performs reference validation. + + :param mnkl: Problem size (M, N, K, L) + :type mnkl: Tuple[int, int, int, int] + :param mma_inst_mn: MMA instruction shape. + :type mma_inst_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :param ab_dtype: Data type for input tensors A and B + :type ab_dtype: Type[Numeric] + :param sf_dtype: Data type for scale factors (SFA/SFB) + :type sf_dtype: Type[Numeric] + :param sf_vec_size: Vector size for the scale factor + :type sf_vec_size: int + :param c_dtype: Data type for output tensor D + :type c_dtype: Type[Numeric] + :param acc_dtype: Accumulator data type (precision) + :type acc_dtype: Type[Numeric] + :param a_major: Major-ness of A tensor (m or k) + :type a_major: str + :param b_major: Major-ness of B tensor (n or k) + :type b_major: str + :param d_major: Major-ness of D tensor (m or n) + :type d_major: str + """ + print("Running Blackwell Dense Block Scaled GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"A: {ab_dtype}, B: {ab_dtype}, D: {d_dtype}, Acc dtype: {acc_dtype}") + print(f"Block scaled MMA with SF: {sf_dtype}, vector size: {sf_vec_size}") + print(f"Matrix majors - A: {a_major}-major, B: {b_major}-major, D: {d_major}-major") + print( + f"Mma Tiler (M, N): {mma_inst_mn}, Cluster Shape: {cluster_shape_mn[0]}x{cluster_shape_mn[1]}x1" + ) + import torch + + # TODO: add can_implement to exclude unsupported/un-implemented test cases + if cluster_shape_mn != (1, 1): + raise RuntimeError("Only 1x1x1 cluster shapes are supported right now.") + if mma_inst_mn != (128, 128): + raise RuntimeError("MMA instruction shape not supported yet.") + if ab_dtype not in (cutlass.Float8E4M3FN, cutlass.Float8E5M2): + raise RuntimeError("Input data type not supported.") + if sf_dtype not in (cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN): + raise RuntimeError("Scale factor data type not supported.") + + if not torch.cuda.is_available(): + raise RuntimeError("A GPU is required to run this example!") + + # Manual seed + torch.manual_seed(111) + + # Create tensors + tb = BlockScaledGemmTestbed( + mnkl, + (ab_dtype, ab_dtype, acc_dtype), + d_dtype, + sf_dtype, + sf_vec_size, + a_major, + b_major, + d_major, + ) + + # JIT-Compile the device kernel + block_scaled_gemm = BlockScaledDenseGemmKernel( + mma_inst_mn=mma_inst_mn, + mma_dtype=(ab_dtype, acc_dtype), + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + ) + + compiled_kernel = cute.experimental.compile( + block_scaled_gemm, + tb.a_tensor, + tb.sfa_tensor, + tb.b_tensor, + tb.sfb_tensor, + tb.d_tensor, + ) + + # Launch the device kernel + compiled_kernel( + tb.a_tensor, + tb.sfa_tensor, + tb.b_tensor, + tb.sfb_tensor, + tb.d_tensor, + ) + + tb.reference_check() + + +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 Sm100 Dense BlockScaled GEMM." + ) + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(512, 256, 256, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_inst_mn", + type=parse_comma_separated_ints, + default=(128, 128), + help="Mma instruction 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.Float8E4M3FN) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) + parser.add_argument("--sf_vec_size", type=int, default=32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--d_dtype", type=cutlass.dtype, default=cutlass.Float16) + 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("--d_major", choices=["n", "m"], type=str, default="n") + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_inst_mn) != 2: + parser.error("--mma_inst_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + run( + args.mnkl, + args.mma_inst_mn, + args.cluster_shape_mn, + args.ab_dtype, + args.sf_dtype, + args.sf_vec_size, + args.d_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.d_major, + ) diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm.py b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm.py new file mode 100644 index 00000000..b717cb76 --- /dev/null +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm.py @@ -0,0 +1,1389 @@ +# Copyright (c) 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 argparse + +import torch +from typing import Type, Tuple + +import cutlass +from cutlass.cute import experimental as cute_ext +from cutlass.base_dsl.typing import Numeric, Constexpr +from cutlass import cute as cute +from cutlass import utils +from cutlass import torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils + +import cutlass.cute.testing as testing + +# ==================================================================================================== +# +# This kernel implements a batched dense GEMM operation: D = A @ B +# where: +# - A has shape (M, K, L) and is stored in global memory +# - B has shape (N, K, L) and is stored in global memory +# - D has shape (M, N, L) and is the output in global memory +# - L is the batch dimension +# +# The kernel uses the LIR (Low-level Intermediate Representation) DSL which is a Python DSL +# for writing high-performance, Blackwell (SM100)-targeted kernels on top of CuTe abstractions. +# +# KEY CONCEPTS: +# - TMA (Tensor Memory Accelerator): Hardware feature for high-bandwidth GMEM <-> SMEM transfers +# - UMMA/MMA: Unified Matrix Multiply-Accumulate hardware units on SM100 +# - TMEM: Tensor Memory - Blackwell's specialized memory for MMA accumulators +# - SMEM: Shared Memory - CTA-local memory for staging data +# - RMEM: Register Memory - Per-thread registers +# +# DATA FLOW: +# GMEM (A,B) --TMA--> SMEM (bufferA, bufferB) --MMA--> TMEM (accumulators) +# TMEM --copy--> RMEM (bufferRAcc) --epilogue--> RMEM (bufferRD) --copy--> SMEM (bufferC) --TMA--> GMEM (D) +# +# WARP SPECIALIZATION: +# This kernel uses 6 warps (192 threads) with specialized roles: +# - Warp 5: TMA load producer (loads A, B tiles from GMEM to SMEM) +# - Warp 4: MMA compute (performs matrix multiply-accumulate) +# - Warps 0-3: Epilogue (TMEM->RMEM->SMEM) and TMA store (warp 0 only) +# +# PIPELINE ARCHITECTURE: +# The kernel uses software pipelining to overlap memory transfers with compute: +# - mainloop_pipe: TMAToUMMAPipeline - synchronizes TMA loads with MMA operations +# - acc_pipe: UMMAtoAsyncPipeline - synchronizes MMA with TMEM->RMEM copies +# - tma_store_pipe: TMAStorePipeline - synchronizes SMEM writes with TMA stores +# +# ==================================================================================================== + + +# ==================================================================================================== +# KERNEL CLASS DEFINITION +# ==================================================================================================== +class DenseGemmKernel: + """ + Dense GEMM kernel class for Blackwell (SM100) GPUs. + + This class encapsulates all the configuration and logic for a high-performance + batched matrix multiplication: D = A @ B (with optional epilogue operation). + + The design follows LIR conventions: + 1. __init__: Store configuration parameters + 2. __call__: JIT-decorated host launcher that computes grid and calls kernel + 3. kernel: Device kernel that performs the actual computation + + Attributes: + mn_tiler (tuple[int, int]): Tile sizes for M and N dimensions (e.g., (128, 256)) + ab_dtype (Type[Numeric]): Data type for input matrices A and B (e.g., Float16) + acc_dtype (Type[Numeric]): Data type for accumulators (typically Float32) + tmem_output_dtype (Type[Numeric]): Data type for TMEM->RMEM copy output + use_2cta_instrs (bool): Whether to use 2-CTA MMA instructions (False = 1-CTA mode) + TMA_STORE_STAGE (int): Number of pipeline stages for TMA store operations + epilogue_op (callable): Optional epilogue function applied to output (default: identity) + """ + + def __init__( + self, + mn_tiler: tuple[int, int], + mma_dtype: tuple[Type[Numeric], Type[Numeric]], + tmem_output_dtype: Type[Numeric], + epilogue_op=lambda x: x, + ): + """ + Initialize the Dense GEMM kernel configuration. + + Args: + mn_tiler: Tuple (M_tile, N_tile) specifying the tile dimensions. + CONSTRAINT: M must be 64 or 128 (SM100 hardware requirement). + Common configurations: (128, 256), (128, 128), (64, 128) + + mma_dtype: Tuple (input_dtype, accumulator_dtype) + - input_dtype: Element type for A and B (e.g., Float16, Float8E4M3FN) + - accumulator_dtype: Precision for accumulation (typically Float32) + + tmem_output_dtype: Element type for TMEM output during epilogue. + Typically matches the output matrix type. + + epilogue_op: Optional function applied to accumulator values before store. + Default is identity (lambda x: x). + Examples: relu, sigmoid, GELU approximations using cute.exp/cute.where + """ + self.mn_tiler = mn_tiler + self.ab_dtype, self.acc_dtype = mma_dtype + self.tmem_output_dtype = tmem_output_dtype + self.use_2cta_instrs = False + + # Number of pipeline stages for TMA store operations. + # More stages = better latency hiding, but more SMEM usage. + self.TMA_STORE_STAGE = 4 + + # Epilogue operation applied in registers before storing output. + self.epilogue_op = epilogue_op + + # ================================================================================================ + # JIT-DECORATED HOST LAUNCHER + # ================================================================================================ + @cute.experimental.jit + def __call__(self, mA: cute.Tensor, mB: cute.Tensor, mD: cute.Tensor): + """ + Host-side JIT-compiled launcher function. + + The @cute.experimental.jit decorator indicates this function: + - Runs on the HOST (CPU) + - Is JIT-compiled when first called + - Computes launch configuration and invokes the GPU kernel + + This function performs two key tasks: + 1. Compute the grid dimensions based on output tensor shape and tile size + 2. Launch the kernel with appropriate grid/block/cluster/smem configuration + + Args: + mA: Input tensor A in global memory, shape (M, K, L) where L is batch + mB: Input tensor B in global memory, shape (N, K, L) + mD: Output tensor D in global memory, shape (M, N, L) + + CUTE ALGEBRA EXPLANATION - tiled_divide: + ----------------------------------------- + cute.tiled_divide(tensor, tiler) divides a tensor into tiles, producing a tensor + with shape: ((Tile), Rest_M, Rest_N, ...) + + Unlike zipped_divide which groups rest dimensions: ((Tile), (Rest_M, Rest_N, ...)) + tiled_divide keeps rest dimensions SEPARATE, making it ideal for grid computation. + + For example, if mD has shape (1024, 1024, 2) and tile_mn = (128, 128, 1): + - div.shape[0] = (128, 128, 1) - the tile shape + - div.shape[1] = 8 - number of tiles in M dimension (1024/128) + - div.shape[2] = 8 - number of tiles in N dimension (1024/128) + - div.shape[3] = 2 - batch dimension L + + The grid is then (8, 8, 2) = 128 CTAs total, each processing one (128, 128) tile. + """ + # Create a packed tile shape for division. The _pack_shape helper handles + # creating the proper CuTe shape representation. + # (*self.mn_tiler, 1) = (M_tile, N_tile, 1) - the 1 handles the batch dimension + tile_mn = cute.core._pack_shape((*self.mn_tiler, 1)) + + # tiled_divide produces shape: ((tile_M, tile_N, 1), num_M_tiles, num_N_tiles, batch_L) + # This is used to compute the grid dimensions. + div = cute.tiled_divide(mD, tile_mn) + + # Grid dimensions: (num_tiles_M, num_tiles_N, batch_size) + # Each CTA (Cooperative Thread Array / thread block) processes one tile. + grid = (div.shape[1], div.shape[2], div.shape[3]) + + # Launch the kernel with Blackwell-specific configuration: + # - block=(192, 1, 1): 6 warps × 32 threads/warp = 192 threads + # Warp assignment: warps 0-3 (epilogue), warp 4 (MMA), warp 5 (TMA load) + # - cluster=(1, 1, 1): Single-CTA mode (no cluster cooperation) + # - smem: Request maximum shared memory capacity for SM100 (~232KB) + self.kernel(mA, mB, mD).launch( + grid=grid, + block=(192, 1, 1), # 6 warps for warp-specialized GEMM + cluster=(1, 1, 1), # Single CTA per cluster + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + ) + + # ================================================================================================ + # DEVICE KERNEL + # ================================================================================================ + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mD: cute.Tensor, + ): + """ + Device-side kernel function - the actual GPU computation. + + The @cute.experimental.kernel decorator indicates this function: + - Runs on the DEVICE (GPU) + - Contains all SMEM/TMEM/RMEM allocations, pipeline setup, and compute logic + - Is compiled to PTX and executed by each thread in the grid + + This kernel follows the standard LIR GEMM structure: + 1. Create tiled_mma configuration + 2. Compute tiler and divide tensors + 3. Allocate SMEM, TMEM, and RMEM buffers + 4. Create pipelines for producer/consumer synchronization + 5. Assign warps to specialized roles + 6. Execute TMA load, MMA compute, and epilogue/store phases + + Args: + mA: Input A tensor (GMEM), shape (M, K, L) + mB: Input B tensor (GMEM), shape (N, K, L) + mD: Output D tensor (GMEM), shape (M, N, L) + """ + + # ======================================================================================== + # STEP 1: CREATE TILED MMA CONFIGURATION + # ======================================================================================== + # The tiled_mma object encapsulates the MMA instruction configuration for Blackwell. + # It defines: + # - The MMA atom shape (the hardware instruction's native tile size) + # - Thread-to-data mapping for the MMA operation + # - Layout requirements for operands + # + # make_trivial_tiled_mma creates a basic tiled MMA configuration: + # - ab_dtype: Element type for A and B operands + # - mma_major_mode(): Returns the major mode for MMA (K-major or MN-major) + # - acc_dtype: Accumulator precision (typically Float32) + # - CtaGroup.ONE: Single-CTA MMA (vs TWO for cooperative 2-CTA) + # - mn_tiler: The (M, N) tile dimensions + # + # The mma_major_mode() is derived from the tensor layout: + # - K-major A: stride(A)[1] < stride(A)[0] (K is the fast dimension) + # - M-major A: stride(A)[0] < stride(A)[1] (M is the fast dimension) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + cute.nvgpu.tcgen05.CtaGroup.ONE, # Single CTA mode + self.mn_tiler, + ) + + # ======================================================================================== + # STEP 2: COMPUTE TILER DIMENSIONS (MNK) + # ======================================================================================== + # The MMA instruction operates on tiles. We need to compute the full MNK tiler + # which includes the K dimension (reduction dimension). + # + # cute.size(tiled_mma.shape_mnk, mode=[2]): + # - tiled_mma.shape_mnk is the (M, N, K) shape of the MMA instruction + # - mode=[2] extracts the K dimension (0=M, 1=N, 2=K) + # - For SM100, this is typically 16 (the instruction's native K) + # + # mma_inst_tile_k (=4) is the number of MMA instructions per K-tile iteration. + # This is a tuning parameter: + # - Higher values (8): Larger K-tile, better MMA utilization, but more SMEM + # - Lower values (2): Smaller K-tile, less SMEM, but more loop iterations + # - 4 is a safe default that balances these tradeoffs + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 # Number of MMA K-tile subdivisions per mainloop iteration + + # Full MNK tiler: (M_tile, N_tile, K_tile) + # K_tile = mma_inst_shape_k * mma_inst_tile_k (e.g., 16 * 4 = 64) + mnk_tiler = ( + self.mn_tiler[0], # M dimension from constructor + self.mn_tiler[1], # N dimension from constructor + mma_inst_shape_k * mma_inst_tile_k, # K dimension + ) + + # Get output tensor layout and type for epilogue configuration + d_layout = utils.LayoutEnum.from_tensor(mD) + d_dtype = mD.element_type + + # Create sub-tilers for each operand: + # - A has shape (M, K, L) → tiler_mk = (M_tile, K_tile) + # - B has shape (N, K, L) → tiler_nk = (N_tile, K_tile) + # - D has shape (M, N, L) → tiler_mn = (M_tile, N_tile) + tiler_mk = (mnk_tiler[0], mnk_tiler[2]) + tiler_nk = (mnk_tiler[1], mnk_tiler[2]) + tiler_mn = (mnk_tiler[0], mnk_tiler[1]) + + # ======================================================================================== + # STEP 3: DIVIDE GLOBAL TENSORS INTO TILES (zipped_divide) + # ======================================================================================== + # cute.zipped_divide is the PRIMARY tiling operation in LIR kernels. + # + # CUTE ALGEBRA EXPLANATION - zipped_divide: + # ------------------------------------------ + # zipped_divide(tensor, tiler) divides a tensor into tiles and produces: + # - Mode 0: The tile shape itself + # - Mode 1: A "zipped" layout of tile coordinates + # + # Result shape: ((TileM, TileK), (RestM, RestK, L)) + # + # For example, if mA has shape (1024, 512, 2) and tiler_mk = (128, 64): + # - gA shape = ((128, 64), (8, 8, 2)) + # - (128, 64): One tile of A + # - (8, 8, 2): 8 tiles in M, 8 tiles in K, 2 batches = 128 total tiles + # + # Key difference from tiled_divide: + # - zipped_divide: ((Tile), (Rest...)) - rest dimensions grouped together + # - tiled_divide: ((Tile), Rest_M, Rest_N, ...) - rest dimensions separate + # + # zipped_divide is preferred for CTA tile selection because the zipped + # rest coordinates can be indexed with a single (cta_m, k, batch) tuple. + gA = cute.zipped_divide(mA, tiler_mk) + gB = cute.zipped_divide(mB, tiler_nk) + gD = cute.zipped_divide(mD, tiler_mn) + + # ======================================================================================== + # STEP 4: PIPELINE CONFIGURATION + # ======================================================================================== + # mainloop_stage: Number of pipeline stages for the TMA load → MMA pipeline. + # More stages allow better overlap of TMA loads with MMA compute. + # - 2 stages: Minimum for double-buffering + # - 4 stages: Good for large GEMMs (better latency hiding) + # Trade-off: More stages = more SMEM usage + # + # acc_stage: Number of accumulator stages in TMEM. + # - For N=256 tiles: use 1 (single accumulator buffer) + # - For N=128 tiles: use 2 (double-buffered) + # Using the correct acc_stage provides measurable performance improvement. + mainloop_stage = 2 + acc_stage = 2 + + # ======================================================================================== + # STEP 5: GET CTA AND THREAD INDICES + # ======================================================================================== + # Each CTA is identified by its position in the 3D grid: (cta_m, cta_n, cta_l) + # - cta_m: Which M-tile this CTA processes + # - cta_n: Which N-tile this CTA processes + # - cta_l: Which batch element this CTA processes + # + # Each thread within a CTA is identified by tid_x (0-191 for 192 threads). + cta_m, cta_n, cta_l = cute.arch.block_idx() + tid_x, _, _ = cute.arch.thread_idx() + + # ======================================================================================== + # STEP 6: SELECT THIS CTA'S TILES FROM GLOBAL TENSORS + # ======================================================================================== + # After zipped_divide, we select the specific tiles for this CTA using slicing. + # + # CUTE SLICING NOTATION: + # - None: Keep this dimension (preserve the mode) + # - integer: Fix this dimension at that index + # + # gA has shape ((M_tile, K_tile), (num_M_tiles, num_K_tiles, batch)) + # gA_tile = gA[(None, None), (cta_m, None, cta_l)] means: + # - (None, None): Keep the tile shape modes (M_tile, K_tile) + # - (cta_m, None, cta_l): Select M-tile cta_m, keep K dimension, select batch cta_l + # + # Result: gA_tile has shape (M_tile, K_tile, num_K_tiles) - one CTA's work + # The K dimension (None) is kept because we iterate over K in the mainloop. + gA_tile = gA[(None, None), (cta_m, None, cta_l)] + gB_tile = gB[(None, None), (cta_n, None, cta_l)] + gD_tile = gD[(None, None), (cta_m, cta_n, cta_l)] + + # ======================================================================================== + # STEP 7: CREATE SMEM LAYOUTS WITH SWIZZLING + # ======================================================================================== + # SMEM layouts must: + # 1. Match the tile dimensions from the tiler + # 2. Include staging for pipeline buffers + # 3. Use swizzle patterns to avoid bank conflicts + # + # make_smem_layout_a/b are helper functions that: + # - Select appropriate swizzle patterns based on major mode and element type + # - Append the stage dimension for pipelining + # - Return a ComposedLayout (layout + swizzle function) + # + # The swizzle pattern interleaves memory addresses across the 32 SMEM banks, + # ensuring that when a warp accesses consecutive elements, they hit different + # banks (avoiding serialization from bank conflicts). + # + # LAYOUT SHAPE: (MMA_ATOM, MMA_TILE, MMA_K, PIPELINE_STAGES) + # For operand A: this encodes how to store M×K tiles with proper bank conflict avoidance + a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + mnk_tiler, + self.ab_dtype, + mainloop_stage, # Number of pipeline stages + ) + b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + mnk_tiler, + self.ab_dtype, + mainloop_stage, + ) + + # ======================================================================================== + # STEP 8: COMPUTE EPILOGUE TILE SHAPE + # ======================================================================================== + # The epilogue processes output tiles in smaller sub-tiles (epi_tile). + # This is necessary because: + # 1. TMEM→RMEM copies have granularity constraints + # 2. TMA stores work on specific tile sizes + # + # cta_tile_shape_mnk: The effective tile shape per CTA after accounting for + # thread-level tiling. This is computed as: + # mnk_tiler / (num_threads_in_mma, 1, 1) + # + # cute.shape_div performs element-wise division of shapes. + # cute.size(tiled_mma.thr_id.shape) gives the number of threads participating in MMA. + cta_tile_shape_mnk = cute.shape_div( + mnk_tiler, (cute.size(tiled_mma.thr_id.shape), 1, 1) + ) + + # compute_epilogue_tile_shape determines the sub-tile size for epilogue operations. + # It considers: + # - CTA tile shape + # - Whether using 1-CTA or 2-CTA instructions + # - Output layout (M-major or N-major) + # - Output data type + epi_tile = sm100_utils.compute_epilogue_tile_shape( + cta_tile_shape_mnk, + self.use_2cta_instrs, + d_layout, + d_dtype, + ) + + # Create epilogue SMEM layout for TMA stores. + # This layout is used for the bufferC staging buffer before TMA store to GMEM. + sc_smem_layout_staged = sm100_utils.make_smem_layout_epi( + d_dtype, + d_layout, + epi_tile, + self.TMA_STORE_STAGE, # Number of TMA store pipeline stages + ) + + # ======================================================================================== + # STEP 9: CREATE TMEM LAYOUT FOR ACCUMULATORS + # ======================================================================================== + # TMEM (Tensor Memory) is Blackwell's specialized memory for MMA accumulators. + # It provides high-bandwidth access for accumulator updates during MMA operations. + # + # TMEM CHARACTERISTICS: + # - Accessible only by the MMA unit within a warpgroup + # - Has a capacity limit of 512 columns + # - Requires specific layout patterns matching MMA instructions + # + # make_tmem_layout_acc: Derives the TMEM accumulator buffer layout from the + # tiled MMA and MNK tiler, with the given number of pipeline stages. + tmem_layout = cute_ext.make_tmem_layout_acc(tiled_mma, mnk_tiler, acc_stage) + + # ======================================================================================== + # STEP 10: ALLOCATE SMEM BUFFERS + # ======================================================================================== + # cute_ext.allocate creates a tensor in the specified address space. + # + # Arguments: + # - type: Element type (e.g., Float16, Float32) + # - address_space: One of smem, tmem, rmem, gmem + # - layout: The layout including staging dimensions + # - alignment: Byte alignment (1024 for SMEM, 16 for TMEM, 32 for RMEM) + # + # ALIGNMENT RATIONALE: + # - SMEM (1024 bytes): Optimal for TMA transfers and swizzle patterns + # - TMEM (16 bytes): Standard tensor memory alignment + # - RMEM (32 bytes): Vectorized register loads/stores + + # Allocate SMEM buffers for A and B operands. + # These buffers hold multiple pipeline stages of tiles loaded from GMEM. + bufferA = cute_ext.allocate( + self.ab_dtype, + cute.AddressSpace.smem, + a_smem_layout_staged, + alignment=1024, + ) + + bufferB = cute_ext.allocate( + self.ab_dtype, + cute.AddressSpace.smem, + b_smem_layout_staged, + alignment=1024, + ) + + # Allocate TMEM buffer for MMA accumulators. + # This stores the running sum: C += A × B across K iterations. + bufferAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.tmem, + tmem_layout, + alignment=16, + ) + + # Allocate SMEM buffer for output (C) - used during epilogue before TMA store. + bufferC = cute_ext.allocate( + d_dtype, + cute.AddressSpace.smem, + sc_smem_layout_staged, + alignment=1024, + ) + + # ======================================================================================== + # STEP 11: CREATE TMEM->RMEM COPY CONFIGURATION + # ======================================================================================== + # The epilogue copies data from TMEM (accumulators) → RMEM (registers) → SMEM → GMEM. + # This section sets up the copy atoms and tiled copies for this path. + # + # get_tmem_load_op: Returns the appropriate tcgen05 load operation for TMEM→RMEM. + # It selects the right instruction based on: + # - CTA tile shape + # - Output layout orientation + # - Data types + # - Epilogue tile size + # - 1-CTA vs 2-CTA mode + copy_atom_t2r = sm100_utils.get_tmem_load_op( + cta_tile_shape_mnk, + d_layout, + self.tmem_output_dtype, + self.acc_dtype, + epi_tile, + self.use_2cta_instrs, + ) + + # ======================================================================================== + # STEP 12: PREPARE ACCUMULATOR FOR EPILOGUE ITERATION + # ======================================================================================== + # The accumulator buffer is divided into epilogue-sized sub-tiles for iteration. + # + # CUTE ALGEBRA EXPLANATION - zipped_divide on accumulators: + # ---------------------------------------------------------- + # We divide bufferAcc by (epi_tile, 1) to create sub-tiles for epilogue processing. + # The "1" preserves the stage dimension. + # + # accumulators = cute.zipped_divide(bufferAcc, ((epi_tile), 1)) + # This creates: ((epi_tile_shape), (rest_subtiles, stages)) + # + # acc_epi_div = accumulators[((None, None), 0), 0] + # - (None, None): Keep the epilogue tile shape + # - 0: Select the first rest-mode position + # - 0: Select the first stage (for tiled_copy_t2r creation) + # + # This gives us one epilogue tile's worth of data for configuring the copy. + accumulators = cute.zipped_divide(bufferAcc, ((epi_tile), 1)) + acc_epi_div = accumulators[((None, None), 0), 0] + + # Create the tiled copy operation for TMEM→RMEM. + # make_tmem_copy creates a TiledCopy object that defines: + # - How threads partition the source (TMEM) + # - How threads partition the destination (RMEM) + # - The mapping between source and destination layouts + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div) + + # ======================================================================================== + # STEP 13: DERIVE RMEM LAYOUT FROM COPY PARTITION + # ======================================================================================== + # RMEM layouts must match the thread-value ownership pattern of the copy. + # We derive the RMEM layout by partitioning the destination and extracting + # the per-thread layout. + # + # CUTE ALGEBRA EXPLANATION - flat_divide: + # --------------------------------------- + # flat_divide(tensor, tiler) flattens all dimensions: + # Result shape: (Tile_M, Tile_N, Rest_M, Rest_N, ...) + # + # Unlike zipped_divide which groups tile and rest separately, + # flat_divide keeps everything flat, which is useful for iteration. + # + # make_t2r_rmem_layout: Derives the per-thread RMEM buffer layout + # produced by a TMEM->RMEM copy for a single epilogue iteration. + gC_mnl_epi = cute.flat_divide(gD_tile, epi_tile) + acc_d_rmem_layout = cute_ext.make_t2r_rmem_layout( + tiled_copy_t2r, gC_mnl_epi, tid_x + ) + + # ======================================================================================== + # STEP 14: ALLOCATE RMEM BUFFERS FOR EPILOGUE + # ======================================================================================== + # RMEM (Register Memory) is per-thread storage. Each thread has its own + # private copy of these buffers. + # + # bufferRAcc: Holds accumulator values copied from TMEM (FP32) + # bufferRD: Holds output values after epilogue conversion (output dtype) + bufferRAcc = cute_ext.allocate( + self.acc_dtype, # FP32 for accumulators + cute.AddressSpace.rmem, + acc_d_rmem_layout, + alignment=32, + ) + bufferRD = cute_ext.allocate( + d_dtype, # Output dtype (e.g., FP16) + cute.AddressSpace.rmem, + acc_d_rmem_layout, + alignment=32, + ) + + # ======================================================================================== + # STEP 15: CREATE PIPELINES + # ======================================================================================== + # Pipelines provide producer/consumer synchronization using hardware barriers. + # They enable overlapping of memory operations with compute. + # + # PIPELINE 1: TMAToUMMAPipeline (mainloop_pipe) + # --------------------------------------------- + # Synchronizes TMA loads (producer) with UMMA/MMA operations (consumer). + # - num_stages: Number of pipeline stages (matches mainloop_stage) + # - mma_operation_type: The type of MMA operation being consumed + # SM100_MMA_1SM_SS = Single SM, Single-Stage MMA (1-CTA mode) + mainloop_pipe = cute_ext.TMAToUMMAPipeline.create( + num_stages=mainloop_stage, + mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS, + ) + + # PIPELINE 2: UMMAtoAsyncPipeline (acc_pipe) + # ------------------------------------------ + # Synchronizes UMMA/MMA operations (producer) with TMEM→RMEM copies (consumer). + # - num_stages: Accumulator stages (acc_stage) + # - mma_operation_type: The MMA operation producing data + # - consumer: The operation consuming data (SM100_COPY_T2R = TMEM→RMEM copy) + # - consumer_arv_count: Number of threads participating as consumers (128 = 4 warps) + acc_pipe = cute_ext.UMMAtoAsyncPipeline.create( + num_stages=acc_stage, + mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS, + consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R, + consumer_arv_count=128, # 4 epilogue warps × 32 threads + ) + + # ======================================================================================== + # STEP 16: WARP ASSIGNMENT AND SPECIALIZATION + # ======================================================================================== + # This kernel uses 6 warps (192 threads) with specialized roles: + # + # Warp 0: TMA store (also participates in epilogue) + # Warps 0-3: Epilogue processing (TMEM→RMEM→SMEM) + # Warp 4: MMA compute + # Warp 5: TMA load + # + # cute.arch.warp_idx(): Returns this thread's warp index (0-5) + # make_warp_uniform: Ensures all threads in a warp see the same value + # (important for conditional branching to avoid divergence) + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Assign warp roles + tma_store_warp_id = 0 + mma_warp_id = 4 + tma_load_warp_id = 5 + + # Boolean flags for role-based execution + is_tma_thr = warp_idx == tma_load_warp_id # Only warp 5 + is_mma_thr = warp_idx == mma_warp_id # Only warp 4 + is_epi_thr = warp_idx < 4 # Warps 0, 1, 2, 3 + + # PIPELINE 3: TMAStorePipeline (tma_store_pipe) + # --------------------------------------------- + # Synchronizes RMEM→SMEM writes with TMA stores. + # Uses named barriers (not mbarriers) for synchronization. + # + # - stages: Number of TMA store pipeline stages + # - arv_count: Number of threads participating in barriers (128 = 4 warps) + # - barrier_id: Named barrier ID (must be unique per pipeline) + # - tma_warp_id: Which warp issues TMA stores (warp 0) + tma_store_pipe = cute_ext.TMAStorePipeline( + stages=self.TMA_STORE_STAGE, + arv_count=128, + barrier_id=1, + tma_warp_id=tma_store_warp_id, + ) + + # ======================================================================================== + # STEP 17: COMPUTE K-TILE ITERATION COUNT + # ======================================================================================== + # cute.size(gA, mode=[1, 1]) extracts the size of the K-tile dimension. + # gA shape after zipped_divide: ((M_tile, K_tile), (num_M_tiles, num_K_tiles, batch)) + # mode=[1, 1] accesses the second element of the second mode = num_K_tiles + k_tile_size = cute.size(gA, mode=[1, 1]) + + # ======================================================================================== + # STEP 18: TMA LOAD WARP - PRODUCER PHASE + # ======================================================================================== + # The TMA load warp (warp 5) loads A and B tiles from GMEM to SMEM. + # This is the PRODUCER in the mainloop pipeline. + # + # The producer loop iterates over K-tiles, loading data ahead of consumption. + # Pipeline stages allow loads to overlap with MMA operations. + if is_tma_thr: + # cutlass.range: A loop construct that supports unrolling. + # unroll=1 means don't unroll (iterate normally). + # This iterates over K-tiles: k = 0, 1, 2, ... k_tile_size-1 + for k in cutlass.range(0, k_tile_size, 1, unroll=1): + # Select the K-tile from the CTA's tile view. + # gA_tile has shape (M_tile, K_tile, num_K_tiles) + # gA_tile[None, None, k] selects the k-th K-tile: shape (M_tile, K_tile) + gA_k = gA_tile[None, None, k] + gB_k = gB_tile[None, None, k] + + # ============================================================================ + # PIPELINE PRODUCER PROTOCOL + # ============================================================================ + # 1. Acquire a pipeline stage (wait for it to be empty) + # 2. Get the mbarrier for TMA synchronization + # 3. Issue TMA loads + # 4. Commit and advance to the next stage + # + # producer_acquire_and_get_stage(): + # - Waits for the next pipeline stage to be empty (consumer released it) + # - Returns (stage_token, idx) where: + # - stage_token: Handle for getting the mbarrier + # - idx: Integer index (0 to num_stages-1) for buffer slicing + ( + producer_stage_token, + idx, + ) = mainloop_pipe.producer_acquire_and_get_stage() + + # get_mbarrier: Retrieves the hardware mbarrier pointer for this stage. + # The mbarrier is signaled by TMA hardware when the load completes. + mbar = cute_ext.get_mbarrier(producer_stage_token) + + ## producer_body begin ## + + # Slice SMEM buffers to the current pipeline stage. + # bufferA has shape (atoms, M, K, stages) + # bufferA[None, None, None, idx] selects stage idx: shape (atoms, M, K) + bufferA_sliced = bufferA[None, None, None, idx] + bufferB_sliced = bufferB[None, None, None, idx] + + # ============================================================================ + # CTA-TO-VALUE MAPS FOR TMA + # ============================================================================ + # cta_v_map (CTA-to-Value map) tells TMA which portion of the global tensor + # this CTA should load. It encodes the mapping from CTA coordinates to + # tensor indices. + # + # get_cta_v_map_ab: Computes the CTA-to-value map for operands A or B. + # Arguments: + # - mA/mB: The global tensor + # - mnk_tiler: The MNK tiler dimensions + # - tiled_mma: The MMA configuration + # - "A"/"B": Which operand this is for + a_cta_v_map = cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A") + b_cta_v_map = cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B") + + # ============================================================================ + # TMA LOAD OPERATIONS + # ============================================================================ + # tma_load: Asynchronous TMA load from GMEM to SMEM. + # + # Arguments: + # - src: Source tensor in GMEM (the K-tile slice) + # - dst: Destination buffer in SMEM (the stage-sliced buffer) + # - mbar: Mbarrier for completion signaling + # - cta_v_map: CTA-to-value mapping layout + # + # The TMA hardware: + # 1. Reads from GMEM at the location specified by cta_v_map + # 2. Writes to SMEM at dst + # 3. Signals mbar when complete + # + # IMPORTANT: src and dst must have matching shapes! + # This is a common source of "source/destination size mismatch" errors. + cute_ext.tma_load( + gA_k, # Source: K-tile from global A + bufferA_sliced, # Destination: SMEM buffer stage + mbar, # Mbarrier for synchronization + cta_v_map=a_cta_v_map, + ) + cute_ext.tma_load( + gB_k, + bufferB_sliced, + mbar, + cta_v_map=b_cta_v_map, + ) + + ## producer_body end ## + + # producer_commit_and_advance: + # - Signals that producer work is complete (mbarrier will be triggered by TMA) + # - Advances internal pipeline state to the next stage + mainloop_pipe.producer_commit_and_advance() + + # ======================================================================================== + # STEP 19: MMA WARP - COMPUTE PHASE + # ======================================================================================== + # The MMA warp (warp 4) performs matrix multiply-accumulate operations. + # It consumes data from SMEM (loaded by TMA warp) and produces results in TMEM. + # + # The MMA warp is both: + # - CONSUMER of mainloop_pipe (waits for TMA loads to complete) + # - PRODUCER of acc_pipe (signals when accumulation is complete) + if is_mma_thr: + # Acquire accumulator pipeline stage before starting MMA operations. + # This reserves a TMEM accumulator buffer for this K-reduction. + producer_stage_token, idx = acc_pipe.producer_acquire_and_get_stage() + + ## acc_producer_body begin ## + + # Select the TMEM accumulator for this stage. + # bufferAcc has shape (MMA_shape, stages) + accumulators_sliced = bufferAcc[None, None, None, idx] + + # ============================================================================ + # MMA ATOM CONFIGURATION + # ============================================================================ + # cute.make_mma_atom: Creates an MMA atom from the tiled_mma operation. + # The MMA atom represents the hardware MMA instruction configuration. + # + # ACCUMULATE field controls whether to: + # - False: Overwrite accumulator (C = A × B) - used for first iteration + # - True: Accumulate into existing value (C += A × B) - used after first + mma_atom = cute.make_mma_atom(tiled_mma.op) + mma_atom.set( + cute.nvgpu.tcgen05.Field.ACCUMULATE, False + ) # First iteration: overwrite + + # Iterate over K-tiles (same loop as TMA load warp) + for k in cutlass.range(0, k_tile_size, 1, unroll=1): + # ============================================================================ + # PIPELINE CONSUMER PROTOCOL + # ============================================================================ + # Wait for TMA load to complete before reading from SMEM. + # consumer_wait_and_get_stage(): + # - Waits for the producer (TMA) to signal the mbarrier + # - Returns (stage_token, mainloop_idx) where mainloop_idx is the stage to read + ( + _, # Stage token not needed for consumer + mainloop_idx, + ) = mainloop_pipe.consumer_wait_and_get_stage() + + ## tma_consumer_body begin ## + + # cute.core.slice_: An alternative slicing function that creates a view. + # This slices the SMEM buffers to the current pipeline stage. + # Equivalent to bufferA[None, None, None, mainloop_idx] + bufferA_sliced_stage = cute.core.slice_( + bufferA, (None, None, None, mainloop_idx) + ) + bufferB_sliced_stage = cute.core.slice_( + bufferB, (None, None, None, mainloop_idx) + ) + + # ============================================================================ + # INNER K-TILE LOOP (MMA INSTRUCTION LOOP) + # ============================================================================ + # Within each K-tile, we execute multiple MMA instructions. + # mma_inst_tile_k (=4) MMA instructions are executed per K-tile. + # + # unroll_full=True: Fully unroll this loop (generate 4 copies of the body) + # This is important for MMA instruction scheduling. + for k_tile in cutlass.range(mma_inst_tile_k, unroll_full=True): + # Select the k_tile-th sub-slice for this MMA instruction. + # bufferA_sliced_stage has shape (MMA_atom, M_tile, K_tile) + # After slicing [None, None, k_tile]: shape (MMA_atom, M_tile) + bufferA_sliced = bufferA_sliced_stage[None, None, k_tile] + bufferB_sliced = bufferB_sliced_stage[None, None, k_tile] + + # ======================================================================== + # CUTE.DOT - MATRIX MULTIPLY-ACCUMULATE + # ======================================================================== + # cute_ext.dot: Performs MMA operation C = A × B (or C += A × B) + # + # Arguments: + # - mma_atom: The MMA instruction configuration + # - a: Input tensor A (must be rank-3) + # - b: Input tensor B (must be rank-3) + # - c: Accumulator tensor C (in TMEM) + # + # CUTE ALGEBRA EXPLANATION - append_ones: + # --------------------------------------- + # cute.append_ones(tensor, up_to_rank=3): + # The MMA instruction expects rank-3 operands. If bufferA_sliced + # is rank-2 after slicing, append_ones pads it to rank-3 by + # appending singleton dimensions: shape (M, K) → (M, K, 1) + # + # This is necessary because the MMA instruction operates on + # 3D tiles even when the logical operation is 2D. + cute_ext.dot( + mma_atom, + cute.append_ones(bufferA_sliced, up_to_rank=3), + cute.append_ones(bufferB_sliced, up_to_rank=3), + accumulators_sliced, + ) + + # After the first MMA instruction, enable accumulation mode. + # Subsequent instructions add to the existing accumulator value. + mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True) + + ## tma_consumer_body end ## + + # Release the mainloop pipeline stage for TMA to reuse. + # consumer_release_and_advance(): + # - Signals that consumer has finished reading this stage + # - Advances internal state to the next stage + mainloop_pipe.consumer_release_and_advance() + + ## acc_producer_body end ## + + # Signal that MMA computation is complete for this tile. + # The epilogue warps will consume this data. + acc_pipe.producer_commit_and_advance() + + # ======================================================================================== + # STEP 20: EPILOGUE WARPS - CONSUME AND STORE PHASE + # ======================================================================================== + # Warps 0-3 handle the epilogue: copying results from TMEM to GMEM. + # This involves: TMEM → RMEM → apply epilogue op → SMEM → TMA store to GMEM + # + # The epilogue is both: + # - CONSUMER of acc_pipe (waits for MMA to complete) + # - PRODUCER/CONSUMER of tma_store_pipe (coordinates SMEM→GMEM stores) + if is_epi_thr: + # Wait for accumulator data to be ready. + _, idx = acc_pipe.consumer_wait_and_get_stage() + + ## acc_consume_body begin ## + + # Select the accumulator stage and reshape for epilogue iteration. + # accumulators_sliced: shape (M_epi, N_epi) after removing stage dimension + accumulators_sliced = bufferAcc[(None, None), 0, 0, idx] + + # Divide the accumulator into epilogue-sized sub-tiles. + # flat_divide creates a flat iteration space over sub-tiles. + # acc_epi_div_tiled: allows iteration with index mn over sub-tiles + acc_epi_div_tiled = cute.flat_divide(accumulators_sliced, epi_tile) + + # Get the number of sub-tiles to process. + # mode=[3] accesses the sub-tile count dimension + subtile_cnt = cute.size(acc_epi_div_tiled.shape, mode=[3]) + + # Iterate over epilogue sub-tiles + for mn in range(subtile_cnt): + # ============================================================================ + # TMEM → RMEM COPY + # ============================================================================ + # partition_and_copy: High-level function that combines partitioning and copying. + # It handles: + # 1. Partitioning source/destination according to the tiled copy layout + # 2. Selecting the appropriate copy method based on memory spaces + # 3. Executing the copy + # + # For TMEM→RMEM, this uses specialized tcgen05 load instructions. + # + # Arguments: + # - tiled_copy.get_slice(tid_x): Per-thread copy configuration + # - source: TMEM accumulator sub-tile + # - destination: RMEM buffer (per-thread, not partitioned) + cute_ext.partition_and_copy( + tiled_copy_t2r.get_slice(tid_x), + acc_epi_div_tiled[None, None, 0, mn], + bufferRAcc, + ) + + # ============================================================================ + # APPLY EPILOGUE OPERATION IN REGISTERS + # ============================================================================ + # bufferRAcc.load(): Reads all values from the RMEM tensor into a register + # .to(d_dtype): Converts from accumulator type (FP32) to output type (FP16) + # self.epilogue_op: Applies user-specified transformation (default: identity) + # bufferRD.store(): Writes the result back to RMEM + # + # Common epilogue operations: + # - Identity: lambda x: x (default) + # - ReLU: cute.where(x > 0, x, cute.full_like(x, 0)) + # - GELU: Uses cute.exp for tanh approximation + # - Sigmoid: 1 / (1 + cute.exp(-x)) + bufferRD.store(self.epilogue_op(bufferRAcc.load().to(d_dtype))) + + # ============================================================================ + # TMA STORE PIPELINE PROTOCOL + # ============================================================================ + # The TMA store pipeline coordinates multiple warps writing to SMEM + # before a single warp (warp 0) issues the TMA store. + # + # acquire_sync(): + # - TMA warp waits for any in-flight TMA ops to complete + # - All warps synchronize via a named barrier + tma_store_pipe.acquire_sync() + + # Get the current pipeline stage index for buffer access + idx = tma_store_pipe.get_index() + + # ============================================================================ + # RMEM → SMEM COPY + # ============================================================================ + # Create a tiled copy for RMEM→SMEM using the same layout as TMEM→RMEM. + # make_tiled_copy_D creates a copy with destination-oriented partitioning. + # + # CopyUniversalOp: A generic copy operation that works for any memory pair. + # The partition_and_copy function will select appropriate vectorization. + tiled_copy_r2s = cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), d_dtype), + tiled_copy_t2r, + ) + + # Copy from RMEM to the current SMEM stage buffer + cute_ext.partition_and_copy( + tiled_copy_r2s.get_slice(tid_x), + bufferRD, + bufferC[None, None, idx], + ) + + # commit_sync(): + # - Fences SMEM writes to ensure visibility for TMA + # - All warps synchronize before TMA store + # This is CRITICAL - TMA must see committed SMEM writes! + tma_store_pipe.commit_sync() + + # ============================================================================ + # TMA STORE (SINGLE WARP) + # ============================================================================ + # Only the designated TMA store warp (warp 0) issues the actual TMA store. + # Other warps skip this but still participate in synchronization. + if warp_idx == tma_store_warp_id: + # get_cta_v_map_c: CTA-to-value map for the output tensor. + # Arguments: + # - mD: Global output tensor + # - epi_tile: Epilogue tile shape + c_cta_v_map = cute_ext.get_cta_v_map_c(mD, epi_tile) + + # tma_store: Asynchronous TMA store from SMEM to GMEM. + # + # Arguments: + # - src: Source buffer in SMEM (current stage) + # - dst: Destination in GMEM (sub-tile at position mn) + # - cta_v_map: CTA-to-value mapping + # + # The store is added to an async bulk group managed by the pipeline. + cute_ext.tma_store( + bufferC[None, None, idx], + gC_mnl_epi[None, None, 0, mn], + cta_v_map=c_cta_v_map, + ) + + # release_advance(): + # - TMA warp commits TMA ops to bulk group + # - All warps advance to the next pipeline stage + tma_store_pipe.release_advance() + + # ============================================================================ + # PIPELINE CLEANUP + # ============================================================================ + # tail(): Called at the end of the pipeline to ensure all TMA stores complete. + # This waits for all in-flight TMA operations before the kernel exits. + # Without this, the kernel might exit before stores are globally visible! + tma_store_pipe.tail() + + # Release the accumulator pipeline stage + acc_pipe.consumer_release_and_advance() + + +# ==================================================================================================== +# HOST-SIDE UTILITY FUNCTIONS +# ==================================================================================================== + + +def create_tensors(l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype): + """ + Create input and output tensors for GEMM operation. + + This function creates: + 1. CPU tensors with proper layouts (for reference computation) + 2. GPU tensors wrapped as CuTe tensors (for kernel execution) + + Args: + l: Batch size (L dimension) + m: M dimension (rows of A, rows of D) + n: N dimension (columns of B, columns of D) + k: K dimension (columns of A, rows of B - the reduction dimension) + a_major: "m" for M-major (column-major in M), "k" for K-major + b_major: "n" for N-major, "k" for K-major + d_major: "m" for M-major, "n" for N-major + ab_dtype: Data type for A and B matrices + d_dtype: Data type for output matrix + + Returns: + Tuple of (a_tensor, b_tensor, d_tensor, a_cpu, b_cpu, d_cpu, d_gpu) + - *_tensor: CuTe tensor wrappers for kernel input + - *_cpu: PyTorch CPU tensors for reference + - d_gpu: PyTorch GPU tensor for result extraction + + TENSOR LAYOUT CONVENTIONS: + - cutlass_torch.matrix(l, m, k, m_major, dtype) creates a tensor of shape (m, k, l) + - m_major=True: M is the fast (stride-1) dimension + - m_major=False: K is the fast dimension + + CUTE TENSOR CREATION: + - cute_tensor_like wraps a PyTorch tensor as a CuTe tensor + - is_dynamic_layout=True: Allows variable problem sizes + - assumed_align=16: Assumes 16-byte alignment for TMA + """ + torch.manual_seed(1111) # For reproducibility + + # Create PyTorch CPU tensors with specified layouts. + # cutlass_torch.matrix(l, m, k, m_major, dtype) creates (m, k, l) tensor + a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype) + d_torch_cpu = cutlass_torch.matrix(l, m, n, d_major == "m", d_dtype) + + # Wrap as CuTe tensors for kernel input. + # cute_tensor_like returns (cute_tensor, pytorch_gpu_tensor) + 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 + ) + d_tensor, d_torch_gpu = cutlass_torch.cute_tensor_like( + d_torch_cpu, d_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + return ( + a_tensor, + b_tensor, + d_tensor, + a_torch_cpu, + b_torch_cpu, + d_torch_cpu, + d_torch_gpu, + ) + + +def compare(a_torch_cpu, b_torch_cpu, d_torch_gpu, d_dtype, tolerance): + """ + Compare kernel output against PyTorch reference. + + The reference computation uses torch.einsum with the pattern "mkl,nkl->mnl": + - A has shape (m, k, l): indices m, k, l + - B has shape (n, k, l): indices n, k, l + - Output has shape (m, n, l): indices m, n, l + - The 'k' index is summed (contraction) + + This computes: D[m,n,l] = sum_k A[m,k,l] * B[n,k,l] + + Args: + a_torch_cpu: Input A tensor on CPU + b_torch_cpu: Input B tensor on CPU + d_torch_gpu: Kernel output tensor on GPU + d_dtype: Output data type (for reference tensor creation) + tolerance: Absolute tolerance for comparison + + Raises: + AssertionError: If kernel output doesn't match reference within tolerance + """ + # Compute reference using einsum + ref = torch.einsum("mkl,nkl->mnl", a_torch_cpu, b_torch_cpu) + + # Wrap reference as CuTe tensor (for consistent comparison) + _, ref_torch_gpu = cutlass_torch.cute_tensor_like( + ref, d_dtype, is_dynamic_layout=True, assumed_align=16 + ) + ref_result = ref_torch_gpu.cpu() + + # Compare with tolerance + torch.testing.assert_close( + d_torch_gpu.cpu(), ref_result, atol=tolerance, rtol=1e-05 + ) + + +def run( + mnkl: Tuple[int, int, int, int], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ab_dtype: Type[Numeric], + c_dtype: Type[Numeric], + acc_dtype: Type[Numeric], + a_major: str, + b_major: str, + c_major: str, + warmup_iterations: int = 0, + iterations: int = 1, + use_cold_l2: bool = False, + tolerance: float = 1e-02, + skip_ref_check: bool = False, + **kwargs, +): + """Execute a batched dense GEMM operation on Blackwell architecture with performance benchmarking. + + This function: + 1. Creates input tensors + 2. Instantiates and compiles the kernel + 3. Executes the kernel + 4. Validates correctness against PyTorch reference + 5. Benchmarks performance + + COMPILATION PATTERN: + ------------------- + CRITICAL: Always use explicit compilation to avoid JIT overhead! + + WRONG (recompiles every call, ~1000x slower): + kernel = DenseGemmKernel(...) + kernel(a, b, d) # JIT compilation happens here every time! + + CORRECT (compile once, run many times): + kernel = DenseGemmKernel(...) + compiled = cute_ext.compile(kernel, a, b, d) # Compile once + compiled(a, b, d) # Fast execution + + Args: + mnkl: Problem size tuple (M, N, K, L) + mma_tiler_mn: MMA tile shape (M_tile, N_tile) + cluster_shape_mn: Cluster shape (currently unused in 1-CTA mode) + ab_dtype: Input data type + d_dtype: Output data type + acc_dtype: Accumulator data type + a_major, b_major, d_major: Layout specifications ("m"/"k"/"n") + warmup_iterations: Warmup iterations before timing + iterations: Timed iterations + use_cold_l2: Whether to use cold L2 cache (requires fresh tensors) + tolerance: Tolerance for numerical comparison + skip_ref_check: Skip reference validation + + Returns: + exec_time: Execution time in microseconds per iteration + """ + print("Running Blackwell Dense GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, D dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, D: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + 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'}") + + m, n, k, l = mnkl + + ab_dtype = ab_dtype + d_dtype = c_dtype + d_major = c_major + + # Create tensors + a_tensor, b_tensor, d_tensor, a_torch_cpu, b_torch_cpu, d_torch_cpu, d_torch_gpu = ( + create_tensors(l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype) + ) + + # Instantiate kernel with configuration + dense_gemm = DenseGemmKernel( + mn_tiler=mma_tiler_mn, + mma_dtype=(ab_dtype, acc_dtype), + tmem_output_dtype=d_dtype, + ) + + # compile() pre-compiles the kernel for the given tensor shapes/types + compiled_dense_gemm = cute_ext.compile(dense_gemm, a_tensor, b_tensor, d_tensor) + + # Execute the kernel (now fast - no recompilation) + compiled_dense_gemm(a_tensor, b_tensor, d_tensor) + + # Validate correctness + if not skip_ref_check: + compare(a_torch_cpu, b_torch_cpu, d_torch_gpu, d_dtype, tolerance) + print("check reference: PASS") + + # Tensor generator for benchmarking + 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 + ) + d_tensor, _ = cutlass_torch.cute_tensor_like( + d_torch_cpu, d_dtype, is_dynamic_layout=True, assumed_align=16 + ) + return testing.JitArguments(a_tensor, b_tensor, d_tensor) + + # For cold L2 benchmarking, we need enough tensor copies to flush the cache + 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() + + d_torch_cpu.numel() * d_torch_cpu.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + # Run benchmark + exec_time = testing.benchmark( + compiled_dense_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time + + +# ==================================================================================================== +# COMMAND-LINE INTERFACE +# ==================================================================================================== + +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 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.Float32) + parser.add_argument("--d_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + + 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("--d_major", choices=["n", "m"], type=str, default="n") + + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", type=int, default=1, help="Number of iterations" + ) + parser.add_argument("--use_cold_l2", action="store_true", help="Use cold L2") + parser.add_argument( + "--tolerance", type=float, default=1e-02, help="Tolerance for validation" + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + + 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") + + exec_time = run( + args.mnkl, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.ab_dtype, + args.d_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.d_major, + args.warmup_iterations, + args.iterations, + args.use_cold_l2, + args.tolerance, + args.skip_ref_check, + ) + + print(f"Execution time: {exec_time} microseconds per iteration") diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_2sm.py b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_2sm.py new file mode 100644 index 00000000..13561b0d --- /dev/null +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_2sm.py @@ -0,0 +1,519 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +""" +2SM Dense GEMM example using cute_ext decorators. +""" + +import torch +import math +import cutlass +from cutlass import cute +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils as utils +from cutlass.base_dsl.typing import Numeric +from typing import Type + + +def create_gemm_tensors_torch( + M, + N, + K, + majors: tuple[ + cute.nvgpu.tcgen05.OperandMajorMode, + cute.nvgpu.tcgen05.OperandMajorMode, + cute.nvgpu.tcgen05.OperandMajorMode, + ], + dtypes: tuple[torch.dtype, torch.dtype, torch.dtype], +): + A = None + B = None + D = None + + if majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.MN: + A = torch.empty(K, M).random_(-4, 4).permute(1, 0).to(dtypes[0]).cuda() + elif majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.K: + A = torch.empty(M, K).random_(-4, 4).permute(0, 1).to(dtypes[0]).cuda() + if majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.MN: + B = torch.empty(K, N).random_(-4, 4).permute(1, 0).to(dtypes[1]).cuda() + elif majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.K: + B = torch.empty(N, K).random_(-4, 4).permute(0, 1).to(dtypes[1]).cuda() + if majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.MN: + D = torch.empty(N, M).random_(-4, 4).permute(1, 0).to(dtypes[2]).cuda() + elif majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.K: + D = torch.empty(M, N).random_(-4, 4).permute(0, 1).to(dtypes[2]).cuda() + + return A, B, D + + +def get_gemm_tensors( + M, + N, + K, + majors: tuple[ + cute.nvgpu.tcgen05.OperandMajorMode, + cute.nvgpu.tcgen05.OperandMajorMode, + cute.nvgpu.tcgen05.OperandMajorMode, + ], + dtypes: tuple[torch.dtype, torch.dtype, torch.dtype], +): + A, B, D = create_gemm_tensors_torch(M, N, K, majors, dtypes) + + A_cute = from_dlpack(A, assumed_align=16).mark_layout_dynamic( + leading_dim=1 if majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0 + ) + B_cute = from_dlpack(B, assumed_align=16).mark_layout_dynamic( + leading_dim=1 if majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0 + ) + D_cute = from_dlpack(D, assumed_align=16).mark_layout_dynamic( + leading_dim=1 if majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0 + ) + + return A, B, D, A_cute, B_cute, D_cute + + +def sm100_4x4x1_kernel_builder( + use_tma_multicast: bool, + use_2cta_instrs: bool, + acc_dtype: Type[Numeric], + M: int, + N: int, +): + CLUSTER_SHAPE = (2, 1, 1) + GRID_SHAPE = ( + math.ceil(M / 128), + math.ceil(N / 256), + 1, + ) # TODO (xpbowler): remove hard-code + NUM_WARPS_PER_CTA = 6 + TMA_STORE_PIPE_DEPTH = 4 + MAINLOOP_STAGE_DEPTH = 4 # pipeline depth of TMA->MMA + # pipeline depth of mainloop->epilogue. only useful if using persistent CTA + EPILOGUE_STAGE_DEPTH = 1 + + # m256n256k16 2SM MMA / m128n256k16 1SM MMA + mma_inst_shape_mnk = (256, 256, 16) if use_2cta_instrs else (128, 256, 16) + + @cute_ext.kernel + def kernel( + mA: cute.Tensor, + mB: cute.Tensor, + mD: cute.Tensor, + ): + d_layout = utils.LayoutEnum.from_tensor(mD) + d_dtype = mD.element_type + ab_dtype = mA.element_type + + mma_inst_shape_m, mma_inst_shape_n, mma_inst_shape_k = mma_inst_shape_mnk + if cutlass.const_expr(use_2cta_instrs): + cta_group = cute.nvgpu.tcgen05.CtaGroup.TWO + else: + cta_group = cute.nvgpu.tcgen05.CtaGroup.ONE + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + acc_dtype, + cta_group, + (mma_inst_shape_m, mma_inst_shape_n), + ) + + mma_inst_tile_k = ( + 4 # 4 MMAs per MMA tile K. For 16b types, tcgen05.mma has K=16. + ) + mma_inst_tile_m = mma_inst_tile_n = 1 # 1 MMAs per MMA tile M/N + bM = mma_inst_shape_m * mma_inst_tile_m + bN = mma_inst_shape_n * mma_inst_tile_n + bK = mma_inst_shape_k * mma_inst_tile_k + mnk_tiler = (bM, bN, bK) + + cta_m, cta_n, _ = cute.arch.block_idx() + tid_x, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(CLUSTER_SHAPE), + cute.core._pack_shape((cute.size(tiled_mma.thr_id.shape),)), + ) + cluster_layout_v_size = cute.size(cluster_layout_vmnk.shape[0]) + mma_coord_vmnk = ( + cta_m % cluster_layout_v_size, + cta_m // cluster_layout_v_size, + cta_n, + ) + + gA = cute.zipped_divide(mA, (bM, bK)) # ((bM, bK), (M/bM, K/bK)) + gA_tma = cute.zipped_divide( + mA, (bM // cluster_layout_v_size, bK) + ) # ((bM/2, bK), (2*M/bM, K/bK)) + tAgA = gA_tma[(None, None), (cta_m, None)] # ((bM/2, bK), (1, K/bK)) + + gB_tma = cute.zipped_divide( + mB, (bN // cluster_layout_v_size, bK) + ) # ((bN/2, bK), (2*M/bM, K/bK)) + # ((bN/2, bK), (1, K/bK)) + tBgB = gB_tma[ + (None, None), + (cluster_layout_v_size * cta_n + cta_m % cluster_layout_v_size, None), + ] + + gD_tma = cute.zipped_divide( + mD, (bM // cluster_layout_v_size, bN) + ) # ((bM/2, bN), (2*M/bM, N/bN)) + tDgD = gD_tma[(None, None), (cta_m, cta_n)] # ((bM/2, bN), (1, 1)) + + a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + mnk_tiler, + ab_dtype, + MAINLOOP_STAGE_DEPTH, + ) + b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + mnk_tiler, + ab_dtype, + MAINLOOP_STAGE_DEPTH, + ) + + cta_tile_shape_mnk = cute.shape_div(mnk_tiler, (cluster_layout_v_size, 1, 1)) + epi_tile = sm100_utils.compute_epilogue_tile_shape( + cta_tile_shape_mnk, + use_2cta_instrs, + d_layout, + d_dtype, + ) + sc_smem_layout_staged = sm100_utils.make_smem_layout_epi( + d_dtype, + d_layout, + epi_tile, + TMA_STORE_PIPE_DEPTH, + ) + + tmem_layout = cute_ext.make_tmem_layout_acc( + tiled_mma, mnk_tiler, EPILOGUE_STAGE_DEPTH + ) + + bufferA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + a_smem_layout_staged, + alignment=1024, + ) + + bufferB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + b_smem_layout_staged, + alignment=1024, + ) + + bufferAcc = cute_ext.allocate( + acc_dtype, + cute.AddressSpace.tmem, + tmem_layout, + alignment=16, + is2cta=use_2cta_instrs, + ) + + bufferC = cute_ext.allocate( + d_dtype, + cute.AddressSpace.smem, + sc_smem_layout_staged, + alignment=1024, + ) + + copy_atom_t2r = sm100_utils.get_tmem_load_op( + cta_tile_shape_mnk, + d_layout, + d_dtype, + acc_dtype, + epi_tile, + use_2cta_instrs, + ) + + # Take only one stage of the TMEM buffer for the epilogue + accumulators = cute.zipped_divide(bufferAcc, ((epi_tile), 1)) + acc_epi_div = accumulators[((None, None), 0), 0] + + # Create the TMEM copy atom based on the size of transfer within one iteration of epilogue + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div) + + # Calculate the per thread destination size per iteration for output of TMEM and input of SMEM + gC_mnl_epi = cute.flat_divide(tDgD, epi_tile) + acc_d_rmem_layout = cute_ext.make_t2r_rmem_layout( + tiled_copy_t2r, gC_mnl_epi, tid_x + ) + + bufferRAcc = cute_ext.allocate( + acc_dtype, + cute.AddressSpace.rmem, + acc_d_rmem_layout, + alignment=32, + ) + bufferRD = cute_ext.allocate( + d_dtype, + cute.AddressSpace.rmem, + acc_d_rmem_layout, + alignment=32, + ) + + tma_mcast_proj_A = 2 + tma_mcast_proj_B = 1 + + mma_operation_type = tma_operation_type = None + acc_pipe = mainloop_pipe = None + if cutlass.const_expr(use_2cta_instrs): + mma_operation_type = cute_ext.OperationTypeEnum.SM100_MMA_2SM_SS + if cutlass.const_expr(use_tma_multicast): + tma_operation_type = ( + cute_ext.OperationTypeEnum.SM100_TMA_LOAD_2SM_MULTICAST + ) + else: + tma_operation_type = cute_ext.OperationTypeEnum.SM100_TMA_LOAD_2SM + + else: + mma_operation_type = cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS + if cutlass.const_expr(use_tma_multicast): + tma_operation_type = cute_ext.OperationTypeEnum.SM90_TMA_LOAD_MULTICAST + else: + tma_operation_type = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + + # MMA <-> TMEM load pipeline + # if 2CTA MMA, warpgroup from both peer and leader CTA consumer.release + acc_pipe_consumer_arv_count = 256 if use_2cta_instrs else 128 + acc_pipe = cute_ext.UMMAtoAsyncPipeline.create( + num_stages=EPILOGUE_STAGE_DEPTH, + mma_operation_type=mma_operation_type, + consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R, + consumer_arv_count=acc_pipe_consumer_arv_count, + cluster_layout_vmnk=cluster_layout_vmnk, + ) + + if cutlass.const_expr(use_tma_multicast): + # TMA load <-> MMA pipeline + mainloop_pipe = cute_ext.TMAToUMMAPipeline.create_with_mask( + num_stages=MAINLOOP_STAGE_DEPTH, + tma_operation_type=tma_operation_type, + mma_operation_type=mma_operation_type, + cluster_layout_vmnk=cluster_layout_vmnk, + ) + else: + mainloop_pipe = cute_ext.TMAToUMMAPipeline.create( + num_stages=MAINLOOP_STAGE_DEPTH, + mma_operation_type=mma_operation_type, + tma_operation_type=tma_operation_type, + cluster_layout_vmnk=cluster_layout_vmnk, + ) + + tma_store_warp_id = 0 + mma_warp_id = 4 + tma_load_warp_id = 5 + is_tma_thr = warp_idx == tma_load_warp_id + is_mma_thr = warp_idx == mma_warp_id + is_epi_thr = warp_idx < 4 + is_leader_cta = mma_coord_vmnk[0] == 0 + + # SMEM -> GMEM + tma_store_pipe = cute_ext.TMAStorePipeline( + stages=TMA_STORE_PIPE_DEPTH, + arv_count=128, + barrier_id=1, + tma_warp_id=tma_store_warp_id, + ) + + k_tile_count = cute.size(gA, mode=[1, 1]) + if is_tma_thr: + for k_tile in cutlass.range(0, k_tile_count, 1, unroll=1): + gA_k = tAgA[None, None, k_tile] + gB_k = tBgB[None, None, k_tile] + + producer_stage_token, idx = ( + mainloop_pipe.producer_acquire_and_get_stage() + ) + mbar = cute_ext.get_mbarrier(producer_stage_token) + bufferA_sliced = bufferA[None, None, None, idx] + bufferB_sliced = bufferB[None, None, None, idx] + a_cta_v_map = cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A") + b_cta_v_map = cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B") + + if cutlass.const_expr(use_tma_multicast): + cute_ext.tma_load_multicast( + gA_k, + bufferA_sliced, + mbar, + vmnk_layout=cluster_layout_vmnk, + cta_v_map=a_cta_v_map, + tma_operation_type=tma_operation_type, + multicast_mode=tma_mcast_proj_A, + ) + cute_ext.tma_load_multicast( + gB_k, + bufferB_sliced, + mbar, + vmnk_layout=cluster_layout_vmnk, + cta_v_map=b_cta_v_map, + tma_operation_type=tma_operation_type, + multicast_mode=tma_mcast_proj_B, + ) + else: + cute_ext.tma_load( + gA_k, + bufferA_sliced, + mbar, + cta_v_map=a_cta_v_map, + tma_operation_type=tma_operation_type, + ) + cute_ext.tma_load( + gB_k, + bufferB_sliced, + mbar, + cta_v_map=b_cta_v_map, + tma_operation_type=tma_operation_type, + ) + + if is_leader_cta: + mainloop_pipe.producer_commit() + mainloop_pipe.producer_state = cute_ext.pipeline_advance_iterator( + mainloop_pipe.raw_pipeline, mainloop_pipe.producer_state + ) + + if is_mma_thr and is_leader_cta: + producer_stage_token, idx = acc_pipe.producer_acquire_and_get_stage() + accumulators_sliced = bufferAcc[None, None, None, idx] + + mma_atom = cute.make_mma_atom(tiled_mma.op) + mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, k_tile_count, 1, unroll=1): + _, mainloop_idx = mainloop_pipe.consumer_wait_and_get_stage() + bufferA_sliced_stage = cute.core.slice_( + bufferA, (None, None, None, mainloop_idx) + ) + bufferB_sliced_stage = cute.core.slice_( + bufferB, (None, None, None, mainloop_idx) + ) + + for k_block in cutlass.range(mma_inst_tile_k, unroll_full=True): + cute_ext.dot( + mma_atom, + cute.append_ones( + bufferA_sliced_stage[None, None, k_block], up_to_rank=3 + ), + cute.append_ones( + bufferB_sliced_stage[None, None, k_block], up_to_rank=3 + ), + accumulators_sliced, + ) + mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True) + + mainloop_pipe.consumer_release_and_advance() + + acc_pipe.producer_commit_and_advance() + + if is_epi_thr: + _, idx = acc_pipe.consumer_wait_and_get_stage() + accumulators_sliced = bufferAcc[(None, None), 0, 0, idx] + acc_epi_div_tiled = cute.flat_divide(accumulators_sliced, epi_tile) + + tiled_copy_r2s = cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), d_dtype), + tiled_copy_t2r, + ) + c_cta_v_map = cute_ext.get_cta_v_map_c(mD, epi_tile) + + subtile_cnt = cute.size(acc_epi_div_tiled.shape, mode=[3]) + for mn in range(subtile_cnt): + # TMEM -> RMEM + cute_ext.partition_and_copy( + tiled_copy_t2r.get_slice(tid_x), + acc_epi_div_tiled[None, None, 0, mn], + bufferRAcc, + ) + + # RMEM -> RMEM + bufferRD.store(bufferRAcc.load().to(d_dtype)) + + tma_store_pipe.acquire_sync() + store_idx = tma_store_pipe.get_index() + + # RMEM -> SMEM + cute_ext.partition_and_copy( + tiled_copy_r2s.get_slice(tid_x), + bufferRD, + bufferC[None, None, store_idx], + ) + + tma_store_pipe.commit_sync() + + if warp_idx == tma_store_warp_id: + cute_ext.tma_store( + bufferC[None, None, store_idx], + gC_mnl_epi[None, None, 0, mn], + cta_v_map=c_cta_v_map, + ) + + tma_store_pipe.release_advance() + + tma_store_pipe.tail() + acc_pipe.consumer_release_and_advance() + + # Return a callable that launches the kernel with proper grid/block/cluster + @cute_ext.jit + def launch_kernel(mA: cute.Tensor, mB: cute.Tensor, mD: cute.Tensor): + kernel(mA, mB, mD).launch( + grid=GRID_SHAPE, + block=(32 * NUM_WARPS_PER_CTA, 1, 1), + cluster=CLUSTER_SHAPE, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + ) + + return launch_kernel + + +if __name__ == "__main__": + M = 256 + N = 256 + K = 64 + use_tma_multicast = True + use_2cta_instrs = True + acc_dtype = cutlass.Float32 + + majors = ( + cute.nvgpu.tcgen05.OperandMajorMode.K, + cute.nvgpu.tcgen05.OperandMajorMode.K, + cute.nvgpu.tcgen05.OperandMajorMode.K, + ) + dtypes = (torch.float16, torch.float16, torch.float16) + + A_torch, B_torch, D_torch, A_cute, B_cute, D_cute = get_gemm_tensors( + M, N, K, majors, dtypes + ) + + kernel_launcher = sm100_4x4x1_kernel_builder( + use_tma_multicast, use_2cta_instrs, acc_dtype, M, N + ) + + compiled_kernel = cute_ext.compile(kernel_launcher, A_cute, B_cute, D_cute) + + compiled_kernel(A_cute, B_cute, D_cute) + + # Reference check (may fail on simulator/unsupported GPU) + try: + ref = torch.mm(A_torch.float(), B_torch.float().T) + torch.testing.assert_close(D_torch.float(), ref, atol=1e-2, rtol=1e-2) + print("PASS") + except RuntimeError as e: + if "no kernel image is available" in str(e): + print("SKIP: Reference check skipped - GPU not supported by PyTorch") + else: + raise diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py new file mode 100755 index 00000000..b4f7e551 --- /dev/null +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py @@ -0,0 +1,1803 @@ +# 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 argparse +from typing import Optional, Tuple, Type, Union +from functools import lru_cache +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +import cutlass.utils as utils +from cutlass.utils import is_fp8_dtype, create_cute_tensor_for_fp8 +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.cute.experimental as cute_ext + +""" +A high-performance 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") + +This example attempts to show interoperability between cute.experimental and existing CUTE DSL APIs by using +cute.experimental APIs for TMA loading operations for A and B tensors. + +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.py. + +.. code-block:: bash + + python examples/blackwell/dense_gemm_persistent.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 python examples/blackwell/dense_gemm_persistent.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.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 +""" + + +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 = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = utils.sm100.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, b, c, 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, + ): + """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 + """ + + 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.arch = "sm_100" + + 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.threads_per_cta = 32 * len( + (self.mma_warp_id, self.tma_warp_id, *self.epilogue_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 + + def _create_tiled_mma(self): + return utils.sm100.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], + ) + + 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 = self._create_tiled_mma() + + # 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 = utils.sm100.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 = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + self.smem_capacity = utils.get_smem_capacity_in_bytes() + + # 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 = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.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 = utils.sm100.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, self.arch + ) + + @cute.experimental.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, + ): + """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 + 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}") + + tiled_mma = self._create_tiled_mma() + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + + # Setup TMA load for B + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + + 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, + a, + 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, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + return + + # GPU device kernel + @cute.experimental.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + mA: cute.Tensor, # Global A tensor + mB: cute.Tensor, # Global B 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, + ): + """ + 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: + 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 + ) + cluster_layout_v_size = cute.size(cluster_layout_vmnk.shape[0]) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + ## Tiling the global tensors for cute.experimental TMA Loads + num_mma_ctas = cute.size(tiled_mma.thr_id.shape) + cta_tile_shape_mnk = cute.shape_div(self.mma_tiler, (num_mma_ctas, 1, 1)) + # A is tiled (M/2, K) for 2CTA + a_tiler_mk = (cta_tile_shape_mnk[0], cta_tile_shape_mnk[2]) + # B is tiled (N/2, K) for 2CTA + b_tiler_nk = (cta_tile_shape_mnk[1] // num_mma_ctas, cta_tile_shape_mnk[2]) + + gA = cute.zipped_divide(mA, a_tiler_mk) + gB = cute.zipped_divide(mB, b_tiler_nk) + + # Determine pipeline operation types based on 2-CTA mode and TMA multicast + if cutlass.const_expr(self.use_2cta_instrs): + tma_operation_type = cute_ext.OperationTypeEnum.SM100_TMA_LOAD_2SM + else: + tma_operation_type = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + + # + # 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) + bufferA = 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) + bufferB = 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, 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, mode=[1, 1]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(bufferA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(bufferB) + # (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) + + # + # Construct the scheduler + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + ) + work_tile = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + + # + # Slice to per mma tile index + # + gA_tile = gA[(None, None), (cur_tile_coord[0], None, cur_tile_coord[2])] + + # For B loading in 2-CTA mode, compute proper N coordinate + if cutlass.const_expr(self.use_2cta_instrs): + # In 2CTA mode, the cur_tile_coord[1] gives a full MMA tile, but we want a CTA level tile to load + # Each CTA in the pair loads a half of the N tile + gB_tma_coord_n = ( + cluster_layout_v_size * cur_tile_coord[1] + + bidx % cluster_layout_v_size + ) + else: + gB_tma_coord_n = cur_tile_coord[1] + + gB_tile = gB[(None, None), (gB_tma_coord_n, None, cur_tile_coord[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): + gA_k = gA_tile[None, None, k_tile] + gB_k = gB_tile[None, None, k_tile] + # Conditionally wait for AB buffer empty + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + idx = handle.index + bufferA_sliced = bufferA[None, None, None, idx] + a_cta_v_map = cute_ext.get_cta_v_map_ab( + mA, self.mma_tiler, tiled_mma, "A" + ) + b_cta_v_map = cute_ext.get_cta_v_map_ab( + mB, self.mma_tiler, tiled_mma, "B" + ) + bufferB_sliced = bufferB[None, None, None, idx] + + # TMA load A/B + cute_ext.tma_load( + gA_k, + bufferA_sliced, + handle.barrier.value, + cta_v_map=a_cta_v_map, + update_expect_tx=False, # Does not automatically update the mbarrier's transaction bytes + tma_operation_type=tma_operation_type, + ) + + cute_ext.tma_load( + gB_k, + bufferB_sliced, + handle.barrier.value, + cta_v_map=b_cta_v_map, + update_expect_tx=False, # Does not automatically update the mbarrier's transaction bytes + tma_operation_type=tma_operation_type, + ) + + # 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 + # + + 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 + # + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + if cutlass.const_expr(self.use_tma_store): + assert tma_atom_c is not None and sC is not None + 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 + ) + 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], + ) + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete + c_pipeline.producer_tail() + else: + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + @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, + arch: str, + ) -> 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, arch=arch) + + return num_tmem_alloc_cols + + def check_supported_dtypes( + self, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + ): + """ + Check if the dtypes are valid + + :param a_dtype: The data type of the A operands + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operands + :type b_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] + + :raises testing.CantImplementError: If the dtypes are invalid + """ + valid_ab_dtypes = { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Uint8, + cutlass.Int8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + if a_dtype not in valid_ab_dtypes or b_dtype not in valid_ab_dtypes: + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype}" + ) + + if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: + raise testing.CantImplementError( + f"Unsupported accumulator dtype: {self.acc_dtype}" + ) + + # 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 ( + a_dtype not in acc_ab_compatibility[self.acc_dtype] + or b_dtype not in acc_ab_compatibility[self.acc_dtype] + ): + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype} for accumulator dtype: {self.acc_dtype}" + ) + + # 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]: + raise testing.CantImplementError( + f"Unsupported C dtype: {c_dtype} for accumulator dtype: {self.acc_dtype}" + ) + + def check_mma_tiler_and_cluster_shape(self): + """Check if the mma tiler and cluster shape are valid. + + :raises testing.CantImplementError: If the mma tiler and cluster shape are invalid + """ + # 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]) + ): + raise testing.CantImplementError( + f"Invalid mma tiler & use_2cta_instrs: {self.mma_tiler_mn}, {self.use_2cta_instrs}" + ) + if self.mma_tiler_mn[1] not in range(32, 257, 32): + raise testing.CantImplementError( + f"Invalid mma tiler N: {self.mma_tiler_mn[1]}" + ) + # Skip illegal cluster shape + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + raise testing.CantImplementError( + f"Invalid cluster shape M: {self.cluster_shape_mn[0]}" + ) + # 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]) + ): + raise testing.CantImplementError( + f"Invalid cluster shape: {self.cluster_shape_mn}" + ) + + def check_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ): + """ + 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 a_dtype: The data type of the A operands + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operands + :type b_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 + + :raises testing.CantImplementError: If the tensor alignment is invalid + """ + + # 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(a_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(b_dtype, b_major == "n", (n, k, l)) + or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + raise testing.CantImplementError( + f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {a_dtype}, {b_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}" + ) + + def check_epilog_store_option(self, m: int, n: int): + """ + 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 + + :raises testing.CantImplementError: If the epilogue store option is invalid + """ + # 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): + raise testing.CantImplementError( + f"Invalid epilog store option: {m}, {n}" + ) + + def can_implement( + self, + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Determine if the given tensor configuration can be implemented by this kernel. + + :param mnkl: Problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] + :param a_dtype: Data type for input tensors A. + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Data type for input tensors B. + :type b_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C. + :type c_dtype: Type[cutlass.Numeric] + :param a_major: Major dimension of the A tensor layout ("m" or "k"). + :type a_major: str + :param b_major: Major dimension of the B tensor layout ("n" or "k"). + :type b_major: str + :param c_major: Major dimension of the C tensor layout ("m" or "n"). + :type c_major: str + :return: True if the kernel supports the given configuration, False otherwise. + :rtype: bool + """ + + try: + # Skip unsupported types + self.check_supported_dtypes(a_dtype, b_dtype, c_dtype) + + # Skip invalid mma tile shape and cluster shape + self.check_mma_tiler_and_cluster_shape() + + m, n, k, l = mnkl + self.check_tensor_alignment( + m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ) + self.check_epilog_store_option(m, n) + except testing.CantImplementError: + return False + return True + + +@cute.experimental.jit +def bmm( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, # (l, m, k) + b: cute.Tensor, # (l, k, n) + c: cute.Tensor, # (l, m, n) + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, +): + """ + Wrapper API for persistent GEMM kernel to follow the convention of PyTorch's batch matrix-multiply (bmm). + + Internally, the tensors are permuted to match CuTe's convention: + - a: (m, k, l) + - b: (n, k, l) + - c: (m, n, l) + + :param gemm_op: Kernel operation, expects (a, b, c, max_active_clusters, stream, epilogue_op) + :type gemm_op: cutlass.Constexpr + :param a: Input tensor of shape (l, m, k) + :type a: cute.Tensor + :param b: Input tensor of shape (l, k, n) + :type b: cute.Tensor + :param c: Output tensor of shape (l, m, n) + :type c: cute.Tensor + :param max_active_clusters: Maximum number of hardware clusters to launch + :type max_active_clusters: cutlass.Constexpr + :param epilogue_op: Optional elementwise lambda function to apply per output element, defaults to identity + :type epilogue_op: cutlass.Constexpr, optional + """ + # (l,m,k) -> (m,k,l) + a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])) + # (l,k,n) -> (n,k,l) + b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])) + # (l,m,n) -> (m,n,l) + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + + gemm_op(a, b, c, max_active_clusters, stream, epilogue_op) + + +@lru_cache(maxsize=1) +def prepare_tensors( + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + init_random: bool = True, + normal_mean: float = 0.0, + normal_std: float = 1.0, +): + """Prepare tensors for GEMM. + + Returns: + Tuple of (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage): + - *_f32: Float32 tensors with the logical data (for reference and fp8 conversion) + - *_storage: Storage tensors for DLPack (uint8 for fp8, otherwise the target dtype) + """ + import torch + from cutlass.torch import dtype as torch_dtype + + m, n, k, l = mnkl + + if a_major == "k": + a_f32 = torch.empty((l, m, k), dtype=torch.float32, device="cuda") + elif a_major == "m": + a_f32 = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) + + if b_major == "n": + b_f32 = torch.empty((l, k, n), dtype=torch.float32, device="cuda") + elif b_major == "k": + b_f32 = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) + + if c_major == "n": + c_f32 = torch.empty((l, m, n), dtype=torch.float32, device="cuda") + elif c_major == "m": + c_f32 = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) + + if init_random: + # Uniform random initialization in range [-2, 3) + a_f32.random_(-2, 3) + b_f32.random_(-2, 3) + c_f32.random_(-2, 3) + + else: + # Normal (Gaussian) initialization with user-specified mean and std + a_f32.normal_(mean=normal_mean, std=normal_std) + b_f32.normal_(mean=normal_mean, std=normal_std) + c_f32.normal_(mean=normal_mean, std=normal_std) + + # For float8 types, use uint8 as storage type to avoid dlpack limitation + # (dlpack doesn't support float8 types) + # For other types, convert to the target dtype + a_storage_dtype = torch.uint8 if is_fp8_dtype(a_dtype) else torch_dtype(a_dtype) + b_storage_dtype = torch.uint8 if is_fp8_dtype(b_dtype) else torch_dtype(b_dtype) + c_storage_dtype = torch.uint8 if is_fp8_dtype(c_dtype) else torch_dtype(c_dtype) + + a_storage = a_f32.to(dtype=a_storage_dtype) + b_storage = b_f32.to(dtype=b_storage_dtype) + c_storage = c_f32.to(dtype=c_storage_dtype) + + return (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage) + + +@lru_cache(maxsize=1) +def compile_bmm( + mnkl: Tuple[int, int, int, int], + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + 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), + max_active_clusters: cutlass.Constexpr = None, + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + epilogue_op: cutlass.Constexpr = lambda x: x, +): + from cutlass.cute.runtime import make_fake_stream + + gemm = PersistentDenseGemmKernel( + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + ) + # Check if configuration can be implemented + can_implement = gemm.can_implement( + mnkl, a.element_type, b.element_type, c.element_type, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + 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}" + ) + + stream = make_fake_stream() + return cute.compile(bmm, gemm, a, b, c, max_active_clusters, stream, epilogue_op) + + +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, + benchmark: bool = False, + **kwargs, +): + """ + Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + + Prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks execution. + + :param mnkl: Problem size as a tuple (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: Accumulator data type for the matrix multiplication. + :type acc_dtype: Type[cutlass.Numeric] + :param a_major: Memory layout of tensor A. + :type a_major: str + :param b_major: Memory layout of tensor B. + :type b_major: str + :param c_major: Memory layout of tensor C. + :type c_major: str + :param mma_tiler_mn: MMA tiling size (M, N), defaults to (256, 256). + :type mma_tiler_mn: Tuple[int, int], optional + :param cluster_shape_mn: Cluster shape (M, N), defaults to (2, 1). + :type cluster_shape_mn: Tuple[int, int], optional + :param use_2cta_instrs: Whether to use 2CTA MMA instructions, defaults to True. + :type use_2cta_instrs: bool, optional + :param use_tma_store: Whether to use TMA store, defaults to True. + :type use_tma_store: bool, optional + :param tolerance: Tolerance for reference validation, 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 benchmark: Whether to only benchmark the kernel, defaults to False. + :type benchmark: 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 + """ + import torch + from cutlass.torch import dtype as torch_dtype + + 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) + + # Check if configuration can be implemented + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + # Run and verify BMM with torch + a_f32, b_f32, c_f32, a_storage, b_storage, c_storage = prepare_tensors( + mnkl, ab_dtype, ab_dtype, c_dtype, a_major, b_major, c_major + ) + + leading_dim_a = 2 if a_major == "k" else 1 + leading_dim_b = 1 if b_major == "k" else 2 + leading_dim_c = 2 if c_major == "n" else 1 + + # Create CuTe tensors, passing float32 source for fp8 conversion + a_tensor = create_cute_tensor_for_fp8( + a_storage, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_tensor = create_cute_tensor_for_fp8( + b_storage, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_tensor = create_cute_tensor_for_fp8( + c_storage, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) + + compiled_fn = compile_bmm( + mnkl, + a_tensor, + b_tensor, + c_tensor, + acc_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + max_active_clusters, + use_2cta_instrs, + use_tma_store, + epilogue_op=lambda x: x, + ) + + print("Running Blackwell Persistent Dense GEMM test with:") + print(f"mnkl: {mnkl}") + 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'}") + + if not skip_ref_check: + # Use small random number for deterministic result for reference check + compiled_fn(a_tensor, b_tensor, c_tensor, current_stream) + + # Manually quantize to be comparable + # Use float32 source data for reference calculation + ref = ( + torch.bmm(a_f32, b_f32) + .to(dtype=torch_dtype(c_dtype)) + .to(dtype=torch.float32) + ) + # Read back the result from CuTe tensor (c_storage was updated in-place) + torch.testing.assert_close( + c_storage.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 + ) + + if not benchmark: + return 0 + + def generate_tensors(): + a_f32, b_f32, c_f32, a_st, b_st, c_st = prepare_tensors( + mnkl, + ab_dtype, + ab_dtype, + c_dtype, + a_major, + b_major, + c_major, + ) + a_tensor = create_cute_tensor_for_fp8( + a_st, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_tensor = create_cute_tensor_for_fp8( + b_st, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_tensor = create_cute_tensor_for_fp8( + c_st, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) + return testing.JitArguments(a_tensor, b_tensor, c_tensor, current_stream) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_storage.numel() * a_storage.element_size() + + b_storage.numel() * b_storage.element_size() + + c_storage.numel() * c_storage.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + # Return execution time in microseconds + return testing.benchmark( + compiled_fn, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + +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." + ) + + +def prepare_parser(): + 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( + "--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( + "--benchmark", + type=str, + default="default", + choices=[ + "default", + "none", + ], + help="Benchmark the kernel with nsight or default (cute.testing.benchmark) or none", + ) + 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", + ) + + return parser + + +if __name__ == "__main__": + parser = prepare_parser() + parser.add_argument( + "--mma_tiler_mn", + type=_parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + + 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") + + print(f"[DSL INFO] Compiling Blackwell Persistent Dense GEMM with:") + print( + f"[DSL INFO] A dtype: {args.ab_dtype}, B dtype: {args.c_dtype}, C dtype: {args.acc_dtype}, Acc dtype: {args.acc_dtype}" + ) + print( + f"[DSL INFO] Matrix majors - A: {args.a_major}, B: {args.b_major}, C: {args.c_major}" + ) + print(f"[DSL INFO] Mma Tiler (M, N): {args.mma_tiler_mn}") + print(f"[DSL INFO] Cluster Shape (M, N): {args.cluster_shape_mn}") + print( + f"[DSL INFO] 2CTA MMA instructions: {'True' if args.use_2cta_instrs else 'False'}" + ) + print(f"[DSL INFO] Use TMA Store: {'True' if args.use_tma_store else 'False'}") + + exec_time = 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, + args.benchmark == "default", + ) + print(f"Execution time: {exec_time} seconds") + print("PASS") diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_ptr_array.py b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_ptr_array.py new file mode 100755 index 00000000..bd4e09c1 --- /dev/null +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_ptr_array.py @@ -0,0 +1,820 @@ +# Copyright (c) 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 argparse + +import torch +from typing import Type, Tuple, List + +import cutlass +from cutlass.cute import experimental as cute_ext +from cutlass.base_dsl.typing import Numeric +from cutlass import cute as cute +from cutlass import utils +from cutlass import torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils + +import cutlass.cute.testing as testing + +class DenseGemmPtrArrayKernel: + def __init__( + self, + mn_tiler: tuple[int, int], + mma_dtype: tuple[Type[Numeric], Type[Numeric], Type[Numeric]], + tmem_output_dtype: Type[Numeric], + batch_count: int, # Number of batches, each batch will have its own pointer for the matrix + A_shape: tuple, # Shape of the matrix A + A_stride: tuple, # Stride of the matrix A + B_shape: tuple, # Shape of the matrix B + B_stride: tuple, # Stride of the matrix B + D_shape: tuple, # Shape of the matrix D + D_stride: tuple, # Stride of the matrix D + epilogue_op=lambda x: x, + ): + self.mn_tiler = mn_tiler + self.ab_dtype, self.acc_dtype, self.d_dtype = mma_dtype + self.tmem_output_dtype = tmem_output_dtype + self.use_2cta_instrs = False + self.TMA_STORE_STAGE = 4 + self.epilogue_op = epilogue_op + self.batch_count = batch_count + self.A_shape = A_shape + self.A_stride = A_stride + self.B_shape = B_shape + self.B_stride = B_stride + self.D_shape = D_shape + self.D_stride = D_stride + + """ + Helper function to convert an int64 to a cute.ptr of a certain type. + The cute.ptr is always located in Gmem. + This is used to load the pointers for A/B/D from the Ptr array. + """ + + @cute.experimental.jit + def _get_pointer(self, address_as_int, cute_type): + cute_ptr = cute.make_ptr( + cute_type, + address_as_int, + mem_space=cute.AddressSpace.gmem, + assumed_align=16, + ) + return cute_ptr + + @cute.experimental.jit + def __call__( + self, mA_tensor: cute.Tensor, mB_tensor: cute.Tensor, mD_tensor: cute.Tensor + ): + # Get the pointer to the first batch of D + d_ptr = self._get_pointer(mD_tensor[0], self.d_dtype) + d_ptr_base_tensor = cute.make_tensor( + d_ptr, layout=cute.make_layout(self.D_shape, stride=self.D_stride) + ) + tile_mn = cute.core._pack_shape((*self.mn_tiler, 1)) + div = cute.tiled_divide(d_ptr_base_tensor, tile_mn) + grid = (div.shape[1], div.shape[2], div.shape[3]) + self.kernel(mA_tensor, mB_tensor, mD_tensor).launch( + grid=grid, + block=(192, 1, 1), + cluster=(1, 1, 1), + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + ) + + @cute.experimental.kernel + def kernel( + self, + mA_tensor: cute.Tensor, + mB_tensor: cute.Tensor, + mD_tensor: cute.Tensor, + ): + # Get pointers for the first batch to perform shape and stage calculations + A_0_ptr = self._get_pointer(mA_tensor[0], self.ab_dtype) + B_0_ptr = self._get_pointer(mB_tensor[0], self.ab_dtype) + D_0_ptr = self._get_pointer(mD_tensor[0], self.d_dtype) + + mA = cute.make_tensor( + A_0_ptr, layout=cute.make_layout(self.A_shape, stride=self.A_stride) + ) + + mB = cute.make_tensor( + B_0_ptr, layout=cute.make_layout(self.B_shape, stride=self.B_stride) + ) + + mD = cute.make_tensor( + D_0_ptr, layout=cute.make_layout(self.D_shape, stride=self.D_stride) + ) + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mn_tiler, + ) + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + mnk_tiler = ( + self.mn_tiler[0], + self.mn_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + d_layout = utils.LayoutEnum.from_tensor(mD) + d_dtype = mD.element_type + + tiler_mk = (mnk_tiler[0], mnk_tiler[2]) + tiler_nk = (mnk_tiler[1], mnk_tiler[2]) + tiler_mn = (mnk_tiler[0], mnk_tiler[1]) + + gA = cute.zipped_divide(mA, tiler_mk) + gB = cute.zipped_divide(mB, tiler_nk) + gD = cute.zipped_divide(mD, tiler_mn) + + mainloop_stage = 2 + acc_stage = 2 + + cta_m, cta_n, cta_l = cute.arch.block_idx() + tid_x, _, _ = cute.arch.thread_idx() + + gA_tile = gA[(None, None), (cta_m, None, cta_l)] + gB_tile = gB[(None, None), (cta_n, None, cta_l)] + gD_tile = gD[(None, None), (cta_m, cta_n, cta_l)] + + # Compute A/B/C shared memory layout + a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + mnk_tiler, + self.ab_dtype, + mainloop_stage, + ) + b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + mnk_tiler, + self.ab_dtype, + mainloop_stage, + ) + + cta_tile_shape_mnk = cute.shape_div( + mnk_tiler, (cute.size(tiled_mma.thr_id.shape), 1, 1) + ) + epi_tile = sm100_utils.compute_epilogue_tile_shape( + cta_tile_shape_mnk, + self.use_2cta_instrs, + d_layout, + d_dtype, + ) + sc_smem_layout_staged = sm100_utils.make_smem_layout_epi( + d_dtype, + d_layout, + epi_tile, + self.TMA_STORE_STAGE, + ) + + # UMMA ACC TMEM Layout + tmem_layout = cute_ext.make_tmem_layout_acc(tiled_mma, mnk_tiler, acc_stage) + + # Allocate UMMA Buffers + bufferA = cute_ext.allocate( + self.ab_dtype, + cute.AddressSpace.smem, + a_smem_layout_staged, + alignment=1024, + ) + + bufferB = cute_ext.allocate( + self.ab_dtype, + cute.AddressSpace.smem, + b_smem_layout_staged, + alignment=1024, + ) + + bufferAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.tmem, + tmem_layout, + alignment=16, + ) + + # Allocate SMEM buffer for C + bufferC = cute_ext.allocate( + d_dtype, + cute.AddressSpace.smem, + sc_smem_layout_staged, + alignment=1024, + ) + + # Create the TMEM load atom + copy_atom_t2r = sm100_utils.get_tmem_load_op( + cta_tile_shape_mnk, + d_layout, + self.tmem_output_dtype, + self.acc_dtype, + epi_tile, + self.use_2cta_instrs, + ) + + # Take only one stage of the TMEM buffer + accumulators = cute.zipped_divide(bufferAcc, ((epi_tile), 1)) + acc_epi_div = accumulators[((None, None), 0), 0] + + # Create the TMEM copy atom based on the size of transfer within one iteration of epilogue + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div) + + # Calculate the per thread destination size per iteration for output of TMEM and input of SMEM + gC_mnl_epi = cute.flat_divide(gD_tile, epi_tile) + acc_d_rmem_layout = cute_ext.make_t2r_rmem_layout( + tiled_copy_t2r, gC_mnl_epi, tid_x + ) + + # Allocate RMEM buffers + bufferRAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + acc_d_rmem_layout, + alignment=32, + ) + bufferRD = cute_ext.allocate( + d_dtype, + cute.AddressSpace.rmem, + acc_d_rmem_layout, + alignment=32, + ) + + # TMA -> UMMA + mainloop_pipe = cute_ext.TMAToUMMAPipeline.create( + num_stages=mainloop_stage, + mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS, + ) + + # UMMA -> TMEM + acc_pipe = cute_ext.UMMAtoAsyncPipeline.create( + num_stages=acc_stage, + mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS, + consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R, + consumer_arv_count=128, + ) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + # warp assignment: [0]-tma_store, [0-3]-epi, [4]-mma, [5]-tma_load + tma_store_warp_id = 0 + mma_warp_id = 4 + tma_load_warp_id = 5 + is_tma_thr = warp_idx == tma_load_warp_id + is_mma_thr = warp_idx == mma_warp_id + is_epi_thr = warp_idx < 4 + + # SMEM -> GMEM + tma_store_pipe = cute_ext.TMAStorePipeline( + stages=self.TMA_STORE_STAGE, + arv_count=128, + barrier_id=1, + tma_warp_id=tma_store_warp_id, + ) + + k_tile_size = cute.size(gA, mode=[1, 1]) + + # Outer loop over batches and perform GEMM for each batch as usual + # This is a dynamic for loop that lowers to an scf.for + # Note that the tensor loading is done in the if `thread` warp specialized + # sections. This is essential to ensure proper synchronization of tma loads + # and tma updates across batches. + for batch_idx in range(0, self.batch_count): + # Load pointers for the current batch + ptr_A = self._get_pointer(mA_tensor[batch_idx], self.ab_dtype) + ptr_B = self._get_pointer(mB_tensor[batch_idx], self.ab_dtype) + ptr_D = self._get_pointer(mD_tensor[batch_idx], self.d_dtype) + + gALayout = cute.zipped_divide(mA, tiler_mk) + k_tile_size = cute.size(gALayout, mode=[1, 1]) + + if is_tma_thr: + mA = cute.make_tensor( + ptr_A, layout=cute.make_layout(self.A_shape, stride=self.A_stride) + ) + mB = cute.make_tensor( + ptr_B, layout=cute.make_layout(self.B_shape, stride=self.B_stride) + ) + gA = cute.zipped_divide(mA, tiler_mk) + gB = cute.zipped_divide(mB, tiler_nk) + gA_tile = gA[(None, None), (cta_m, None, cta_l)] + gB_tile = gB[(None, None), (cta_n, None, cta_l)] + for k in cutlass.range(0, k_tile_size, 1, unroll=1): + gA_k = gA_tile[None, None, k] + gB_k = gB_tile[None, None, k] + + # Scoped state management - pipeline object manages state internally + ( + producer_stage_token, + idx, + ) = mainloop_pipe.producer_acquire_and_get_stage() + mbar = cute_ext.get_mbarrier(producer_stage_token) + ## producer_body begin ## + bufferA_sliced = bufferA[None, None, None, idx] + bufferB_sliced = bufferB[None, None, None, idx] + a_cta_v_map = cute_ext.get_cta_v_map_ab( + mA, mnk_tiler, tiled_mma, "A" + ) + b_cta_v_map = cute_ext.get_cta_v_map_ab( + mB, mnk_tiler, tiled_mma, "B" + ) + cute_ext.tma_load( + gA_k, + bufferA_sliced, + mbar, + cta_v_map=a_cta_v_map, + ) + cute_ext.tma_load( + gB_k, + bufferB_sliced, + mbar, + cta_v_map=b_cta_v_map, + ) + ## producer_body end ## + mainloop_pipe.producer_commit_and_advance() + + # MMA section remains same as a regular GEMM + if is_mma_thr: + producer_stage_token, idx = acc_pipe.producer_acquire_and_get_stage() + ## acc_producer_body begin ## + accumulators_sliced = bufferAcc[None, None, None, idx] + + mma_atom = cute.make_mma_atom(tiled_mma.op) + mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, k_tile_size, 1, unroll=1): + # Scoped state management - pipeline object manages consumer state internally + ( + _, + mainloop_idx, + ) = mainloop_pipe.consumer_wait_and_get_stage() + ## tma_consumer_body begin ## + + bufferA_sliced_stage = cute.core.slice_( + bufferA, (None, None, None, mainloop_idx) + ) + bufferB_sliced_stage = cute.core.slice_( + bufferB, (None, None, None, mainloop_idx) + ) + + for k_block in cutlass.range(mma_inst_tile_k, unroll_full=True): + bufferA_sliced = bufferA_sliced_stage[None, None, k_block] + bufferB_sliced = bufferB_sliced_stage[None, None, k_block] + + cute_ext.dot( + mma_atom, + cute.append_ones(bufferA_sliced, up_to_rank=3), + cute.append_ones(bufferB_sliced, up_to_rank=3), + accumulators_sliced, + ) + mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True) + + ## tma_consumer_body end ## + mainloop_pipe.consumer_release_and_advance() + + ## acc_producer_body end ## + acc_pipe.producer_commit_and_advance() + + if is_epi_thr: + # Load the D tensor in the warp specialized section + mD = cute.make_tensor( + ptr_D, layout=cute.make_layout(self.D_shape, stride=self.D_stride) + ) + gD = cute.zipped_divide(mD, tiler_mn) + gD_tile = gD[(None, None), (cta_m, cta_n, cta_l)] + gC_mnl_epi = cute.flat_divide(gD_tile, epi_tile) + _, idx = acc_pipe.consumer_wait_and_get_stage() + ## acc_consume_body begin ## + accumulators_sliced = bufferAcc[(None, None), 0, 0, idx] + acc_epi_div_tiled = cute.flat_divide(accumulators_sliced, epi_tile) + + subtile_cnt = cute.size(acc_epi_div_tiled.shape, mode=[3]) + for mn in range(subtile_cnt): + # TMEM -> RMEM + cute_ext.partition_and_copy( + tiled_copy_t2r.get_slice(tid_x), + acc_epi_div_tiled[None, None, 0, mn], + bufferRAcc, + ) + + # RMEM -> RMEM + bufferRD.store(self.epilogue_op(bufferRAcc.load().to(self.d_dtype))) + + # Acquire pipeline stage and synchronize before RMEM->SMEM copy + tma_store_pipe.acquire_sync() + idx = tma_store_pipe.get_index() + + # RMEM -> SMEM + tiled_copy_r2s = cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.d_dtype), + tiled_copy_t2r, + ) + cute_ext.partition_and_copy( + tiled_copy_r2s.get_slice(tid_x), + bufferRD, + bufferC[None, None, idx], + ) + + # Fence SMEM writes and synchronize before TMA store + tma_store_pipe.commit_sync() + + # SMEM -> GMEM (only designated TMA store warp performs TMA store) + if warp_idx == tma_store_warp_id: + c_cta_v_map = cute_ext.get_cta_v_map_c(mD, epi_tile) + cute_ext.tma_store( + bufferC[None, None, idx], + gC_mnl_epi[None, None, 0, mn], + cta_v_map=c_cta_v_map, + ) + + # Release pipeline stage and advance + tma_store_pipe.release_advance() + + tma_store_pipe.tail() + acc_pipe.consumer_release_and_advance() + + +def create_tensors(l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype): + torch.manual_seed(1111) + + a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype) + d_torch_cpu = cutlass_torch.matrix(l, m, n, d_major == "m", d_dtype) + + a_tensor, a_torch_gpu = cutlass_torch.cute_tensor_like( + a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, b_torch_gpu = cutlass_torch.cute_tensor_like( + b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + d_tensor, d_torch_gpu = cutlass_torch.cute_tensor_like( + d_torch_cpu, d_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + return ( + a_tensor, + b_tensor, + d_tensor, + a_torch_cpu, + b_torch_cpu, + d_torch_cpu, + a_torch_gpu, + b_torch_gpu, + d_torch_gpu, + ) + + +# Helper creates a cute.Tensor from a List of device pointers +def make_tensor_of_ptrs(torch_tensor_array: List): + tensor_of_ptrs_torch = torch.tensor( + [t.data_ptr() for t in torch_tensor_array], + dtype=torch.int64, + device="cuda", + requires_grad=False, + ) + tensor_of_ptrs_cute, backing_torch_tensor = cutlass_torch.cute_tensor_like( + tensor_of_ptrs_torch, + cutlass.Int64, + is_dynamic_layout=False, + assumed_align=16, + ) + return tensor_of_ptrs_cute, backing_torch_tensor + + +def create_tensors_for_ptr_array( + l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype +): + # Store torch gpu pointers + As_torch_gpu = [] + Bs_torch_gpu = [] + Ds_torch_gpu = [] + # Store cute tensors + A_cutes = [] + B_cutes = [] + D_cutes = [] + + for batch_idx in range(l): + torch.manual_seed(111 + batch_idx) + + ( + A_tensor, + B_tensor, + D_tensor, + A_torch_cpu, + B_torch_cpu, + D_torch_cpu, + A_torch_gpu, + B_torch_gpu, + D_torch_gpu, + ) = create_tensors( + 1, # outer loop creates a new tensor for each batch + m, + n, + k, + a_major, + b_major, + d_major, + ab_dtype, + d_dtype, + ) + + A_cutes.append(A_tensor) + B_cutes.append(B_tensor) + D_cutes.append(D_tensor) + As_torch_gpu.append(A_torch_gpu) + Bs_torch_gpu.append(B_torch_gpu) + Ds_torch_gpu.append(D_torch_gpu) + + # Create cute tensors of pointers + a_tensor, a_backing_torch_tensor = make_tensor_of_ptrs(As_torch_gpu) + b_tensor, b_backing_torch_tensor = make_tensor_of_ptrs(Bs_torch_gpu) + d_tensor, d_backing_torch_tensor = make_tensor_of_ptrs(Ds_torch_gpu) + + return ( + a_tensor, + b_tensor, + d_tensor, + a_backing_torch_tensor, + b_backing_torch_tensor, + d_backing_torch_tensor, + A_cutes, + B_cutes, + D_cutes, + As_torch_gpu, + Bs_torch_gpu, + Ds_torch_gpu, + ) + + +def compare(a_torch_cpu, b_torch_cpu, d_torch_gpu, d_dtype, tolerance): + ref = torch.einsum("mkl,nkl->mnl", a_torch_cpu, b_torch_cpu) + + _, ref_torch_gpu = cutlass_torch.cute_tensor_like( + ref, d_dtype, is_dynamic_layout=True, assumed_align=16 + ) + ref_result = ref_torch_gpu.cpu() + torch.testing.assert_close( + d_torch_gpu.cpu(), ref_result, atol=tolerance, rtol=1e-05 + ) + + +def run( + mnkl: Tuple[int, int, int, int], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ab_dtype: Type[Numeric], + c_dtype: Type[Numeric], + acc_dtype: Type[Numeric], + a_major: str, + b_major: str, + c_major: str, + warmup_iterations: int = 0, + iterations: int = 1, + use_cold_l2: bool = False, + tolerance: float = 1e-02, + skip_ref_check: bool = False, + **kwargs, +): + """Execute a Pointer array batched dense GEMM operation on Blackwell architecture with performance benchmarking. + The main difference between this and a regular bathced GEMM is that the inputs to the kernel are arrays of pointers. + Every batch of each operand (A/B/D) has its own pointer. Thus, the size of the array of pointers is the batch size. + These pointers NEED NOT be stored contiguously in memory. + This example also demonstrates how cute_ext.tma_load/cute_ext.tma_store performs automatic device side TMA updates. + Note that the dimensions of the operand for each batch are the same across all batches. That is, all batches of A have the same shape and stride, same for B and D. + + This function prepares input tensors, configures and launches the 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 mma_tiler_mn: MMA tiling size. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :param ab_dtype: Data type for input tensors A and B + :type ab_dtype: Type[Numeric] + :param d_dtype: Data type for output tensor D + :type d_dtype: Type[Numeric] + """ + print("Running Blackwell Dense GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, D dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, D: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + 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'}") + + m, n, k, l = mnkl + + ab_dtype = ab_dtype + d_major = c_major + d_dtype = c_dtype + + # a_tensor, b_tensor, d_tensor are cute Tensors where each element is an Int64 pointer to global memory + # A_cutes, B_cutes, D_cutes are lists of cute Tensors for each batch of A/B/D + ( + a_tensor, + b_tensor, + d_tensor, + a_backing_torch_tensor, + b_backing_torch_tensor, + d_backing_torch_tensor, + A_cutes, + B_cutes, + D_cutes, + As_torch_gpu, + Bs_torch_gpu, + Ds_torch_gpu, + ) = create_tensors_for_ptr_array( + l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype + ) + + ptr_array_dense_gemm = DenseGemmPtrArrayKernel( + mn_tiler=mma_tiler_mn, + mma_dtype=(ab_dtype, acc_dtype, d_dtype), + tmem_output_dtype=d_dtype, + batch_count=l, + A_shape=A_cutes[0].shape, + A_stride=A_cutes[0].stride, + B_shape=B_cutes[0].shape, + B_stride=B_cutes[0].stride, + D_shape=D_cutes[0].shape, + D_stride=D_cutes[0].stride, + ) + + compiled_dense_gemm = cute_ext.compile( + ptr_array_dense_gemm, a_tensor, b_tensor, d_tensor + ) + compiled_dense_gemm(a_tensor, b_tensor, d_tensor) + + if not skip_ref_check: + for batch_idx in range(l): + compare( + As_torch_gpu[batch_idx].cpu(), + Bs_torch_gpu[batch_idx].cpu(), + Ds_torch_gpu[batch_idx], + d_dtype, + tolerance, + ) + print("check reference: PASS") + + def generate_tensors(): + ( + a_tensor, + b_tensor, + d_tensor, + a_backing_torch_tensor, + b_backing_torch_tensor, + d_backing_torch_tensor, + A_cutes, + B_cutes, + D_cutes, + As_torch_gpu, + Bs_torch_gpu, + Ds_torch_gpu, + ) = create_tensors_for_ptr_array( + l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype + ) + args = testing.JitArguments(a_tensor, b_tensor, d_tensor) + args.add_to_scope([A_cutes, B_cutes, D_cutes]) + return args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + sum( + As_torch_gpu[batch_idx].numel() * As_torch_gpu[batch_idx].element_size() + for batch_idx in range(l) + ) + + sum( + Bs_torch_gpu[batch_idx].numel() * Bs_torch_gpu[batch_idx].element_size() + for batch_idx in range(l) + ) + + sum( + Ds_torch_gpu[batch_idx].numel() * Ds_torch_gpu[batch_idx].element_size() + for batch_idx in range(l) + ) + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_dense_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time + + +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 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.Float32) + parser.add_argument("--d_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + + 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("--d_major", choices=["n", "m"], type=str, default="n") + + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", type=int, default=1, help="Number of iterations" + ) + parser.add_argument("--use_cold_l2", action="store_true", help="Use cold L2") + parser.add_argument( + "--tolerance", type=float, default=1e-02, help="Tolerance for validation" + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + + 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") + + exec_time = run( + args.mnkl, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.ab_dtype, + args.d_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.d_major, + args.warmup_iterations, + args.iterations, + args.use_cold_l2, + args.tolerance, + args.skip_ref_check, + ) + + print(f"Execution time: {exec_time} microseconds per iteration") diff --git a/examples/python/CuTeDSL/helpers/sparse_utils.py b/examples/python/CuTeDSL/helpers/sparse_utils.py deleted file mode 100644 index 24b3f791..00000000 --- a/examples/python/CuTeDSL/helpers/sparse_utils.py +++ /dev/null @@ -1,457 +0,0 @@ -import numpy as np -import cutlass -import cutlass.cute as cute -from cutlass.cute.runtime import from_dlpack -import torch - - -@cute.jit -def print_tensor_dlpack(src: cute.Tensor): - print(src) - cute.print_tensor(src) - - -# Sparse emulation -class SparseEmulation: - def __init__(self, M: int, N: int, K: int, L: int): - self.M = M - self.N = N - self.K = K - self.L = L - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor, e: cute.Tensor): - """Sparse emulation""" - num_threads = 128 - grid = (cute.ceil_div(self.M, num_threads), 1, 1) - block = (num_threads, 1, 1) - self.kernel(a, b, d, e).launch(grid=grid, block=block) - return - - @cute.kernel - def kernel(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor, e: cute.Tensor): - """CUDA kernel to emulate sparse tensor core""" - tidx, tidy, tidz = cute.arch.thread_idx() - bidx, bidy, bidz = cute.arch.block_idx() - - row_idx = tidx + bidx * self.M - meta_idx = self.K // 4 // 8 - if row_idx < self.M: - # each thread process 1 row - for col in range(self.N): - # each meta_idx stands for 32 elements - for e_idx in range(meta_idx): - meta_val = e[(row_idx, e_idx)] - for k in range(8): - # each k stands for 4 elements - meta_row = (meta_val >> (k * 4)) & 0xF - idx0 = meta_row & 0x3 - idx1 = (meta_row >> 2) & 0x3 - # calculate the idx in b tensor which has value in A tensor - km = e_idx * 16 + k * 2 - km_1 = km + 1 - kn = e_idx * 32 + k * 4 + idx0 - kn_1 = e_idx * 32 + k * 4 + idx1 - d[row_idx, col] += a[row_idx, km] * b[col, kn] - d[row_idx, col] += a[row_idx, km_1] * b[col, kn_1] - return - - -# Compressor -# compress a sparse tensor to a dense tensor && generate metadata -class Compressor: - def __init__(self, M: int, K: int, L: int): - self.M = M - self.K = K - self.L = L - self.pos_map = { - 0x4: [0, 1], - 0x8: [0, 2], - 0xC: [0, 3], - 0x9: [1, 2], - 0xD: [1, 3], - 0xE: [2, 3], - } - - @cute.jit - def _init__(self, a: cute.Tensor): - self.__init__(a.shape[0], a.shape[1], a.shape[2]) - - def compress(self, a, a_compressed, meta, run_on_cpu: bool): - if run_on_cpu: - if a.device.type != "cpu": - raise ValueError("a must be on cpu") - return self.__compress_on_cpu(a, a_compressed, meta) - else: - if a.device.type != "cuda": - raise ValueError("a must be on cuda") - return self.__compress_on_cuda(a, a_compressed, meta) - - def __compress_on_cpu(self, a, a_compressed, meta): - """ - compress the tensor on cpu - # Convert to 4-bit metadata value - # The metadata value represents which 2 elements are non-zero - # 0x4: [1,1,0,0] - first two elements are non-zero - # 0x8: [1,0,1,0] - first and third elements are non-zero - # 0xC: [1,0,0,1] - first and fourth elements are non-zero - # 0x9: [0,1,1,0] - second and third elements are non-zero - # 0xD: [0,1,0,1] - second and fourth elements are non-zero - # 0xE: [0,0,1,1] - third and fourth elements are non-zero - # special case: - # [0,0,0,0] == [0,0,1,1] - # [1,0,0,0] == [1,0,0,1] - # [0,1,0,0] == [0,1,0,1] - # [0,0,1,0] == [0,0,1,1] - # [0,0,0,1] == [0,0,1,1] - """ - M, K = a.shape - assert a_compressed.shape == ( - M, - K // 2, - ), f"Expected a_compressed shape {(M, K // 2)}, got {a_compressed.shape}" - assert meta.shape == ( - M, - K // 4 // 8, - ), f"Expected meta shape {(M, K // 4 // 8)}, got {meta.shape}" - for m in range(M): - k_meta = 0 - for k in range(0, K, 4): - chunk = a[m, k : k + 4] - - non_zero_indices = torch.nonzero(chunk).squeeze() - meta_val = 0xE - if torch.equal(non_zero_indices, torch.tensor([0, 1])): - meta_val = 0x4 - elif torch.equal(non_zero_indices, torch.tensor([0, 2])): - meta_val = 0x8 - elif torch.equal(non_zero_indices, torch.tensor([0, 3])) or torch.equal( - non_zero_indices, torch.tensor(0) - ): - meta_val = 0xC - elif torch.equal(non_zero_indices, torch.tensor([1, 2])): - meta_val = 0x9 - elif torch.equal(non_zero_indices, torch.tensor([1, 3])) or torch.equal( - non_zero_indices, torch.tensor(1) - ): - meta_val = 0xD - elif torch.equal(non_zero_indices, torch.tensor([2, 3])) or torch.equal( - non_zero_indices, torch.tensor(2) - ): - meta_val = 0xE - elif torch.equal(non_zero_indices, torch.tensor([])) or torch.equal( - non_zero_indices, torch.tensor(3) - ): - meta_val = 0xE - else: - raise ValueError(f"Invalid non-zero pattern: {non_zero_indices}") - meta_idx = k // 4 // 8 - meta_bit_pos = (k // 4) % 8 - if k_meta == meta_idx: - k_meta = meta_idx + 1 - meta[m, meta_idx] = 0 - meta[m, meta_idx] |= meta_val << (meta_bit_pos * 4) - compressed_idx = k // 2 - index = self.pos_map[meta_val] - a_compressed[m, compressed_idx] = chunk[index[0]] - a_compressed[m, compressed_idx + 1] = chunk[index[1]] - - def __compress_on_cuda(self, a, a_compressed, meta): - """ - compress the tensor on cuda - """ - a_tensor = from_dlpack(a) - a_compressed_tensor = from_dlpack(a_compressed) - meta_tensor = from_dlpack(meta) - self.compress_on_cuda_impl(a_tensor, a_compressed_tensor, meta_tensor) - return - - @cute.jit - def compress_on_cuda_impl( - self, a: cute.Tensor, a_compressed: cute.Tensor, meta: cute.Tensor - ): - """Compress the input tensor using the metadata""" - num_threads = 128 - grid = (cute.ceil_div(self.M, num_threads), 1, 1) - block = (num_threads, 1, 1) - self.compressor_impl(a, a_compressed, meta).launch(grid=grid, block=block) - - @cute.kernel - def compressor_impl( - self, a: cute.Tensor, a_compressed: cute.Tensor, meta: cute.Tensor - ): - """CUDA kernel to compress the tensor""" - tidx, tidy, tidz = cute.arch.thread_idx() - bidx, bidy, bidz = cute.arch.block_idx() - m = a.shape[0] - k = a.shape[1] - - # each thread process 1 row - row_idx = tidx + bidx * self.M - meta_idx = self.K // 4 // 8 - if row_idx < self.M: - # each meta_idx stands for 32 elements - for i in range(meta_idx): - meta[row_idx, i] = 0 - # each k stands for 4 elements - for j in range(8): - val = a[row_idx, i * 32 + j * 4] - val_1 = a[row_idx, i * 32 + j * 4 + 1] - val_2 = a[row_idx, i * 32 + j * 4 + 2] - val_3 = a[row_idx, i * 32 + j * 4 + 3] - value_idx = 0 - value_idx_1 = 0 - value_idx_2 = 0 - value_idx_3 = 0 - pos0 = 0 - pos1 = 0 - if val != 0: - value_idx = 1 - pos0 = 0 - if val_1 != 0: - value_idx_1 = 1 - if val_2 != 0: - value_idx_2 = 1 - if val_3 != 0: - value_idx_3 = 1 - pos = [value_idx, value_idx_1, value_idx_2, value_idx_3] - tmp = 0 - if pos == [0, 0, 0, 0]: - tmp = 0xE - pos0 = 2 - pos1 = 3 - elif pos == [1, 0, 0, 0]: - tmp = 0xC - pos0 = 0 - pos1 = 3 - elif pos == [0, 1, 0, 0]: - tmp = 0xD - pos0 = 1 - pos1 = 3 - elif pos == [0, 0, 1, 0]: - tmp = 0xE - pos0 = 2 - pos1 = 3 - elif pos == [0, 0, 0, 1]: - tmp = 0xE - pos0 = 2 - pos1 = 3 - elif pos == [1, 1, 0, 0]: - tmp = 0x4 - pos0 = 0 - pos1 = 1 - elif pos == [1, 0, 1, 0]: - tmp = 0x8 - pos0 = 0 - pos1 = 2 - elif pos == [1, 0, 0, 1]: - tmp = 0xC - pos0 = 0 - pos1 = 3 - elif pos == [0, 1, 1, 0]: - tmp = 0x9 - pos0 = 1 - pos1 = 2 - elif pos == [0, 1, 0, 1]: - tmp = 0xD - pos0 = 1 - pos1 = 3 - elif pos == [0, 0, 1, 1]: - tmp = 0xE - pos0 = 2 - pos1 = 3 - # cute.printf(row_idx, cutlass.Float32(val), cutlass.Float32(val_1), cutlass.Float32(val_2), cutlass.Float32(val_3), tmp) - meta[row_idx, i] |= tmp << (j * 4) - - a_compressed[row_idx, i * 16 + j * 2] = a[ - row_idx, i * 32 + j * 4 + pos0 - ] - a_compressed[row_idx, i * 16 + j * 2 + 1] = a[ - row_idx, i * 32 + j * 4 + pos1 - ] - - return - - -# SparseUtils is used to generate sparse tensor -# format torch.Tensor -class SparseUtils: - #!brief: SparseUtils is used to generate sparse tensor - #!param: M: int, K: int, L: int, dtype: cutlass.DataType - def __init__(self, M: int, K: int, L: int, dtype): - self.M = M - self.K = K - self.L = L - self.dtype = dtype - self.meta_data = self._generate_meta_data_4_2() - self._use_specific_meta_data = False - - #!brief: cast cutlass.DataType to torch.Tensor - def _get_type(self): - if self.dtype == cutlass.Float16: - return torch.float16 - elif self.dtype == cutlass.Float32: - return torch.float32 - elif self.dtype == cutlass.Int8: - return torch.int8 - else: - raise ValueError(f"Unsupported dtype: {self.dtype}") - - def _generate_meta_data_4_2(self): - # metadata for 4:2 sparse will in range( 4,8,9,c,d,e) - # represents - # 0: [1,1,0,0] no zero pos 00,01 -> 0100 = 4 - # 1: [1,0,1,0] no zero pos 00,10 -> 1000 = 8 - # 2: [1,0,0,1] no zero pos 00,11 -> 1100 = c - # 3: [0,1,1,0] no zero pos 01,10 -> 1001 = 9 - # 4: [0,1,0,1] no zero pos 01,11 -> 1101 = d - # 5: [0,0,1,1] no zero pos 10,11 -> 1011 = e - meta_value = [0x4, 0x8, 0x9, 0xC, 0xD, 0xE] - # 4:2 sparse, so each chunk is 4 elements, map to 4 bits - K_NumChunk = self.K // 4 - meta_data = np.random.choice( - meta_value, size=(self.M, K_NumChunk), replace=True - ) - meta_data = torch.from_numpy( - np.array(meta_data).astype(np.uint8).reshape(self.M, K_NumChunk) - ) - return meta_data - - #!brief: pack meta data - def _pack_meta_data(self): - tmp = [] - K_NumChunk = self.K // 4 - for i in range(self.M): - for j in range(K_NumChunk // 8): - v = 0 - for k in range(8): - vv = int(self.meta_data[i, j * 8 + k] & 0xF) - tt = vv << (k * 4) - v = v | tt - tmp.append(v) - # debug print - # print([hex(vt) for vt in tmp]) - result = torch.from_numpy( - np.array(tmp).astype(np.uint32).reshape(self.M, K_NumChunk // 8) - ) - return result - - #!brief: use specific meta data - def use_specific_meta_data(self, meta_data: torch.Tensor = None): - if meta_data is not None: - self.meta_data = meta_data - self._use_specific_meta_data = True - - #!brief: generate sparse tensor with tensor - #!param: a: torch.Tensor - #!param: run_on_cpu: bool - #!return: torch.Tensor - def generate_sparse_4_2_tensor_with_tensor(self, a, run_on_cpu): - if run_on_cpu: - if a.device.type != "cpu": - raise ValueError("a must be on cpu") - return self.__generate_sparse_tensor_cpu(a) - else: - if a.device.type != "cuda": - raise ValueError("a must be on cuda") - a_tensor = from_dlpack(a) - packed_meta_data = self._pack_meta_data() - meta_tensor = from_dlpack(packed_meta_data.cuda()) - self.__generate_sparse_tensor_cuda(a_tensor, meta_tensor) - return a - - #!brief: generate sparse tensor - #!param: run_on_cpu: bool - #!return: torch.Tensor - def generate_4_2_sparse_tensor(self, run_on_cpu): - dtype = self._get_type() - a = torch.empty(self.M, self.K).random_(-5, 5).to(dtype) - if run_on_cpu: - return self.generate_sparse_4_2_tensor_with_tensor(a, run_on_cpu) - else: - return self.generate_sparse_4_2_tensor_with_tensor(a.cuda(), run_on_cpu) - - #!brief: generate sparse tensor on cpu - #!param: a: torch.Tensor - #!return: torch.Tensor - def __generate_sparse_tensor_cpu(self, a): - if not self._use_specific_meta_data: - for m in range(self.M): - for k in range(0, self.K, 4): - # random choose 2 zero positions - zero_indices = torch.randperm(4)[:2] - a[m, k + zero_indices[0]] = 0 - a[m, k + zero_indices[1]] = 0 - return a - else: - # use specific meta data - tensor_mask = [] - for i in range(self.M): - for j in range(self.K // 4): - meta_val = self.meta_data[i, j] - tmp = [] - if meta_val == 0x4: - tmp = [1, 1, 0, 0] - elif meta_val == 0x8: - tmp = [1, 0, 1, 0] - elif meta_val == 0xC: - tmp = [1, 0, 0, 1] - elif meta_val == 0x9: - tmp = [0, 1, 1, 0] - elif meta_val == 0xD: - tmp = [0, 1, 0, 1] - elif meta_val == 0xE: - tmp = [0, 0, 1, 1] - tensor_mask.extend(tmp) - a = torch.reshape(a, (-1,)) - mask = torch.tensor(tensor_mask) - a = a * mask - a = torch.reshape(a, (self.M, self.K)) - return a - - @cute.jit - def __generate_sparse_tensor_cuda(self, a: cute.Tensor, meta: cute.Tensor): - """Generate a sparse tensor from a dense tensor using metadata""" - assert a.shape[0] == self.M and a.shape[1] == self.K - assert meta.shape[0] == self.M and meta.shape[1] == self.K // 4 // 8 - num_threads = 128 - grid = (cute.ceil_div(self.M, num_threads), 1, 1) - block = (num_threads, 1, 1) - self.kernel(a, meta).launch(grid=grid, block=block) - - @cute.kernel - def kernel(self, a: cute.Tensor, meta: cute.Tensor): - """Apply sparsity mask to input tensor using metadata""" - tidx, tidy, tidz = cute.arch.thread_idx() - bidx, bidy, bidz = cute.arch.block_idx() - - # each thread process 1 ro - row_idx = tidx + bidx * self.M - meta_idx = self.K // 4 // 8 - # each thread process 1 row - if row_idx < self.M: - # iterate over each chunk(32 elements) - for i in range(meta_idx): - meta_val = meta[(row_idx, i)] - # iterate over each sparse pattern(4 elements) - for j in range(8): - meta_row = (meta_val >> (j * 4)) & 0xF - idx0 = meta_row & 0x3 - idx1 = (meta_row >> 2) & 0x3 - r_id0 = 0 - r_id1 = 0 - # r_id is the idx that value is 0 - if idx0 >= 2 and idx1 >= 2: - r_id0 = 0 - r_id1 = 1 - elif idx0 <= 1 and idx1 <= 1: - r_id0 = 2 - r_id1 = 3 - else: - r_id0 = idx0 ^ 0b1 - r_id1 = idx1 ^ 0b1 - row_id0 = r_id0 + i * 32 + j * 4 - row_id1 = r_id1 + i * 32 + j * 4 - a[row_idx, row_id0] = self.dtype(0.0) - a[row_idx, row_id1] = self.dtype(0.0) - return diff --git a/examples/python/CuTeDSL/helpers/test_sparse_utils.py b/examples/python/CuTeDSL/helpers/test_sparse_utils.py deleted file mode 100644 index 3264f191..00000000 --- a/examples/python/CuTeDSL/helpers/test_sparse_utils.py +++ /dev/null @@ -1,104 +0,0 @@ -import sparse_utils as su -import cutlass -import torch -from cutlass.cute.runtime import from_dlpack -import numpy as np -import pytest - - -@pytest.mark.L0 -def test_sparse_cpu(): - M = 128 - N = 32 - K = 32 - L = 1 - debug = False - # generate sparse tensor - a = torch.empty(M, K).random_(-5, 5).to(torch.float16) - sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16) - if debug: - sparse_utils.use_specific_meta_data() - a_gen_from_cpu = sparse_utils.generate_sparse_4_2_tensor_with_tensor(a, True) - # print(a_gen_from_cpu) - # generate compressed tensor and meta data - a_compressed_cpu = torch.empty(M, K // 2).to(torch.float16) - meta_data_cpu = torch.empty(M, K // 4 // 8).to(torch.uint32) - compressor = su.Compressor(M, K, L) - compressor.compress(a_gen_from_cpu, a_compressed_cpu, meta_data_cpu, True) - # # test with gemm - b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda() - d = torch.empty(M, N).zero_().to(torch.float16).cuda() - b_tensor = from_dlpack(b) - d_tensor = from_dlpack(d) - a_compressed_cpu_tensor = from_dlpack(a_compressed_cpu.cuda()) - meta_data_cpu_tensor = from_dlpack(meta_data_cpu.cuda()) - sparse_emulation = su.SparseEmulation(M, N, K, 1) - sparse_emulation(a_compressed_cpu_tensor, b_tensor, d_tensor, meta_data_cpu_tensor) - - ref = torch.einsum("mk,nk->mn", a_gen_from_cpu.cpu(), b.cpu()) - if debug: - a_ori = a_gen_from_cpu.cpu().numpy() - np.savetxt("a.txt", a_ori, fmt="%f") - a_compressed_cpu_ori = a_compressed_cpu.cpu().numpy() - np.savetxt("a_compressed_cpu.txt", a_compressed_cpu_ori, fmt="%f") - meta_data_cpu_ori = meta_data_cpu.cpu().numpy() - np.savetxt("meta_data_cpu.txt", meta_data_cpu_ori, fmt="%f") - d_ori = d.cpu().numpy() - np.savetxt("d.txt", d_ori, fmt="%f") - ref_ori = ref.cpu().numpy() - np.savetxt("ref.txt", ref_ori, fmt="%f") - torch.testing.assert_close(d.cpu(), ref) - print("cpu d == ref") - - -@pytest.mark.L0 -def test_sparse_cuda(): - M = 128 - N = 32 - K = 32 - L = 1 - debug = False - sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16) - if debug: - sparse_utils.use_specific_meta_data() - # generate sparse tensor - a = torch.empty(M, K).random_(-5, 5).to(torch.float16).cuda() - a_gen_from_cuda = sparse_utils.generate_4_2_sparse_tensor(False) - # print(a_gen_from_cuda) - # generate compressed tensor and meta data - a_compressed_cuda = torch.empty(M, K // 2).to(torch.float16).cuda() - meta_data_cuda = torch.empty(M, K // 4 // 8).to(torch.uint32).cuda() - compressor = su.Compressor(M, K, L) - compressor.compress(a_gen_from_cuda, a_compressed_cuda, meta_data_cuda, False) - # test with gemm - b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda() - d = torch.empty(M, N).zero_().to(torch.float16).cuda() - b_tensor = from_dlpack(b) - d_tensor = from_dlpack(d) - a_compressed_cuda_tensor = from_dlpack(a_compressed_cuda) - meta_data_cuda_tensor = from_dlpack(meta_data_cuda) - sparse_emulation = su.SparseEmulation(M, N, K, 1) - sparse_emulation( - a_compressed_cuda_tensor, b_tensor, d_tensor, meta_data_cuda_tensor - ) - - ref = torch.einsum("mk,nk->mn", a_gen_from_cuda.cpu(), b.cpu()) - if debug: - a_ori = a_gen_from_cuda.cpu().numpy() - np.savetxt("a.txt", a_ori, fmt="%f") - a_compressed_cuda_ori = a_compressed_cuda.cpu().numpy() - np.savetxt("a_compressed_cuda.txt", a_compressed_cuda_ori, fmt="%f") - meta_data_cuda_ori = meta_data_cuda.cpu().numpy() - np.savetxt("meta_data_cuda.txt", meta_data_cuda_ori, fmt="%f") - d_ori = d.cpu().numpy() - np.savetxt("d.txt", d_ori, fmt="%f") - ref_ori = ref.cpu().numpy() - np.savetxt("ref.txt", ref_ori, fmt="%f") - torch.testing.assert_close(d.cpu(), ref) - print("cuda d == ref") - - -if __name__ == "__main__": - cutlass.cuda.initialize_cuda_context() - test_sparse_cpu() - test_sparse_cuda() diff --git a/examples/python/CuTeDSL/hopper/dense_gemm.py b/examples/python/CuTeDSL/hopper/dense_gemm.py index c00ead6f..96ad8338 100644 --- a/examples/python/CuTeDSL/hopper/dense_gemm.py +++ b/examples/python/CuTeDSL/hopper/dense_gemm.py @@ -31,15 +31,12 @@ from typing import Tuple, Type import math import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -import cutlass.torch as cutlass_torch from cutlass.cute.runtime import from_dlpack import cutlass.utils.hopper_helpers as sm90_utils @@ -1006,7 +1003,10 @@ class HopperWgmmaGemmKernel: tiled_copy_r2s, tRS_rD_out, tRS_sD[(None, None, None, epi_buffer)] ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # barrier for sync pipeline.sync(barrier_id=1) @@ -1431,6 +1431,9 @@ def run( :rtype: float """ + import torch + import cutlass.torch as cutlass_torch + print("Running Hopper Dense GEMM with:") print(f"mnkl: {mnkl}") print( @@ -1519,7 +1522,7 @@ def run( gemm = HopperWgmmaGemmKernel(acc_dtype, tile_shape_mn, cluster_shape_mn) - torch_stream = torch.cuda.Stream() + torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) # compile gemm kernel compiled_gemm = cute.compile(gemm, mA, mB, mC, stream) diff --git a/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py b/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py index 5b999c0b..a8119c8a 100644 --- a/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py +++ b/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py @@ -31,14 +31,11 @@ from typing import Optional, Tuple, Type import math import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.utils.hopper_helpers as sm90_utils @@ -952,7 +949,10 @@ class HopperWgmmaGemmPersistentKernel: tRS_sD[(None, None, None, epi_buffer)], ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() gmem_coord = epi_tile_layout.get_hier_coord(epi_idx) @@ -1465,6 +1465,8 @@ def run( :return: Execution time of the GEMM kernel in microseconds :rtype: float """ + import torch + import cutlass.torch as cutlass_torch print("Running Hopper Persistent Dense GEMM with:") print(f"mnkl: {mnkl}") diff --git a/examples/python/CuTeDSL/hopper/fmha.py b/examples/python/CuTeDSL/hopper/fmha.py index 9efce75b..e384c53b 100644 --- a/examples/python/CuTeDSL/hopper/fmha.py +++ b/examples/python/CuTeDSL/hopper/fmha.py @@ -83,9 +83,6 @@ import sys import time from typing import Type, Tuple, Optional -import torch - - import cuda.bindings.driver as cuda import cutlass @@ -96,7 +93,6 @@ import cutlass.cute.nvgpu.warpgroup as warpgroup import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -import cutlass.torch as cutlass_torch from cutlass._mlir.dialects import math as _math import cutlass.utils.hopper_helpers as sm90_utils @@ -598,6 +594,7 @@ class HopperFusedMultiHeadAttentionForward: k_smem_layout_staged.outer, swizzle=k_smem_layout_staged.inner ) # (MMA, MMA_K, MMA_D, PIPE) + # Adjust swizzle info to reuse smem sV_ptr = cute.recast_ptr(sK.iterator, v_smem_layout_staged.inner) sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer) @@ -648,11 +645,19 @@ class HopperFusedMultiHeadAttentionForward: producer_warp_role = warp_idx % 4 # self.num_warps_per_warp_group + # Fence the mbarrier init to ensure all mbarrier initializations are visible + # to all threads. This is critical for FP8 performance - without this fence, + # the compiler may generate software polling loops instead of hardware waits. + cute.arch.mbarrier_init_fence() + # We need this to guarantee that the Pipeline init is visible # To all producers and consumer blocks in the Cluster # and to finish smem init - pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) - pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + if cute.size(self.cluster_shape_mnk) > 1: + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + else: + cute.arch.sync_threads() if warp_idx == 0: cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) @@ -1164,7 +1169,10 @@ class HopperFusedMultiHeadAttentionForward: tRS_sD[(None, None, None, epi_buffer, warp_group_idx - 1)], ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) pipeline.arrive_and_wait( barrier_id=warp_group_idx, num_threads=self.num_threads_per_warp_group, @@ -1935,6 +1943,9 @@ def run( :return: Execution time of the FMHA kernel in microseconds :rtype: float """ + import torch + import cutlass.torch as cutlass_torch + print("Running Hopper SM90 FMHA test with:") print(f" q_shape: {q_shape}") print(f" k_shape: {k_shape}") diff --git a/examples/python/CuTeDSL/jax/cutlass_call_sharding.py b/examples/python/CuTeDSL/jax/cutlass_call_sharding.py index a2d03c34..e40687ff 100644 --- a/examples/python/CuTeDSL/jax/cutlass_call_sharding.py +++ b/examples/python/CuTeDSL/jax/cutlass_call_sharding.py @@ -31,7 +31,7 @@ import argparse import jax import jax.numpy as jnp -from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from jax.sharding import NamedSharding, PartitionSpec as P, AxisType from jax.experimental.custom_partitioning import custom_partitioning import cutlass @@ -42,8 +42,8 @@ import cuda.bindings.driver as cuda """ -Examples of combining jax.jit and jax.shard_map for sharding and executing kernels -across multiple GPU devices. +Examples of combining jax.jit, jax.shard_map and custom_partitioning for sharding +and executing kernels across multiple GPU devices. To run this example: @@ -84,25 +84,57 @@ def launch( ) +def sharded_cutlass_call_impl(a_block, b_block): + """The sharded implementation that operates on a single device.""" + call = cjax.cutlass_call( + launch, + use_static_tensors=True, + output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype), + ) + ref_result = a_block + b_block + return call(a_block, b_block), ref_result + + +@custom_partitioning +def custom_shared_call(a, b): + return sharded_cutlass_call_impl(a, b) + + +def custom_shared_call_partitioner(mesh, arg_shapes, result_shape): + arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes) + result_shardings = tuple([arg_shardings[0]] * len(result_shape)) + + def lower_fn(*args): + return sharded_cutlass_call_impl(*args) + + return mesh, lower_fn, result_shardings, arg_shardings + + +custom_shared_call.def_partition(custom_shared_call_partitioner) + + def run_example(): # Create a device mesh with one axis b ngpu = jax.device_count() - mesh = jax.make_mesh((ngpu,), "b") + mesh = jax.make_mesh((ngpu,), "b", axis_types=(AxisType.Explicit,)) if ngpu == 1: print("Note: only 1 GPU was detected.") # We will shard our 3D tensors over b sharding = P("b", None, None) + named_sharding = NamedSharding(mesh, sharding) - @partial(jax.jit, static_argnums=[0, 1]) + print("Testing shard_map...") + + @partial( + jax.jit, static_argnums=[0, 1], out_shardings=(named_sharding, named_sharding) + ) def allocate_sharded_tensors(shape, dtype): key = jax.random.key(1123) - a_key, b_keys = jax.random.split(key, 2) + a_key, b_key = jax.random.split(key, 2) a = create_tensor(shape, dtype, a_key) - b = create_tensor(shape, dtype, b_keys) - a = jax.lax.with_sharding_constraint(a, NamedSharding(mesh, sharding)) - b = jax.lax.with_sharding_constraint(b, NamedSharding(mesh, sharding)) + b = create_tensor(shape, dtype, b_key) return a, b @jax.jit @@ -115,13 +147,7 @@ def run_example(): out_specs=(sharding, sharding), ) def sharded_call(a_block, b_block): - call = cjax.cutlass_call( - launch, - use_static_tensors=True, - output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype), - ) - ref_result = a_block + b_block - return call(a_block, b_block), ref_result + return sharded_cutlass_call_impl(a_block, b_block) return sharded_call(a, b) @@ -134,6 +160,17 @@ def run_example(): assert jnp.allclose(c, c_ref) + print("Testing custom_partitioning...") + + # Test custom_partitioning implementation which should produce identical results + @jax.jit + def compute_cp(a, b): + return custom_shared_call(a, b) + + c, c_ref = compute_cp(a, b) + + assert jnp.allclose(c, c_ref) + if __name__ == "__main__": run_example() diff --git a/examples/python/CuTeDSL/notebooks/hello_world.ipynb b/examples/python/CuTeDSL/notebooks/hello_world.ipynb index 218378ad..6bf35b76 100644 --- a/examples/python/CuTeDSL/notebooks/hello_world.ipynb +++ b/examples/python/CuTeDSL/notebooks/hello_world.ipynb @@ -55,7 +55,7 @@ " # Get the x component of the thread index (y and z components are unused)\n", " tidx, _, _ = cute.arch.thread_idx()\n", " # Only the first thread (thread 0) prints the message\n", - " if tidx == 0:\n", + " if cutlass.dynamic_expr(tidx == 0):\n", " cute.printf(\"Hello world\")" ] }, @@ -142,8 +142,6 @@ "from cutlass.cute import KeepPTX, KeepCUBIN\n", "\n", "print(\"Compiling with PTX/CUBIN dumped...\")\n", - "# Alternatively, compile with string based options like\n", - "# cute.compile(hello_world, options=\"--keep-ptx --keep-cubin\") would also work.\n", "hello_world_compiled_ptx_on = cute.compile[KeepPTX, KeepCUBIN](hello_world)\n", "\n", "# Run the pre-compiled version\n", diff --git a/examples/python/CuTeDSL/notebooks/tensorssa.ipynb b/examples/python/CuTeDSL/notebooks/tensorssa.ipynb index f60e0365..62804a10 100644 --- a/examples/python/CuTeDSL/notebooks/tensorssa.ipynb +++ b/examples/python/CuTeDSL/notebooks/tensorssa.ipynb @@ -167,7 +167,7 @@ "outputs": [], "source": [ "@cute.jit\n", - "def binary_op_1(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n", + "def binary_op_1(a: cute.Tensor, b: cute.Tensor):\n", " a_vec = a.load()\n", " b_vec = b.load()\n", "\n", @@ -184,7 +184,7 @@ " cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n", "\n", " floor_div_res = a_vec // b_vec\n", - " cute.print_tensor(res) # prints [0.000000, 0.000000, 0.000000]\n", + " cute.print_tensor(floor_div_res) # prints [0.000000, 0.000000, 0.000000]\n", "\n", " mod_res = a_vec % b_vec\n", " cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n", @@ -194,8 +194,7 @@ "a.fill(1.0)\n", "b = np.empty((3,), dtype=np.float32)\n", "b.fill(2.0)\n", - "res = np.empty((3,), dtype=np.float32)\n", - "binary_op_1(from_dlpack(res), from_dlpack(a), from_dlpack(b))" + "binary_op_1(from_dlpack(a), from_dlpack(b))" ] }, { @@ -205,7 +204,7 @@ "outputs": [], "source": [ "@cute.jit\n", - "def binary_op_2(res: cute.Tensor, a: cute.Tensor, c: cutlass.Constexpr):\n", + "def binary_op_2(a: cute.Tensor, c: cutlass.Constexpr):\n", " a_vec = a.load()\n", "\n", " add_res = a_vec + c\n", @@ -230,8 +229,7 @@ "a = np.empty((3,), dtype=np.float32)\n", "a.fill(1.0)\n", "c = 2.0\n", - "res = np.empty((3,), dtype=np.float32)\n", - "binary_op_2(from_dlpack(res), from_dlpack(a), c)" + "binary_op_2(from_dlpack(a), c)" ] }, { diff --git a/include/cutlass/arch/grid_dependency_control.h b/include/cutlass/arch/grid_dependency_control.h index 912b635a..63a7714e 100644 --- a/include/cutlass/arch/grid_dependency_control.h +++ b/include/cutlass/arch/grid_dependency_control.h @@ -62,6 +62,8 @@ (defined(__CUDA_ARCH_FEAT_SM100_ALL) || CUDA_ARCH_FAMILY(1000))) || \ (__CUDA_ARCH__ == 1010 &&\ (defined(__CUDA_ARCH_FEAT_SM101_ALL) || CUDA_ARCH_FAMILY(1010))) || \ + (__CUDA_ARCH__ == 1100 &&\ + (defined(__CUDA_ARCH_FEAT_SM110_ALL) || CUDA_ARCH_FAMILY(1100))) || \ (__CUDA_ARCH__ == 1030 &&\ (defined(__CUDA_ARCH_FEAT_SM103_ALL) || CUDA_ARCH_FAMILY(1030))) || \ (__CUDA_ARCH__ == 1200 &&\ diff --git a/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp b/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp index fbf88553..494ffe7a 100644 --- a/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp +++ b/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp @@ -276,7 +276,11 @@ public: static constexpr int MaxClusterSize = 16; implementable &= size(args.hw_info.cluster_shape) <= MaxClusterSize; implementable &= size(args.hw_info.cluster_shape_fallback) <= MaxClusterSize; - implementable &= cutlass::detail::preferred_cluster_can_implement(args.hw_info.cluster_shape, args.hw_info.cluster_shape_fallback); + // Early return if cluster shape validation failed to avoid division by zero below + if (not cutlass::detail::preferred_cluster_can_implement(args.hw_info.cluster_shape, args.hw_info.cluster_shape_fallback)) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Invalid dynamic cluster shape\n"); + return false; + } } auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, args.hw_info.cluster_shape); diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp index bdeae396..15dd91bc 100644 --- a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp @@ -181,6 +181,11 @@ struct CollectiveMma< static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 3; constexpr static int NumBandsToCompute = DispatchPolicy::NumBandsToCompute; diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp index d9275cf0..45b9cb43 100644 --- a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp @@ -189,6 +189,11 @@ public: static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 3; constexpr static int ConjSwapMode = 2; diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp index 9e0a6c5e..33507c9b 100644 --- a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp @@ -173,6 +173,11 @@ public: static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 152; + static constexpr uint32_t TransformRegisterRequirement = 200; + static constexpr uint32_t AccumRegisterRequirement = 152; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 2; constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount; diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp index 57bc1cca..1be80601 100644 --- a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp @@ -192,6 +192,11 @@ struct CollectiveMma< static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 3; constexpr static int NumBandsToCompute = DispatchPolicy::NumBandsToCompute; diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp index 140b7854..a8fe8c4a 100644 --- a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp @@ -185,6 +185,11 @@ public: static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 3; constexpr static int ConjSwapMode = 2; diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp index 1dc7b02c..a1f25017 100644 --- a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp @@ -184,6 +184,11 @@ public: static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 152; + static constexpr uint32_t TransformRegisterRequirement = 200; + static constexpr uint32_t AccumRegisterRequirement = 152; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 2; constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount; diff --git a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp index 492371c5..18d1cd61 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp @@ -832,6 +832,12 @@ public: mainloop_pipeline.init_masks(cluster_shape, block_id_in_cluster); accumulator_pipeline.init_masks(cluster_shape, block_id_in_cluster); + // Ensure that the prefetched kernel does not touch + // unflushed global memory prior to this instruction. + // For the static grouped scheduler, the problem shapes + // might be produced by a previous kernel in global memory. + cutlass::arch::wait_on_dependent_grids(); + // TileID scheduler TileScheduler scheduler( (!IsTensorMapUpdateAsync || is_participant.sched || is_participant.tensor_map_updater) @@ -842,12 +848,6 @@ public: ); auto work_tile_info = [&] () { - // Ensure that the prefetched kernel does not touch - // unflushed global memory prior to this instruction. - // For the static grouped scheduler, the problem shapes - // might be produced by a previous kernel in global memory. - cutlass::arch::wait_on_dependent_grids(); - if constexpr (IsTensorMapUpdateAsync) { return scheduler.initial_work_tile_info(cluster_shape, [] (typename TileScheduler::CLCResponse response) { CLCResponseWithAdditionalInformation response_with_additional_info = response; diff --git a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp index 48b1f738..25d5c6e3 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp @@ -148,9 +148,10 @@ public: static constexpr uint32_t NumFixupBarriers = 1; static constexpr uint32_t CLCResponseSize = sizeof(typename TileScheduler::CLCResponse); - // Transfer registers from regular warps to Accum warps - static constexpr uint32_t GenericRegisterRequirement = 152; - static constexpr uint32_t AccumRegisterRequirement = 200; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = CollectiveMainloop::GenericRegisterRequirement; + static constexpr uint32_t TransformRegisterRequirement = CollectiveMainloop::TransformRegisterRequirement; + static constexpr uint32_t AccumRegisterRequirement = CollectiveMainloop::AccumRegisterRequirement; // Pipeline and pipeline state types using Load2TransformPipeline = typename CollectiveMainloop::Load2TransformPipeline; @@ -412,6 +413,22 @@ public: return dim3(MaxThreadsPerBlock, 1, 1); } + // Register alloc/dealloc behavior might change according to the underlying collective used + template + CUTLASS_DEVICE + static constexpr void + warpgroup_reg_reconfig() { + // Compute default-allocated registers per thread: round_down((512 / NumWG), 8) + constexpr int32_t MaxWarpGroupsPerBlock = ceil_div(MaxThreadsPerBlock, NumThreadsPerWarpGroup); + constexpr int32_t NumRegsPerThread = (512 / MaxWarpGroupsPerBlock) / 8 * 8; + if constexpr (NReg < NumRegsPerThread) { + arch::warpgroup_reg_dealloc(); + } + else if constexpr (NReg > NumRegsPerThread) { + arch::warpgroup_reg_alloc(); + } + } + CUTLASS_DEVICE void operator() (Params const& params, char* smem_buf) { @@ -652,12 +669,12 @@ public: // Allocate accumulators auto acc_shape = collective_mainloop.partition_accumulator_shape(); - // TileID scheduler - TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster); - // Ensure memory ops in this kernel are not done prior to completion of dependent grids. cutlass::arch::wait_on_dependent_grids(); + // TileID scheduler + TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster); + typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape); auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info); @@ -677,7 +694,7 @@ public: if (is_participant.main_load) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Ensure that the prefetched kernel does not touch // unflushed global memory prior to this instruction @@ -791,7 +808,7 @@ public: else if (is_participant.transformation) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -833,7 +850,7 @@ public: else if (is_participant.sched) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -898,7 +915,7 @@ public: else if (is_participant.mma) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Allocate all tmem tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr); @@ -966,7 +983,7 @@ public: else if (is_participant.epi_load) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Ensure that the prefetched kernel does not touch // unflushed global memory prior to this instruction @@ -1051,7 +1068,7 @@ public: else if (is_participant.epilogue) { // Register reconfiguration - arch::warpgroup_reg_alloc(); + warpgroup_reg_reconfig(); // Throttle the epilogue warps to improve prologue performance static constexpr int epilogue_throttle_phase_bit = 0; @@ -1182,7 +1199,7 @@ public: else { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); } } }; diff --git a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp index 0025e63b..83ae76ac 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp @@ -735,12 +735,12 @@ public: Tensor accumulators = cutlass::detail::make_sm100_accumulator( tiled_mma, acc_shape, EpilogueTile{}); - // TileID scheduler - TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster); - // Ensure memory ops in this kernel are not done prior to completion of dependent grids. cutlass::arch::wait_on_dependent_grids(); + // TileID scheduler + TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster); + typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape); auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info); diff --git a/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp b/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp index 67853638..c82e084f 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp @@ -143,10 +143,10 @@ public: static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent; - // Transfer registers from regular warps to Accum warps - static constexpr uint32_t GenericRegisterRequirement = 64; - static constexpr uint32_t TransformRegisterRequirement = 184; - static constexpr uint32_t AccumRegisterRequirement = 256; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = CollectiveMainloop::GenericRegisterRequirement; + static constexpr uint32_t TransformRegisterRequirement = CollectiveMainloop::TransformRegisterRequirement; + static constexpr uint32_t AccumRegisterRequirement = CollectiveMainloop::AccumRegisterRequirement; // Pipeline and pipeline state types using Load2TransformPipeline = typename CollectiveMainloop::Load2TransformPipeline; @@ -389,6 +389,22 @@ public: return dim3(MaxThreadsPerBlock, 1, 1); } + // Register alloc/dealloc behavior might change according to the underlying collective used + template + CUTLASS_DEVICE + static constexpr void + warpgroup_reg_reconfig() { + // Compute default-allocated registers per thread: round_down((512 / NumWG), 8) + constexpr int32_t MaxWarpGroupsPerBlock = ceil_div(MaxThreadsPerBlock, NumThreadsPerWarpGroup); + constexpr int32_t NumRegsPerThread = (512 / MaxWarpGroupsPerBlock) / 8 * 8; + if constexpr (NReg < NumRegsPerThread) { + arch::warpgroup_reg_dealloc(); + } + else if constexpr (NReg > NumRegsPerThread) { + arch::warpgroup_reg_alloc(); + } + } + CUTLASS_DEVICE void operator() (Params const& params, char* smem_buf) { @@ -638,7 +654,7 @@ public: if (is_participant.main_load) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Ensure that the prefetched kernel does not touch // unflushed global memory prior to this instruction @@ -716,7 +732,7 @@ public: else if (is_participant.sched) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -770,7 +786,7 @@ public: else if (is_participant.transformation) { // Register reconfiguration - arch::warpgroup_reg_alloc(); + warpgroup_reg_reconfig(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -813,7 +829,7 @@ public: else if (is_participant.mma) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Tmem allocation sequence tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr); @@ -880,7 +896,7 @@ public: else if (is_participant.epi_load) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Ensure that the prefetched kernel does not touch // unflushed global memory prior to this instruction @@ -943,7 +959,7 @@ public: else if (is_participant.epilogue) { // Register reconfiguration - arch::warpgroup_reg_alloc(); + warpgroup_reg_reconfig(); // Throttle the epilogue warps to improve prologue performance static constexpr int epilogue_throttle_phase_bit = 0; @@ -1067,7 +1083,7 @@ public: else { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); } } }; diff --git a/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp b/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp index 9c821773..2d8728a9 100755 --- a/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp +++ b/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp @@ -116,12 +116,15 @@ public: CUTLASS_DEVICE PersistentTileSchedulerSm100Group() { } - + + // Note: constructing this tile scheduler can touch global memory that was + // written to by the prior kernel. CUTLASS_DEVICE PersistentTileSchedulerSm100Group(CLCResponse* clc_response_ptr, Params const& params) : scheduler_params(params), scheduler_sm90(params.params_sm90_, clc_response_ptr) { } - + // Note: constructing this tile scheduler can touch global memory that was + // written to by the prior kernel. CUTLASS_DEVICE PersistentTileSchedulerSm100Group(CLCResponse* clc_response_ptr, Params const& params, dim3 /* block_id_in_cluster */) : scheduler_params(params), @@ -161,9 +164,6 @@ public: // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently Arguments args{}; - if constexpr (!std::is_const_v) { - args.max_swizzle_size = 1 << params.params_sm90_.log_swizzle_size_; - } args.raster_order = params.params_sm90_.raster_order_ == RasterOrder::AlongN ? RasterOrderOptions::AlongN : RasterOrderOptions::AlongM; return Params::get_grid_shape( diff --git a/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp b/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp index c90b3e3d..7416f417 100644 --- a/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp @@ -752,12 +752,12 @@ public: Tensor accumulators = cutlass::detail::make_sm100_accumulator( tiled_mma, acc_shape, EpilogueTile{}); - // TileID scheduler - TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster); - // Ensure memory ops in this kernel are not done prior to completion of dependent grids. cutlass::arch::wait_on_dependent_grids(); + // TileID scheduler + TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster); + typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape); auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info); diff --git a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp index 67639bf7..3a5149d6 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp @@ -387,9 +387,6 @@ public: get_grid_shape(Params const& params) { // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently TileSchedulerArguments args{}; - if constexpr (!std::is_const_v) { - args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_; - } args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM; dim3 grid_shape; if constexpr (IsGroupedGemmKernel) { @@ -454,16 +451,6 @@ public: // Kernel level shared memory storage SharedStorage& shared_storage = *reinterpret_cast(smem_buf); - auto scheduler = [&] () { - // Group scheduler requires a different constructor that takes a response ptr - if constexpr (cute::is_same_v) { - return TileScheduler{params.scheduler, shared_storage.scheduler_response}; - } - else { - return TileScheduler{params.scheduler}; - } - } (); - // In a warp specialized kernel, collectives expose data movement and compute operations separately CollectiveMainloop collective_mainloop; CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue); @@ -585,6 +572,16 @@ public: // Ensure memory ops in this kernel are not done prior to completion of dependent grids. cutlass::arch::wait_on_dependent_grids(); + auto scheduler = [&] () { + // Group scheduler requires a different constructor that takes a response ptr + if constexpr (cute::is_same_v) { + return TileScheduler{params.scheduler, shared_storage.scheduler_response}; + } + else { + return TileScheduler{params.scheduler}; + } + } (); + auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{}); if (not work_tile_info.is_valid()) { diff --git a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp index 295a44dd..e1fa1c86 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp @@ -399,9 +399,6 @@ public: get_grid_shape(Params const& params) { // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently TileSchedulerArguments args{}; - if constexpr (!std::is_const_v) { - args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_; - } args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM; dim3 grid_shape; if constexpr (IsGroupedGemmKernel) { @@ -463,16 +460,6 @@ public: // Kernel level shared memory storage SharedStorage& shared_storage = *reinterpret_cast(smem_buf); - auto scheduler = [&] () { - // Group scheduler requires a different constructor that takes a response ptr - if constexpr (cute::is_same_v) { - return TileScheduler{params.scheduler, shared_storage.scheduler_response}; - } - else { - return TileScheduler{params.scheduler}; - } - } (); - // In a warp specialized kernel, collectives expose data movement and compute operations separately CollectiveMainloop collective_mainloop; CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue); @@ -600,6 +587,16 @@ public: // Ensure memory ops in this kernel are not done prior to completion of dependent grids. cutlass::arch::wait_on_dependent_grids(); + auto scheduler = [&] () { + // Group scheduler requires a different constructor that takes a response ptr + if constexpr (cute::is_same_v) { + return TileScheduler{params.scheduler, shared_storage.scheduler_response}; + } + else { + return TileScheduler{params.scheduler}; + } + } (); + auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{}); if (not work_tile_info.is_valid()) { diff --git a/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp b/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp index edab79bb..d746ca73 100644 --- a/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp +++ b/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp @@ -59,6 +59,7 @@ private: uint64_t start_linear_idx = 0; uint64_t total_tiles = 0; uint64_t problem_blocks_along_raster_order = 0; + int32_t log_swizzle_size = 0; } current_group_info_; public: @@ -244,8 +245,29 @@ public: return true; } + // Calculate the log of the swizzle size based on the problem CTAs and the max swizzle size + CUTLASS_DEVICE + static int32_t + get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) { + int min_cta_dim = platform::min(problem_ctas_m, problem_ctas_n); + if (max_swizzle_size >= 8 && min_cta_dim >= 6) { + return 3; + } + else if (max_swizzle_size >= 4 && min_cta_dim >= 3) { + return 2; + } + else if (max_swizzle_size >= 2 && min_cta_dim >= 2) { + return 1; + } + else { + return 0; + } + } + PersistentTileSchedulerSm90Group() = default; + // Note: constructing this tile scheduler can touch global memory that was + // written to by the prior kernel. CUTLASS_DEVICE explicit PersistentTileSchedulerSm90Group(Params const& params_, SchedulerResponse* response_ptr) : scheduler_params(params_), response_ptr_(response_ptr) { // MSVC requires protecting use of CUDA-specific nonstandard syntax, // like blockIdx and gridDim, with __CUDA_ARCH__. @@ -274,8 +296,9 @@ public: ctas_along_m = scheduler_params.divmod_cta_shape_m_.divide(cute::shape<0>(problem_shape) + scheduler_params.divmod_cta_shape_m_.divisor - 1); ctas_along_n = scheduler_params.divmod_cta_shape_n_.divide(cute::shape<1>(problem_shape) + scheduler_params.divmod_cta_shape_n_.divisor - 1); } - auto problem_blocks_m = round_up(ctas_along_m, (1 << params_.log_swizzle_size_) * params_.cluster_shape_.m()); - auto problem_blocks_n = round_up(ctas_along_n, (1 << params_.log_swizzle_size_) * params_.cluster_shape_.n()); + current_group_info_.log_swizzle_size = get_log_swizzle_size(ctas_along_m, ctas_along_n, params_.max_swizzle_size_); + auto problem_blocks_m = round_up(ctas_along_m, (1 << current_group_info_.log_swizzle_size) * params_.cluster_shape_.m()); + auto problem_blocks_n = round_up(ctas_along_n, (1 << current_group_info_.log_swizzle_size) * params_.cluster_shape_.n()); current_group_info_.total_tiles = problem_blocks_m * problem_blocks_n; current_group_info_.problem_blocks_along_raster_order = params_.raster_order_ == RasterOrder::AlongN ? problem_blocks_n : problem_blocks_m; @@ -300,7 +323,7 @@ public: FastDivmodU64Pow2 const& divmod_cluster_shape_minor, FastDivmodU64 const& divmod_cta_shape_m, FastDivmodU64 const& divmod_cta_shape_n, - int32_t log_swizzle_size, + int32_t max_swizzle_size, RasterOrder raster_order) { uint8_t valid_tile = 1; @@ -308,7 +331,6 @@ public: // Use a warp to "speculatively" check if the work tile maps to the next 32 groups int lane_idx = canonical_lane_idx(); int total_problem_groups = problem_shapes.groups(); - if (linear_idx >= group_info.total_tiles + group_info.start_linear_idx) { group_info.group_idx += lane_idx; for ( ; ; group_info.group_idx += NumThreadsPerWarp) { @@ -327,8 +349,9 @@ public: ctas_along_m = divmod_cta_shape_m.divide(cute::shape<0>(cached_problem_shapes[0]) + divmod_cta_shape_m.divisor - 1); ctas_along_n = divmod_cta_shape_n.divide(cute::shape<1>(cached_problem_shapes[0]) + divmod_cta_shape_n.divisor - 1); } - auto problem_blocks_m = round_up(ctas_along_m, (1 << log_swizzle_size) * cluster_shape.m()); - auto problem_blocks_n = round_up(ctas_along_n, (1 << log_swizzle_size) * cluster_shape.n()); + group_info.log_swizzle_size = get_log_swizzle_size(ctas_along_m, ctas_along_n, max_swizzle_size); + auto problem_blocks_m = round_up(ctas_along_m, (1 << group_info.log_swizzle_size) * cluster_shape.m()); + auto problem_blocks_n = round_up(ctas_along_n, (1 << group_info.log_swizzle_size) * cluster_shape.n()); group_info.problem_blocks_along_raster_order = raster_order == RasterOrder::AlongN ? problem_blocks_n : problem_blocks_m; group_info.total_tiles = problem_blocks_m * problem_blocks_n; } @@ -354,6 +377,7 @@ public: group_info.start_linear_idx = __shfl_sync(0xffffffff, group_info.start_linear_idx, first_succeeding_thread); group_info.total_tiles = __shfl_sync(0xffffffff, group_info.total_tiles, first_succeeding_thread); group_info.problem_blocks_along_raster_order = __shfl_sync(0xffffffff, group_info.problem_blocks_along_raster_order, first_succeeding_thread); + group_info.log_swizzle_size = __shfl_sync(0xffffffff, group_info.log_swizzle_size, first_succeeding_thread); if (group_info.group_idx + lane_idx < total_problem_groups) { cached_problem_shapes[1] = problem_shapes.get_problem_shape(group_info.group_idx + lane_idx); } @@ -388,15 +412,15 @@ public: uint64_t cluster_idx_minor_div_swizzle, extra, offset; - offset = cluster_id & ((1 << log_swizzle_size) - 1); - extra = cluster_id >> log_swizzle_size; + offset = cluster_id & ((1 << group_info.log_swizzle_size) - 1); + extra = cluster_id >> group_info.log_swizzle_size; uint64_t curr_group_cluster_blk_major = divmod_cluster_shape_major.divide(group_info.problem_blocks_along_raster_order); cluster_idx_minor_div_swizzle = extra / curr_group_cluster_blk_major; cluster_idx_major = extra % curr_group_cluster_blk_major; - cluster_idx_minor = cluster_idx_minor_div_swizzle * (1 << log_swizzle_size) + offset; + cluster_idx_minor = cluster_idx_minor_div_swizzle * (1 << group_info.log_swizzle_size) + offset; auto minor_work_idx = static_cast(cluster_idx_minor * divmod_cluster_shape_minor.divisor + cluster_minor_offset); @@ -428,7 +452,7 @@ public: scheduler_params.divmod_cluster_shape_minor_, scheduler_params.divmod_cta_shape_m_, scheduler_params.divmod_cta_shape_n_, - scheduler_params.log_swizzle_size_, + scheduler_params.max_swizzle_size_, scheduler_params.raster_order_); } diff --git a/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp b/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp index 1e206aa6..c874d638 100644 --- a/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp +++ b/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp @@ -246,8 +246,9 @@ public: static bool can_implement(Arguments const& args) { - // Split count > 1 is only valid for heuristic and split-K decomposition modes - return (args.splits == 1 || + // Split count must be positive, and > 1 is only valid for heuristic and split-K decomposition modes + return args.splits >= 1 && + (args.splits == 1 || args.decomposition_mode == DecompositionMode::Heuristic || args.decomposition_mode == DecompositionMode::SplitK); } diff --git a/include/cutlass/gemm/kernel/tile_scheduler_params.h b/include/cutlass/gemm/kernel/tile_scheduler_params.h index bc09688d..09fda77e 100644 --- a/include/cutlass/gemm/kernel/tile_scheduler_params.h +++ b/include/cutlass/gemm/kernel/tile_scheduler_params.h @@ -1635,7 +1635,7 @@ struct PersistentTileSchedulerSm90GroupParams { uint64_t blocks_across_problem_ = 0; bool pre_processed_problem_shapes = true; - int32_t log_swizzle_size_ = 0; + int32_t max_swizzle_size_ = 0; RasterOrder raster_order_ = RasterOrder::AlongN; GroupProblemShape problem_shapes_; @@ -1658,10 +1658,8 @@ struct PersistentTileSchedulerSm90GroupParams { CUTLASS_UNUSED(hw_info); - // Round up to nearest multiple of swizzle_size along each mode - auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size); - auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m()); - auto problem_blocks_n = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n()); + auto problem_blocks_m = round_up(problem_blocks.x, cluster_shape.m()); + auto problem_blocks_n = round_up(problem_blocks.y, cluster_shape.n()); RasterOrder raster_order = get_rasterization_order( problem_blocks_m, @@ -1678,7 +1676,7 @@ struct PersistentTileSchedulerSm90GroupParams { blocks_across_problem_ = problem_blocks.x * problem_blocks.y * problem_blocks.z; pre_processed_problem_shapes = problem_shapes.is_host_problem_shape_available(); - log_swizzle_size_ = log_swizzle_size; + max_swizzle_size_ = max_swizzle_size; raster_order_ = raster_order; if (raster_order == RasterOrder::AlongN) { @@ -1727,10 +1725,8 @@ struct PersistentTileSchedulerSm90GroupParams { int const sm_count = hw_info.sm_count; int const max_active_clusters = hw_info.max_active_clusters; - // Round up to nearest multiple of swizzle_size along each mode - auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size); - auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m()); - auto problem_blocks_n = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n()); + auto problem_blocks_m = round_up(problem_blocks.x, cluster_shape.m()); + auto problem_blocks_n = round_up(problem_blocks.y, cluster_shape.n()); int problem_blocks_total = problem_blocks_m * problem_blocks_n * problem_blocks.z; @@ -1803,24 +1799,6 @@ struct PersistentTileSchedulerSm90GroupParams { return launch_grid; } - CUTLASS_HOST_DEVICE - static int32_t - get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) { - int min_cta_dim = platform::min(problem_ctas_m, problem_ctas_n); - if (max_swizzle_size >= 8 && min_cta_dim >= 6) { - return 3; - } - else if (max_swizzle_size >= 4 && min_cta_dim >= 3) { - return 2; - } - else if (max_swizzle_size >= 2 && min_cta_dim >= 2) { - return 1; - } - else { - return 0; - } - } - CUTLASS_HOST_DEVICE static RasterOrder get_rasterization_order( @@ -2496,10 +2474,8 @@ struct PersistentTileSchedulerSm100GroupParams { int const sm_count = hw_info.sm_count; int const max_active_clusters = hw_info.max_active_clusters; - // Round up to nearest multiple of swizzle_size along each mode - auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size); - auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m()); - auto problem_blocks_n = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n()); + auto problem_blocks_m = round_up(problem_blocks.x, cluster_shape.m()); + auto problem_blocks_n = round_up(problem_blocks.y, cluster_shape.n()); int problem_blocks_total = problem_blocks_m * problem_blocks_n * problem_blocks.z; @@ -2581,12 +2557,6 @@ struct PersistentTileSchedulerSm100GroupParams { return launch_grid; } - CUTLASS_HOST_DEVICE - static int32_t - get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) { - return UnderlyingSm90Params::get_log_swizzle_size(problem_ctas_m, problem_ctas_n, max_swizzle_size); - } - CUTLASS_HOST_DEVICE static RasterOrder get_rasterization_order( diff --git a/include/cutlass/platform/platform.h b/include/cutlass/platform/platform.h index ee58b054..46e1adf4 100644 --- a/include/cutlass/platform/platform.h +++ b/include/cutlass/platform/platform.h @@ -582,6 +582,7 @@ template struct alignment_of : std::alignment_of {}; #endif + #if CUDA_VERSION >= 11080 /* 16B specializations where 32-bit Win32 host compiler disagrees with device compiler */ template <> @@ -609,12 +610,7 @@ struct alignment_of { enum { value = 16 }; }; - -#if !defined(CUDA_VECTOR_TYPE_ALIGNMENT_16_32_ENABLED) -#define CUDA_VECTOR_TYPE_ALIGNMENT_16_32_ENABLED (__CUDACC_VER_MAJOR__ >= 13) -#endif - -#if (CUDA_VECTOR_TYPE_ALIGNMENT_16_32_ENABLED) +#if CUDA_VERSION >= 13000 template <> struct alignment_of { enum { value = 16 }; @@ -655,7 +651,9 @@ template <> struct alignment_of { enum { value = 32 }; }; + #else + template <> struct alignment_of { enum { value = 16 }; @@ -677,7 +675,7 @@ struct alignment_of { enum { value = 16 }; }; -#endif +#endif // CUDA_VERSION >= 13000 #endif // CUDA_VERSION >= 11080 // Specializations for volatile/const qualified types diff --git a/media/docs/pythonDSL/cute_dsl.rst b/media/docs/pythonDSL/cute_dsl.rst index a05c434e..5ead90f0 100644 --- a/media/docs/pythonDSL/cute_dsl.rst +++ b/media/docs/pythonDSL/cute_dsl.rst @@ -20,3 +20,4 @@ CuTe DSL Deprecation Policy Compile with TVM FFI Ahead-of-Time (AOT) Compilation + Talks and Presentations diff --git a/media/docs/pythonDSL/cute_dsl_general/debugging.rst b/media/docs/pythonDSL/cute_dsl_general/debugging.rst index 0e6319de..759086e3 100644 --- a/media/docs/pythonDSL/cute_dsl_general/debugging.rst +++ b/media/docs/pythonDSL/cute_dsl_general/debugging.rst @@ -112,9 +112,6 @@ For compiled kernels, the generated PTX/CUBIN/IR can be accessed programmaticall - ``__cubin__``: The generated CUBIN data of the compiled kernel. - ``__mlir__``: The generated IR code of the compiled kernel. -These attributes are populated only when the corresponding ``CUTE_DSL_KEEP_*`` environment variable is enabled; -otherwise they return ``None``. - .. code:: python compiled_foo = cute.compile(foo, ...) diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst index 67b4354f..4433f65a 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst @@ -236,4 +236,4 @@ For more information, see the section "Exporting Compiled Module" in :doc:`compi The primary distinction is that, when TVM FFI is enabled, |DSL| generates a dedicated wrapper function on top of the underlying CuTe ABI. This wrapper adheres to the calling conventions defined by TVM FFI. In contrast, the CuTe ABI entry function is specified directly in the generated header file, which affects how arguments must be provided. -For instance, with the TVM FFI wrapper function, users are able to pass in arguments such as ``torch.Tensor`` directly. However, when calling the CuTe ABI entry function, arguments should be provided as ``cute.Tensor`` types. +For instance, with the TVM FFI wrapper function, users are able to pass in arguments such as ``torch.Tensor`` directly. However, when calling the CuTe ABI entry function, arguments should be provided as ``cute.Tensor`` types. \ No newline at end of file diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst index 67180c4d..baeb6a7c 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst @@ -7,11 +7,24 @@ End-to-End Code Generation ========================== -1. Techniques for Turning Python into |IR| ------------------------------------------- +1. Hybrid DSL: Python Metaprogramming, Structured GPU Code +---------------------------------------------------------- -1.1 AST rewrite -^^^^^^^^^^^^^^^^ +|DSL| is a **hybrid DSL** that combines two compilation techniques: *AST rewrite* +and *tracing*. This combination gives you the best of both worlds: + +* **Program structure is preserved** — control flow (loops, branches) is + captured via AST rewrite, compiling to proper structured code instead of + flattened traces. +* **Python stays Python** — arithmetic and tensor operations are captured via + tracing, so dynamic shapes, metaprogramming, and Python's rich expression + language work naturally. + +To understand why this matters, let's look at each technique. + + +1.1 AST Rewrite +^^^^^^^^^^^^^^^ The function’s abstract-syntax tree is analysed **before** execution. Python control-flow (``for``/``while``, ``if``/``else``) and built-ins are converted to structured |IR| constructs. Computation inside each region is left untouched at this stage. @@ -47,11 +60,206 @@ trace that is lowered to |IR|. * Data-dependent control-flow freezes to a single execution path. -2. |DSL| Code-Generation Modes +1.3 The Hybrid Solution +^^^^^^^^^^^^^^^^^^^^^^^ + +As shown above, neither technique alone is sufficient—but together they +complement each other perfectly. + +**Why this works: GPU kernels are simple at runtime** + +High-performance GPU kernels are structurally simple at runtime: they avoid +deep call hierarchies, complex branching, and dynamic dispatch. However, +*authoring* such kernels benefits greatly from Python's abstractions—classes, +metaprogramming, and polymorphic patterns improve readability and +maintainability. + +The hybrid approach resolves this tension by evaluating Python abstractions at +compile time while emitting simple, optimized code for runtime execution. + +**How |DSL| divides the work:** + +1. **AST rewrite handles structure** — loops (``for``, ``while``) and branches + (``if``/``else``) are converted to structured |IR| *before* execution. + This solves tracing's control-flow problem. + +2. **Tracing handles arithmetic** — inside each structured region, the tracer + records tensor operations exactly as they execute. No need to model Python's + complex semantics—just run Python and record what happens. This solves AST + rewriting's complexity problem. + +The result: + +* Loops compile to real loops, not unrolled traces. +* All branches are preserved, even if not taken during tracing. +* Dynamic shapes, metaprogramming, and Python idioms work naturally. +* The rewriter only needs to understand control flow, not all of Python. + + +2. |DSL| Compilation Flow: Meta-Stage to Object-Stage +------------------------------------------------------ + +|DSL| bridges Python and GPU hardware through a three-stage pipeline. + +.. _fig-dsl-modes: + +.. figure:: dsl_modes.png + :width: 400 + :align: center + + *Left*: tracing mode records only the path that executed. + *Right*: preprocessor mode emits structured |IR| for every branch and loop + before tracing the arithmetic. + + + The default |DSL| compilation pipeline (mode 2): Python source flows through AST preprocessing + and interpreter-driven tracing to produce |IR|, which is then lowered and + compiled to device code. + +**Stage 1: Pre-Staging (Python AST)** + +Before any code executes, the AST preprocessor rewrites the decorated function. +It inserts *callbacks* around control-flow constructs—loops, branches, and +function boundaries—so that program structure is captured explicitly rather than +lost during execution. + +**Stage 2: Meta-Stage (Python Interpreter)** + +The rewritten function runs in the Python interpreter with proxy tensor +arguments. As execution proceeds: + +* Callbacks fire at control-flow boundaries, emitting structured |IR| (loops, + branches, etc.). +* Tensor operations are traced: each operator invocation records the + corresponding operation. +* Compile-time constants are *partially evaluated*—values known at JIT time + fold directly into the |IR|, enabling aggressive specialization. + +The result is a complete representation of the kernel, with both high-level +structure and low-level arithmetic intact. + +**Stage 3: Object-Stage (Compiler Backend)** + +The internal representation passes through a lowering pipeline: + +1. High-level operations are progressively lowered toward hardware-specific + representations. +2. Optimization passes (tiling, vectorization, memory promotion) reshape the + code for the target architecture. +3. The final code is translated to PTX/SASS (for NVIDIA GPUs) and assembled + into a device binary. + +At runtime, the compiled kernel is loaded and launched on the accelerator. + + +3. Meta-Programming vs Runtime: Two Worlds in One Function +---------------------------------------------------------- + +A key insight for understanding |DSL| is that **your Python code runs twice**, +in two very different contexts: + +1. **Meta-programming time (compilation)** — Python executes to *build* the + kernel. This happens on the host CPU when you call a ``@jit`` function. +2. **Runtime (execution)** — The compiled kernel runs on the GPU with actual + tensor data. + +This distinction determines what you can observe and when. + +``print()`` vs ``cute.printf()``: Meta-Stage vs Object-Stage Output +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +|DSL| provides two ways to print values, each operating at a different stage: + +* **Python's** ``print()`` — executes during the **meta-stage** (compilation). + Use it to inspect what the compiler sees. +* ``cute.printf()`` — compiles into the kernel and executes at **runtime** on + the GPU. Use it to observe actual tensor values during execution. + +The following examples demonstrate how the same ``result`` variable appears +differently depending on when and how you print it. + +**Example 1: Dynamic variables (both** ``a`` **and** ``b`` **are runtime values)** + +.. code-block:: python + + @cute.jit + def add_dynamicexpr(b: cutlass.Float32): + a = cutlass.Float32(2.0) + result = a + b + print("[meta-stage] result =", result) # runs at compile time + cute.printf("[object-stage] result = %f\n", result) # runs on GPU + + add_dynamicexpr(5.0) + +.. code-block:: text + + $> python myprogram.py + [meta-stage] result = + [object-stage] result = 7.000000 + +At meta-stage, ``result`` is a proxy—its value is unknown until the kernel runs. +At runtime, ``cute.printf()`` prints the actual GPU-computed value. + +**Example 2: Compile-time constants (both** ``a`` **and** ``b`` **are Constexpr)** + +.. code-block:: python + + @cute.jit + def add_constexpr(b: cutlass.Constexpr): + a = 2.0 + result = a + b + print("[meta-stage] result =", result) # runs at compile time + cute.printf("[object-stage] result = %f\n", result) # runs on GPU + + add_constexpr(5.0) + +.. code-block:: text + + $> python myprogram.py + [meta-stage] result = 7.0 + [object-stage] result = 7.000000 + +Both values are known at compile time, so Python evaluates ``2.0 + 5.0 = 7.0`` +during tracing. The constant is baked into the compiled kernel. + +**Example 3: Hybrid (** ``a`` **is dynamic,** ``b`` **is Constexpr)** + +.. code-block:: python + + @cute.jit + def add_hybrid(b: cutlass.Constexpr): + a = cutlass.Float32(2.0) + result = a + b + print("[meta-stage] result =", result) # runs at compile time + cute.printf("[object-stage] result = %f\n", result) # runs on GPU + + add_hybrid(5.0) + +.. code-block:: text + + $> python myprogram.py + [meta-stage] result = + [object-stage] result = 7.000000 + +The constant ``b = 5.0`` is folded in, but since ``a`` is dynamic, the result +remains a proxy at meta-stage. The GPU computes the final answer at runtime. + + +Practical Implications +^^^^^^^^^^^^^^^^^^^^^^ + +* **Use** ``print()`` **to debug your meta-program** — inspect shapes, strides, + tile sizes, and compile-time decisions. +* **Constexpr parameters enable specialization** — the compiler can generate + tighter code when values are known at JIT time. +* **Dynamic parameters preserve generality** — a single compiled kernel can + handle varying input sizes without recompilation. + +4. |DSL| Code-Generation Modes ------------------------------ -CuTe’s Python front-end combines the techniques above into **two mutually -exclusive modes**, selectable with the ``preprocessor`` flag of the +CuTe's Python front-end combines the techniques above into **two mutually +exclusive modes** (see :ref:`fig-dsl-modes`), selectable with the ``preprocessor`` flag of the ``@jit`` decorator: 1. Tracing mode ``@jit(preprocess=False)`` – tracing only. @@ -64,23 +272,3 @@ optimisation problems of pure tracing; tracing then fills in the arithmetic. This hybrid “preprocessor” pipeline is unique to |DSL| and was designed specifically to overcome the disadvantages identified above. -.. figure:: dsl_modes.png - :width: 400 - :align: center - - *Left*: tracing mode records only the path that executed. - *Right*: preprocessor mode emits structured |IR| for every branch and loop - before tracing the arithmetic. - - -Why Tracing-Only Is Insufficient for Control-Flow -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -* **Branch loss** – The untaken side of an ``if``/``else`` is never lowered. -* **Loop unrolling** – Loops are flattened to the iteration count observed, - destroying structure needed for parallel mapping and tiling. -* **Data-dependent paths** – Control-flow that depends on tensor values freezes - to a single execution path at trace time. - -The preprocessor mode fixes all of these by lowering control-flow first and delegating -only the arithmetic to the tracer. diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_compilation.png b/media/docs/pythonDSL/cute_dsl_general/dsl_compilation.png new file mode 100644 index 00000000..4689a1b2 Binary files /dev/null and b/media/docs/pythonDSL/cute_dsl_general/dsl_compilation.png differ diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst index af487563..73b7e236 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst @@ -5,20 +5,30 @@ Introduction -====================== +============ Overview -------- -|DSL| is a Python-based domain-specific language (DSL) designed for |DC| of numeric and GPU-oriented code. Its primary goals are: +|DSL| is a Python-based domain-specific language (DSL) designed for |DC| of +high-performance GPU kernels. It evolved from the C++ CUTLASS library and is +now available as a decorator-based DSL. -- **Consistent with CuTe C++**, allowing users to express GPU kernels with full control of the hardware. +Its primary goals are: + +- **Zero-cost abstraction**, DSL is a zero-cost abstraction thanks to Hybrid DSL approach. +- **Consistent with CuTe C++**, allowing users to express GPU kernels with full + control of the hardware. - **JIT compilation** for both host and GPU execution. -- `DLPack `_ **integration**, enabling seamless interop with frameworks (e.g., PyTorch, JAX). -- **JIT caching**, so that repeated calls to the same function benefit from cached |IR| modules. -- **Native types and type inference** to reduce boilerplate and improve performance. -- **Optional lower-level control**, offering direct access to GPU backends or specialized |IR| dialects. +- `DLPack `_ **integration**, enabling seamless + interop with frameworks (e.g., PyTorch, JAX). +- **JIT caching**, so that repeated calls to the same function benefit from + cached |IR| modules. +- **Native types and type inference** to reduce boilerplate and improve + performance. +- **Optional lower-level control**, offering direct access to GPU backends or + specialized |IR| dialects. Decorators ---------- diff --git a/media/docs/pythonDSL/cute_dsl_general/resources.rst b/media/docs/pythonDSL/cute_dsl_general/resources.rst new file mode 100644 index 00000000..0a63a87a --- /dev/null +++ b/media/docs/pythonDSL/cute_dsl_general/resources.rst @@ -0,0 +1,28 @@ +.. _talks_and_presentations: +.. |DSL| replace:: CuTe DSL + +Talks and Presentations +======================= + +This page collects talks, presentations, and other resources related to |DSL| +and CUTLASS Python infrastructure. + +Conference Talks +---------------- + +**CuTeDSL: CUTLASS Python DSL Infrastructure** — *LLVM 2025* + +An introduction to the |DSL| architecture, covering the hybrid AST-rewrite and +tracing approach, MLIR code generation, and integration with CUTLASS. + +* `LLVM Video `_ +* `Slides (PDF) `_ + +---- + +**Enable Tensor Core Programming in Python with CUTLASS 4.0** — *GTC 2025* + +Learn how to leverage Tensor Cores directly from Python using CUTLASS 4.0's +new DSL front-end, enabling rapid kernel development without writing CUDA C++. + +* `GTC Video `_ diff --git a/media/docs/pythonDSL/overview.rst b/media/docs/pythonDSL/overview.rst index c98bea59..fbd3abd8 100644 --- a/media/docs/pythonDSL/overview.rst +++ b/media/docs/pythonDSL/overview.rst @@ -105,4 +105,4 @@ You can: - Propose support for additional data types or kernel variants - Help prioritize roadmap features by upvoting GitHub issues -Thank you for helping shape the future of CUTLASS DSLs! +Thank you for helping shape the future of CUTLASS DSLs! \ No newline at end of file diff --git a/media/docs/pythonDSL/quick_start.rst b/media/docs/pythonDSL/quick_start.rst index cd6dfc2f..e97e21a7 100644 --- a/media/docs/pythonDSL/quick_start.rst +++ b/media/docs/pythonDSL/quick_start.rst @@ -8,6 +8,14 @@ The CUTLASS DSL 4.4 release currently supports **Linux** and **Python 3.10 - 3.1 Installation ----------------------- +Before installing the latest version, you need to uninstall any previous CUTLASS DSL Installation. + +.. code-block:: bash + + pip uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu13 -y + + + To ensure compatibility with the examples and code on `GitHub `_, use the `setup.sh `_ file from the corresponding commit in the repository. diff --git a/python/CuTeDSL/cutlass/__init__.py b/python/CuTeDSL/cutlass/__init__.py index 3df1c5c2..8836db84 100644 --- a/python/CuTeDSL/cutlass/__init__.py +++ b/python/CuTeDSL/cutlass/__init__.py @@ -14,6 +14,16 @@ from ._mlir._mlir_libs import _cutlass_ir _cutlass_ir.populate(_cutlass_ir) __version__ = "@CUTLASS_IR_WHEEL_RELEASE_VERSION@" +# Monkey patch CUDA version query function +from ._mlir._mlir_libs._cutlass_ir._base_dsl import ( + get_cuda_version as _get_cuda_version, +) +from .base_dsl import common as _common + +_common._get_cuda_version = _get_cuda_version + +# Import CUDA version from base_dsl +from .base_dsl.version_info import CUDA_VERSION from .cutlass_dsl import ( Constexpr, @@ -47,6 +57,7 @@ from .cutlass_dsl import ( extract_mlir_values, new_from_mlir_values, DSLCudaVersion, + target_version, ) from .cute.typing import * @@ -54,7 +65,6 @@ from .cute.typing import * # Utilities not belonging to CuTe from . import utils as utils from . import pipeline as pipeline -from .utils.version_info import CUDA_VERSION # Used as internal symbol from . import cutlass_dsl as _dsl @@ -67,5 +77,3 @@ cuda = _dsl.cuda_helpers # Jax Framework support from . import jax as jax - -CACHE_FILE = "compiled_cache.db" diff --git a/python/CuTeDSL/cutlass/base_dsl/__init__.py b/python/CuTeDSL/cutlass/base_dsl/__init__.py index 2b4927fa..567d4da7 100644 --- a/python/CuTeDSL/cutlass/base_dsl/__init__.py +++ b/python/CuTeDSL/cutlass/base_dsl/__init__.py @@ -24,4 +24,4 @@ from .utils.tree_utils import ( DSLTreeFlattenError, ) -from .common import DSLCudaVersion +from .common import DSLCudaVersion, target_version diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py index c434c982..cda5137d 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py @@ -52,6 +52,7 @@ class Executor: self._any_executor = None self._all_executor = None self._builtin_redirector = None + self._ifexp_dynamic = None def set_functions( self, @@ -64,6 +65,7 @@ class Executor: any_executor: Callable = None, all_executor: Callable = None, builtin_redirector: Callable = None, + ifexp_dynamic: Callable = None, ): self._is_dynamic_expression = is_dynamic_expression self._loop_execute_range_dynamic = loop_execute_range_dynamic @@ -73,6 +75,7 @@ class Executor: self._any_executor = any_executor self._all_executor = all_executor self._builtin_redirector = builtin_redirector + self._ifexp_dynamic = ifexp_dynamic @staticmethod def convert_to_list(x): @@ -173,6 +176,16 @@ class Executor: write_args_names, ) + def ifexp_execute( + self, + pred, + block_args: tuple, + then_block: Callable, + else_block: Callable, + ): + assert self._ifexp_dynamic, "Functions must be set before execution." + return self._ifexp_dynamic(pred, block_args, then_block, else_block) + # ============================================================================= # Decorator @@ -293,6 +306,19 @@ def if_executor( ) +def ifExp_executor( + *, + pred, + block_args: tuple, + then_block: Callable, + else_block: Callable, +): + if not executor._is_dynamic_expression(pred): + return then_block(*block_args) if pred else else_block(*block_args) + else: + return executor.ifexp_execute(pred, block_args, then_block, else_block) + + # ============================================================================= # Range # ============================================================================= @@ -552,18 +578,19 @@ def cf_symbol_check(symbol): name = symbol.__name__ self_module = _get_self_module() if inspect.ismodule(symbol): - name = "range" - if not self_module.__name__.startswith(symbol.__name__): + if not self_module.__name__.startswith(name): failed = True else: owning_module = inspect.getmodule(symbol) - if owning_module != self_module: + root_module = owning_module.__name__.split(".")[0] + self_root_module = self_module.__name__.split(".")[0] + if root_module != self_root_module: failed = True if failed: raise DSLRuntimeError( - f"Incorrect {symbol.__name__} is used.", - suggestion=f"Please avoid overriding `{symbol.__name__}` from DSL package.", + f"Incorrect `{name}` is used.", + suggestion=f"Please avoid overriding `{name}` from DSL package.", ) diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py index 4afc61b6..5eea3018 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py @@ -45,6 +45,7 @@ from typing import List, Set, Dict, Any, Callable, Optional from types import ModuleType from collections import OrderedDict from copy import deepcopy +from itertools import chain from .common import * from .utils.logger import log @@ -140,25 +141,150 @@ class ScopeManager: """ scopes: List[Set[str]] + callables: List[Set[str]] @classmethod def create(cls) -> "ScopeManager": - return cls([]) + return cls([], []) def add_to_scope(self, name: str) -> None: if name == "_": return self.scopes[-1].add(name) + def add_to_callables(self, name: str) -> None: + if not self.callables: + return + self.callables[-1].add(name) + def get_active_symbols(self) -> List[Set[str]]: return self.scopes.copy() - def __enter__(self) -> "ScopeManager": + def get_active_callables(self) -> List[Set[str]]: + return self.callables.copy() + + @contextlib.contextmanager + def enter_local_scope(self): + """ + Context manager for entering a new local variable and callable scope. + + This is conceptually Python's local scope, such as within a function or class definition. + + Use this in a ``with`` statement to temporarily push a new, empty set for both variable and callable + tracking onto the respective ScopeManager stacks. These sets accumulate any new symbols + introduced within the local context. When the context manager exits, the local sets are popped, + restoring the previous scope state. + + **Example** + .. code-block:: python + + with scope_manager.enter_local_scope(): + # Symbols defined here are local to this scope + ... + + :yields: None + """ self.scopes.append(set()) + self.callables.append(set()) + yield + self.scopes.pop() + self.callables.pop() + + @contextlib.contextmanager + def enter_control_flow_scope(self): + """ + Context manager for entering a new dynamic control-flow scope. + + This scope rule diverge from Python's local scope, variables defined here are discarded after exiting the block, but callables are kept in parent scope. + + This context manager pushes a new, empty variable scope onto the stack for the + duration of a control-flow block (such as within loops or if/else blocks). Variables + introduced inside this block are tracked separately and discarded after exiting the block. + Callable symbol scopes are not affected. + + :yields: None + + **Example** + .. code-block:: python + + with scope_manager.enter_control_flow_scope(): + # Variables defined here are local to this control-flow scope + ... + """ + self.scopes.append(set()) + yield + self.scopes.pop() + + +class Region: + """ + Context manager for handling regions during AST transformations. + + This class is used to manage region-scoped state during DSL preprocessing. + It is responsible for tracking and collecting new statements generated while + visiting and transforming regions, such as the bodies of AST nodes representing + constructs like loops or conditional blocks. + + Upon entering a region (using a ``with`` statement), the region is pushed onto + the session's ``region_stack``, and prepares a place for new statements to be collected. + On exit, the region is popped from the stack and any temporary state is cleaned up. + + Parameters + ---------- + session_data : SessionData + The shared session context for the AST preprocessor, which holds the region stack. + owning_node : Optional[ast.stmt], default=None + If provided, the AST statement node that owns this region; new statements will be append to _new_value of this new node. + new_value : Optional[list[ast.stmt]], default=None + If provided, a list for collecting new statements for this region. + + Methods + ------- + __enter__() + Enter the region context, mutate state as needed. + __exit__(exc_type, exc_value, traceback) + Exit the context, clean up state. + append_new_stmts(stmts) + Append new AST statements to the region's collection. + """ + + def __init__( + self, + session_data: "SessionData", + *, + owning_node: ast.stmt = None, + new_value: list[ast.stmt] = None, + ): + self.session_data = session_data + self.owning_node = owning_node + self.new_value = new_value + + def __enter__(self): + if self.new_value is not None or isinstance(self.owning_node, ast.stmt): + self.session_data.region_stack.append(self) + if self.owning_node is not None: + self.owning_node._new_value = [] return self - def __exit__(self, exc_type, exc_value, traceback) -> None: - self.scopes.pop() + def __exit__(self, exc_type, exc_value, traceback): + if self.new_value is not None or isinstance(self.owning_node, ast.stmt): + self.session_data.region_stack.pop() + if self.owning_node is not None: + delattr(self.owning_node, "_new_value") + + def append_new_stmts(self, stmts: list[ast.stmt]): + """ + Append a list of statements to the region's collection. + + Parameters + ---------- + stmts : list[ast.stmt] + The AST statements to append to this region. + """ + if self.owning_node is not None: + self.owning_node._new_value.extend(stmts) + else: + self.new_value.extend(stmts) @dataclass @@ -173,10 +299,25 @@ class SessionData: function_name: str = "" class_name: Optional[str] = None file_name: str = "" - function_depth: int = 0 - local_closures: set[str] = field(default_factory=set) function_globals: Optional[dict[str, Any]] = None import_top_module: bool = False + region_stack: list[Region] = field(default_factory=list) + generator_targets: list[str] = field(default_factory=list) + lambda_args: list[str] = field(default_factory=list) + + @contextlib.contextmanager + def set_current_class_name(self, class_name: str): + old_class_name = self.class_name + self.class_name = class_name + yield + self.class_name = old_class_name + + @contextlib.contextmanager + def set_current_function_name(self, function_name: str): + old_function_name = self.function_name + self.function_name = function_name + yield + self.function_name = old_function_name def _create_module_attribute( @@ -225,54 +366,6 @@ def _create_module_attribute( set_location(node, lineno, col_offset) return node - -class DSLPreprocessorSession: - """Context manager for managing a DSL preprocessor session. - - This context manager is used to ensure that each preprocessing operation - (typically a transformation of a Python AST via the DSL preprocessor) - is performed within a well-defined session. When entering the context, - it initializes session-specific resources or state by calling - `_start_session()` on the provided DSL object. Upon exit, it performs - appropriate cleanup by calling `_end_session()`. - - Example usage:: - - with DSLPreprocessorSession(dsl_object): - # perform AST transformations or other preprocessing actions - - :param dsl_object: An instance of a DSL object that implements - `_start_session()` and `_end_session()` methods to manage the - session state - :type dsl_object: DSLPreprocessor - """ - - def __init__(self, dsl_object): - self.dsl_object = dsl_object - - def __enter__(self): - """Starts the DSL preprocessor session. - - :return: The DSL object for use within the context - :rtype: DSLPreprocessor - """ - self.dsl_object._start_session() - # Let `with preprocessor.get_session() as p:` keep using `p` as the preprocessor. - return self.dsl_object - - def __exit__(self, exc_type, exc_value, traceback): - """Ends the DSL preprocessor session. - - :param exc_type: The exception type if an exception was raised in the context - :type exc_type: type, optional - :param exc_value: The exception value if an exception was raised in the context - :type exc_value: Exception, optional - :param traceback: The traceback if an exception was raised in the context - :type traceback: traceback, optional - """ - self.dsl_object._end_session() - - class DSLPreprocessor(ast.NodeTransformer): """ A preprocessor for transforming Python ASTs. It supports: @@ -286,15 +379,52 @@ class DSLPreprocessor(ast.NodeTransformer): DECORATOR_IF_STATEMENT = "if_selector" DECORATOR_WHILE_STATEMENT = "while_selector" IF_EXECUTOR = "if_executor" + IFEXP_EXECUTOR = "ifExp_executor" WHILE_EXECUTOR = "while_executor" ASSERT_EXECUTOR = "assert_executor" BOOL_CAST = "bool_cast" IMPLICIT_DOWNCAST_NUMERIC_TYPE = "implicitDowncastNumericType" SUPPORTED_FOR_RANGE_STATEMENTS = {"range", "range_dynamic", "range_constexpr"} + CONST_EXPR_NAME = {"const_expr", "target_version"} COMPARE_EXECUTOR = "compare_executor" ANY_EXECUTOR = "any_executor" ALL_EXECUTOR = "all_executor" + def generic_visit(self, node): + """ + Copy of :meth:`ast.NodeTransformer.generic_visit` with support for inserting statements during expression visits. + + This version provides the same recursive traversal and transformation as the standard + ``generic_visit``, but extends it to allow statement insertion when visiting expressions. + This is particularly useful for DSL AST processing that needs to emit new statements within + regions associated with expression nodes (e.g., using the ``Region`` context manager). + + :param node: The AST node to process. + :type node: ast.AST + :return: The transformed AST node. + :rtype: ast.AST + """ + for field, old_value in ast.iter_fields(node): + if isinstance(old_value, list): + with Region(self.session_data, owning_node=node): + for value in old_value: + if isinstance(value, ast.AST): + value = self.visit(value) + if value is None: + continue + elif not isinstance(value, ast.AST): + node._new_value.extend(value) + continue + node._new_value.append(value) + old_value[:] = node._new_value + elif isinstance(old_value, ast.AST): + new_node = self.visit(old_value) + if new_node is None: + delattr(node, field) + else: + setattr(node, field, new_node) + return node + def __init__(self, client_module_name): super().__init__() # Persistent state @@ -303,29 +433,14 @@ class DSLPreprocessor(ast.NodeTransformer): self.module_cache = {} self._session_data = None - def _start_session(self): - """ - Starts a new preprocessing session by initializing session data. - - This method sets up a fresh SessionData instance for use during - AST transformations. It must be called before performing any - preprocessing actions that require access to context-specific - information during the transformation of a function's AST. - """ - self._session_data = SessionData() - - def _end_session(self): - """ - Ends the current preprocessing session and clears session data. - - This method resets the session-specific data, marking the end of - a preprocessing context. It should be called after all necessary - AST processing is complete to ensure no stale context remains. - """ - self._session_data = None + @contextlib.contextmanager def get_session(self): - return DSLPreprocessorSession(dsl_object=self) + try: + self._session_data = SessionData() + yield self + finally: + self._session_data = None @property def session_data(self): @@ -733,10 +848,13 @@ class DSLPreprocessor(ast.NodeTransformer): if isinstance(node.test, ast.Call): func = node.test.func - if isinstance(func, ast.Attribute) and func.attr == "const_expr": + if ( + isinstance(func, ast.Attribute) + and func.attr in self.CONST_EXPR_NAME + ): return True - elif isinstance(func, ast.Name) and func.id == "const_expr": + elif isinstance(func, ast.Name) and func.id in self.CONST_EXPR_NAME: return True return False @@ -775,7 +893,10 @@ class DSLPreprocessor(ast.NodeTransformer): return unified_tree def analyze_region_variables( - self, node: Union[ast.For, ast.If, ast.While], active_symbols: List[Set[str]] + self, + node: Union[ast.For, ast.If, ast.While], + active_symbols: List[Set[str]], + active_callables: List[Set[str]], ): """ Analyze variables in different code regions to identify read-only, write-only, @@ -785,8 +906,7 @@ class DSLPreprocessor(ast.NodeTransformer): # we need orderedset to keep the insertion order the same. otherwise generated IR is different each time write_args = OrderedSet() invoked_args = OrderedSet() - local_closure = self.session_data.local_closures - called_closures = OrderedSet() + called_functions = OrderedSet() class RegionAnalyzer(ast.NodeVisitor): force_store = False @@ -853,8 +973,7 @@ class DSLPreprocessor(ast.NodeTransformer): if isinstance(node.func, ast.Name): func_name = node.func.id - if func_name in local_closure: - called_closures.add(func_name) + called_functions.add(func_name) # Classes are mutable by default. Mark them as write. If they are # dataclass(frozen=True), treat them as read in runtime. @@ -878,8 +997,8 @@ class DSLPreprocessor(ast.NodeTransformer): write_args = list(write_args.intersections(active_symbols)) invoked_args = list(invoked_args.intersections(active_symbols)) - - return write_args + invoked_args, len(write_args), called_closures + called_functions = list(called_functions.intersections(active_callables)) + return write_args + invoked_args, len(write_args), called_functions def extract_range_args(self, iter_node): args = iter_node.args @@ -958,12 +1077,15 @@ class DSLPreprocessor(ast.NodeTransformer): # Create the loop body transformed_body = [] - for stmt in node.body: - transformed_stmt = self.visit(stmt) # Recursively visit inner statements - if isinstance(transformed_stmt, list): - transformed_body.extend(transformed_stmt) - else: - transformed_body.append(transformed_stmt) + with Region(self.session_data, new_value=transformed_body): + for stmt in node.body: + transformed_stmt = self.visit( + stmt + ) # Recursively visit inner statements + if isinstance(transformed_stmt, list): + transformed_body.extend(transformed_stmt) + else: + transformed_body.append(transformed_stmt) # Handle the return for a single iterated argument correctly if len(write_args) == 0: @@ -1186,8 +1308,9 @@ class DSLPreprocessor(ast.NodeTransformer): return node active_symbols = self.session_data.scope_manager.get_active_symbols() + active_callables = self.session_data.scope_manager.get_active_callables() - with self.session_data.scope_manager: + with self.session_data.scope_manager.enter_control_flow_scope(): if isinstance(node.target, ast.Name): self.session_data.scope_manager.add_to_scope(node.target.id) @@ -1210,7 +1333,9 @@ class DSLPreprocessor(ast.NodeTransformer): # Get toplevel module check_call = self._insert_cf_symbol_check(node.iter.func.value) - new_for_node = self.transform_for_loop(node, active_symbols) + new_for_node = self.transform_for_loop( + node, active_symbols, active_callables + ) if check_call is not None: new_for_node = [check_call] + new_for_node @@ -1234,7 +1359,6 @@ class DSLPreprocessor(ast.NodeTransformer): ), location, ) - self.generic_visit(node) return node def _handle_negative_step(self, node, start_expr, stop_expr, step_expr): @@ -1309,11 +1433,12 @@ class DSLPreprocessor(ast.NodeTransformer): location=node, ) - extra_exprs.append(isNegative) - extra_exprs.append(start) - extra_exprs.append(stop) - extra_exprs.append(step) - extra_exprs.append(offset) + with Region(self.session_data, new_value=extra_exprs): + extra_exprs.append(self.generic_visit(isNegative)) + extra_exprs.append(self.generic_visit(start)) + extra_exprs.append(self.generic_visit(stop)) + extra_exprs.append(self.generic_visit(step)) + extra_exprs.append(self.generic_visit(offset)) # Add this to begining of loop body # for i in range(start, stop, step): @@ -1360,7 +1485,7 @@ class DSLPreprocessor(ast.NodeTransformer): ) ) - def transform_for_loop(self, node, active_symbols): + def transform_for_loop(self, node, active_symbols, active_callables): # Check for early exit and raise exception self.check_early_exit(node, "for") if node.orelse: @@ -1409,7 +1534,7 @@ class DSLPreprocessor(ast.NodeTransformer): prefetch_stages = self.extract_prefetch_stages_args(node.iter) vectorize = self.extract_vectorize_args(node.iter) write_args, full_write_args_count, called_closures = ( - self.analyze_region_variables(node, active_symbols) + self.analyze_region_variables(node, active_symbols, active_callables) ) if has_step and self.client_module_name[0] == "cutlass": @@ -1678,10 +1803,8 @@ class DSLPreprocessor(ast.NodeTransformer): return node def visit_ClassDef(self, node): - self.session_data.class_name = node.name - self.generic_visit(node) - self.session_data.class_name = None - return node + with self.session_data.set_current_class_name(node.name): + return self.generic_visit(node) def _visit_target(self, target): if isinstance(target, ast.Name): @@ -1798,13 +1921,13 @@ class DSLPreprocessor(ast.NodeTransformer): return new_decorator_list def visit_FunctionDef(self, node): - with self.session_data.scope_manager: - self.session_data.function_counter += 1 - self.session_data.function_name = node.name - if self.session_data.function_depth > 0: - self.session_data.local_closures.add(node.name) + # Add self to active symbols of parent scope + self.session_data.scope_manager.add_to_callables(node.name) - self.session_data.function_depth += 1 + with self.session_data.scope_manager.enter_local_scope(), self.session_data.set_current_function_name( + node.name + ): + self.session_data.function_counter += 1 # Add function name and arguments self.session_data.scope_manager.add_to_scope(node.name) @@ -1822,7 +1945,6 @@ class DSLPreprocessor(ast.NodeTransformer): self.generic_visit(node) - self.session_data.function_depth -= 1 # Remove .jit and .kernel decorators node.decorator_list = self.remove_dsl_decorator(node.decorator_list) @@ -1832,13 +1954,10 @@ class DSLPreprocessor(ast.NodeTransformer): return node def visit_With(self, node): - with self.session_data.scope_manager: - for item in node.items: - if isinstance(item.optional_vars, ast.Name): - self.session_data.scope_manager.add_to_scope(item.optional_vars.id) - self.generic_visit(node) - - return node + for item in node.items: + if isinstance(item.optional_vars, ast.Name): + self.session_data.scope_manager.add_to_scope(item.optional_vars.id) + return self.generic_visit(node) def visit_While(self, node): # Constexpr doesn't get preprocessed @@ -1848,12 +1967,14 @@ class DSLPreprocessor(ast.NodeTransformer): return [check, node] active_symbols = self.session_data.scope_manager.get_active_symbols() - with self.session_data.scope_manager: + active_callables = self.session_data.scope_manager.get_active_callables() + + with self.session_data.scope_manager.enter_control_flow_scope(): # Check for early exit and raise exception self.check_early_exit(node, "while") write_args, full_write_args_count, called_closures = ( - self.analyze_region_variables(node, active_symbols) + self.analyze_region_variables(node, active_symbols, active_callables) ) exprs = [] if called_closures: @@ -1869,18 +1990,6 @@ class DSLPreprocessor(ast.NodeTransformer): return exprs + [func_def] + assign - def visit_Try(self, node): - with self.session_data.scope_manager: - self.generic_visit(node) - return node - - def visit_ExceptHandler(self, node): - with self.session_data.scope_manager: - if node.name: # Exception variable - self.session_data.scope_manager.add_to_scope(node.name) - self.generic_visit(node) - return node - def create_cf_call(self, func_name, yield_args, node): """Creates the assignment statement for the if function call""" if not yield_args: @@ -1928,42 +2037,164 @@ class DSLPreprocessor(ast.NodeTransformer): else: return [ast.copy_location(assign, node)] + def _visit_Comprehension(self, node, ele_visitor): + node.generators = [self.visit(generator) for generator in node.generators] + + targets = [] + + class NameCollector(ast.NodeVisitor): + def visit_Name(self, node): + if isinstance(node.ctx, ast.Store): + targets.append(node.id) + + # Collect generator targets + collector = NameCollector() + [collector.visit(generator) for generator in node.generators] + + self.session_data.generator_targets = targets + + ele_visitor(node) + + self.session_data.generator_targets = [] + return node + + def visit_DictComp(self, node): + def key_value_visitor(n): + n.key = self.visit(n.key) + n.value = self.visit(n.value) + + return self._visit_Comprehension(node, key_value_visitor) + + def visit_Lambda(self, node): + current_lambda_args = len(self.session_data.lambda_args) + for arg in node.args.args: + self.session_data.lambda_args.append(arg.arg) + + node.body = self.visit(node.body) + + self.session_data.lambda_args = self.session_data.lambda_args[ + :current_lambda_args + ] + + return node + + def visit_ListComp(self, node): + return self._visit_Comprehension( + node, lambda n: setattr(n, "elt", self.visit(n.elt)) + ) + + def visit_GeneratorExp(self, node): + return self._visit_Comprehension( + node, lambda n: setattr(n, "elt", self.visit(n.elt)) + ) + + def visit_SetComp(self, node): + return self._visit_Comprehension( + node, lambda n: setattr(n, "elt", self.visit(n.elt)) + ) + def visit_IfExp(self, node): """ - Visits an inline if-else expression (ternary operator). - This is the Python equivalent of `x if condition else y`. + Transforms an inline if-else (ternary) expression into runtime-dispatched + control flow using synthesized function definitions for each branch. + + This converts an expression of the form ``x if cond else y`` into two local + function blocks (for the ``then`` and ``else`` branches), inserts those blocks + just before the current statement, and produces a call to the conditional executor. + + This lets the DSL infrastructure analyze and dispatch dynamic inline conditionals + in a uniform way at runtime. + + Parameters + ---------- + node : ast.IfExp + The AST node representing the inline if-else expression. + + Returns + ------- + ast.Call + An AST node that calls the conditional expression executor, referencing + the synthesized blocks and the predicate. """ - self.generic_visit(node) - # Emit - # node if type(pred) == bool else select_(pred, body, orelse) - # so if pred is a python bool, use python to short-circuit and avoid emit arith.select - self.session_data.import_top_module = True - return ast.copy_location( - ast.IfExp( - test=ast.Compare( - left=ast.Call( - func=ast.Name(id="type", ctx=ast.Load()), - args=[node.test], - keywords=[], - ), - ops=[ast.Eq()], - comparators=[ast.Name(id="bool", ctx=ast.Load())], - ), - body=node, # Original ternary expression - orelse=ast.Call( - func=_create_module_attribute( - "select_", use_base_dsl=False, submodule_name=None - ), - args=[ - node.test, - node.body, - node.orelse, - ], - keywords=[], - ), + # Create unique names for the then and else branch function blocks + then_block_name = f"ifexp_then_block_{self.session_data.counter}" + else_block_name = f"ifexp_else_block_{self.session_data.counter}" + self.session_data.counter += 1 + + # Define the then-block function, with no arguments and returning the visited body + then_block_def = ast.FunctionDef( + name=then_block_name, + args=ast.arguments( + posonlyargs=[], + args=[ + ast.arg(arg=target, annotation=None) + for target in chain( + self.session_data.generator_targets, + self.session_data.lambda_args, + ) + ], + kwonlyargs=[], + kw_defaults=[], + defaults=[], ), - node, + body=[ast.Return(value=self.visit(node.body))], + decorator_list=[], ) + # Define the else-block function, with no arguments and returning the visited orelse + else_block_def = ast.FunctionDef( + name=else_block_name, + args=ast.arguments( + posonlyargs=[], + args=[ + ast.arg(arg=target, annotation=None) + for target in chain( + self.session_data.generator_targets, + self.session_data.lambda_args, + ) + ], + kwonlyargs=[], + kw_defaults=[], + defaults=[], + ), + body=[ast.Return(value=self.visit(node.orelse))], + decorator_list=[], + ) + + # Insert the block definitions into the most recent (innermost) region before the statement + self.session_data.region_stack[-1].append_new_stmts( + [then_block_def, else_block_def] + ) + + # Create the executor call node, wiring up the predicate and newly synthesized blocks + executor_call = ast.Call( + func=_create_module_attribute(self.IFEXP_EXECUTOR), + args=[], + keywords=[ + ast.keyword(arg="pred", value=self.visit(node.test)), + ast.keyword( + arg="block_args", + value=ast.Tuple( + elts=[ + ast.Name(id=name, ctx=ast.Load()) + for name in chain( + self.session_data.generator_targets, + self.session_data.lambda_args, + ) + ], + ctx=ast.Load(), + ), + ), + ast.keyword( + arg="then_block", value=ast.Name(id=then_block_name, ctx=ast.Load()) + ), + ast.keyword( + arg="else_block", value=ast.Name(id=else_block_name, ctx=ast.Load()) + ), + ], + ) + + # Return the transformed executor call node at the original location in the AST + return ast.copy_location(executor_call, node) cmpops = { "Eq": "==", @@ -2016,12 +2247,14 @@ class DSLPreprocessor(ast.NodeTransformer): return [check, node] active_symbols = self.session_data.scope_manager.get_active_symbols() - with self.session_data.scope_manager: + active_callables = self.session_data.scope_manager.get_active_callables() + + with self.session_data.scope_manager.enter_control_flow_scope(): # Check for early exit and raise exception self.check_early_exit(node, "if") yield_args, full_write_args_count, called_closures = ( - self.analyze_region_variables(node, active_symbols) + self.analyze_region_variables(node, active_symbols, active_callables) ) exprs = [] if called_closures: @@ -2060,12 +2293,18 @@ class DSLPreprocessor(ast.NodeTransformer): func_args_then_else = [ast.arg(arg=var, annotation=None) for var in write_args] then_body = [] - for stmt in node.body: - transformed_stmt = self.visit(stmt) # Recursively visit inner statements - if isinstance(transformed_stmt, list): - then_body.extend(transformed_stmt) - else: - then_body.append(transformed_stmt) + with ( + Region(self.session_data, new_value=then_body), + self.session_data.scope_manager.enter_control_flow_scope(), + ): + for stmt in node.body: + transformed_stmt = self.visit( + stmt + ) # Recursively visit inner statements + if isinstance(transformed_stmt, list): + then_body.extend(transformed_stmt) + else: + then_body.append(transformed_stmt) # Create common return list for all blocks return_list = ast.List( @@ -2210,14 +2449,18 @@ class DSLPreprocessor(ast.NodeTransformer): ) else: else_body = [] - for stmt in node.orelse: - transformed_stmt = self.visit( - stmt - ) # Recursively visit inner statements - if isinstance(transformed_stmt, list): - else_body.extend(transformed_stmt) - else: - else_body.append(transformed_stmt) + with ( + Region(self.session_data, new_value=else_body), + self.session_data.scope_manager.enter_control_flow_scope(), + ): + for stmt in node.orelse: + transformed_stmt = self.visit( + stmt + ) # Recursively visit inner statements + if isinstance(transformed_stmt, list): + else_body.extend(transformed_stmt) + else: + else_body.append(transformed_stmt) # Regular else block else_block = ast.FunctionDef( @@ -2302,7 +2545,6 @@ class DSLPreprocessor(ast.NodeTransformer): cond, write_args = while_before_block(write_args) return write_args """ - test_expr = self.visit(node.test) # Section: decorator construction decorator_keywords = [ @@ -2343,11 +2585,15 @@ class DSLPreprocessor(ast.NodeTransformer): ) # Section: while_before_block FunctionDef, which contains condition + while_before_stmts = [] + with Region(self.session_data, new_value=while_before_stmts): + test_expr = ast.copy_location(self.visit(node.test), node.test) + while_before_return_list = ast.List( elts=[test_expr, yield_args_ast_name_list], ctx=ast.Load(), ) - while_before_stmts = [ast.Return(value=while_before_return_list)] + while_before_stmts.append(ast.Return(value=while_before_return_list)) while_before_block = ast.copy_location( ast.FunctionDef( name=while_before_block_name, @@ -2360,12 +2606,15 @@ class DSLPreprocessor(ast.NodeTransformer): # Section: while_after_block FunctionDef, which contains loop body while_after_stmts = [] - for stmt in node.body: - transformed_stmt = self.visit(stmt) # Recursively visit inner statements - if isinstance(transformed_stmt, list): - while_after_stmts.extend(transformed_stmt) - else: - while_after_stmts.append(transformed_stmt) + with Region(self.session_data, new_value=while_after_stmts): + for stmt in node.body: + transformed_stmt = self.visit( + stmt + ) # Recursively visit inner statements + if isinstance(transformed_stmt, list): + while_after_stmts.extend(transformed_stmt) + else: + while_after_stmts.append(transformed_stmt) while_after_stmts.append(ast.Return(value=yield_args_ast_name_list)) while_after_block = ast.copy_location( diff --git a/python/CuTeDSL/cutlass/base_dsl/common.py b/python/CuTeDSL/cutlass/base_dsl/common.py index 202288ad..be85fd99 100644 --- a/python/CuTeDSL/cutlass/base_dsl/common.py +++ b/python/CuTeDSL/cutlass/base_dsl/common.py @@ -10,7 +10,9 @@ # is strictly prohibited. import os -from typing import Any, Dict, Iterable, Optional, Union, Sequence +from typing import Any, Dict, Optional, Union +from functools import total_ordering +from dataclasses import dataclass """ This module provides a Exception classes DSL class for any Dialect. @@ -325,29 +327,143 @@ This error typically occurs when: ) +def _get_cuda_version() -> str: + # Client of this module should implement this function + """ + Placeholder for CUDA version query. + + This function should be implemented by the client of this module. + When implemented, it must return the CUDA version as a string, e.g. "12.2". + + Raises: + NotImplementedError: Always, unless overridden by the package initializer or client. + """ + raise NotImplementedError("_get_cuda_version is not implemented") + + +@total_ordering +@dataclass(frozen=True) class DSLCudaVersion: """ Class to represent the CUDA version used to build the DSL. """ - def __init__(self, version: str): - self.version_tuple = tuple(int(part) for part in version.split(".")) + major: int + minor: int - def __str__(self): - return f"{self.major}.{self.minor}" + def __init__(self, version: str): + parts = version.split(".") + object.__setattr__(self, "major", int(parts[0])) + object.__setattr__(self, "minor", int(parts[1])) def __eq__(self, other): - if isinstance(other, DSLCudaVersion): - return self.version_tuple == other.version_tuple - elif isinstance(other, str): - return self == DSLCudaVersion(other) - else: - return False + return self.major == other.major and self.minor == other.minor - @property - def major(self): - return self.version_tuple[0] + def __lt__(self, other): + return [self.major, self.minor] < [other.major, other.minor] - @property - def minor(self): - return self.version_tuple[1] + +def _coerce_to_cuda_version( + value: Optional[Union[DSLCudaVersion, str]], param_name: str +) -> Optional[DSLCudaVersion]: + """ + Coerce a value to DSLCudaVersion. + + :param value: The value to coerce (DSLCudaVersion, str, or None). + :param param_name: The parameter name for error messages. + :returns: DSLCudaVersion or None if value is None. + :raises DSLRuntimeError: If value is not a supported type. + """ + if value is None: + return None + if isinstance(value, DSLCudaVersion): + return value + if isinstance(value, str): + return DSLCudaVersion(value) + raise DSLRuntimeError( + f"{param_name} must be a DSLCudaVersion or str, got {type(value).__name__}" + ) + + +def target_version( + exact_version: Optional[Union[DSLCudaVersion, str]] = None, + *, + min_version: Optional[Union[DSLCudaVersion, str]] = None, + max_version: Optional[Union[DSLCudaVersion, str]] = None, +) -> bool: + """ + Check if the current CUDA version used to build the DSL matches an exact version + or falls within specified bounds at compile-time. + + Only one of ``exact_version`` *or* ``min_version``/``max_version`` may be specified. + At least one must be provided. + + :param exact_version: The required CUDA version (e.g., "12.3"). + :type exact_version: Optional[Union[DSLCudaVersion, str]] + :param min_version: The minimum CUDA version required (inclusive, e.g., "12.0"). + :type min_version: Optional[Union[DSLCudaVersion, str]] + :param max_version: The maximum CUDA version allowed (inclusive, e.g., "13.2"). + :type max_version: Optional[Union[DSLCudaVersion, str]] + + :returns: ``True`` if the CUDA version matches the requirement(s) specified. + :rtype: bool + + :raises DSLRuntimeError: + - If neither an ``exact_version`` nor version range is given. + - If both an exact version and a range are provided. + - If ``min_version`` > ``max_version``. + - If any version parameter is not a DSLCudaVersion or str. + + **Examples** + + .. code-block:: python + + target_version(exact_version="12.3") + # True if CUDA_VERSION == 12.3 + + target_version(min_version="12.0") + # True if CUDA_VERSION >= 12.0 + + target_version(max_version="13.2") + # True if CUDA_VERSION <= 13.2 + + target_version(min_version="12.0", max_version="13.2") + # True if 12.0 <= CUDA_VERSION <= 13.2 + """ + # Avoid circular dependency + from .version_info import CUDA_VERSION + + # Coerce all version parameters to DSLCudaVersion at the start + exact_v = _coerce_to_cuda_version(exact_version, "exact_version") + min_v = _coerce_to_cuda_version(min_version, "min_version") + max_v = _coerce_to_cuda_version(max_version, "max_version") + + # Sanity check + is_range_check = min_v is not None or max_v is not None + is_exact_version_check = exact_v is not None + if is_range_check and is_exact_version_check: + raise DSLRuntimeError( + "Cannot use exact_version and [min_version, max_version] check at the same time" + ) + + if is_range_check: + if min_v is None and max_v is None: + raise DSLRuntimeError( + "min_version and max_version cannot be None at the same time" + ) + if min_v is not None and max_v is not None: + if min_v > max_v: + raise DSLRuntimeError("min_version must be less than max_version") + + result = True + if min_v is not None: + result = result and CUDA_VERSION >= min_v + if max_v is not None: + result = result and CUDA_VERSION <= max_v + return result + elif is_exact_version_check: + return CUDA_VERSION == exact_v + else: + raise DSLRuntimeError( + "either exact_version, min_version, or max_version must be provided" + ) diff --git a/python/CuTeDSL/cutlass/base_dsl/compiler.py b/python/CuTeDSL/cutlass/base_dsl/compiler.py index 339f695e..10afaa4c 100644 --- a/python/CuTeDSL/cutlass/base_dsl/compiler.py +++ b/python/CuTeDSL/cutlass/base_dsl/compiler.py @@ -340,21 +340,6 @@ class LinkLibraries(StringCompileOption): class GPUArch(StringCompileOption): option_name = "cubin-chip" - def __init__(self, val): - if isinstance(val, str) and val.startswith("sm_110"): - val = val.replace("sm_110", "sm_101") - super().__init__(val) - - @property - def value(self) -> bool: - return self._value - - @value.setter - def value(self, value: bool): - if isinstance(value, str) and value.startswith("sm_110"): - value = value.replace("sm_110", "sm_101") - self._value = value - class EnableTVMFFI(EmptyCompileOption): pass diff --git a/python/CuTeDSL/cutlass/base_dsl/dsl.py b/python/CuTeDSL/cutlass/base_dsl/dsl.py index 118e406b..0305ba30 100644 --- a/python/CuTeDSL/cutlass/base_dsl/dsl.py +++ b/python/CuTeDSL/cutlass/base_dsl/dsl.py @@ -50,7 +50,11 @@ from .jit_executor import JitCompiledFunction, JitFunctionArtifacts from .utils.timer import timer from .utils.logger import log from .utils.stacktrace import filter_exception, walk_to_top_module, filter_stackframe -from .runtime.jit_arg_adapters import is_argument_constexpr, JitArgAdapterRegistry +from .runtime.jit_arg_adapters import ( + is_argument_constexpr, + is_arg_spec_constexpr, + JitArgAdapterRegistry, +) from .ast_preprocessor import DSLPreprocessor from .common import * @@ -846,6 +850,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): use_pdl: bool = False auto_smem: bool = False cooperative: bool = False + @staticmethod def _check_and_canonicalize_dim(dim, name): if not isinstance(dim, (list, tuple)): @@ -967,18 +972,13 @@ class BaseDSL(metaclass=DSLSingletonMeta): sys.stderr = redirect_stderr = io.StringIO() sys.stdout = redirect_stdout = io.StringIO() - compile_gpu_arch = ( - self.envar.arch - if not self.compile_options.gpu_arch - else self.compile_options.gpu_arch - ) try: kernel = self.compiler_provider.compile_and_jit( module, pipeline, shared_libs=shared_libs, cuda_toolkit=self.envar.cuda_toolkit, - arch=compile_gpu_arch, + arch=self.envar.arch, ) finally: @@ -1314,8 +1314,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): dynamic_args = [] dynamic_kwargs = OrderedDict() for i, arg in enumerate(args): - if not is_argument_constexpr( - arg, + if not is_arg_spec_constexpr( args_spec.annotations.get(args_spec.args[i], None), args_spec.args[i], i, @@ -1323,7 +1322,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): ): dynamic_args.append(arg) for i, (k, v) in enumerate(kwargs.items()): - if not is_argument_constexpr(v, args_spec.kwonlyargs[i], k, i, funcBody): + if not is_arg_spec_constexpr(args_spec.kwonlyargs[i], k, i, funcBody): dynamic_kwargs[k] = v return dynamic_args, dynamic_kwargs diff --git a/python/CuTeDSL/cutlass/base_dsl/env_manager.py b/python/CuTeDSL/cutlass/base_dsl/env_manager.py index f135258b..19c66322 100644 --- a/python/CuTeDSL/cutlass/base_dsl/env_manager.py +++ b/python/CuTeDSL/cutlass/base_dsl/env_manager.py @@ -125,8 +125,6 @@ def detect_gpu_arch(prefix): suffix = "" if major >= 9: suffix = "a" - if major == 11 and minor == 0: - major, minor = 10, 1 return f"sm_{major}{minor}{suffix}" @@ -367,8 +365,6 @@ class EnvironmentVarManager(LogEnvironmentManager): # Other options self.dryrun = get_bool_env_var(f"{prefix}_DRYRUN", False) self.arch = get_str_env_var(f"{prefix}_ARCH", detect_gpu_arch(prefix)) - if self.arch.startswith("sm_110"): - self.arch = self.arch.replace("sm_110", "sm_101") self.warnings_as_errors = get_bool_env_var( f"{prefix}_WARNINGS_AS_ERRORS", False ) diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py index 2d5ab527..9fc1e02b 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py @@ -371,63 +371,6 @@ class MLIRBuilder(MLIRTypeBuilder): self.const_str_table[content] = symbol return symbol - def get_or_load_global_func_ptr_from_text( - self, - current_block: ir.Block, - function_name: str, - ) -> ir.Value: - """Get or create a function pointer global in .text section and load it. - - This creates a constant global function pointer in the .text section - (for AArch64 ADRP range compatibility) and performs a volatile load - to prevent optimization. - - This forces the function pointer to be local to the code, bypassing GOT entry - ADRP lookup issues on AArch64 when GOT and .text section are more than 4GB - apart which can happen when ASLR is applied. - """ - # Check if we've already created this global - if function_name not in self.const_func_ptr_table: - symbol = f"__func_ptr_{function_name}" - - module_body = self.module.body - with ir.InsertionPoint(module_body): - # 1. Create the global constant - # We use 'private' linkage so it doesn't conflict across modules - global_ptr = llvm.GlobalOp( - self.ptr_type, - symbol, - ir.Attribute.parse("#llvm.linkage"), - # Initialization via block below - ) - - # 2. Set the necessary attributes for JIT safety and AArch64 range - # We use 'constant' to mark it as immutable - # We use 'section = ".text"' to force it into the code block - global_ptr.attributes["constant"] = ir.UnitAttr.get() - global_ptr.attributes["section"] = ir.StringAttr.get(".text") - - # 3. Add a constructor block to the GlobalOp to initialize it - # with the address of the target function - initializer_block = global_ptr.initializer.blocks.append() - with ir.InsertionPoint(initializer_block): - # Get the address of the external function - func_addr = llvm.AddressOfOp(self.ptr_type, function_name).res - # Return the address as the initial value of the global - llvm.return_(arg=func_addr) - - self.const_func_ptr_table[function_name] = symbol - else: - symbol = self.const_func_ptr_table[function_name] - - # Load it with volatile semantics in the current block - with ir.InsertionPoint(current_block): - symbol_addr = self.address_of(symbol, self.ptr_type) - # Perform a volatile load to prevent optimization - load_op = llvm.load(self.ptr_type, symbol_addr) - # Set volatile attribute to prevent optimization - load_op.owner.attributes["volatile_"] = ir.UnitAttr.get() - return load_op # function def function( @@ -477,9 +420,7 @@ class MLIRBuilder(MLIRTypeBuilder): ) func_op.attributes["llvm.linkage"] = ir.StringAttr.get("external") - def create_alloca( - self, entry_block: ir.Block, alloca_type: ir.Type, array_size: int - ) -> ir.Value: + def create_alloca(self, entry_block: ir.Block, alloca_type: ir.Type, array_size: int) -> ir.Value: """Create an alloca operation.""" with ir.InsertionPoint(entry_block.operations[0]): # declare the struct type diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py index 9948d500..5b13b213 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py @@ -1277,7 +1277,10 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): return cond return self.check_condition( - current_block, check_value_mismatch, error_kind, error_msg_mismatch + current_block, + check_value_mismatch, + error_kind, + error_msg_mismatch, ) def set_or_check_matched_var_binding_from_shape( diff --git a/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py b/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py index e447f353..5939d2aa 100644 --- a/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py +++ b/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py @@ -290,7 +290,7 @@ def set_dataclass_attributes( for field, value in zip(fields, values): setattr(instance, field, value) - return instance + return instance def default_dataclass_from_iterable( diff --git a/python/CuTeDSL/cutlass/utils/version_info.py b/python/CuTeDSL/cutlass/base_dsl/version_info.py similarity index 62% rename from python/CuTeDSL/cutlass/utils/version_info.py rename to python/CuTeDSL/cutlass/base_dsl/version_info.py index 231856ec..203ddfb7 100644 --- a/python/CuTeDSL/cutlass/utils/version_info.py +++ b/python/CuTeDSL/cutlass/base_dsl/version_info.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # Use of this software is governed by the terms and conditions of the @@ -9,14 +9,15 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from ..cutlass_dsl import DSLCudaVersion, DSLRuntimeError +from typing import Callable + +from .common import DSLCudaVersion, DSLRuntimeError, _get_cuda_version try: - from .._mlir._mlir_libs._cutlass_ir._base_dsl import get_cuda_version - CUDA_VERSION = DSLCudaVersion(get_cuda_version()) + CUDA_VERSION = DSLCudaVersion(_get_cuda_version()) except Exception as e: raise DSLRuntimeError( "💥💥💥 Failed to get CUDA version 💥💥💥", cause=e, - suggestion="Consider re-installing the package." + suggestion="Consider re-installing the package.", ) from e diff --git a/python/CuTeDSL/cutlass/cute/__init__.py b/python/CuTeDSL/cutlass/cute/__init__.py index 5522bd67..d044b025 100644 --- a/python/CuTeDSL/cutlass/cute/__init__.py +++ b/python/CuTeDSL/cutlass/cute/__init__.py @@ -181,19 +181,21 @@ from .atom import ( make_tiled_copy_C_atom, make_cotiled_copy, copy_atom_call, + mma_atom_call, ) from .algorithm import gemm, copy, basic_copy, basic_copy_if, autovec_copy, prefetch from . import typing as typing_module from . import core from . import arch - from . import export + from . import nvgpu from . import testing from . import runtime from . import math + # Export all math ops without "math." from .math import * @@ -212,7 +214,6 @@ GenerateLineInfo = _dsl.GenerateLineInfo KeepCUBIN = _dsl.KeepCUBIN KeepPTX = _dsl.KeepPTX GPUArch = _dsl.GPUArch -LinkLibraries = _dsl.LinkLibraries EnableTVMFFI = _dsl.EnableTVMFFI # attach the TVM FFI ABI interface postprocessor to the DSL @@ -222,16 +223,52 @@ _tvm_ffi_args_spec_converter.attach_args_spec_converter(_dsl.CuTeDSL._get_dsl()) # Explicitly export all symbols for documentation generation __all__ = [ - # Core types - *core.__all__, + # ==================== cutlass._mlir.dialects.cute ==================== "AddressSpace", "CacheEvictionPriority", + # ==================== .typing ==================== "Tensor", "Layout", "ComposedLayout", - "Swizzle", - "E", - "ScaledBasis", + "SymInt", + "is_integer", + "is_int_tuple", + # ==================== .core ==================== + *core.__all__, + # ==================== .tuple ==================== + "transform_leaf", + "find_if", + "find", + "flatten_to_tuple", + "unflatten", + "product", + "product_like", + "product_each", + "elem_less", + "tuple_cat", + "transform_apply", + "filter_tuple", + # ==================== .tensor ==================== + "TensorSSA", + "ReductionOp", + "make_tensor", + "make_identity_tensor", + "make_fragment", + "make_fragment_like", + "make_rmem_tensor_like", + "make_rmem_tensor", + "recast_tensor", + "domain_offset", + "print_tensor", + "full", + "full_like", + "empty_like", + "ones_like", + "zeros_like", + "where", + "any_", + "all_", + # ==================== .atom ==================== "Atom", "MmaAtom", "CopyAtom", @@ -239,106 +276,6 @@ __all__ = [ "TiledMma", "ThrMma", "ThrCopy", - "TensorSSA", - "ReductionOp", - "SymInt", - # Basic utility functions - "assume", - "is_integer", - "is_int_tuple", - "is_static", - "has_underscore", - "shape", - "printf", - "print_tensor", - "pretty_str", - # Layout functions - "make_layout", - "recast_layout", - "make_identity_layout", - "make_ordered_layout", - "make_layout_like", - "make_composed_layout", - "make_layout_tv", - "make_layout_image_mask", - "get_nonswizzle_portion", - "get_swizzle_portion", - # Tensor functions - "make_ptr", - "make_tensor", - "make_identity_tensor", - "make_fragment", - "make_fragment_like", - "make_rmem_tensor", - "make_rmem_tensor_like", - "recast_ptr", - "recast_tensor", - # Tensor manipulation - "get", - "select", - "front", - "is_major", - "leading_dim", - "find", - "find_if", - "transform_leaf", - "basis_value", - "basis_get", - "coalesce", - "group_modes", - "cosize", - "size_in_bytes", - # Tuple operations - "flatten_to_tuple", - "flatten", - "unflatten", - "product", - "product_like", - "product_each", - "prepend", - "append", - "prepend_ones", - "append_ones", - "elem_less", - "tuple_cat", - "transform_apply", - "filter_tuple", - # Math operations - "ceil_div", - "round_up", - # Layout operations - "slice_and_offset", - "crd2idx", - "domain_offset", - "filter_zeros", - "filter", - "tile_to_shape", - "shape_div", - "dice", - # Layout algebra - "composition", - "complement", - "right_inverse", - "left_inverse", - "max_common_layout", - "max_common_vector", - "is_congruent", - "is_weakly_congruent", - # Product operations - "logical_product", - "zipped_product", - "tiled_product", - "flat_product", - "raked_product", - "blocked_product", - # Division operations - "flat_divide", - "logical_divide", - "zipped_divide", - "tiled_divide", - "local_partition", - "local_tile", - # MMA and Copy atom operations "make_atom", "make_mma_atom", "make_tiled_mma", @@ -353,39 +290,24 @@ __all__ = [ "make_tiled_copy_C_atom", "make_cotiled_copy", "copy_atom_call", - # Algorithm operations + "mma_atom_call", + # ==================== .algorithm ==================== + "gemm", + "copy", "basic_copy", "basic_copy_if", "autovec_copy", - "copy", "prefetch", - "gemm", - # Tensor creation - "full", - "full_like", - "empty_like", - "ones_like", - "zeros_like", - "where", - "any_", - "all_", - "repeat_as_tuple", - "repeat", - "repeat_like", - # User defined struct - "struct", - # FastDivmod operations - "FastDivmodDivisor", - "fast_divmod_create_divisor", - # Modules + # ==================== .extension ==================== + # ==================== .math ==================== + *math.__all__, + # ==================== submodules ==================== "arch", "export", "nvgpu", "testing", "runtime", - # Math utils - *math.__all__, - # Decorators and code generation + # ==================== DSL (cutlass_dsl) ==================== "jit", "kernel", "register_jit_arg_adapter", diff --git a/python/CuTeDSL/cutlass/cute/algorithm.py b/python/CuTeDSL/cutlass/cute/algorithm.py index 6be6f7fd..bc03231c 100644 --- a/python/CuTeDSL/cutlass/cute/algorithm.py +++ b/python/CuTeDSL/cutlass/cute/algorithm.py @@ -10,7 +10,7 @@ # is strictly prohibited. import math -from typing import Optional, Dict, Any, List, Tuple +from typing import Optional, Dict, Any, List, Tuple, Union from cutlass._mlir import ir from cutlass.cutlass_dsl import for_generate, yield_out, if_generate, dsl_user_op @@ -29,15 +29,35 @@ from .core import ( append_ones, group_modes, ) -from .atom import MmaAtom, CopyAtom, make_atom +from .atom import ( + MmaAtom, + CopyAtom, + make_atom, + _normalize_variadic_tensor_operand, + copy_atom_call, +) +from .nvgpu.common import CacheEvictionPriority + +def _normalize_gemm_operand_list( + x: Union["Tensor", List["Tensor"], Tuple["Tensor", ...]], name: str +) -> List["Tensor"]: + if isinstance(x, Tensor): + return [x] + if isinstance(x, (list, tuple)): + if len(x) == 0: + raise ValueError(f"`{name}` must contain at least one Tensor") + if not all(isinstance(t, Tensor) for t in x): + raise TypeError(f"All elements of `{name}` must be Tensor") + return list(x) # type: ignore + raise TypeError(f"`{name}` must be a Tensor or a sequence of Tensors") @dsl_user_op def gemm( atom: MmaAtom, d: Tensor, - a: Tensor, - b: Tensor, + a: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], + b: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], c: Tensor, *, loc=None, @@ -62,14 +82,17 @@ def gemm( - Dispatch [4]: (V,M) x (V,N) => (V,M,N) => (V,M,1) x (V,N,1) => (V,M,N) - Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N) + Operand flexibility: + - `a` and `b` can be a single Tensor (regular GEMM) or a sequence `[operand, scale_factor]` for block-scaled GEMM. + :param atom: MMA atom :type atom: MmaAtom :param d: Destination tensor :type d: Tensor - :param a: First source tensor - :type a: Tensor - :param b: Second source tensor - :type b: Tensor + :param a: First source tensor or sequence for advanced modes (e.g., `[a, sfa]`) + :type a: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] + :param b: Second source tensor or sequence for advanced modes (e.g., `[b, sfb]`) + :type b: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] :param c: Third source tensor :type c: Tensor :param loc: Source location for MLIR, defaults to None @@ -82,8 +105,13 @@ def gemm( :rtype: None """ - a_rank = rank(a.shape) - b_rank = rank(b.shape) + # Normalize A/B to lists for variadic IR operands, while keeping old API working. + a_list = _normalize_gemm_operand_list(a, "a") + b_list = _normalize_gemm_operand_list(b, "b") + + # Rank validations based on the primary A/B tensors (guaranteed non-empty) + a_rank = rank(a_list[0].shape) + b_rank = rank(b_list[0].shape) c_rank = rank(c.shape) d_rank = rank(d.shape) @@ -104,7 +132,9 @@ def gemm( raise ValueError("`c` must have rank 3 when `a` has rank 3") value = atom._unpack(loc=loc, ip=ip, **kwargs) - return _cute_ir.gemm(value, d.value, a.value, b.value, c.value, loc=loc, ip=ip) + a_vals = [t.value for t in a_list] + b_vals = [t.value for t in b_list] + return _cute_ir.gemm(value, d.value, a_vals, b_vals, c.value, loc=loc, ip=ip) @dsl_user_op @@ -132,7 +162,7 @@ def basic_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None: src.element_type.mlir_type, src.element_type.width ) simt_copy = make_atom(simt_copy_ty, loc=loc, ip=ip) - return _cute_ir.copy(simt_copy, src.value, dst.value, loc=loc, ip=ip) + return _cute_ir.copy(simt_copy, [src.value], [dst.value], loc=loc, ip=ip) s = size(dst, loc=loc, ip=ip) # Always generate an scf.for Op when one of the tensors is dynamic @@ -186,7 +216,14 @@ def _basic_copy_if_static( @dsl_user_op -def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None: +def autovec_copy( + src: Tensor, + dst: Tensor, + *, + l1c_evict_priority: CacheEvictionPriority = CacheEvictionPriority.EVICT_NORMAL, + loc=None, + ip=None, +) -> None: """ Auto-vectorization SIMT copy policy. @@ -239,11 +276,15 @@ def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None: # Dispatch to copy with atom simt_type = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get( - src.element_type.mlir_type, num_bits_per_copy + src.element_type.mlir_type, + num_bits_per_copy, + 0, + 0, + l1c_evict_priority._to_ir(), ) simt_copy = make_atom(simt_type, loc=loc, ip=ip) return _cute_ir.copy( - simt_copy, tiled_src.value, tiled_dst.value, loc=loc, ip=ip + simt_copy, [tiled_src.value], [tiled_dst.value], loc=loc, ip=ip ) # Failed to vectorize, use a basic copy @@ -258,19 +299,21 @@ def _parse_auto_multicast_args( This function consumes the following key from kwargs if present: - 'auto_multicast': dict - dict: { 'multicast_layout': str, 'use_2cta': bool } + dict: { 'multicast_layout': str, 'use_2cta': bool, 'from_block_api': bool } Returns: List of (attr_name, ir.Attribute) pairs to be attached to the op. Recognized attributes: - ('multicast_layout', #cute.layout<...>) when a layout string is provided - ('use_2cta', unit) when use_2cta is True + - ('from_block_api', unit) when from_block_api is True """ attr_pairs: List[Tuple[str, ir.Attribute]] = [] # Pop known keys to avoid leaking to trait unpack auto_multicast = kwargs.pop("auto_multicast", None) + from_block_api: bool = False use_2cta: bool = False layout_str: Optional[str] = None @@ -281,6 +324,7 @@ def _parse_auto_multicast_args( ) layout_str = auto_multicast.get("multicast_layout", None) use_2cta = bool(auto_multicast.get("use_2cta", False)) + from_block_api = bool(auto_multicast.get("from_block_api", False)) if layout_str is not None: if not isinstance(layout_str, str): @@ -293,7 +337,8 @@ def _parse_auto_multicast_args( ir.Attribute.parse(f'#cute.layout<"{layout_str}">'), ) ) - + if from_block_api: + attr_pairs.append(("from_block_api", ir.UnitAttr.get())) if use_2cta: attr_pairs.append(("use_2cta", ir.UnitAttr.get())) @@ -303,8 +348,8 @@ def _parse_auto_multicast_args( @dsl_user_op def copy( atom: CopyAtom, - src: Tensor, - dst: Tensor, + src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], + dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], *, pred: Optional[Tensor] = None, loc=None, @@ -315,10 +360,10 @@ def copy( :param atom: Copy atom specifying the transfer operation :type atom: CopyAtom - :param src: Source tensor with layout profile ``(V, Rest...)`` - :type src: Tensor - :param dst: Destination tensor with layout profile ``(V, Rest...)`` - :type dst: Tensor + :param src: Source tensor or list of tensors with layout profile ``(V, Rest...)`` + :type src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] + :param dst: Destination tensor or list of tensors with layout profile ``(V, Rest...)`` + :type dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] :param pred: Optional predication tensor for conditional transfers, defaults to None :type pred: Optional[Tensor], optional :param loc: Source location information, defaults to None @@ -346,6 +391,12 @@ def copy( Source and destination tensors must be partitioned in accordance with the Copy Atom specifications. Post-partitioning, both tensors will exhibit a ``(V, Rest...)`` layout profile. + The operands `src` and `dst` are variadic, each containing a variable number of tensors: + + - For regular copy, `src` and `dst` contain single source and destination tensors respectively. + - For copy with auxiliary operands, `src` and `dst` contain the primary tensors followed by + their respective auxiliary tensors. + **Precondition:** The size of mode 1 must be equal for both source and destination tensors: ``size(src, mode=[1]) == size(dst, mode=[1])`` @@ -371,41 +422,54 @@ def copy( for future releases. """ - if isinstance(src.type, _cute_ir.MemRefType) and isinstance( - dst.type, _cute_ir.MemRefType + # Normalize src/dst to lists for variadic IR operands + src_list = _normalize_variadic_tensor_operand(src, "src") + dst_list = _normalize_variadic_tensor_operand(dst, "dst") + + # Validate primary tensors (first element) + src_primary = src_list[0] + dst_primary = dst_list[0] + + if isinstance(src_primary.type, _cute_ir.MemRefType) and isinstance( + dst_primary.type, _cute_ir.MemRefType ): - if src.element_type.width != dst.element_type.width: + if src_primary.element_type.width != dst_primary.element_type.width: raise TypeError( "`copy` currently only supports equal source and destination " "element type bit width" ) - if rank(src) != rank(dst): + if rank(src_primary) != rank(dst_primary): raise ValueError( "Expected source and destination tensors to have the same rank, " - f"but got {rank(src)} and {rank(dst)}" + f"but got {rank(src_primary)} and {rank(dst_primary)}" ) - # Canonicalize to at least rank-2 tensors - src = group_modes(append_ones(src, up_to_rank=2), 1) - dst = group_modes(append_ones(dst, up_to_rank=2), 1) + # Canonicalize all tensors to at least rank-2 + src_list = [group_modes(append_ones(t, up_to_rank=2), 1) for t in src_list] + dst_list = [group_modes(append_ones(t, up_to_rank=2), 1) for t in dst_list] if pred is not None: pred = group_modes(append_ones(pred, up_to_rank=2), 1) - if is_static(src.shape[1]) and is_static(dst.shape[1]): - if size(src, mode=[1]) != size(dst, mode=[1]): + # Recompute primary references after canonicalization + src_primary = src_list[0] + dst_primary = dst_list[0] + + if is_static(src_primary.shape[1]) and is_static(dst_primary.shape[1]): + if size(src_primary, mode=[1]) != size(dst_primary, mode=[1]): raise ValueError( "Expected source and destination tensors to have the same size in mode-1, " - f"but got {size(src, mode=[1])} and {size(dst, mode=[1])}" + f"but got {size(src_primary, mode=[1])} and {size(dst_primary, mode=[1])}" ) multicast_attr_pairs = _parse_auto_multicast_args(kwargs) value = atom._unpack(loc=loc, ip=ip, **kwargs) - if isinstance(pred, Tensor): - pred = pred.value + pred_value = pred.value if isinstance(pred, Tensor) else pred - op = _cute_ir.copy(value, src.value, dst.value, pred=pred, loc=loc, ip=ip) + src_vals = [t.value for t in src_list] + dst_vals = [t.value for t in dst_list] + op = _cute_ir.copy(value, src_vals, dst_vals, pred=pred_value, loc=loc, ip=ip) for name, attr in multicast_attr_pairs: op.attributes[name] = attr diff --git a/python/CuTeDSL/cutlass/cute/arch/__init__.py b/python/CuTeDSL/cutlass/cute/arch/__init__.py index 6d10c66e..5213b00e 100644 --- a/python/CuTeDSL/cutlass/cute/arch/__init__.py +++ b/python/CuTeDSL/cutlass/cute/arch/__init__.py @@ -11,7 +11,6 @@ from .elect import * from .mbar import * -from .numeric_conversion import * from .nvvm_wrappers import * from .smem import * from .tmem import * @@ -74,6 +73,8 @@ __all__ = [ "vote_any_sync", "vote_all_sync", "vote_uni_sync", + "warp_redux_sync", + "atomic_max_float32", "atomic_add", "atomic_and", "atomic_or", @@ -95,15 +96,19 @@ __all__ = [ "fma_packed_f32x2", "mul_packed_f32x2", "add_packed_f32x2", + "sub_packed_f32x2", "fmax", "rcp_approx", "exp2", + "cvt_i8x4_to_f32x4", + "cvt_i8x2_to_f32x2", + "cvt_i8_bf16", + "cvt_i8x2_to_bf16x2", + "cvt_i8x4_to_bf16x4", + "cvt_f32x2_bf16x2", + "warp_redux_sync", # Constants "WARP_SIZE", - # Forward from auto-generated nvvm python - "ProxyKind", - "SharedSpace", - "RoundingModeKind", # # smem.py # diff --git a/python/CuTeDSL/cutlass/cute/arch/clc.py b/python/CuTeDSL/cutlass/cute/arch/clc.py index 29af7348..b2f7015f 100644 --- a/python/CuTeDSL/cutlass/cute/arch/clc.py +++ b/python/CuTeDSL/cutlass/cute/arch/clc.py @@ -10,8 +10,11 @@ # is strictly prohibited. from typing import Tuple + from cutlass.cutlass_dsl import T, dsl_user_op -from cutlass._mlir.dialects import nvvm, vector + +from cutlass._mlir import ir +from cutlass._mlir.dialects import nvvm, llvm, vector, arith from ..typing import Int32, Pointer, Int128 @@ -20,6 +23,7 @@ from ..typing import Int32, Pointer, Int128 def issue_clc_query( mbar_ptr: Pointer, clc_response_ptr: Pointer, + multicast: bool = True, loc=None, ip=None, ) -> None: @@ -36,12 +40,20 @@ def issue_clc_query( """ mbar_llvm_ptr = mbar_ptr.llvm_ptr clc_response_llvm_ptr = clc_response_ptr.llvm_ptr - nvvm.clusterlaunchcontrol_try_cancel_multicast( - clc_response_llvm_ptr, - mbar_llvm_ptr, - loc=loc, - ip=ip, - ) + if multicast: + nvvm.clusterlaunchcontrol_try_cancel_multicast( + clc_response_llvm_ptr, + mbar_llvm_ptr, + loc=loc, + ip=ip, + ) + else: + nvvm.clusterlaunchcontrol_try_cancel( + clc_response_llvm_ptr, + mbar_llvm_ptr, + loc=loc, + ip=ip, + ) @dsl_user_op @@ -78,7 +90,6 @@ def clc_response( ) # Query if the cluster was canceled pred = nvvm.clusterlaunchcontrol_query_cancel_is_canceled( - T.bool(), clc_result_i128, loc=loc, ip=ip, @@ -87,7 +98,6 @@ def clc_response( # Get first CTA ID x component m_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_x( - T.i32(), clc_result_i128, loc=loc, ip=ip, @@ -95,7 +105,6 @@ def clc_response( # Get first CTA ID y component n_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_y( - T.i32(), clc_result_i128, loc=loc, ip=ip, @@ -103,7 +112,6 @@ def clc_response( # Get first CTA ID z component l_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_z( - T.i32(), clc_result_i128, loc=loc, ip=ip, diff --git a/python/CuTeDSL/cutlass/cute/arch/elect.py b/python/CuTeDSL/cutlass/cute/arch/elect.py index 9f51484e..abd213b1 100644 --- a/python/CuTeDSL/cutlass/cute/arch/elect.py +++ b/python/CuTeDSL/cutlass/cute/arch/elect.py @@ -9,7 +9,6 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from cutlass.base_dsl.arch import Arch from cutlass.cutlass_dsl import BaseDSL, T, dsl_user_op import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir @@ -72,6 +71,6 @@ def elect_one(*, loc=None, ip=None) -> IfOpRegion: from cutlass.base_dsl.arch import Arch BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) - is_thread_leader = nvvm.elect_sync(T.bool()) + is_thread_leader = nvvm.elect_sync() if_op = scf.IfOp(is_thread_leader, loc=loc, ip=ip) return IfOpRegion(if_op.then_block, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/cute/arch/mbar.py b/python/CuTeDSL/cutlass/cute/arch/mbar.py index 3f9f2ec5..17e541e8 100644 --- a/python/CuTeDSL/cutlass/cute/arch/mbar.py +++ b/python/CuTeDSL/cutlass/cute/arch/mbar.py @@ -13,7 +13,7 @@ from typing import Optional from cutlass.base_dsl.arch import Arch from cutlass.cutlass_dsl import BaseDSL, T, if_generate, dsl_user_op -from cutlass._mlir.dialects import nvvm +from cutlass._mlir.dialects import nvvm, llvm from ..typing import Pointer, Int, Boolean, Int32, AddressSpace @@ -35,10 +35,7 @@ def mbarrier_init(mbar_ptr: Pointer, cnt: Int, *, loc=None, ip=None) -> None: :type cnt: Int """ nvvm.mbarrier_init_shared( - mbar_ptr.to_llvm_ptr(loc=loc, ip=ip), - Int32(cnt).ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, + mbar_ptr.llvm_ptr, Int32(cnt).ir_value(loc=loc, ip=ip), loc=loc, ip=ip ) @@ -68,15 +65,18 @@ def mbarrier_arrive_and_expect_tx( """ BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) - mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip) + mbar_llvm_ptr = mbar_ptr.llvm_ptr if peer_cta_rank_in_cluster is not None: - mbar_llvm_ptr = nvvm.mapa_shared_cluster( - mbar_llvm_ptr.type, + mbar_cluster_type = llvm.PointerType.get(AddressSpace.dsmem) + mbar_llvm_ptr = nvvm.mapa( + mbar_cluster_type, mbar_llvm_ptr, Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip), loc=loc, ip=ip, ) + mbar_shared_type = llvm.PointerType.get(AddressSpace.smem) + mbar_llvm_ptr = llvm.addrspacecast(mbar_shared_type, mbar_llvm_ptr) space = nvvm.MBarrierSpaceKind.CLUSTER else: space = nvvm.MBarrierSpaceKind.CTA @@ -108,15 +108,18 @@ def mbarrier_expect_tx( """ BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) - mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip) + mbar_llvm_ptr = mbar_ptr.llvm_ptr if peer_cta_rank_in_cluster is not None: + mbar_cluster_type = llvm.PointerType.get(AddressSpace.dsmem) mbar_llvm_ptr = nvvm.mapa( - mbar_llvm_ptr.type, + mbar_cluster_type, mbar_llvm_ptr, Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip), loc=loc, ip=ip, ) + mbar_shared_type = llvm.PointerType.get(AddressSpace.smem) + mbar_llvm_ptr = llvm.addrspacecast(mbar_shared_type, mbar_llvm_ptr) space = nvvm.MBarrierSpaceKind.CLUSTER else: space = nvvm.MBarrierSpaceKind.CTA @@ -147,7 +150,7 @@ def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None: # This NVVM Op is a spin-loop wrapping the mbarrier.try_wait.parity.shared.b64 PTX # The timeout in ns only applies to the latter and this call is truly blocking nvvm.mbarrier_try_wait_parity_shared( - mbar_ptr.to_llvm_ptr(loc=loc, ip=ip), + mbar_ptr.llvm_ptr, Int32(phase).ir_value(loc=loc, ip=ip), Int32(timeout_ns).ir_value(loc=loc, ip=ip), loc=loc, @@ -171,8 +174,7 @@ def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Bo return Boolean( nvvm.mbarrier_wait_parity( - T.bool(), - mbar_ptr.to_llvm_ptr(loc=loc, ip=ip), + mbar_ptr.llvm_ptr, Int32(phase).ir_value(loc=loc, ip=ip), nvvm.MBarrierWaitKind.TRY, loc=loc, @@ -226,17 +228,20 @@ def mbarrier_arrive( the mbarrier is converted to a remote address in the peer CTA's SMEM. """ - mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip) + mbar_llvm_ptr = mbar_ptr.llvm_ptr if peer_cta_rank_in_cluster is not None: BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) - mbar_llvm_ptr = nvvm.mapa_shared_cluster( - mbar_llvm_ptr.type, + mbar_cluster_type = llvm.PointerType.get(AddressSpace.dsmem) + mbar_llvm_ptr = nvvm.mapa( + mbar_cluster_type, mbar_llvm_ptr, Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip), loc=loc, ip=ip, ) + mbar_shared_type = llvm.PointerType.get(AddressSpace.smem) + mbar_llvm_ptr = llvm.addrspacecast(mbar_shared_type, mbar_llvm_ptr) space = nvvm.MBarrierSpaceKind.CLUSTER else: space = nvvm.MBarrierSpaceKind.CTA @@ -264,5 +269,10 @@ def cp_async_mbarrier_arrive_noinc(mbar_ptr: Pointer, *, loc=None, ip=None) -> N """ BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) - mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip) - nvvm.cp_async_mbarrier_arrive_shared(mbar_llvm_ptr, noinc=True, loc=loc, ip=ip) + mbar_llvm_ptr = mbar_ptr.llvm_ptr + nvvm.cp_async_mbarrier_arrive_shared( + mbar_llvm_ptr, + noinc=True, + loc=loc, + ip=ip, + ) diff --git a/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py b/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py index 6484136f..ba9faa4e 100644 --- a/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py +++ b/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py @@ -9,16 +9,17 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. - from cutlass.base_dsl.arch import Arch from cutlass.base_dsl.common import DSLRuntimeError from cutlass.cutlass_dsl import BaseDSL, dsl_user_op from cutlass._mlir import ir -from cutlass._mlir.dialects import builtin, arith, llvm, vector +from cutlass._mlir.dialects import arith, llvm, vector from .nvvm_wrappers import ( cvt_i8_bf16, + cvt_i8x2_to_bf16x2, + cvt_i8x4_to_bf16x4, cvt_f32x2_bf16x2, cvt_i8x4_to_f32x4, cvt_i8x2_to_f32x2, @@ -26,22 +27,11 @@ from .nvvm_wrappers import ( cvt_i4x4_to_bf16x4, cvt_i4x2_to_bf16x2, cvt_i4_bf16, - cvt_f4e2m1x8_to_f16x8, - cvt_f4e2m1x4_to_f16x4, - cvt_f4e2m1x2_to_f16x2, - cvt_f4e2m1_f16, + cvt_f32_bf16, sext_unpacked_i4x4_to_i8x4, ) +from ..typing import Int4, Int8, Float32, BFloat16, Int32 -from ..typing import ( - Int4, - Int8, - Int32, - Float16, - Float32, - BFloat16, - Float32, -) @dsl_user_op def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None): @@ -64,6 +54,7 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None): vec_f32x2_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc) vec_dst_type = ir.VectorType.get([length], BFloat16.mlir_type, loc=loc) vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip) + arch = BaseDSL._get_dsl().get_arch_enum() # try to use vectorized version if length >= 4: num_vec4 = length // 4 @@ -71,45 +62,66 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None): vec_i8x4 = vector.extract_strided_slice( vec_i8x4_type, vec_i8, [src_pos], [4], [1], loc=loc, ip=ip ) - vec_f32x4 = cvt_i8x4_to_f32x4(vec_i8x4, loc=loc, ip=ip) - vec_f32x2_lo = vector.extract_strided_slice( - vec_f32x2_type, vec_f32x4, [0], [2], [1], loc=loc, ip=ip - ) - vec_f32x2_hi = vector.extract_strided_slice( - vec_f32x2_type, vec_f32x4, [2], [2], [1], loc=loc, ip=ip - ) - vec_bf16x2_lo = cvt_f32x2_bf16x2(vec_f32x2_lo, loc=loc, ip=ip) - vec_bf16x2_hi = cvt_f32x2_bf16x2(vec_f32x2_hi, loc=loc, ip=ip) - vec_dst = vector.insert_strided_slice( - vec_bf16x2_lo, vec_dst, [src_pos], [1], loc=loc, ip=ip - ) - vec_dst = vector.insert_strided_slice( - vec_bf16x2_hi, vec_dst, [src_pos + 2], [1], loc=loc, ip=ip - ) + if arch in cvt_i8_bf16_intrinsic.s26_bf16_supported_archs: + vec_bf16x4 = cvt_i8x4_to_bf16x4(vec_i8x4, loc=loc, ip=ip) + vec_dst = vector.insert_strided_slice( + vec_bf16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip + ) + else: + vec_f32x4 = cvt_i8x4_to_f32x4(vec_i8x4, loc=loc, ip=ip) + vec_f32x2_lo = vector.extract_strided_slice( + vec_f32x2_type, vec_f32x4, [0], [2], [1], loc=loc, ip=ip + ) + vec_f32x2_hi = vector.extract_strided_slice( + vec_f32x2_type, vec_f32x4, [2], [2], [1], loc=loc, ip=ip + ) + vec_bf16x2_lo = cvt_f32x2_bf16x2(vec_f32x2_lo, loc=loc, ip=ip) + vec_bf16x2_hi = cvt_f32x2_bf16x2(vec_f32x2_hi, loc=loc, ip=ip) + vec_dst = vector.insert_strided_slice( + vec_bf16x2_lo, vec_dst, [src_pos], [1], loc=loc, ip=ip + ) + vec_dst = vector.insert_strided_slice( + vec_bf16x2_hi, vec_dst, [src_pos + 2], [1], loc=loc, ip=ip + ) + src_pos += 4 length -= 4 if length >= 2: vec_i8x2 = vector.extract_strided_slice( vec_i8x2_type, vec_i8, [src_pos], [2], [1], loc=loc, ip=ip ) - vec_f32x2 = cvt_i8x2_to_f32x2(vec_i8x2, loc=loc, ip=ip) - vec_bf16x2 = cvt_f32x2_bf16x2(vec_f32x2, loc=loc, ip=ip) + if arch in cvt_i8_bf16_intrinsic.s26_bf16_supported_archs: + vec_bf16x2 = cvt_i8x2_to_bf16x2(vec_i8x2, loc=loc, ip=ip) + else: + vec_f32x2 = cvt_i8x2_to_f32x2(vec_i8x2, loc=loc, ip=ip) + vec_bf16x2 = cvt_f32x2_bf16x2(vec_f32x2, loc=loc, ip=ip) vec_dst = vector.insert_strided_slice( vec_bf16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip ) src_pos += 2 length -= 2 if length >= 1: - val_bf16 = cvt_i8_bf16( - vector.extractelement( + if arch in cvt_i8_bf16_intrinsic.s26_bf16_supported_archs: + val_bf16 = cvt_i8_bf16( + vector.extractelement( + vec_i8, + position=arith.constant(Int32.mlir_type, src_pos), + loc=loc, + ip=ip, + ), + loc=loc, + ip=ip, + ) + else: + src_i8 = vector.extractelement( vec_i8, position=arith.constant(Int32.mlir_type, src_pos), loc=loc, ip=ip, - ), - loc=loc, - ip=ip, - ) + ) + src_i32 = llvm.sext(Int32.mlir_type, src_i8, loc=loc, ip=ip) + src_f32 = llvm.sitofp(Float32.mlir_type, src_i32, loc=loc, ip=ip) + val_bf16 = cvt_f32_bf16(src_f32, loc=loc, ip=ip) vec_dst = vector.insertelement( val_bf16, vec_dst, @@ -121,7 +133,7 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None): @dsl_user_op -def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None): +def cvt_i4_bf16_intrinsic(vec_i4, length, *, with_shuffle=False, loc=None, ip=None): """ Fast conversion from int4 to bfloat16. It converts a vector of int4 to a vector of bfloat16. @@ -129,6 +141,13 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None): :type vec_i4: 1D vector of int4 :param length: The length of the input vector. :type length: int + :param with_shuffle: Whether the input vec_i4 follows a specific shuffle pattern. + If True, for consecutive 8 int4 values with indices of (0, 1, 2, 3, 4, 5, 6, 7), + the input elements are shuffled to (0, 2, 1, 3, 4, 6, 5, 7). For tailing elements less than 8, + the shuffle pattern is (0, 2, 1, 3) for 4 elements. No shuffle is needed for less than 4 elements. + Shuffle could help to produce converted bf16 values in the natural order of (0, 1, 2 ,3 ,4 ,5 ,6 ,7) + without extra prmt instructions and thus better performance. + :type with_shuffle: bool :return: The output 1D vector of bfloat16 with the same length as the input vector. :rtype: 1D vector of bfloat16 """ @@ -141,6 +160,7 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None): vec_i4x2_type = ir.VectorType.get([2], Int4.mlir_type, loc=loc) vec_dst_type = ir.VectorType.get([length], BFloat16.mlir_type, loc=loc) vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip) + # try to use vectorized version if length >= 8: num_vec8 = length // 8 @@ -148,7 +168,9 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None): vec_i4x8 = vector.extract_strided_slice( vec_i4x8_type, vec_i4, [src_pos], [8], [1], loc=loc, ip=ip ) - vec_bf16x8 = cvt_i4x8_to_bf16x8(vec_i4x8, loc=loc, ip=ip) + vec_bf16x8 = cvt_i4x8_to_bf16x8( + vec_i4x8, with_shuffle=with_shuffle, loc=loc, ip=ip + ) vec_dst = vector.insert_strided_slice( vec_bf16x8, vec_dst, [src_pos], [1], loc=loc, ip=ip ) @@ -158,7 +180,9 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None): vec_i4x4 = vector.extract_strided_slice( vec_i4x4_type, vec_i4, [src_pos], [4], [1], loc=loc, ip=ip ) - vec_bf16x4 = cvt_i4x4_to_bf16x4(vec_i4x4, loc=loc, ip=ip) + vec_bf16x4 = cvt_i4x4_to_bf16x4( + vec_i4x4, with_shuffle=with_shuffle, loc=loc, ip=ip + ) vec_dst = vector.insert_strided_slice( vec_bf16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip ) @@ -168,7 +192,9 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None): vec_i4x2 = vector.extract_strided_slice( vec_i4x2_type, vec_i4, [src_pos], [2], [1], loc=loc, ip=ip ) - vec_bf16x2 = cvt_i4x2_to_bf16x2(vec_i4x2, loc=loc, ip=ip) + vec_bf16x2 = cvt_i4x2_to_bf16x2( + vec_i4x2, with_shuffle=with_shuffle, loc=loc, ip=ip + ) vec_dst = vector.insert_strided_slice( vec_bf16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip ) @@ -195,84 +221,6 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None): return vec_dst -@dsl_user_op -def cvt_f4e2m1_f16_intrinsic(vec_f4e2m1, length, *, loc=None, ip=None): - """ - Convert a vector of float4e2m1 to a vector of float16. - - :param vec_f4e2m1: The input vector of float4e2m1. - :type vec_f4e2m1: 1D vector of float4e2m1 - :param length: The length of the input vector. - :type length: int - :return: The output 1D vector of float16 with the same length as the input vector. - :rtype: 1D vector of float16 - """ - src_pos = 0 - vec_src_i4 = builtin.unrealized_conversion_cast( - [ir.VectorType.get([length], Int4.mlir_type, loc=loc)], - [vec_f4e2m1], - loc=loc, - ip=ip, - ) - vec_i4x8_type = ir.VectorType.get([8], Int4.mlir_type, loc=loc) - vec_i4x4_type = ir.VectorType.get([4], Int4.mlir_type, loc=loc) - vec_i4x2_type = ir.VectorType.get([2], Int4.mlir_type, loc=loc) - vec_dst_type = ir.VectorType.get([length], Float16.mlir_type, loc=loc) - vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip) - # try to use vectorized version - if length >= 8: - num_vec8 = length // 8 - for _ in range(num_vec8): - vec_f4e2m1x8 = vector.extract_strided_slice( - vec_i4x8_type, vec_src_i4, [src_pos], [8], [1], loc=loc, ip=ip - ) - vec_f16x8 = cvt_f4e2m1x8_to_f16x8(vec_f4e2m1x8, loc=loc, ip=ip) - vec_dst = vector.insert_strided_slice( - vec_f16x8, vec_dst, [src_pos], [1], loc=loc, ip=ip - ) - src_pos += 8 - length -= 8 - if length >= 4: - vec_f4e2m1x4 = vector.extract_strided_slice( - vec_i4x4_type, vec_src_i4, [src_pos], [4], [1], loc=loc, ip=ip - ) - vec_f16x4 = cvt_f4e2m1x4_to_f16x4(vec_f4e2m1x4, loc=loc, ip=ip) - vec_dst = vector.insert_strided_slice( - vec_f16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip - ) - src_pos += 4 - length -= 4 - if length >= 2: - vec_f4e2m1x2 = vector.extract_strided_slice( - vec_i4x2_type, vec_src_i4, [src_pos], [2], [1], loc=loc, ip=ip - ) - vec_f16x2 = cvt_f4e2m1x2_to_f16x2(vec_f4e2m1x2, loc=loc, ip=ip) - vec_dst = vector.insert_strided_slice( - vec_f16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip - ) - src_pos += 2 - length -= 2 - if length >= 1: - val_f16 = cvt_f4e2m1_f16( - vector.extractelement( - vec_src_i4, - position=arith.constant(Int32.mlir_type, src_pos), - loc=loc, - ip=ip, - ), - loc=loc, - ip=ip, - ) - vec_dst = vector.insertelement( - val_f16, - vec_dst, - position=arith.constant(Int32.mlir_type, src_pos), - loc=loc, - ip=ip, - ) - return vec_dst - - @dsl_user_op def sext_unpacked_i4_i8_intrinsic(vec_unpacked_i4, length, *, loc=None, ip=None): """ @@ -295,9 +243,7 @@ def sext_unpacked_i4_i8_intrinsic(vec_unpacked_i4, length, *, loc=None, ip=None) vec_unpacked_i4x4 = vector.extract_strided_slice( vec_i8x4_type, vec_unpacked_i4, [pos], [4], [1], loc=loc, ip=ip ) - vec_i8x4 = sext_unpacked_i4x4_to_i8x4( - vec_unpacked_i4x4, loc=loc, ip=ip - ) + vec_i8x4 = sext_unpacked_i4x4_to_i8x4(vec_unpacked_i4x4, loc=loc, ip=ip) vec_i8 = vector.insert_strided_slice( vec_i8x4, vec_i8, [pos], [1], loc=loc, ip=ip ) @@ -312,6 +258,12 @@ cvt_i8_bf16_intrinsic.supported_archs = ( *Arch.HopperArchs(), *Arch.BlackwellArchs(), ) +cvt_i8_bf16_intrinsic.s26_bf16_supported_archs = ( + Arch.sm_100a, + Arch.sm_110a, + Arch.sm_120a, + Arch.sm_121a, +) cvt_i4_bf16_intrinsic.supported_archs = ( Arch.sm_100a, Arch.sm_110a, diff --git a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py index 78f68639..ac109718 100644 --- a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py +++ b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py @@ -10,25 +10,16 @@ # is strictly prohibited. from functools import partial -from typing import Optional, Tuple, Union, Callable, Literal +from typing import Any, Optional, Tuple, Union, Callable, Literal from typing_extensions import deprecated -from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.cutlass_dsl import T, dsl_user_op, target_version import cutlass.cutlass_dsl as cutlass_dsl from cutlass._mlir import ir from cutlass._mlir.dialects import arith, llvm, nvvm, vector -# Forward nvvm enums -from cutlass._mlir.dialects.nvvm import ( - ProxyKind, - SharedSpace, - Tcgen05WaitKind, - SetMaxRegisterAction, - RoundingModeKind, -) - from ..core import size from ..typing import ( @@ -95,25 +86,23 @@ def _enhance_enum_with_str_mapping(enum_class): """ Convert a string literal to the corresponding enum member. - :param s: String representation of the enum member, or an enum member itself (deprecated) + :param s: String representation of the enum member :return: The enum member (or None if s is None) :raises ValueError: If the string is not a valid enum member + :raises TypeError: If an enum is passed instead of a string """ - import warnings - if s is None: return None + # Check if user passed an enum (should be a string literal instead) + # This catches cases where user passes e.g., RoundingModeKind.RN instead of "rn" + from enum import Enum - # Check if s is already an enum member of the correct type - if isinstance(s, cls): - warnings.warn( - f"Passing enum member directly to {cls.__name__}.from_str() is deprecated. " - f"Please use string literals instead (e.g., '{str(s)}' instead of {cls.__name__}.{s.name}).", - DeprecationWarning, - stacklevel=2, + if isinstance(s, Enum): + raise TypeError( + f"Expected a string literal for {cls.__name__}, but got enum '{type(s).__name__}.{s.name}'. " + f"Please pass a string instead (e.g., '{str(s)}' instead of {type(s).__name__}.{s.name}). " + f"Valid string options are: {sorted(str_to_enum_map.keys())}" ) - return s - if s not in str_to_enum_map: valid_options = sorted(str_to_enum_map.keys()) raise ValueError( @@ -446,6 +435,7 @@ def warp_reduction( offset = offset // 2 return val + warp_reduction_max = partial( warp_reduction, op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else cutlass_dsl.max(x, y), @@ -460,34 +450,13 @@ def barrier(*, barrier_id=None, number_of_threads=None, loc=None, ip=None) -> No """ if barrier_id is not None: barrier_id = Int32(barrier_id).ir_value(loc=loc, ip=ip) - else: - barrier_id = Int32(0).ir_value(loc=loc, ip=ip) if number_of_threads is not None: number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip) - llvm.inline_asm( - None, - [barrier_id, number_of_threads], - "bar.sync $0, $1;", - "r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - else: - llvm.inline_asm( - None, - [barrier_id], - "bar.sync $0;", - "r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) + + nvvm.barrier( + barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip + ) @dsl_user_op @@ -496,8 +465,6 @@ def barrier_arrive( ) -> None: if barrier_id is not None: barrier_id = Int32(barrier_id).ir_value(loc=loc, ip=ip) - else: - barrier_id = Int32(0).ir_value(loc=loc, ip=ip) if number_of_threads is None: raise ValueError( @@ -505,14 +472,8 @@ def barrier_arrive( ) number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip) - llvm.inline_asm( - None, - [barrier_id, number_of_threads], - "bar.arrive $0, $1;", - "r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, + nvvm.barrier_arrive( + barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip ) @@ -638,19 +599,73 @@ def cluster_arrive_relaxed(*, aligned=None, loc=None, ip=None) -> None: @dsl_user_op def fence_proxy( - kind: ProxyKind, + kind: Literal[ + "alias", "async", "async.global", "async.shared", "tensormap", "generic" + ], *, - space: Optional[SharedSpace] = None, + space: Optional[Literal["cta", "cluster"]] = None, use_intrinsic=None, loc=None, ip=None, ) -> None: + """ + Fence operation to ensure memory consistency between proxies. + + :param kind: Proxy kind string literal: + - "alias" : Alias proxy + - "async" : Async proxy + - "async.global" : Async global proxy + - "async.shared" : Async shared proxy + - "tensormap" : Tensormap proxy + - "generic" : Generic proxy + :type kind: Literal["alias", "async", "async.global", "async.shared", "tensormap", "generic"] + :param space: Shared memory space scope string literal (optional): + - "cta" : CTA (Cooperative Thread Array) scope + - "cluster" : Cluster scope + :type space: Optional[Literal["cta", "cluster"]] + :param use_intrinsic: Whether to use intrinsic version + """ + from cutlass._mlir.dialects.nvvm import ( + SharedSpace, + ProxyKind, + ) + + # Enhance enum with str mapping + SharedSpace = _enhance_enum_with_str_mapping(SharedSpace) + ProxyKind = _enhance_enum_with_str_mapping(ProxyKind) + + kind = ProxyKind.from_str(kind) + space = SharedSpace.from_str(space) + nvvm.fence_proxy( - kind=kind, space=space, use_intrinsic=use_intrinsic, loc=loc, ip=ip + kind=kind, + space=space, + use_intrinsic=use_intrinsic, + loc=loc, + ip=ip, ) @dsl_user_op +def vote_sync_op( + pred: Boolean, kind: nvvm.VoteSyncKind, mask: Int = FULL_MASK, *, loc=None, ip=None +) -> Union[Int32, Boolean]: + """ + Performs a vote operation across the warp. + """ + return_type = Int32 if kind == nvvm.VoteSyncKind.ballot else Boolean + return return_type( + nvvm.vote_sync( + T.i32() if kind == nvvm.VoteSyncKind.ballot else T.bool(), + Int32(mask).ir_value(loc=loc, ip=ip), + Boolean(pred).ir_value(loc=loc, ip=ip), + kind, + loc=loc, + ip=ip, + ) + ) + + def vote_ballot_sync( pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None ) -> Int32: @@ -668,45 +683,7 @@ def vote_ballot_sync( See the `PTX documentation `__. """ - return Int32( - nvvm.vote_ballot_sync( - T.i32(), - Int32(mask).ir_value(loc=loc, ip=ip), - Boolean(pred).ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def vote_sync_op( - pred: Boolean, kind: str, mask: Int = FULL_MASK, *, loc=None, ip=None -) -> Union[Int32, Boolean]: - return_type = Boolean - return_type_str = "pred" - return return_type( - llvm.inline_asm( - T.bool(), - [ - Boolean(pred).ir_value(loc=loc, ip=ip), - Int32(mask).ir_value(loc=loc, ip=ip), - ], - f"""{{\n\t - .reg .pred ps;\n\t - .reg .pred pd;\n\t - setp.ne.b32 ps, $1, 0;\n\t - vote.sync.{kind}.{return_type_str} pd, ps, $2;\n\t - selp.b32 $0, 1, 0, pd;\n\t - }}""", - "=r,r,i", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) + return vote_sync_op(pred, nvvm.VoteSyncKind.ballot, mask, loc=loc, ip=ip) @dsl_user_op @@ -714,7 +691,7 @@ def vote_any_sync( pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None ) -> Boolean: """True if source predicate is True for any non-exited threads in mask. Negate the source - predicate to compute .not_all. + predicate to compute .none. :param pred: The predicate value for the current thread :type pred: Boolean @@ -727,7 +704,7 @@ def vote_any_sync( See the `PTX documentation `__. """ - return vote_sync_op(pred, "any", mask, loc=loc, ip=ip) + return vote_sync_op(pred, nvvm.VoteSyncKind.any, mask, loc=loc, ip=ip) @dsl_user_op @@ -748,7 +725,7 @@ def vote_all_sync( See the `PTX documentation `__. """ - return vote_sync_op(pred, "all", mask, loc=loc, ip=ip) + return vote_sync_op(pred, nvvm.VoteSyncKind.all, mask, loc=loc, ip=ip) @dsl_user_op @@ -767,7 +744,7 @@ def vote_uni_sync( threads in mask :rtype: Boolean """ - return vote_sync_op(pred, "uni", mask, loc=loc, ip=ip) + return vote_sync_op(pred, nvvm.VoteSyncKind.uni, mask, loc=loc, ip=ip) @dsl_user_op @@ -821,8 +798,8 @@ def fence_view_async_tmem_op( from cutlass._mlir.dialects.nvvm import Tcgen05WaitKind # Enhance enum and convert string literal to enum type - Tcgen05WaitKind_enhanced = _enhance_enum_with_str_mapping(Tcgen05WaitKind) - kind = Tcgen05WaitKind_enhanced.from_str(kind) + Tcgen05WaitKind = _enhance_enum_with_str_mapping(Tcgen05WaitKind) + kind = Tcgen05WaitKind.from_str(kind) nvvm.tcgen05_wait(kind=kind, loc=loc, ip=ip) @@ -847,9 +824,8 @@ def fence_view_async_shared( This function is usually used for async execution unit (like TMA, UMMA) after the load/store operations. """ - nvvm.fence_proxy( - nvvm.ProxyKind.async_shared, space=nvvm.SharedSpace.shared_cta, loc=loc, ip=ip - ) + # Use the fence_proxy wrapper function with string literals + fence_proxy(kind="async.shared", space="cta", loc=loc, ip=ip) @dsl_user_op @@ -859,6 +835,7 @@ def setmaxregister_increase( loc=None, ip=None, ): + from cutlass._mlir.dialects.nvvm import SetMaxRegisterAction return nvvm.setmaxregister(reg_count, SetMaxRegisterAction.increase, loc=loc, ip=ip) @@ -869,6 +846,7 @@ def setmaxregister_decrease( loc=None, ip=None, ): + from cutlass._mlir.dialects.nvvm import SetMaxRegisterAction return nvvm.setmaxregister(reg_count, SetMaxRegisterAction.decrease, loc=loc, ip=ip) @@ -880,6 +858,7 @@ def warpgroup_reg_alloc( loc=None, ip=None, ) -> None: + from cutlass._mlir.dialects.nvvm import SetMaxRegisterAction nvvm.setmaxregister(reg_count, SetMaxRegisterAction.increase, loc=loc, ip=ip) @@ -891,8 +870,10 @@ def warpgroup_reg_dealloc( loc=None, ip=None, ) -> None: + from cutlass._mlir.dialects.nvvm import SetMaxRegisterAction nvvm.setmaxregister(reg_count, SetMaxRegisterAction.decrease, loc=loc, ip=ip) + @dsl_user_op def calc_packed_f32x2_op( src_a: Tuple[Float32, Float32], @@ -908,8 +889,8 @@ def calc_packed_f32x2_op( from cutlass._mlir.dialects.nvvm import RoundingModeKind # Enhance enum and convert string literal to enum type - RoundingModeKind_enhanced = _enhance_enum_with_str_mapping(RoundingModeKind) - rnd = RoundingModeKind_enhanced.from_str(rnd) + RoundingModeKind = _enhance_enum_with_str_mapping(RoundingModeKind) + rnd = RoundingModeKind.from_str(rnd) vec_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc) vec_src_a = vector.from_elements( @@ -959,15 +940,17 @@ mul_packed_f32x2 = partial( add_packed_f32x2 = partial( calc_packed_f32x2_op, src_c=None, calc_func=nvvm.add_packed_f32x2 ) +sub_packed_f32x2 = partial( + calc_packed_f32x2_op, src_c=None, calc_func=nvvm.sub_packed_f32x2 +) + @dsl_user_op def fmax( a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None ) -> Float32: - return Float32( nvvm.fmax( - T.f32(), Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip), loc=loc, @@ -975,12 +958,25 @@ def fmax( ) ) + +@dsl_user_op +def fmin( + a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None +) -> Float32: + return Float32( + nvvm.fmin( + Float32(a).ir_value(loc=loc, ip=ip), + Float32(b).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + ) + + @dsl_user_op def rcp_approx(a: Union[float, Float32], *, loc=None, ip=None): return Float32( - nvvm.rcp_approx_ftz_f( - T.f32(), Float32(a).ir_value(loc=loc, ip=ip), loc=loc, ip=ip - ) + nvvm.rcp_approx_ftz_f(Float32(a).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) ) @@ -1024,6 +1020,68 @@ def cvt_i8_bf16(src_i8, *, loc=None, ip=None): return val_bf16 +@dsl_user_op +def cvt_i8x2_to_bf16x2(src_vec2, *, loc=None, ip=None): + # pack 2 int8 into 1 int16 value + src_i16 = llvm.bitcast(Int16.mlir_type, src_vec2, loc=loc, ip=ip) + val_i32 = llvm.inline_asm( + Int32.mlir_type, + [ + src_i16, + ], + """{\n\t + .reg .b16 scale;\n\t + mov.b16 scale, 0x8585;\n\t + cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, $1, scale;\n\t + }""", + "=r,h", + ) + + vec_bf16x2_type = ir.VectorType.get([2], BFloat16.mlir_type, loc=loc) + vec_bf16x2 = llvm.bitcast(vec_bf16x2_type, val_i32, loc=loc, ip=ip) + return vec_bf16x2 + + +@dsl_user_op +def cvt_i8x4_to_bf16x4(src_vec4, *, loc=None, ip=None): + # pack 4 int8 into 1 int32 value + src_i32 = llvm.bitcast(Int32.mlir_type, src_vec4, loc=loc, ip=ip) + rst01 = llvm.inline_asm( + Int32.mlir_type, + [ + src_i32, + ], + """{\n\t + .reg .b16 pair<2>;\n\t + .reg .b16 scale;\n\t + mov.b32 {pair0, pair1}, $1;\n\t + mov.b16 scale, 0x8585;\n\t + cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair0, scale;\n\t + }""", + "=r,r", + ) + + rst23 = llvm.inline_asm( + Int32.mlir_type, + [ + src_i32, + ], + """{\n\t + .reg .b16 pair<2>;\n\t + .reg .b16 scale;\n\t + mov.b32 {pair0, pair1}, $1;\n\t + mov.b16 scale, 0x8585;\n\t + cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair1, scale;\n\t + }""", + "=r,r", + ) + vec_type = ir.VectorType.get([2], Int32.mlir_type, loc=loc) + rst_i32 = vector.from_elements(vec_type, [rst01, rst23], loc=loc, ip=ip) + vec_bf16x4_type = ir.VectorType.get([4], BFloat16.mlir_type, loc=loc) + vec_bf16x4 = llvm.bitcast(vec_bf16x4_type, rst_i32, loc=loc, ip=ip) + return vec_bf16x4 + + # Convert vector of 2 float values to vector of 2 bfloat16 values with satfinite rounding @dsl_user_op def cvt_f32x2_bf16x2(src_vec2, *, loc=None, ip=None): @@ -1263,12 +1321,116 @@ def prmt(src, src_reg_shifted, prmt_indices, *, loc=None, ip=None): @dsl_user_op def cvt_i4_bf16(src_i4, *, loc=None, ip=None): # i4 -> i32 -> f32 -> bf - src_i32 = llvm.zext(Int32.mlir_type, src_i4, loc=loc, ip=ip) + src_i32 = llvm.sext(Int32.mlir_type, src_i4, loc=loc, ip=ip) src_f32 = llvm.sitofp(Float32.mlir_type, src_i32, loc=loc, ip=ip) bf16_val = cvt_f32_bf16(src_f32, loc=loc, ip=ip) return bf16_val +# Convert multiple shuffled int4 values to bfloat16 values. +# The input elements are assumed to be already shuffled following a specific shuffle pattern. +# Specifically, for consecutive 8 int4 values with indices of (0, 1, 2, 3, 4, 5, 6, 7), +# they are shuffled to (0, 2, 1, 3, 4, 6, 5, 7). For tailing elements less than 8, the +# shuffle pattern is (0, 2, 1, 3) for 4 elements. No shuffle is needed for less than 4 elements. +# Shuffle could help to produce converted bf16 values in the natural order of (0, 1, 2 ,3 ,4 ,5 ,6 ,7) +# without extra prmt instructions and thus better performance. +# The number of elements to be converted must be be even as specified by num_elts. +# Int4 values are packed into int32 values with upper bits filled with 0 if there are less than 4 int4 values. +# Results bfloat16 values are also packed into int32 values. +@dsl_user_op +def cvt_i4_to_bf16_with_shuffle_impl(src_i32, num_elts, *, loc=None, ip=None): + from cutlass import CUDA_VERSION + if CUDA_VERSION.major < 13: + raise cutlass_dsl.DSLCudaVerNotImplemented( + feature="cvt_i4_to_bf16_with_shuffle_impl", required_version="13.1" + ) + + num_i32_elts = num_elts // 2 + mask_odd = arith.constant(Int32.mlir_type, 0xF0F0F0F0, loc=loc, ip=ip) + mask_even = arith.constant(Int32.mlir_type, 0x0F0F0F0F, loc=loc, ip=ip) + src_odd = arith.andi(src_i32, mask_odd, loc=loc, ip=ip) + src_even = arith.andi(src_i32, mask_even, loc=loc, ip=ip) + c4 = arith.constant(Int32.mlir_type, 4, loc=loc, ip=ip) + src_even = arith.shli(src_even, c4, loc=loc, ip=ip) + rst13 = llvm.inline_asm( + Int32.mlir_type, + [ + src_odd, + ], + """{\n\t + .reg .b16 pair<2>;\n\t + .reg .b16 scale;\n\t + mov.b32 {pair0, pair1}, $1;\n\t + mov.b16 scale, 0x8181;\n\t + cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair0, scale;\n\t + }""", + "=r,r", + ) + rst57 = llvm.inline_asm( + Int32.mlir_type, + [ + src_odd, + ], + """{\n\t + .reg .b16 pair<2>;\n\t + .reg .b16 scale;\n\t + mov.b32 {pair0, pair1}, $1;\n\t + mov.b16 scale, 0x8181;\n\t + cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair1, scale;\n\t + }""", + "=r,r", + ) + rst02 = llvm.inline_asm( + Int32.mlir_type, + [ + src_even, + ], + """{\n\t + .reg .b16 pair<2>;\n\t + .reg .b16 scale;\n\t + mov.b16 scale, 0x8181;\n\t + mov.b32 {pair0, pair1}, $1;\n\t + cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair0, scale;\n\t + }""", + "=r,r", + ) + rst46 = llvm.inline_asm( + Int32.mlir_type, + [ + src_even, + ], + """{\n\t + .reg .b16 pair<2>;\n\t + .reg .b16 scale;\n\t + mov.b16 scale, 0x8181;\n\t + mov.b32 {pair0, pair1}, $1;\n\t + cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair1, scale;\n\t + }""", + "=r,r", + ) + vec_type = ir.VectorType.get([num_i32_elts], Int32.mlir_type, loc=loc) + if num_elts == 2: + prmt_index = arith.constant(Int32.mlir_type, 0x00005410, loc=loc, ip=ip) + rst = llvm.inline_asm( + Int32.mlir_type, + [ + rst02, + rst13, + prmt_index, + ], + "prmt.b32 $0, $1, $2, $3;", + "=r,r,r,r", + ) + vec_rsts = vector.from_elements(vec_type, [rst], loc=loc, ip=ip) + elif num_elts == 4: + vec_rsts = vector.from_elements(vec_type, [rst02, rst13], loc=loc, ip=ip) + else: + vec_rsts = vector.from_elements( + vec_type, [rst02, rst13, rst46, rst57], loc=loc, ip=ip + ) + return vec_rsts + + # Convert multiple int4 values to bfloat16 values. # The number of elements to be converted must be be even as specified by num_elts. # Int4 values are packed into int32 values with upper bits filled with 0 if there are less than 4 int4 values. @@ -1357,11 +1519,12 @@ def cvt_i4_to_bf16_impl(src_i32, num_elts, *, loc=None, ip=None): # Convert 2 int4 values to 2 bfloat16 values @dsl_user_op -def cvt_i4x2_to_bf16x2(src_vec2, *, loc=None, ip=None): +def cvt_i4x2_to_bf16x2(src_vec2, *, with_shuffle=False, loc=None, ip=None): + cvt_func = cvt_i4_to_bf16_with_shuffle_impl if with_shuffle else cvt_i4_to_bf16_impl # pack 2 int4 into 1 int32 value and fill upper bits with 0 src_i8 = llvm.bitcast(Int8.mlir_type, src_vec2, loc=loc, ip=ip) src_i32 = llvm.zext(Int32.mlir_type, src_i8, loc=loc, ip=ip) - rst_i32 = cvt_i4_to_bf16_impl(src_i32, 2, loc=loc, ip=ip) + rst_i32 = cvt_func(src_i32, 2, loc=loc, ip=ip) vec_bf16x2_type = ir.VectorType.get([2], BFloat16.mlir_type, loc=loc) vec_bf16x2 = llvm.bitcast(vec_bf16x2_type, rst_i32, loc=loc, ip=ip) return vec_bf16x2 @@ -1369,11 +1532,12 @@ def cvt_i4x2_to_bf16x2(src_vec2, *, loc=None, ip=None): # Convert 4 int4 values to 4 bfloat16 values @dsl_user_op -def cvt_i4x4_to_bf16x4(src_vec4, *, loc=None, ip=None): +def cvt_i4x4_to_bf16x4(src_vec4, *, with_shuffle=False, loc=None, ip=None): + cvt_func = cvt_i4_to_bf16_with_shuffle_impl if with_shuffle else cvt_i4_to_bf16_impl # pack 4 int4 into 1 int32 value and fill upper bits with 0 src_i16 = llvm.bitcast(Int16.mlir_type, src_vec4, loc=loc, ip=ip) src_i32 = llvm.zext(Int32.mlir_type, src_i16, loc=loc, ip=ip) - rst_i32 = cvt_i4_to_bf16_impl(src_i32, 4, loc=loc, ip=ip) + rst_i32 = cvt_func(src_i32, 4, loc=loc, ip=ip) vec_bf16x4_type = ir.VectorType.get([4], BFloat16.mlir_type, loc=loc) vec_bf16x4 = llvm.bitcast(vec_bf16x4_type, rst_i32, loc=loc, ip=ip) return vec_bf16x4 @@ -1381,14 +1545,16 @@ def cvt_i4x4_to_bf16x4(src_vec4, *, loc=None, ip=None): # Convert 8 int4 values to 8 bfloat16 values @dsl_user_op -def cvt_i4x8_to_bf16x8(src_vec8, *, loc=None, ip=None): +def cvt_i4x8_to_bf16x8(src_vec8, *, with_shuffle=False, loc=None, ip=None): + cvt_func = cvt_i4_to_bf16_with_shuffle_impl if with_shuffle else cvt_i4_to_bf16_impl # pack 8 int4 into 1 int32 value and fill upper bits with 0 src_i32 = llvm.bitcast(Int32.mlir_type, src_vec8, loc=loc, ip=ip) - rst_i32 = cvt_i4_to_bf16_impl(src_i32, 8, loc=loc, ip=ip) + rst_i32 = cvt_func(src_i32, 8, loc=loc, ip=ip) vec_bf16x8_type = ir.VectorType.get([8], BFloat16.mlir_type, loc=loc) vec_bf16x8 = llvm.bitcast(vec_bf16x8_type, rst_i32, loc=loc, ip=ip) return vec_bf16x8 + # Sign extend 4 int4 unpacked in 8b containers @dsl_user_op def sext_unpacked_i4x4_to_i8x4(src_vec4, *, loc=None, ip=None): @@ -1485,6 +1651,196 @@ def griddepcontrol_launch_dependents(*, loc=None, ip=None) -> None: +@dsl_user_op +def _warp_redux_sync_nvvm( + value: Numeric, + kind: Literal[ + "fmax", + "fmin", + "max", + "min", + "add", + "xor", + "or", + "and", + ], + mask_and_clamp: Int = FULL_MASK, + abs: bool = False, + nan: bool = None, + *, + loc=None, + ip=None, +) -> Numeric: + from cutlass._mlir.dialects.nvvm import ReduxKind + + # Enhance enum and convert string literal to enum type + ReduxKind = _enhance_enum_with_str_mapping(ReduxKind) + kind = ReduxKind.from_str(kind) + + value_type = type(value) + value_ir = value.ir_value(loc=loc, ip=ip) + + return value_type( + nvvm.redux_sync( + res=value_ir.type, + val=value_ir, + kind=kind, + mask_and_clamp=Int32(mask_and_clamp).ir_value(loc=loc, ip=ip), + abs=abs, + nan=nan, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _warp_redux_sync_ptx( + value: Numeric, + kind: Literal[ + "fmax", + "fmin", + "max", + "min", + ], + mask_and_clamp: Int = FULL_MASK, + abs: bool = None, + nan: bool = None, + *, + loc=None, + ip=None, +) -> Numeric: + value_type = type(value) + value_ir = value.ir_value(loc=loc, ip=ip) + mlir_type = value_type.mlir_type + mask_ir = Int32(mask_and_clamp).ir_value(loc=loc, ip=ip) + + kind_ptx_str = kind + if kind == "fmax": + kind_ptx_str = "max" + elif kind == "fmin": + kind_ptx_str = "min" + + modifiers = [] + if nan is True: + modifiers.append("NaN") + if abs is True: + modifiers.append("abs") + + modifier_str = "." + ".".join(modifiers) if modifiers else "" + ptx_instr = f"redux.sync.{kind_ptx_str}{modifier_str}.f32 $0, $1, $2;" + + return value_type( + llvm.inline_asm( + mlir_type, + [value_ir, mask_ir], + f"{ptx_instr}", + f"=f,f,i", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def warp_redux_sync( + value: Numeric, + kind: Literal[ + "fmax", + "fmin", + "max", + "min", + "add", + "xor", + "or", + "and", + ], + mask_and_clamp: Int = FULL_MASK, + *, + abs: bool = None, + nan: bool = None, + loc=None, + ip=None, +) -> Numeric: + """ + Perform warp-level reduction operation across threads. + + Reduces values from participating threads in a warp according to the specified operation. + All threads in the mask receive the same result. + + :param value: Input value to reduce + :type value: Numeric + :param kind: Reduction operation. Supported operations: + - Integer types (Int32/Uint32): "add", "and", "max", "min", "or", "xor" + - Float types (Float32): "fmax", "fmin" (or "max"/"min" which auto-convert to "fmax"/"fmin") + :type kind: Literal["add", "and", "max", "min", "or", "xor", "fmin", "fmax"] + :param mask_and_clamp: Warp participation mask (default: FULL_MASK = 0xFFFFFFFF) + :type mask_and_clamp: Int + :param abs: Apply absolute value before reduction (float types only) + :type abs: bool + :param nan: Enable NaN propagation for fmax/fmin operations (float types only) + :type nan: Optional[bool] + :return: Reduced value (same for all participating threads) + :rtype: Numeric + """ + # Convert value to Numeric type if needed + if not isinstance(value, Numeric): + value = as_numeric(value) + + # Determine value type and choose appropriate implementation + value_type = type(value) + mlir_type = value_type.mlir_type + + # Use inline PTX for float types, NVVM for integer types + if mlir_type == T.f32(): + return _warp_redux_sync_ptx( + value, kind, mask_and_clamp, abs, nan, loc=loc, ip=ip + ) + else: + return _warp_redux_sync_nvvm( + value, kind, mask_and_clamp, abs, nan, loc=loc, ip=ip + ) + + +@dsl_user_op +def atomic_max_float32( + ptr, + value: Float32, + *, + positive_only: bool = True, + loc=None, + ip=None, +) -> Float32: + """ + Performs an atomic max operation on a float32 value in global memory. + + This implementation works correctly for non-negative values (>= 0) using direct bitcast. + + :param ptr: Pointer to the memory location + :param value: The float32 value to compare and potentially store (should be >= 0 for correct results) + :type value: Float32 + :param positive_only: If True (default), assumes input values are non-negative. + This parameter is provided for API compatibility and future extensions. + :type positive_only: bool + :return: The old value at the memory location + :rtype: Float32 + """ + from cutlass._mlir.dialects.nvvm import AtomicOpKind + + value_int = llvm.bitcast(T.i32(), value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + + old_value_int = nvvm.atomicrmw( + AtomicOpKind.MAX, + ptr, + value_int, + loc=loc, + ip=ip, + ) + + return Float32(llvm.bitcast(T.f32(), old_value_int, loc=loc, ip=ip)) + + def _normalize_ptr(addr, *, loc=None, ip=None) -> ir.Value: """ Helper function to normalize pointer types to MLIR ir.Value. @@ -1549,7 +1905,7 @@ def _atomic( :rtype: Union[Numeric, ir.Value] """ from cutlass._mlir.dialects.nvvm import AtomicOpKind, MemOrderKind, MemScopeKind - from cutlass.utils.version_info import CUDA_VERSION + from cutlass import CUDA_VERSION # Enhance enums and convert string literals to enum types AtomicOpKind = _enhance_enum_with_str_mapping(AtomicOpKind) @@ -1848,7 +2204,7 @@ def atomic_cas( :rtype: Numeric """ from cutlass._mlir.dialects.nvvm import AtomicOpKind, MemOrderKind, MemScopeKind - from cutlass.utils.version_info import CUDA_VERSION + from cutlass import CUDA_VERSION # Enhance enums and convert string literals to enum types MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind) @@ -1882,6 +2238,17 @@ def atomic_cas( loc=loc, ip=ip, ) + elif CUDA_VERSION.major == 13 and CUDA_VERSION.minor == 1: + result = nvvm.atomicrmw( + op=AtomicOpKind.CAS, + ptr=ptr, + a=val_ir, + b=cmp_ir, + mem_order=sem, + syncscope=scope, + loc=loc, + ip=ip, + ) else: result = nvvm.atomicrmw( op=AtomicOpKind.CAS, @@ -2190,3 +2557,42 @@ def cvt_f4e2m1x8_to_f16x8(src_vec8, *, loc=None, ip=None): vec_f16x8_type = ir.VectorType.get([8], Float16.mlir_type, loc=loc) vec_f16x8 = llvm.bitcast(vec_f16x8_type, vec_f32x4, loc=loc, ip=ip) return vec_f16x8 + + +@dsl_user_op +def mapa(ptr, cta_rank_in_cluster=0, *, loc=None, ip=None): + """ + Map a pointer to distributed shared memory across cluster. + + Portable wrapper that uses the appropriate NVVM API based on CUDA version: + - CUDA 13.1+: Uses nvvm.mapa with dsmem address space + - CUDA 12.9: Uses nvvm.mapa_shared_cluster + + Args: + ptr: Pointer to shared memory (llvm_ptr attribute will be used) + cta_rank_in_cluster: CTA rank within the cluster (default 0) + + Returns: + Mapped LLVM pointer to shared memory + """ + if target_version(min_version="13.1"): + dsmem_ptr_ty = llvm.PointerType.get(7) # dsmem + smem_ptr_ty = llvm.PointerType.get(3) # smem + + llvm_ptr = nvvm.mapa( + dsmem_ptr_ty, + ptr.llvm_ptr, + Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return llvm.addrspacecast(smem_ptr_ty, llvm_ptr, loc=loc, ip=ip) + else: + llvm_ptr = ptr.llvm_ptr + return nvvm.mapa_shared_cluster( + llvm_ptr.type, + llvm_ptr, + Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) diff --git a/python/CuTeDSL/cutlass/cute/arch/smem.py b/python/CuTeDSL/cutlass/cute/arch/smem.py index 35b420ea..1a89db71 100644 --- a/python/CuTeDSL/cutlass/cute/arch/smem.py +++ b/python/CuTeDSL/cutlass/cute/arch/smem.py @@ -17,7 +17,7 @@ import cutlass._mlir.dialects.cute as _cute_ir import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir from cutlass._mlir import ir -from ..typing import Pointer, Numeric, NumericMeta +from ..typing import Pointer, Numeric, NumericMeta, Layout @dsl_user_op diff --git a/python/CuTeDSL/cutlass/cute/arch/tmem.py b/python/CuTeDSL/cutlass/cute/arch/tmem.py index 2e393672..c5f63750 100644 --- a/python/CuTeDSL/cutlass/cute/arch/tmem.py +++ b/python/CuTeDSL/cutlass/cute/arch/tmem.py @@ -55,7 +55,6 @@ def get_max_tmem_alloc_cols(compute_capability: str) -> int: return TMEM_MAX_ALLOC_COLUMNS_MAP[compute_capability] - def get_min_tmem_alloc_cols(compute_capability: str) -> int: """Get the minimum TMEM allocation columns for a given compute capability. @@ -179,11 +178,9 @@ def dealloc_tmem( :param num_columns: The number of columns in the TMEM allocation :type num_columns: Int :param is_two_cta: Optional boolean parameter for 2-CTA MMAs - :param arch: The architecture of the GPU. - :type arch: str """ - tmem_max_alloc_cols = get_max_tmem_alloc_cols(arch) tmem_min_alloc_cols = get_min_tmem_alloc_cols(arch) + tmem_max_alloc_cols = get_max_tmem_alloc_cols(arch) if isinstance(num_columns, int): if ( num_columns < tmem_min_alloc_cols diff --git a/python/CuTeDSL/cutlass/cute/atom.py b/python/CuTeDSL/cutlass/cute/atom.py index ba453317..0d4ec36e 100644 --- a/python/CuTeDSL/cutlass/cute/atom.py +++ b/python/CuTeDSL/cutlass/cute/atom.py @@ -10,7 +10,7 @@ # is strictly prohibited. from abc import ABC, ABCMeta, abstractmethod -from typing import Type, Union, Optional, Any, overload +from typing import Type, Union, Optional, Any, List, Tuple, overload from .typing import Shape, Layout, Tile, Tensor, Numeric, Int32 from .core import ( @@ -285,6 +285,8 @@ class MmaAtom(Atom): if self.op is not None: self.op._verify_fragment_B(input, loc=loc, ip=ip) input = input.value + if isinstance(input, tuple): + input = _pack_shape(input, loc=loc, ip=ip) return _cute_ir.mma_make_fragment( _cute_ir.MmaOperand.B, self._trait.value, input, loc=loc, ip=ip ) @@ -1111,25 +1113,66 @@ def make_tiled_copy_C_atom(atom: CopyAtom, mma: TiledMma, *, loc=None, ip=None): return _make_tiled_copy(atom, layout_tv, tiler_mn, loc=loc, ip=ip) +def _normalize_variadic_tensor_operand( + x: Union["Tensor", List["Tensor"], Tuple["Tensor", ...]], name: str +) -> List["Tensor"]: + """Normalize a Tensor or sequence of Tensors to a list of Tensors. + + Helper function for operations with variadic operands. + """ + if isinstance(x, Tensor): + return [x] + if isinstance(x, (list, tuple)): + if len(x) == 0: + raise ValueError(f"`{name}` must contain at least one Tensor") + if not all(isinstance(t, Tensor) for t in x): + raise TypeError(f"All elements of `{name}` must be Tensor") + return list(x) # type: ignore + raise TypeError(f"`{name}` must be a Tensor or a sequence of Tensors") + + @dsl_user_op def copy_atom_call( atom: CopyAtom, - src: Tensor, - dst: Tensor, + src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], + dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], *, pred: Optional[Tensor] = None, loc=None, ip=None, **kwargs, ) -> None: - """Executes a single copy atom operation between two tensors. + """ + Execute a single copy atom operation. + + The copy_atom_call operation executes a copy atom with the given operands. + Source and destination tensors have layout profile ``(V)``. + + The ``V-mode`` represents either: + + - A singular mode directly consumable by the provided Copy Atom + - A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``, + + For src/dst layout like ``(V, Rest...)``, the layout profile of ``pred`` must match ``(Rest...)``. + + - Certain Atoms may require additional operation-specific keyword arguments. + - Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned + for future releases. + + Both ``src`` and ``dst`` operands are variadic, containing a variable number of tensors: + + - For regular copy, ``src`` and ``dst`` each contain a single tensor. + - For copy with auxiliary operands, they contain the main tensor followed by + auxiliary tensors. For example: :param atom: Copy atom specifying the transfer operation :type atom: CopyAtom - :param src: Source tensor with layout profile ``(V)`` - :type src: Tensor - :param dst: Destination tensor with layout profile ``(V)`` - :type dst: Tensor + :param src: Source tensor(s) with layout profile ``(V)``. Can be a single Tensor + or a list/tuple of Tensors for operations with auxiliary source operands. + :type src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] + :param dst: Destination tensor(s) with layout profile ``(V)``. Can be a single Tensor + or a list/tuple of Tensors for operations with auxiliary destination operands. + :type dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] :param pred: Optional predication tensor for conditional transfers, defaults to None :type pred: Optional[Tensor], optional :param loc: Source location information, defaults to None @@ -1142,51 +1185,89 @@ def copy_atom_call( :return: None :rtype: None - The copy_atom_call operation executes a single copy atom with the given operands. - Source and destination tensors with layout profile like ``(V)``. - - The ``V-mode`` represents either: - - - A singular mode directly consumable by the provided Copy Atom - - A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``, - - For src/dst layout like ``(V, Rest...)``, the layout profile of ``pred`` must match ``(Rest...)``. - **Examples**: .. code-block:: python - # Basic copy atom operation + # Regular copy atom operation cute.copy_atom_call(copy_atom, src, dst) # Predicated copy atom operation cute.copy_atom_call(copy_atom, src, dst, pred=pred) - .. note:: - - - Certain Atoms may require additional operation-specific keyword arguments. - - Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned - for future releases. - """ - if isinstance(src.type, _cute_ir.MemRefType) and isinstance( - dst.type, _cute_ir.MemRefType + # Normalize src/dst to lists for variadic IR operands, while keeping old API working. + src_list = _normalize_variadic_tensor_operand(src, "src") + dst_list = _normalize_variadic_tensor_operand(dst, "dst") + + # Validate first src/dst for element type width check + if isinstance(src_list[0].type, _cute_ir.MemRefType) and isinstance( + dst_list[0].type, _cute_ir.MemRefType ): - if src.element_type.width != dst.element_type.width: + if src_list[0].element_type.width != dst_list[0].element_type.width: raise TypeError( "`copy_atom_call` currently only supports equal source and destination " "element type bit width" ) - if rank(src, mode=[0]) > 2 or rank(dst, mode=[0]) > 2: + if rank(src_list[0], mode=[0]) > 2 or rank(dst_list[0], mode=[0]) > 2: raise NotImplementedError( "V-mode (mode-0) with rank > 2 is not supported yet, " - f"but got rank(src, mode=[0]) = {rank(src, mode=[0])} and rank(dst, mode=[0]) = {rank(dst, mode=[0])}" + f"but got rank(src, mode=[0]) = {rank(src_list[0], mode=[0])} and rank(dst, mode=[0]) = {rank(dst_list[0], mode=[0])}" ) value = atom._unpack(loc=loc, ip=ip, **kwargs) if isinstance(pred, Tensor): pred = pred.value - return _cute_ir.copy_atom_call( - value, src.value, dst.value, pred=pred, loc=loc, ip=ip + src_vals = [t.value for t in src_list] + dst_vals = [t.value for t in dst_list] + return _cute_ir.copy_atom_call(value, src_vals, dst_vals, pred=pred, loc=loc, ip=ip) + + +@dsl_user_op +def mma_atom_call( + atom: MmaAtom, + d: Tensor, + a: Tensor, + b: Tensor, + c: Tensor, + *, + loc=None, + ip=None, + **kwargs, +) -> None: + """ + Execute a single MMA atom operation. + + The mma_atom_call operation executes an MMA atom with the given operands. + This performs a matrix multiplication and accumulation operation: + D = A * B + C + + Note: The tensors 'd', 'a', 'b', and 'c' must only have a single fragment. + + :param atom: The MMA atom to execute + :type atom: MmaAtom + :param d: Destination tensor (output accumulator) + :type d: Tensor + :param a: First source tensor (matrix A) + :type a: Tensor + :param b: Second source tensor (matrix B) + :type b: Tensor + :param c: Third source tensor (input accumulator C) + :type c: Tensor + :param loc: Source location for MLIR, defaults to None + :type loc: Optional[Location], optional + :param ip: Insertion point, defaults to None + :type ip: Optional[InsertionPoint], optional + + Examples: + + .. code-block:: python + + # Call an MMA atom operation + cute.mma_atom_call(mma_atom, d_tensor, a_tensor, b_tensor, c_tensor) + """ + value = atom._unpack(loc=loc, ip=ip, **kwargs) + return _cute_ir.mma_atom_call( + value, d.value, a.value, b.value, c.value, loc=loc, ip=ip ) diff --git a/python/CuTeDSL/cutlass/cute/core.py b/python/CuTeDSL/cutlass/cute/core.py index 632a3b77..2ab6fd6c 100644 --- a/python/CuTeDSL/cutlass/cute/core.py +++ b/python/CuTeDSL/cutlass/cute/core.py @@ -10,14 +10,14 @@ # is strictly prohibited. from functools import partial, reduce -import inspect from inspect import isclass from typing import Any, Dict, List, Optional, Tuple, Type, Union, overload +from cutlass import const_expr from typing_extensions import deprecated from cutlass._mlir import ir -from cutlass._mlir.dialects import builtin, llvm, vector +from cutlass._mlir.dialects import builtin, llvm, vector, arith, nvvm from cutlass._mlir.dialects import cute as _cute_ir from cutlass._mlir.dialects.cute import ( Ratio as _Ratio, @@ -125,6 +125,7 @@ __all__ = [ "shape", "recast_ptr", "make_ptr", + "get_remote_smem_ptr_in_cluster", "composition", "complement", "right_inverse", @@ -247,7 +248,7 @@ def _unpack_x_tuple(t: Union[ir.Type, ir.Value], *, loc=None, ip=None) -> XTuple vals = [] else: vals = get_leaves(t, loc=loc, ip=ip) - if not isinstance(vals, list): + if not isinstance(vals, ir.OpResultList): vals = [vals] else: raise TypeError(f"expects static type or value, but got {t}") @@ -383,9 +384,9 @@ class IntValue(cutlass_arith.ArithValue): @property def divisibility(self): - assert isinstance( - self.get_typed_value().type, _cute_ir.IntTupleType - ), f"expected self.get_typed_value() to be int_tuple type, but got {self.get_typed_value().type}" + assert isinstance(self.get_typed_value().type, _cute_ir.IntTupleType), ( + f"expected self.get_typed_value() to be int_tuple type, but got {self.get_typed_value().type}" + ) return self.get_typed_value().type.get_divisibility([0]) def __str__(self): @@ -429,7 +430,9 @@ class IntValue(cutlass_arith.ArithValue): @dsl_user_op @_binary_op def __add__(self, other, *, loc=None, ip=None): - return _cute_ir.tuple_add(self.get_typed_value(), other, loc=loc, ip=ip) + return _cute_ir.tuple_add( + self.get_typed_value(loc=loc, ip=ip), other, loc=loc, ip=ip + ) @dsl_user_op @_binary_op @@ -461,8 +464,10 @@ class IntValue(cutlass_arith.ArithValue): @dsl_user_op @_binary_op - def __radd__(self, other, *, loc=None, ip=None): - return _cute_ir.tuple_add(other, self.get_typed_value(), loc=loc, ip=ip) + def __radd__(self, other, *, loc=None, ip=None) -> "IntValue": + return _cute_ir.tuple_add( + other, self.get_typed_value(loc=loc, ip=ip), loc=loc, ip=ip + ) @dsl_user_op @_binary_op @@ -1207,10 +1212,6 @@ class _ComposedLayout(ComposedLayout): @property @dsl_user_op def shape(self, *, loc=None, ip=None) -> Shape: - return self.shape_method(loc=loc, ip=ip) - - @dsl_user_op - def shape_method(self, *, loc=None, ip=None) -> Shape: return _unpack_x_tuple( _cute_ir.get_shape(self.value, loc=loc, ip=ip), loc=loc, ip=ip ) @@ -1262,9 +1263,9 @@ class _ComposedLayout(ComposedLayout): # In this context, a _ComposedLayout instance is an encapsulated ir.Value which is automatically created # by value caster for ComposedLayout typed values assert len(values) == 1, f"Expected 1 value, but got {len(values)}" - assert isinstance( - values[0], (_ComposedLayout, ir.Value) - ), f"Expected _ComposedLayout or ir.Value, but got {type(values[0])}" + assert isinstance(values[0], (_ComposedLayout, ir.Value)), ( + f"Expected _ComposedLayout or ir.Value, but got {type(values[0])}" + ) return _ComposedLayout( values[0] if isinstance(values[0], ir.Value) else values[0].value, ) @@ -1313,9 +1314,9 @@ class _Pointer(Pointer): # In this context, a _Pointer instance is an encapsulated ir.Value which is automatically created # by value caster for cute.ptr typed values assert len(values) == 1, f"Expected 1 value, but got {len(values)}" - assert isinstance( - values[0], (_Pointer, ir.Value) - ), f"Expected _Pointer or ir.Value, but got {type(values[0])}" + assert isinstance(values[0], (_Pointer, ir.Value)), ( + f"Expected _Pointer or ir.Value, but got {type(values[0])}" + ) return _Pointer( values[0] if isinstance(values[0], ir.Value) else values[0].value ) @@ -1359,29 +1360,12 @@ class _Pointer(Pointer): """ Get the LLVM pointer representation of this pointer. - :param loc: Source location for MLIR, defaults to None - :type loc: Optional[Location] - :param ip: Insertion point for MLIR, defaults to None - :type ip: Optional[InsertionPoint] :return: The LLVM pointer representation :rtype: ir.Value """ - return self.to_llvm_ptr(loc=loc, ip=ip) - - @dsl_user_op - @lru_cache_ir() - def to_llvm_ptr(self, *, loc=None, ip=None) -> ir.Value: - """ - Get the LLVM pointer representation of this pointer. (Used by internal API to propagate loc and ip) - - :param loc: Source location for MLIR, defaults to None - :type loc: Optional[Location] - :param ip: Insertion point for MLIR, defaults to None - :type ip: Optional[InsertionPoint] - :return: The LLVM pointer representation - :rtype: ir.Value - """ - llvm_ptr_ty = llvm.PointerType.get(self.memspace.value) + llvm_ptr_ty = llvm.PointerType.get( + self.memspace.value if self.memspace != AddressSpace.rmem else 0 + ) return builtin.unrealized_conversion_cast( [llvm_ptr_ty], [self.value], loc=loc, ip=ip ) @@ -1679,7 +1663,16 @@ def printf(*args, loc=None, ip=None) -> None: elif isinstance(arg0, tuple): # Assume it's a tile return _pack_tile(arg0) - elif isinstance(arg0, (_Tensor, _Pointer, _ComposedLayout)): + elif isinstance(arg0, _Tensor): + arg0._check_can_load_store() + if isinstance(arg0.layout, ComposedLayout) and isinstance( + arg0.layout.inner, Swizzle + ): + raise NotImplementedError( + "tensor with swizzled layout (PISL) is not supported in printf, please use swizzled pointer (PDSL) instead" + ) + return arg0.value + elif isinstance(arg0, (_Pointer, _ComposedLayout)): return arg0.value else: raise TypeError(f"unsupported argument type in printf, got {type(arg)}") @@ -1751,6 +1744,7 @@ def make_swizzle(b, m, s, *, loc=None, ip=None): return Swizzle(static(ty, loc=loc, ip=ip)) + @dsl_user_op def static(value, *, loc=None, ip=None): return _cute_ir.static(value, loc=loc, ip=ip) @@ -3409,39 +3403,90 @@ def make_ptr( loc=None, ip=None, ) -> Pointer: + # Perform checks if dtype is None or not isinstance(dtype, NumericMeta): raise TypeError(f"expects dtype to be a type of Numeric, but got {dtype}") - if not isinstance(mem_space, AddressSpace): raise TypeError(f"expects mem_space to be an AddressSpace, but got {mem_space}") - if isinstance(value, ir.Value) and llvm.PointerType.isinstance(value.type): value = llvm.ptrtoint(T.i64(), value) - if not is_integer(value): raise TypeError(f"expects integer value, but got {type(value)}") - value = Int32(value) if mem_space == AddressSpace.tmem else Int64(value) + + # TMEM addresses are 32b wide + is_tmem = mem_space == AddressSpace.tmem value = Int32(value) if mem_space == AddressSpace.tmem else Int64(value) + # Set the alignment of the pointer bytes_per_elt = max(1, dtype.width // 8) if assumed_align is None: assumed_align = bytes_per_elt - if bytes_per_elt % assumed_align != 0 and assumed_align % bytes_per_elt != 0: raise ValueError( f"{bytes_per_elt=} is not a multiple of {assumed_align=} and vice versa." ) - aligned_ty = _cute_ir.ConstrainedIntType.get(assumed_align, type(value).width) aligned_intptr = _cute_ir.assume( aligned_ty, value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip ) + # Construct the pointer Type data_ty = T.i8() if dtype is None else dtype.mlir_type ptr_ty = _cute_ir.PtrType.get(data_ty, mem_space, assumed_align) return _cute_ir.inttoptr(ptr_ty, aligned_intptr, loc=loc, ip=ip) +@dsl_user_op +def get_remote_smem_ptr_in_cluster( + smem_ptr: Pointer, + cta_rank_in_cluster: Int, + *, + loc=None, + ip=None, +) -> Pointer: + """ + Get the remote shared memory CuTe pointer in a cluster. + + :param smem_ptr: The current shared memory pointer + :type smem_ptr: Pointer + :param cta_rank_in_cluster: The peer CTA rank in cluster to get the remote pointer for + :type cta_rank_in_cluster: Int + :param loc: Source location for MLIR, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point, defaults to None + :type ip: Optional[InsertionPoint] + + :return: The remote shared memory CuTe pointer + :rtype: Pointer + + """ + cur_llvm_ptr = smem_ptr.llvm_ptr + remote_llvm_ptr = nvvm.mapa( + llvm.PointerType.get(7), # LLVM dsmem address space + cur_llvm_ptr, + Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + remote_llvm_ptr_cast = llvm.addrspacecast( + llvm.PointerType.get(AddressSpace.smem), remote_llvm_ptr, loc=loc, ip=ip + ) + remote_ptr = make_ptr( + smem_ptr.dtype, + remote_llvm_ptr_cast, + AddressSpace.smem, + assumed_align=smem_ptr.alignment, + loc=loc, + ip=ip, + ) + if const_expr(smem_ptr.value.type.is_swizzled): + sw = Swizzle(static(smem_ptr.value.type.swizzle_type)) + remote_ptr = recast_ptr( + remote_ptr, swizzle_=sw, dtype=smem_ptr.dtype, loc=loc, ip=ip + ) + return remote_ptr + + # # Layout algebra # @@ -3868,9 +3913,7 @@ def local_tile( return _cute_ir.local_tile( input=input.value, tile=tiler_val, - static_tile=None, coord=coord_val, - static_coord=None, proj=proj, loc=loc, ip=ip, @@ -3907,9 +3950,9 @@ def make_layout_image_mask( sliced_lay, offset = slice_and_offset(slicer, lay, loc=loc, ip=ip) # Given that we replace only one mode with _, the rank of the slice should be 1 assert rank(sliced_lay) == 1 - assert is_static( - sliced_lay - ), "make_layout_image_mask requires the layout to be static" + assert is_static(sliced_lay), ( + "make_layout_image_mask requires the layout to be static" + ) # Create the mask of the image mcast_mask = Int16(0) @@ -3952,6 +3995,7 @@ def leading_dim(shape: Shape, stride: Stride) -> Union[int, Tuple[int, ...], Non return find_if(stride, pred_fn=pred_fn) + @dsl_user_op def make_layout_tv( thr_layout: Layout, val_layout: Layout, *, loc=None, ip=None @@ -4468,9 +4512,9 @@ class struct: """ Return the round-up offset up to the next multiple of align. """ - assert align > 0 and not ( - align & (align - 1) - ), "align should be a strictly positive power of 2." + assert align > 0 and not (align & (align - 1)), ( + "align should be a strictly positive power of 2." + ) return (offset + (align - 1)) & ~(align - 1) @@ -4607,29 +4651,11 @@ class FastDivmodDivisor: new_obj = object.__new__(FastDivmodDivisor) new_obj._divisor = values[0] return new_obj + def __repr__(self): return f"FastDivmodDivisor({self._divisor.type})" -# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator -FastDivmodDivisor.__init__.__signature__ = inspect.Signature( - [ - inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD), - inspect.Parameter( - "divisor", - inspect.Parameter.POSITIONAL_OR_KEYWORD, - annotation=Integer, - ), - inspect.Parameter( - "is_power_of_2", - inspect.Parameter.POSITIONAL_OR_KEYWORD, - default=None, - annotation=bool, - ), - ] -) - - @dsl_user_op def fast_divmod_create_divisor( divisor: Integer, *, loc=None, ip=None diff --git a/python/CuTeDSL/cutlass/cute/experimental/README.md b/python/CuTeDSL/cutlass/cute/experimental/README.md new file mode 100644 index 00000000..914330ab --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/experimental/README.md @@ -0,0 +1,57 @@ +# CuTe Experimental APIs + +> **Note:** APIs in this module are experimental and subject to change. +> +> This module serves as a staging area for new CuTe functionality that is still under active development. Performance, compile time, and interoperability with CuTe are works in progress. API signatures, behavior, and naming conventions may change without notice between releases. +> +> Once these APIs are stabilized, they will be migrated to the main `cute` submodules. +> +> Users are encouraged to experiment with these APIs but should be prepared to update their code as the interfaces evolve. + +## Core APIs (`core.py`) + +- `elect_sync` — Elects one thread within a warp +- `get_mbarrier` — Returns the mbarrier pointer for a given stage token +- `create_pipeline` — Creates a circular buffer of synchronization primitives indexed by stage count +- `create_pipeline_with_mask` — Creates a pipeline with an arrival mask for cluster-scoped synchronization +- `pipeline_advance_iterator` — Advances a pipeline iterator to the next stage +- `producer_acquire` / `producer_commit` — Producer-side pipeline synchronization +- `consumer_wait` / `consumer_release` / `consumer_tail` — Consumer-side pipeline synchronization +- `get_pipeline_produce_stage` / `get_pipeline_consume_stage` — Gets pipeline stage tokens + +## Memory APIs (`memory.py`) + +- `allocate` — Allocate a buffer with given type, layout, and address space +- `tma_load` — Copy tensor from global memory to shared memory using TMA +- `tma_load_multicast` — Copy tensor from global memory to shared memory using TMA with multicast +- `tma_store` — Copy tensor from shared memory to global memory using TMA +- `copy` — Copy tensor from src to dst using a given copy atom + +## Algorithm APIs (`algorithm.py`) + +- `simt_auto_vec_copy` — Copies a tensor between buffers with single thread (auto-vectorized) +- `partition` — Partition a buffer into a given layout and tiler +- `partition_and_copy` — Combines partitioning and copying in a single operation + +## Math APIs (`math.py`) + +- `dot` — Computes a dot product of two tensors using an MMA atom +- `dot_block_scaled` — Computes a block-scaled dot product with scale factors + +## Pipeline Classes (`pipeline.py`) + +- `GenericPipeline` — Generic pipeline for any producer/consumer combination +- `TMAToUMMAPipeline` — Pipeline for TMA load to UMMA consumption +- `TMAToAsyncPipeline` — Pipeline for TMA load to async consumer +- `AsyncToUMMAPipeline` — Pipeline for async producer to UMMA consumption +- `UMMAtoAsyncPipeline` — Pipeline for UMMA producer to async consumer +- `TMAStorePipeline` — Pipeline for SMEM producer to TMA store consumer + +## Utilities (`utils.py`) + +- `get_cta_v_map_ab` — Compute CTA-V map for A/B operands +- `get_cta_v_map_c` — Compute CTA-V map for C operand +- `make_tmem_layout_acc` — Derive TMEM accumulator buffer layout from a tiled MMA +- `make_tmem_layout_a` — Derive TMEM A-operand buffer layout from a tiled MMA +- `make_t2r_rmem_layout` — Derive per-thread RMEM buffer layout for the T2R epilogue copy + diff --git a/python/CuTeDSL/cutlass/cute/experimental/__init__.py b/python/CuTeDSL/cutlass/cute/experimental/__init__.py old mode 100755 new mode 100644 index d913a607..bee9202d --- a/python/CuTeDSL/cutlass/cute/experimental/__init__.py +++ b/python/CuTeDSL/cutlass/cute/experimental/__init__.py @@ -9,6 +9,15 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -raise NotImplementedError( - "CuTe Experimental module is only supported on Cuda toolkit 13.1 and above!" -) +from ... import cutlass_dsl as _dsl + +jit = _dsl.CuteExperimentalDSL.jit +kernel = _dsl.CuteExperimentalDSL.kernel +compile = _dsl.CompileCallable() + +from .algorithm import * +from .core import * +from .math import * +from .memory import * +from .pipeline import * +from .utils import * diff --git a/python/CuTeDSL/cutlass/cute/experimental/algorithm.py b/python/CuTeDSL/cutlass/cute/experimental/algorithm.py new file mode 100644 index 00000000..bf8d7cf4 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/experimental/algorithm.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from cutlass import cute +from cutlass.cutlass_dsl import dsl_user_op +from cutlass._mlir.dialects import lir as cutlass_lir + +from .memory import copy + + +@dsl_user_op +def simt_auto_vec_copy( + src: cute.Tensor, dst: cute.Tensor, async_op=False, loc=None, ip=None +): + """ + Copies a tensor between two cute.memref buffers with single thread. + + :param src: Source tensor + :type src: cute.Tensor + :param dst: Destination tensor + :type dst: cute.Tensor + :param async_op: Whether to use asynchronous operation, defaults to False + :type async_op: bool, optional + """ + if async_op: + cutlass_lir.SimtAutoVecCopyOp( + src.value, dst.value, async_=True, cache="always", loc=loc, ip=ip + ) + else: + cutlass_lir.SimtAutoVecCopyOp(src.value, dst.value, loc=loc, ip=ip) + + +@dsl_user_op +def partition( + buffer: cute.Tensor, agent_id: cute.Int32, *, layout_tv, tiler, loc=None, ip=None +) -> cute.Tensor: + """ + Partition a buffer into a given layout and tiler. + + :param buffer: Buffer to partition + :type buffer: cute.Tensor + :param agent_id: Agent ID + :type agent_id: cute.Int32 + :param layout_tv: Layout tensor + :type layout_tv: cute.Tensor + :param tiler: Tiler + :type tiler: cute.Tensor + """ + assert isinstance(agent_id, cute.Int32), ( + f"Expected agent_id to be cute.Int32, got {type(agent_id)}" + ) + partition_op = cutlass_lir.PartitionOp( + buffer.value, + agent_id.ir_value(), + layout_tv=layout_tv.type.attribute, + tiler=tiler.type.attribute, + loc=loc, + ip=ip, + ) + return partition_op.result + + +@dsl_user_op +def partition_and_copy( + tiled_copy: cute.core.ThrCopy, + src: cute.Tensor, + dst: cute.Tensor, + *, + loc=None, + ip=None, +): + """ + Copies a tensor between two cute.memref buffer + + :param tiled_copy: Tiled copy + :type tiled_copy: cute.core.ThrCopy + :param src: Source tensor + :type src: cute.Tensor + :param dst: Destination tensor + :type dst: cute.Tensor + """ + src_partitioned = src + dst_partitioned = dst + tid_x = tiled_copy.thr_idx + if src.memspace != cute.AddressSpace.rmem: + src_partitioned = partition( + src, + tid_x, + layout_tv=tiled_copy.layout_src_tv_tiled, + tiler=cute.core._pack_tile(tiled_copy.tiler_mn), + ) + if dst.memspace != cute.AddressSpace.rmem: + dst_partitioned = partition( + dst, + tid_x, + layout_tv=tiled_copy.layout_dst_tv_tiled, + tiler=cute.core._pack_tile(tiled_copy.tiler_mn), + ) + + # Handle copy where copy atom is used for both partition and copy during smem to rmem and rmem to smem copies + if type(tiled_copy.op) in [ + cute.nvgpu.warp.LdMatrix8x8x16bOp, + cute.nvgpu.warp.LdMatrix16x16x8bOp, + cute.nvgpu.warp.StMatrix8x8x16bOp, + cute.nvgpu.warp.StMatrix16x8x8bOp, + ]: + copy( + src_partitioned, + dst_partitioned, + copy_atom=tiled_copy, + loc=loc, + ip=ip, + ) + + # The rest handles copy where copy atom is used for partition + elif ( + src.memspace, + dst.memspace, + ) in [ + (cute.AddressSpace.rmem, cute.AddressSpace.smem), + (cute.AddressSpace.smem, cute.AddressSpace.rmem), + (cute.AddressSpace.rmem, cute.AddressSpace.gmem), + (cute.AddressSpace.gmem, cute.AddressSpace.rmem), + ]: + simt_auto_vec_copy(src_partitioned, dst_partitioned, loc=loc, ip=ip) + elif ( + src.memspace == cute.AddressSpace.gmem + and dst.memspace == cute.AddressSpace.smem + ): + simt_auto_vec_copy( + src_partitioned, dst_partitioned, async_op=True, loc=loc, ip=ip + ) + + # Handle copy where copy atom is used for partition and copy + else: + copy( + src_partitioned, + dst_partitioned, + copy_atom=tiled_copy, + loc=loc, + ip=ip, + ) diff --git a/python/CuTeDSL/cutlass/cute/experimental/core.py b/python/CuTeDSL/cutlass/cute/experimental/core.py new file mode 100644 index 00000000..42159d01 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/experimental/core.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +from cutlass.cutlass_dsl import dsl_user_op +from cutlass._mlir.dialects import lir as cutlass_lir_ir, nvvm as _nvvm +from cutlass._mlir import ir +from cutlass.cutlass_dsl import lru_cache_ir +from cutlass._mlir.dialects.core import OperationTypeEnum +from cutlass import cute + + +@dsl_user_op +def elect_sync(loc=None, ip=None): + """ + Elects one predicated thread within a warp. + """ + return _nvvm.elect_sync(loc=loc, ip=ip) + + +@dsl_user_op +def get_mbarrier(stage_token, loc=None, ip=None): + """ + Returns the mbarrier pointer for a given stage token. + """ + return cutlass_lir_ir.GetMbarrierOp(stage_token, loc=loc, ip=ip) + + +@ir.register_value_caster(cutlass_lir_ir.PipelineStateType.get_static_typeid()) +class PipelineState(ir.Value): + def __init__(self, value): + if isinstance(value, ir.Value): + self.value = value + else: + raise TypeError(f"Expected ir.Value, got {type(value)}") + super().__init__(value) + + @property + @lru_cache_ir() + def type(self) -> ir.Type: + return self.value.type + + @classmethod + def __new_from_mlir_values__(cls, values): + assert len(values) == 1, f"Expected 1 value, but got {len(values)}" + return PipelineState(values[0]) + + +@dsl_user_op +def create_pipeline( + stage: cute.Int32, + producer: OperationTypeEnum, + consumer: OperationTypeEnum, + producer_arv_count: cute.Int32, + consumer_arv_count: cute.Int32, + loc=None, + ip=None, +) -> tuple[PipelineState, PipelineState, PipelineState]: + """ + Creates an abstraction for a circular buffer of synchronizatoin primitives + indexed by stage count. + + :param stage: Stage count + :type stage: cute.Int32 + :param producer: Producer operation type + :type producer: OperationTypeEnum + :param consumer: Consumer operation type + :type consumer: OperationTypeEnum + :param producer_arv_count: Producer arrival count + :type producer_arv_count: cute.Int32 + :param consumer_arv_count: Consumer arrival count + :type consumer_arv_count: cute.Int32 + """ + if isinstance(producer_arv_count, int): + producer_arv_count = cute.Int32(producer_arv_count) + if isinstance(consumer_arv_count, int): + consumer_arv_count = cute.Int32(consumer_arv_count) + result = ir.Type.parse(f"!lir.pipeline<{stage}, {producer} -> {consumer}>") + op = cutlass_lir_ir.CreatePipelineOp( + result, + producer_arv_count.ir_value(), + consumer_arv_count.ir_value(), + loc=loc, + ip=ip, + ) + pipeline = op.result + + result = ir.Type.parse(f"!lir.pipeline_state<{stage}>") + op = cutlass_lir_ir.CreatePipelineStateOp(result, pipeline, loc=loc, ip=ip) + producer_state = op.result + + result = ir.Type.parse(f"!lir.pipeline_state<{stage}>") + op = cutlass_lir_ir.CreatePipelineStateOp(result, pipeline, loc=loc, ip=ip) + consumer_state = op.result + + return pipeline, producer_state, consumer_state + + +@dsl_user_op +def create_pipeline_with_mask( + stage: cute.Int32, + producer: OperationTypeEnum, + consumer: OperationTypeEnum, + producer_arv_count: cute.Int32, + consumer_arv_count: cute.Int32, + arrival_mask: cute.Int16, + loc=None, + ip=None, +) -> tuple[PipelineState, PipelineState, PipelineState]: + """ + Creates a pipeline with an arrival mask for cluster-scoped synchronization. + + :param stage: Pipeline stage count. + :param producer: Producer operation type (e.g. SM90_TMA_LOAD_MULTICAST). + :param consumer: Consumer operation type (e.g. SM100_MMA_2SM_SS). + :param producer_arv_count: Producer arrival count for the pipeline barriers. + :param consumer_arv_count: Consumer arrival count for the pipeline barriers. + :param arrival_mask: Bitmask that selects participating peers (e.g. CTAs in a + cluster). This is attached to the pipeline value and is consulted by some + pipeline lowerings to generate cluster-scoped synchronization + """ + if isinstance(producer_arv_count, int): + producer_arv_count = cute.Int32(producer_arv_count) + if isinstance(consumer_arv_count, int): + consumer_arv_count = cute.Int32(consumer_arv_count) + if isinstance(arrival_mask, int): + arrival_mask = cute.Int16(arrival_mask) + + result = ir.Type.parse(f"!lir.pipeline<{stage}, {producer} -> {consumer}>") + op = cutlass_lir_ir.CreatePipelineWithMaskOp( + result, + producer_arv_count.ir_value(), + consumer_arv_count.ir_value(), + arrival_mask.ir_value(), + loc=loc, + ip=ip, + ) + pipeline = op.result + + result = ir.Type.parse(f"!lir.pipeline_state<{stage}>") + op = cutlass_lir_ir.CreatePipelineStateOp(result, pipeline, loc=loc, ip=ip) + producer_state = op.result + + result = ir.Type.parse(f"!lir.pipeline_state<{stage}>") + op = cutlass_lir_ir.CreatePipelineStateOp(result, pipeline, loc=loc, ip=ip) + consumer_state = op.result + + return pipeline, producer_state, consumer_state + + + +@dsl_user_op +def pipeline_advance_iterator(pipe, state, loc=None, ip=None): + """ + Advances a pipeline iterator to the next stage. + """ + op = cutlass_lir_ir.PipelineAdvanceIteratorOp(pipe, state, loc=loc, ip=ip) + return op.result + + +@dsl_user_op +def producer_acquire(pipe, state, loc=None, ip=None): + """ + Acquires exclusive access to a pipeline. + """ + op = cutlass_lir_ir.ProducerAcquireOp(pipe, state, loc=loc, ip=ip) + return op.result + + +@dsl_user_op +def producer_commit(pipe, state, loc=None, ip=None): + """ + Commits results to a pipeline. + """ + op = cutlass_lir_ir.ProducerCommitOp(pipe, state, loc=loc, ip=ip) + return op.result + + +@dsl_user_op +def consumer_wait(pipe, state, loc=None, ip=None): + """ + Waits for a pipeline to transition to `full`. + """ + op = cutlass_lir_ir.ConsumerWaitOp(pipe, state, loc=loc, ip=ip) + return op.result + + +@dsl_user_op +def consumer_release(pipe, state, loc=None, ip=None): + """ + Releases a pipeline that has been consumed. + """ + op = cutlass_lir_ir.ConsumerReleaseOp(pipe, state, loc=loc, ip=ip) + return op.result + + +@dsl_user_op +def consumer_tail(pipe, state, loc=None, ip=None): + """ + Called by the consumer to block until asynchronous tasks have completed. + """ + op = cutlass_lir_ir.ConsumerTailOp(pipe, state, loc=loc, ip=ip) + return op.result + + +@dsl_user_op +def get_pipeline_produce_stage(pipeline, state, loc=None, ip=None): + """ + Gets a pipeline produce stage. + """ + stage_token_type = ir.Type.parse(f"!lir.stage_token<{pipeline.type}>") + stage_idx = ir.IntegerType.get_signless(32) + op = cutlass_lir_ir.GetPipelineProduceStageOp( + stage_token=stage_token_type, + stage_index=stage_idx, + pipeline=pipeline, + pipelineState=state, + loc=loc, + ip=ip, + ) + return op.stage_token, op.stage_index + + +@dsl_user_op +def get_pipeline_consume_stage(pipeline, state, loc=None, ip=None): + """ + Creates a pipeline consume stage. + """ + stage_token_type = ir.Type.parse(f"!lir.stage_token<{pipeline.type}>") + stage_idx = ir.IntegerType.get_signless(32) + op = cutlass_lir_ir.GetPipelineConsumeStageOp( + stage_token=stage_token_type, + stage_index=stage_idx, + pipeline=pipeline, + pipelineState=state, + loc=loc, + ip=ip, + ) + return op.stage_token, op.stage_index diff --git a/python/CuTeDSL/cutlass/cute/experimental/math.py b/python/CuTeDSL/cutlass/cute/experimental/math.py new file mode 100644 index 00000000..92a98471 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/experimental/math.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from cutlass import cute +from cutlass.cutlass_dsl import dsl_user_op +from cutlass._mlir.dialects import lir as cutlass_lir + + +@dsl_user_op +def dot_block_scaled( + mma_atom: cute.MmaAtom, + a: cute.Tensor, + sfa: cute.Tensor, + b: cute.Tensor, + sfb: cute.Tensor, + c: cute.Tensor, + loc=None, + ip=None, +): + """ + Computes the dot product of two tensors with block scaling and accumulates the result into a third tensor. + + :param mma_atom: MMA atom + :type mma_atom: cute.MmaAtom + :param a: First tensor + :type a: cute.Tensor + :param sfa: First scale factor tensor + :type sfa: cute.Tensor + :param b: Second tensor + :type b: cute.Tensor + :param sfb: Second scale factor tensor + :type sfb: cute.Tensor + :param c: Result tensor + :type c: cute.Tensor + """ + cutlass_lir.DotBlockScaledOp( + mma_atom._unpack(), + a.value, + sfa.value, + b.value, + sfb.value, + c.value, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def dot( + mma_atom: cute.MmaAtom, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + loc=None, + ip=None, +): + """ + Computes the dot product of two tensors and accumulates the result into a third tensor. + + :param mma_atom: MMA atom + :type mma_atom: cute.MmaAtom + :param a: First tensor + :type a: cute.Tensor + :param b: Second tensor + :type b: cute.Tensor + :param c: Result tensor + :type c: cute.Tensor + """ + cutlass_lir.DotOp( + mma_atom._unpack(), + a.value, + b.value, + c.value, + loc=loc, + ip=ip, + ) diff --git a/python/CuTeDSL/cutlass/cute/experimental/memory.py b/python/CuTeDSL/cutlass/cute/experimental/memory.py new file mode 100644 index 00000000..5294e09a --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/experimental/memory.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from typing import Type, Optional +from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir +from cutlass._mlir import ir +from cutlass.cutlass_dsl import dsl_user_op +from cutlass._mlir.dialects import ( + lir as cutlass_lir, + cute as _cute_ir, +) +from cutlass._mlir.dialects.core import OperationTypeEnum +from cutlass import cute + + +def _get_tma_load_kind(tma_operation_type: OperationTypeEnum): + """Convert OperationTypeEnum to TiledTmaLoadEnum.""" + if tma_operation_type == OperationTypeEnum.SM100_TMA_LOAD_2SM_MULTICAST: + return _cute_ir.TiledTmaLoadEnum.sm_100_2sm_multicast + if tma_operation_type == OperationTypeEnum.SM90_TMA_LOAD_MULTICAST: + return _cute_ir.TiledTmaLoadEnum.sm_90_multicast + if tma_operation_type == OperationTypeEnum.SM100_TMA_LOAD_2SM: + return _cute_ir.TiledTmaLoadEnum.sm_100_2sm + if tma_operation_type == OperationTypeEnum.SM90_TMA_LOAD: + return _cute_ir.TiledTmaLoadEnum.sm_90 + raise ValueError(f"Unsupported TMA operation type: {tma_operation_type}") + + +@dsl_user_op +def allocate( + type: Type[cute.Numeric], + address_space: cute.AddressSpace, + layout: cute.Layout | cute.ComposedLayout, + alignment: cute.Int32, + is2cta: bool = False, + loc=None, + ip=None, +) -> cute.Tensor: + """ + Allocate a buffer of the given type and layout. + + :param type: The type of the buffer + :type type: cute.Tensor + :param layout: The layout of the buffer + :type layout: cute.Layout + :param address_space: The address space of the buffer + :type address_space: str + :param alignment: The alignment of the buffer + :type alignment: cute.Int32 + :param is2cta: Whether TMEM allocation should span a CTA pair (2CTA TMEM) + :type is2cta: bool + """ + swizzle = None + if isinstance(layout, cute.ComposedLayout): + swizzle = layout.inner + layout = layout.outer + + # Handle SparseElemType (pass through) vs regular types (get mlir_type) + if isinstance(type, _cute_ir.SparseElemType): + pass + else: + type = type.mlir_type + + ptr_ty = _cute_ir.PtrType.get( + type, + address_space, + alignment, + swizzle.type.attribute if swizzle else None, + ) + buffer_type = _cute_ir.MemRefType.get(ptr_ty, layout.type) + + # `is2cta` is a UnitAttr flag in the IR: + # present => true, absent => false. + is2cta_attr = ir.UnitAttr.get() if is2cta else None + buffer_op = cutlass_lir.AllocateBufferOp( + buffer_type, is2cta=is2cta_attr, loc=loc, ip=ip + ) + return buffer_op.result + + +@dsl_user_op +def tma_load( + src: cute.Tensor, + dst: cute.Tensor, + mbar, + *, + cta_v_map, + tma_operation_type: Optional[OperationTypeEnum] = None, + internal_type=None, + update_expect_tx: bool = True, + loc=None, + ip=None, +): + """ + Copies a tensor pointed by a !cute.memref into a Buffer using TMA. + + update_expect_tx (bool): controls whether this operation increments the mbarrier's transaction bytes with the TMA copy size. + When used with Cute DSL pipelines, it must be set to False as the pipeline already initializes the mbarrier's transaction bytes. + tma_operation_type (optional): specifies the TMA operation type (SM90_TMA_LOAD, SM100_TMA_LOAD_2SM, etc.) + internal_type (optional): selects the TMA transfer's internal element encoding used by hardware. + Does not change src/dst memref types. For structured sparsity, use base storage types: + Float16 for 2:4 FP16 sparse element type, Uint8 for 8:1 uint8 sparse element type. + + :param src: Source tensor in global memory + :type src: cute.Tensor + :param dst: Destination tensor in shared memory + :type dst: cute.Tensor + :param mbar: Memory barrier for synchronization + :type mbar: cute.core.Mbarrier + :param cta_v_map: CTA V-map for the tensor + :type cta_v_map: cute.core.CtaVMap + :param tma_operation_type: TMA operation type (e.g., SM90_TMA_LOAD, SM100_TMA_LOAD_2SM, etc.) + :type tma_operation_type: OperationTypeEnum + :param internal_type: Internal type of the TMA transfer + :type internal_type: cute.core.InternalType + :param update_expect_tx: Whether to update expected transaction bytes + :type update_expect_tx: bool + """ + if tma_operation_type is not None: + kind = _get_tma_load_kind(tma_operation_type) + else: + kind = _cute_ir.TiledTmaLoadEnum.sm_90 + + kwargs = { + "cta_v_map": cta_v_map.type.attribute, + "kind": kind, + "loc": loc, + "ip": ip, + } + # Map internal_type to tma_format per updated API + if internal_type is not None: + internal_mlir_ty = ( + internal_type.mlir_type + if hasattr(internal_type, "mlir_type") + else internal_type + ) + kwargs["tma_format"] = _cute_nvgpu_ir.TmaDataFormat( + _cute_nvgpu_ir.get_default_tma_format(internal_mlir_ty, False) + ) + + if update_expect_tx: + kwargs["update_expect_tx"] = True + + cutlass_lir.TmaLoadOp(src.value, dst.value, mbar, **kwargs) + + +@dsl_user_op +def tma_load_multicast( + src: cute.Tensor, + dst: cute.Tensor, + mbar, + *, + vmnk_layout: cute.Layout, + cta_v_map, + tma_operation_type: OperationTypeEnum, + multicast_mode: int, + update_expect_tx: bool = True, + loc=None, + ip=None, +): + """ + Copies a tensor pointed by a !cute.memref into a Buffer using TMA with multicast. + + :param src: Source tensor in global memory + :param dst: Destination tensor in shared memory + :param mbar: Memory barrier for synchronization + :param vmnk_layout: Layout describing the cluster configuration + :param cta_v_map: CTA V-map for the tensor + :param tma_operation_type: TMA operation type (e.g., SM90_TMA_LOAD_MULTICAST, SM100_TMA_LOAD_2SM_MULTICAST) + :param multicast_mode: Multicast projection mode (1=column, 2=row) + :param update_expect_tx: Whether to update expected transaction bytes + """ + kind = _get_tma_load_kind(tma_operation_type) + kwargs = { + "cta_v_map": cta_v_map.type.attribute, + "kind": kind, + "vmnk_layout": vmnk_layout, + "multicast_mode": multicast_mode, + "loc": loc, + "ip": ip, + } + + if update_expect_tx: + kwargs["update_expect_tx"] = True + + cutlass_lir.TmaLoadMulticastOp( + src.value, + dst.value, + mbar, + **kwargs, + ) + + +@dsl_user_op +def tma_store( + src: cute.Tensor, + dst: cute.Tensor, + *, + cta_v_map, + internal_type=None, + loc=None, + ip=None, +): + """ + Copies a tensor from a Buffer to a tensor pointed to by a !cute.memref. + + internal_type (optional): selects the TMA transfer's internal element encoding used by hardware. + Does not change src/dst memref types. For structured sparsity, use base storage types: + Float16 for 2:4 FP16 sparse element type, Uint8 for 8:1 uint8 sparse element type. + + + :param src: Source tensor in shared memory + :type src: cute.Tensor + :param dst: Destination tensor in global memory + :type dst: cute.Tensor + :param cta_v_map: CTA V-map for the tensor + :type cta_v_map: cute.core.CtaVMap + :param internal_type: Internal type of the TMA transfer + :type internal_type: cute.core.InternalType + """ + + kwargs = { + "cta_v_map": cta_v_map.type.attribute, + "loc": loc, + "ip": ip, + } + + # Map internal_type to tma_format per updated API + if internal_type is not None: + internal_mlir_ty = ( + internal_type.mlir_type + if hasattr(internal_type, "mlir_type") + else internal_type + ) + kwargs["tma_format"] = _cute_nvgpu_ir.TmaDataFormat( + _cute_nvgpu_ir.get_default_tma_format(internal_mlir_ty, False) + ) + + cutlass_lir.TmaStoreOp(src.value, dst.value, **kwargs) + + +@dsl_user_op +def copy(src: cute.Tensor, dst: cute.Tensor, *, copy_atom, loc=None, ip=None): + """ + Copy a tensor from src to dst using a given copy atom. + """ + copy_atom = ir.Attribute.parse(f"{copy_atom.type}") + cutlass_lir.CopyOp(src.value, dst.value, copy_atom=copy_atom, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/cute/experimental/pipeline.py b/python/CuTeDSL/cutlass/cute/experimental/pipeline.py new file mode 100644 index 00000000..20c1f3ee --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/experimental/pipeline.py @@ -0,0 +1,684 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +""" +Convenience pipeline classes that hide elect_one synchronization complexity +""" + +from dataclasses import dataclass +from typing import Optional + +import cutlass + +import cutlass.cute as cute +from cutlass._mlir.dialects import lir as cutlass_lir_ir +from cutlass.base_dsl.typing import Int32 +from cutlass._mlir.dialects.core import OperationTypeEnum +from cutlass.cute.experimental.core import ( + create_pipeline, + create_pipeline_with_mask, + producer_acquire, + get_pipeline_produce_stage, + get_pipeline_consume_stage, + producer_commit, + consumer_release, + pipeline_advance_iterator, + consumer_wait, + consumer_tail, +) + +from cutlass.cutlass_dsl import CuteExperimentalDSL + + +class GenericPipelineBase: + """Base class for pipeline convenience wrappers""" + + def __init__( + self, + raw_pipeline, + num_stages, + producer_state, + consumer_state, + ): + self.raw_pipeline = raw_pipeline + self.num_stages = num_stages + # For convenience class, we always manage state internally + self.producer_state = producer_state + self.consumer_state = consumer_state + + def __extract_mlir_values__(self): + """Extract MLIR values for DynamicExpression protocol.""" + # raw_pipeline is always ir.OpResult from create_pipeline (no __extract_mlir_values__) + pipeline_values = [self.raw_pipeline] + + # Create DSL types and extract their underlying MLIR values + num_stages_dsl = Int32(self.num_stages) + + # Pipeline states are already MLIR values (PipelineState objects) + producer_state_values = [self.producer_state] + consumer_state_values = [self.consumer_state] + + return ( + pipeline_values + + [ + num_stages_dsl.__extract_mlir_values__()[0], + ] + + producer_state_values + + consumer_state_values + ) + + @classmethod + def __new_from_mlir_values__(cls, values): + """Reconstruct object from MLIR values.""" + # Parse the known structure: [pipeline] + [num_stages, producer_flag, consumer_flag] + [producer_state] + [consumer_state] + # All lir_* objects are single MLIR values + raw_pipeline = values[0] # Always single ir.OpResult + num_stages_val = values[1] + producer_state = values[2] # Always single PipelineState + consumer_state = values[3] # Always single PipelineState + + # Create temporary DSL objects and extract Python values + temp_num_stages = Int32(0) + + num_stages_dsl = temp_num_stages.__new_from_mlir_values__([num_stages_val]) + + return cls( + raw_pipeline, + ( + num_stages_dsl.value + if hasattr(num_stages_dsl, "value") + else int(num_stages_dsl) + ), + producer_state, + consumer_state, + ) + + def producer_acquire(self): + """Acquire producer state.""" + producer_acquire(self.raw_pipeline, self.producer_state) + return self + + def get_producer_stage(self): + """Get producer stage.""" + return get_pipeline_produce_stage(self.raw_pipeline, self.producer_state) + + def get_consumer_stage(self): + """Get consumer stage.""" + return get_pipeline_consume_stage(self.raw_pipeline, self.consumer_state) + + # Instance methods that can now be used directly in kernel context + def producer_acquire_and_get_stage(self): + """Combined producer acquire + get_stage with automatic elect_one using internal state.""" + + self.producer_acquire() + return get_pipeline_produce_stage(self.raw_pipeline, self.producer_state) + + def producer_commit(self): + """Commit producer state.""" + producer_commit(self.raw_pipeline, self.producer_state) + return self + + def consumer_release(self): + """Release consumer state.""" + consumer_release(self.raw_pipeline, self.consumer_state) + return self + + def producer_commit_and_advance(self): + """Combined producer commit + advance with automatic elect_one using internal state.""" + self.producer_commit() + # Update internal state in-place for better performance + self.producer_state = pipeline_advance_iterator( + self.raw_pipeline, self.producer_state + ) + return self + + def consumer_wait_and_get_stage(self): + """Combined consumer wait + get_stage with automatic elect_one using internal state.""" + self.consumer_wait() + return get_pipeline_consume_stage(self.raw_pipeline, self.consumer_state) + + def consumer_wait(self): + """Wait for consumer to be ready.""" + consumer_wait(self.raw_pipeline, self.consumer_state) + return self + + def consumer_release_and_advance(self): + """Combined consumer release + advance with automatic elect_one using internal state.""" + self.consumer_release() + # Update internal state in-place for better performance + self.consumer_state = pipeline_advance_iterator( + self.raw_pipeline, self.consumer_state + ) + return self + + def consumer_tail(self): + """Combined consumer tail with automatic elect_one using internal state.""" + consumer_tail(self.raw_pipeline, self.consumer_state) + return self + + +class GenericPipeline(GenericPipelineBase): + """ + Generic pipeline for any combination of producer and consumer. + """ + + @staticmethod + def create( + *, + producer: OperationTypeEnum, + consumer: OperationTypeEnum, + producer_arv_count: cute.Int32, + consumer_arv_count: cute.Int32, + num_stages: cute.Int32, + ): + """ + Create a generic pipeline with parameterized producer and consumer. + + Args: + producer: Producer operation type + consumer: Consumer operation type + producer_arv_count: Producer arrival count + consumer_arv_count: Consumer arrival count + num_stages: Number of pipeline stages + """ + raw_pipeline, producer_state, consumer_state = create_pipeline( + num_stages, + producer, + consumer, + producer_arv_count=producer_arv_count, + consumer_arv_count=consumer_arv_count, + ) + + return GenericPipeline( + raw_pipeline, + num_stages, + producer_state, + consumer_state, + ) + + +def _validate_umma_operation_type(operation_type: OperationTypeEnum): + if operation_type not in [ + OperationTypeEnum.SM100_MMA_1SM_SS, + OperationTypeEnum.SM100_MMA_1SM_TS, + OperationTypeEnum.SM100_MMA_2SM_SS, + OperationTypeEnum.SM100_MMA_2SM_TS, + OperationTypeEnum.SM100_MMA_SCALED_1SM_SS, + OperationTypeEnum.SM100_MMA_SCALED_1SM_TS, + OperationTypeEnum.SM100_MMA_SCALED_2SM_SS, + OperationTypeEnum.SM100_MMA_SCALED_2SM_TS, + ]: + raise ValueError(f"Invalid UMMA operation type: {operation_type}") + + +def _is_2sm_umma_operation_type(operation_type: OperationTypeEnum) -> bool: + """Check if the operation type is a 2SM UMMA operation.""" + return operation_type in [ + OperationTypeEnum.SM100_MMA_2SM_SS, + OperationTypeEnum.SM100_MMA_2SM_TS, + OperationTypeEnum.SM100_MMA_SCALED_2SM_SS, + OperationTypeEnum.SM100_MMA_SCALED_2SM_TS, + ] + + +class TMAToUMMAPipeline(GenericPipelineBase): + """ + Pipeline for TMA to UMMA. + """ + + @staticmethod + def create( + *, + num_stages: cute.Int32, + mma_operation_type: OperationTypeEnum, + tma_operation_type: Optional[OperationTypeEnum] = None, + cluster_layout_vmnk: Optional[cute.Layout] = None, + ): + """ + Create a TMA to UMMA pipeline. + + For 2SM MMA with TMA_LOAD_2SM, provide cluster_layout_vmnk for proper mask computation. + """ + _validate_umma_operation_type( + mma_operation_type, + ) + + # Default to SM90_TMA_LOAD if not specified + if tma_operation_type is None: + tma_operation_type = OperationTypeEnum.SM90_TMA_LOAD + + if tma_operation_type == OperationTypeEnum.SM100_TMA_LOAD_2SM: + if cluster_layout_vmnk is None: + raise ValueError( + "cluster_layout_vmnk is required if using 2CTA MMA with TMA" + ) + + # If using 2CTA MMA, need consumer_mask == local_cta | peer_cta + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + cta_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + arrival_mask = cute.make_layout_image_mask( + cluster_layout_vmnk, cta_in_cluster_coord_vmnk, mode=0 + ) + + raw_pipeline, producer_state, consumer_state = create_pipeline_with_mask( + num_stages, + tma_operation_type, + mma_operation_type, + producer_arv_count=1, + consumer_arv_count=1, + arrival_mask=arrival_mask, + ) + else: + raw_pipeline, producer_state, consumer_state = create_pipeline( + num_stages, + tma_operation_type, + mma_operation_type, + producer_arv_count=1, + consumer_arv_count=1, + ) + return TMAToUMMAPipeline( + raw_pipeline, + num_stages, + producer_state, + consumer_state, + ) + + @staticmethod + def create_with_mask( + *, + num_stages: cute.Int32, + tma_operation_type: OperationTypeEnum, + mma_operation_type: OperationTypeEnum, + cluster_layout_vmnk: cute.Layout, + ): + """ + Create a TMA to UMMA pipeline with multicast mask for 2CTA operations. + """ + _validate_umma_operation_type( + mma_operation_type, + ) + + # Calculate TMA multicasting masks + tma_mcast_proj_A = 2 # multicast across CTAs in same row + tma_mcast_proj_B = 1 # multicast across CTAs in same column + + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + cta_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + + # For 2CTA MMA (v-size==2), the peer CTA is the other v-slice (xor 1). + # For 1CTA MMA (v-size==1), the peer is the local CTA (no flip). + v_size = cute.size(cluster_layout_vmnk.shape[0]) + peer_v = ( + (cta_in_cluster_coord_vmnk[0] ^ 1) + if cutlass.const_expr(v_size > 1) + else cta_in_cluster_coord_vmnk[0] + ) + cta_in_cluster_coord_vmnk_peer = ( + peer_v, + *cta_in_cluster_coord_vmnk[1:], + ) + + arrival_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, cta_in_cluster_coord_vmnk, tma_mcast_proj_A + ) + arrival_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, cta_in_cluster_coord_vmnk, tma_mcast_proj_B + ) + + arrival_mask_a_peer = cute.nvgpu.cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, + cta_in_cluster_coord_vmnk_peer, + mcast_mode=tma_mcast_proj_A, + ) + arrival_mask_b_peer = cute.nvgpu.cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, + cta_in_cluster_coord_vmnk_peer, + mcast_mode=tma_mcast_proj_B, + ) + + # if 1SM MMA, arrival_mask_a_peer==arrival_mask_a && arrival_mask_b==arrival_mask_b_peer + arrival_mask_c = ( + arrival_mask_a | arrival_mask_a_peer | arrival_mask_b | arrival_mask_b_peer + ) + + num_mcast_ctas_a = cute.size(cluster_layout_vmnk.shape[2]) + num_mcast_ctas_b = cute.size(cluster_layout_vmnk.shape[1]) + num_mcast_participants = num_mcast_ctas_a + num_mcast_ctas_b - 1 + + raw_pipeline, producer_state, consumer_state = create_pipeline_with_mask( + num_stages, + tma_operation_type, + mma_operation_type, + producer_arv_count=1, + consumer_arv_count=num_mcast_participants, + arrival_mask=arrival_mask_c, + ) + return TMAToUMMAPipeline( + raw_pipeline, num_stages, producer_state, consumer_state + ) + + def producer_commit(self): + """Commit producer state.""" + with cute.arch.elect_one(): + super().producer_commit() + return self + + def consumer_release(self): + """Release consumer state.""" + with cute.arch.elect_one(): + super().consumer_release() + return self + + +class TMAToAsyncPipeline(GenericPipelineBase): + """ + Pipeline for TMA to * (except UMMA). + """ + + @staticmethod + def create( + *, + num_stages: cute.Int32, + consumer: OperationTypeEnum, + consumer_arv_count: cute.Int32, + ): + """ + Create a TMA to * (except UMMA) pipeline. + """ + + raw_pipeline, producer_state, consumer_state = create_pipeline( + num_stages, + OperationTypeEnum.SM90_TMA_LOAD, + consumer, + producer_arv_count=1, + consumer_arv_count=consumer_arv_count, + ) + return TMAToAsyncPipeline( + raw_pipeline, + num_stages, + producer_state, + consumer_state, + ) + + def producer_commit(self): + """Commit producer state.""" + with cute.arch.elect_one(): + super().producer_commit() + return self + + +class AsyncToUMMAPipeline(GenericPipelineBase): + """ + Pipeline for * (except TMA) to UMMA. + """ + + @staticmethod + def create( + *, + num_stages: cute.Int32, + producer: OperationTypeEnum, + producer_arv_count: cute.Int32, + mma_operation_type: OperationTypeEnum, + ): + """ + Create a * (except TMA) to UMMA pipeline. + """ + _validate_umma_operation_type( + mma_operation_type, + ) + + if producer == OperationTypeEnum.SM90_TMA_LOAD: + raise ValueError("TMA to UMMA is not supported.") + + raw_pipeline, producer_state, consumer_state = create_pipeline( + num_stages, + producer, + mma_operation_type, + producer_arv_count=producer_arv_count, + consumer_arv_count=1, + ) + return AsyncToUMMAPipeline( + raw_pipeline, + num_stages, + producer_state, + consumer_state, + ) + + def consumer_release(self): + """Release consumer state.""" + with cute.arch.elect_one(): + super().consumer_release() + return self + + +class UMMAtoAsyncPipeline(GenericPipelineBase): + """ + Pipeline for UMMA to * (except TMA). + """ + + @staticmethod + def create( + *, + num_stages: cute.Int32, + consumer: OperationTypeEnum, + consumer_arv_count: cute.Int32, + mma_operation_type: OperationTypeEnum, + cluster_layout_vmnk: Optional[cute.Layout] = None, + ): + """ + Create a UMMA to * (except TMA) pipeline. + + For 2SM MMA, provide cluster_layout_vmnk for proper mask computation. + """ + _validate_umma_operation_type( + mma_operation_type, + ) + + if consumer == OperationTypeEnum.SM90_TMA_LOAD: + raise ValueError("UMMA to TMA is not supported.") + + if _is_2sm_umma_operation_type(mma_operation_type): + if cluster_layout_vmnk is None: + raise ValueError("cluster_layout_vmnk cannot be None if using 2SM MMA") + return UMMAtoAsyncPipeline.create_with_mask( + num_stages=num_stages, + consumer_type=consumer, + consumer_arv_count=consumer_arv_count, + mma_operation_type=mma_operation_type, + cluster_layout_vmnk=cluster_layout_vmnk, + ) + else: # 1SM MMA + raw_pipeline, producer_state, consumer_state = create_pipeline( + num_stages, + mma_operation_type, + consumer, + producer_arv_count=1, + consumer_arv_count=consumer_arv_count, + ) + return UMMAtoAsyncPipeline( + raw_pipeline, + num_stages, + producer_state, + consumer_state, + ) + + @staticmethod + def create_with_mask( + *, + num_stages: cute.Int32, + consumer_type: OperationTypeEnum, + consumer_arv_count: cute.Int32, + mma_operation_type: OperationTypeEnum, + cluster_layout_vmnk: cute.Layout, + ): + """ + Create a UMMA to * pipeline with arrival mask for 2CTA operations. + """ + tmem_sync_mask = cutlass.pipeline.PipelineUmmaAsync._compute_tmem_sync_mask( + cta_layout_vmnk=cluster_layout_vmnk + ) + raw_pipeline, producer_state, consumer_state = create_pipeline_with_mask( + num_stages, + mma_operation_type, + consumer_type, + producer_arv_count=1, + consumer_arv_count=consumer_arv_count, + arrival_mask=tmem_sync_mask, + ) + return UMMAtoAsyncPipeline( + raw_pipeline, + num_stages, + producer_state, + consumer_state, + ) + + def producer_commit(self): + """Commit producer state.""" + with cute.arch.elect_one(): + super().producer_commit() + return self + + +@dataclass +class TMAStorePipeline: + """ + TMA Store Pipeline modeling SMEM producer to TMA consumer pipeline. + A number of epilogue warps participate in the pipeline as producers, and one of them is designated as the consumer to perform TMA store. + Named barrier is used to synchronize all warps so that producers write SMEM after the pipeline stage is available, and the consumer waits for all producers before issuing TMA store. + The canonical pipeline flow is: + 1. acquire_sync(): wait for pipeline stage availability + barrier + 2. Each producer performs SMEM writes + 3. commit_sync(): fence SMEM writes + barrier + 4. Consumer performs TMA store + 5. release_advance(): commit TMA store + advance stage + + Args: + stages: Number of pipeline stages (type parameter) + arv_count: Number of threads participating in barriers + barrier_id: Barrier ID for synchronization + tma_warp_id: Which warp issues TMA stores (None = no TMA operations) + index: Initial stage index + """ + + stages: cutlass.Constexpr[int] + arv_count: int + barrier_id: int + tma_warp_id: int + index: int = 0 + + def get_num_stages(self): + return self.stages + + def acquire_sync(self): + """ + Acquire pipeline stage and synchronize all warps. + + TMA warp waits for previous TMA operation to the same stage to complete (allowing writes to other stages to be in flight). + All warps then synchronize before producers write to SMEM. + """ + + @CuteExperimentalDSL.jit + def acquire_sync_impl(): + # Only TMA warp needs to wait for bulk async operations + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Use Python if with @Cutlass_LIR.jit preprocessor + if warp_idx == self.tma_warp_id: + # Allow N-1 TMA operations in flight for pipelining + # Now we can use the compile-time constant from type parameter + num_stages = self.get_num_stages() + wait_count = num_stages - 1 if num_stages > 1 else 0 + cute.arch.cp_async_bulk_wait_group(wait_count, read=True) + + # All warps must synchronize before producers write to SMEM + self._barrier() + return self + + return acquire_sync_impl() + + def commit_sync(self): + """ + Fence SMEM writes and synchronize all warps. + + All warps fence their SMEM writes to make them visible to consumer + All warps then synchronize before TMA store operation. + """ + # All warps fence their SMEM writes for TMA visibility + cute.arch.fence_proxy("async.shared", space="cta") + + # All warps synchronize before TMA store + self._barrier() + return self + + def release_advance(self): + """ + Release current stage and advance to next stage. + + TMA warp commits the TMA store operations to a bulk group. + All warps advance to the next pipeline stage. + """ + + @CuteExperimentalDSL.jit + def release_advance_impl(): + # Only TMA warp commits the TMA operations + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Use Python if with @Cutlass_LIR.jit preprocessor + if warp_idx == self.tma_warp_id: + cute.arch.cp_async_bulk_commit_group() + + # All warps advance to next stage + self.index = (self.index + 1) % self.get_num_stages() + return self + + return release_advance_impl() + + def get_index(self): + """Get current pipeline stage index.""" + return self.index + + def tail(self): + """ + Wait for all remaining TMA operations to complete. + + Should be called at the end of the pipeline to ensure all TMA stores finish. + """ + + @CuteExperimentalDSL.jit + def tail_impl(): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Use Python if with @Cutlass_LIR.jit preprocessor + if warp_idx == self.tma_warp_id: + # Wait for all TMA operations to complete + cute.arch.cp_async_bulk_wait_group(0, read=True) + + self._barrier() + return self + + return tail_impl() + + def _barrier(self): + """Internal barrier synchronization.""" + cute.arch.barrier( + barrier_id=self.barrier_id, + number_of_threads=self.arv_count, + ) diff --git a/python/CuTeDSL/cutlass/cute/experimental/utils.py b/python/CuTeDSL/cutlass/cute/experimental/utils.py new file mode 100644 index 00000000..5a9c72be --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/experimental/utils.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from cutlass import cute + + +def get_cta_v_map_ab( + gmem_tensor, + mma_tiler_mnk, + tiled_mma, + input_operand, + *, + loc=None, + ip=None, +): + """ + Build the **CTA-to-value map** (aka **CTA V-map**) layout for a TMA load of A/B + (and scale-factor variants SFA/SFB). + + In practice, `cta_v_map` is a `cute.Layout` that tells TMA how this CTA’s + portion of a global tensor tile maps onto the values being transferred into + shared memory. + + :param gmem_tensor: Global-memory tensor being loaded by TMA. + :type gmem_tensor: cute.Tensor + :param mma_tiler_mnk: The (M,N,K,...) tiler describing the CTA tile shape. + :type mma_tiler_mnk: tuple + :param tiled_mma: The tiled MMA object used to derive the per-operand thread/value mapping. + :type tiled_mma: cute.core.TiledMma + :param input_operand: One of {"A","B","SFA","SFB"} selecting which operand mapping to use. + :type input_operand: str + :returns: A layout suitable to pass as `cta_v_map=...` to `tma_load` / `tma_load_multicast`. + :rtype: cute.Layout + """ + ident = cute.core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip) + mode = 0 if (input_operand in ("A", "SFA")) else 1 + mma_tiler_mk = (mma_tiler_mnk[mode], *mma_tiler_mnk[2:]) + g_tile = cute.core.composition(ident, mma_tiler_mk, loc=loc, ip=ip) + if input_operand in ("A", "SFA"): + cta_v_map = tiled_mma._thrfrg_A(g_tile) + if input_operand in ("B", "SFB"): + cta_v_map = tiled_mma._thrfrg_B(g_tile) + cta_v_map = cute.core.get(cta_v_map, mode=[1]) + cta_v_map = cute.core.dice(cta_v_map, (1, (1,) * cute.core.rank(g_tile))) + return cta_v_map + + +def get_cta_v_map_c( + gmem_tensor, + epi_tile, + *, + loc=None, + ip=None, +): + """ + Build the **CTA-to-value map** (aka **CTA V-map**) layout for a TMA store/load + of the output tensor C/D. + + This returns an identity layout over the global tensor composed with the + epilogue tile, yielding a `cute.Layout` that describes which global indices + this CTA is responsible for. + + :param gmem_tensor: Global-memory tensor being stored/loaded by TMA. + :type gmem_tensor: cute.Tensor + :param epi_tile: Epilogue tile layout describing the CTA's output tile shape. + :type epi_tile: cute.Layout + :returns: A layout suitable to pass as `cta_v_map=...` to `tma_store` / `tma_load`. + :rtype: cute.Layout + """ + ident = cute.core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip) + return cute.core.composition(ident, epi_tile, loc=loc, ip=ip) + + +def make_tmem_layout_acc( + tiled_mma, + mnk_tiler, + acc_stage, + *, + loc=None, + ip=None, +): + """Return TMEM accumulator buffer layout for a tiled MMA. + + This is a small helper around ``tiled_mma.make_fragment_C(...).layout`` to + keep example code fragment-free at the call site. + + :param tiled_mma: The MMA tiler (``cute.TiledMma``). + :type tiled_mma: cute.TiledMma + :param mnk_tiler: Full MNK tiler; only the MN components are used for C. + :type mnk_tiler: tuple + :param acc_stage: Accumulator pipeline stages. + :param loc: Optional location for DSL ops. + :param ip: Optional insertion point for DSL ops. + :return: Layout for the accumulator TMEM buffer. + :rtype: cute.Layout + """ + acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2], loc=loc, ip=ip) + acc_shape_staged = cute.append(acc_shape, acc_stage, loc=loc, ip=ip) + return tiled_mma.make_fragment_C(acc_shape_staged, loc=loc, ip=ip).layout + + +def make_tmem_layout_a( + tiled_mma, + mk_tiler, + stage, + *, + loc=None, + ip=None, +): + """Return TMEM A operand buffer layout for a tiled MMA. + + :param tiled_mma: The MMA tiler (``cute.TiledMma``). + :type tiled_mma: cute.TiledMma + :param mk_tiler: MK tiler used to shape the A operand. + :type mk_tiler: tuple + :param stage: Pipeline stages for the A operand buffer. + :param loc: Optional location for DSL ops. + :param ip: Optional insertion point for DSL ops. + :return: Layout for the A operand TMEM buffer. + :rtype: cute.Layout + """ + a_shape = tiled_mma.partition_shape_A(mk_tiler, loc=loc, ip=ip) + a_shape_staged = cute.append(a_shape, stage, loc=loc, ip=ip) + return tiled_mma.make_fragment_A(a_shape_staged, loc=loc, ip=ip).layout + + +def make_t2r_rmem_layout( + tiled_copy_t2r, + gC_mnl_epi, + tidx, + *, + loc=None, + ip=None, +): + """Return RMEM buffer layout for the T2R epilogue destination. + + Computes the per-thread RMEM buffer layout produced by a TMEM->RMEM copy + for a single epilogue iteration. + + :param tiled_copy_t2r: The TMEM->RMEM tiled copy op (``cute.TiledCopy``). + :type tiled_copy_t2r: cute.TiledCopy + :param gC_mnl_epi: Global C tensor partitioned by epilogue tile. + :type gC_mnl_epi: cute.Tensor + :param tidx: Thread index for the copy slice. + :param loc: Optional location for DSL ops. + :param ip: Optional insertion point for DSL ops. + :return: Layout for the RMEM buffer. + :rtype: cute.Layout + """ + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi, loc=loc, ip=ip) + return cute.make_fragment_like( + tTR_gC[(None, None, None, 0, 0)].layout, loc=loc, ip=ip + ) diff --git a/python/CuTeDSL/cutlass/cute/export/aot_config.py b/python/CuTeDSL/cutlass/cute/export/aot_config.py index c143fac9..5922b91c 100644 --- a/python/CuTeDSL/cutlass/cute/export/aot_config.py +++ b/python/CuTeDSL/cutlass/cute/export/aot_config.py @@ -169,4 +169,3 @@ Examples: if __name__ == "__main__": main() - diff --git a/python/CuTeDSL/cutlass/cute/math.py b/python/CuTeDSL/cutlass/cute/math.py index fe6f9bf0..16a511a6 100644 --- a/python/CuTeDSL/cutlass/cute/math.py +++ b/python/CuTeDSL/cutlass/cute/math.py @@ -16,8 +16,6 @@ from .tensor import TensorSSA from cutlass._mlir.dialects import math, arith -from typing import Callable, Union - def _math_op(func: Callable, fastmath: bool, *args, **kwargs): """Dispatch the function to either a TensorSSA or a Numeric(Float). diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/common.py b/python/CuTeDSL/cutlass/cute/nvgpu/common.py index 3fa1a321..053667e9 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/common.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/common.py @@ -24,6 +24,7 @@ from ..typing import Float16, Float32, Float64, Numeric __all__ = [ "OpError", + "normalize_field_to_ir_name", "MmaUniversalOp", "MmaUniversalTrait", "CopyUniversalOp", @@ -33,6 +34,33 @@ __all__ = [ "CacheEvictionPriority", ] + +def normalize_field_to_ir_name(field, admissible_fields) -> str: + """ + Normalize a field specifier to its IR logical field name. + + Accepted inputs: + + - Enum value present in admissible_fields (must expose _to_ir_field_name()). + - Exact string IR name (e.g., "accum_c", "neg_a", "sf_a"). + + Any other form is rejected. + """ + # Enum path + if any(field is f for f in admissible_fields): + return field._to_ir_field_name() + # String path (must match exactly one of the IR names exposed by admissible_fields) + if isinstance(field, str): + allowed = {f._to_ir_field_name() for f in admissible_fields} + if field in allowed: + return field + # Otherwise, reject + allowed_pretty = [f._to_ir_field_name() for f in admissible_fields] + raise ValueError( + f"invalid field, must be one of {allowed_pretty} or their enum counterparts, but got {field}" + ) + + class OpError(DSLBaseError): """ An exception class for Op construction errors. @@ -178,8 +206,8 @@ class CopyUniversalOp(atom.CopyOp): op = cute.nvgpu.CopyUniversalOp() atom = cute.make_copy_atom( - op, - tensor_dtype, + op, + tensor_dtype, num_bits_per_copy=64, l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.EVICT_NORMAL ) @@ -195,7 +223,6 @@ class CopyUniversalOp(atom.CopyOp): - ``invariant`` is a kw argument specifying whether the load is invariant (read-only data \ that never changes). This enables compiler optimizations like instruction reordering. \ Defaults to ``False`` if not provided. - """ def __str__(self) -> str: diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/__init__.py b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/__init__.py index d59685fd..5dd3f3bc 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/__init__.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/__init__.py @@ -24,6 +24,7 @@ __all__ = [ "CopyBulkTensorTileG2SMulticastOp", "CopyBulkTensorTileS2GOp", "CopyReduceBulkTensorTileS2GOp", + "CopyDsmemStoreOp", # # helpers.py # @@ -36,5 +37,4 @@ __all__ = [ "fence_tma_desc_acquire", "cp_fence_tma_desc_release", "fence_tma_desc_release", - "group_bulk_copy_modes", ] diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py index fbf9d863..b9c167b2 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py @@ -21,7 +21,7 @@ from cutlass._mlir.dialects.cute import ReductionOp as ReductionOp from cutlass._mlir import ir from ...atom import CopyOp, Trait, make_atom -from ...typing import Int16, Int64, Pointer, Integer, Numeric +from ...typing import Int16, Int32, Int64, Pointer, Integer, Numeric from ..common import OpError from ..tcgen05.mma import CtaGroup @@ -112,6 +112,7 @@ TMA_MBAR_PTR_FIELD_NAME = "tma_bar" TMA_MCAST_MASK_FIELD_NAME = "mcast_mask" TMA_DESC_PTR_FIELD_NAME = "tma_descriptor_ptr" TMA_BYTE_MASK_FIELD_NAME = "byte_mask" +TMA_CTA_RANK_FIELD_NAME = "cta_rank" TMA_CACHE_POLICY_FIELD_NAME = "cache_policy" @@ -249,6 +250,7 @@ class CopyBulkTensorTileG2SNonExecTrait(Trait): class CopyBulkTensorTileG2STrait(Trait): pass + # # TMA GMEM -> SMEM multicast copies # @@ -374,6 +376,7 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait): ) return exec_value + class CopyBulkTensorTileG2SMulticastTrait(Trait): pass @@ -457,10 +460,6 @@ class CopyBulkTensorTileS2GTrait(Trait): pass -class CopyBulkTensorTileS2GTrait(Trait): - pass - - @dataclass class CopyReduceBulkTensorTileS2GOp(TmaCopyOp): """ @@ -800,7 +799,7 @@ class CopyBulkS2GByteMaskOp(CopyOp): def __post_init__(self) -> None: # Arch verification - arch: Arch = CuTeDSL._get_dsl().get_arch_enum() + arch: Arch = BaseDSL._get_dsl().get_arch_enum() if not arch >= Arch.sm_100: raise OpError( self, @@ -874,7 +873,7 @@ class CopyBulkS2SOp(CopyOp): def __post_init__(self) -> None: # Arch verification - arch: Arch = CuTeDSL._get_dsl().get_arch_enum() + arch: Arch = BaseDSL._get_dsl().get_arch_enum() if not arch >= Arch.sm_90: raise OpError( self, @@ -958,7 +957,7 @@ class CopyDsmemStoreOp(CopyOp): def __post_init__(self) -> None: # Arch verification - arch: Arch = CuTeDSL._get_dsl().get_arch_enum() + arch: Arch = BaseDSL._get_dsl().get_arch_enum() if not arch >= Arch.sm_90: raise OpError( self, @@ -984,6 +983,11 @@ class CopyDsmemStoreOp(CopyOp): "expects a 'num_bits_per_copy' kw argument of type int that is non-negative " f"when creating a copy Atom for {self.__class__.__name__}" ) + if num_bits_per_copy not in [0, 32, 64, 128]: + raise ValueError( + "expects a 'num_bits_per_copy' kw argument that is one of {0, 32, 64, 128} " + f"when creating a copy Atom for {self.__class__.__name__}" + ) ty = _cute_nvgpu_ir.CopyAtomDsmemStoreType.get( copy_internal_type.mlir_type, num_bits_per_copy ) diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py index b724c2cd..ec678977 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py @@ -10,7 +10,6 @@ # is strictly prohibited. from typing import Optional, Tuple, Type, Union -from typing_extensions import deprecated from cutlass.cutlass_dsl import dsl_user_op @@ -47,11 +46,12 @@ TMAOp = Union[ CopyReduceBulkTensorTileS2GOp, ] + @dsl_user_op def make_tiled_tma_atom( op: TMAOp, gmem_tensor: Tensor, - smem_layout: Union[Layout, ComposedLayout], + smem_layout_: Union[Layout, ComposedLayout], cta_tiler: Tiler, num_multicast: int = 1, *, @@ -84,7 +84,7 @@ def make_tiled_tma_atom( :type op: TMAOp :param gmem_tensor: The GMEM tensor involved in the Copy :type gmem_tensor: Tensor - :param smem_layout: The SMEM layout to construct the Copy Atom + :param smem_layout: The SMEM layout to construct the Copy Atom, either w/ or w/o the stage mode :type smem_layout: Union[Layout, ComposedLayout] :param cta_tiler: The CTA Tiler to use :type cta_tiler: Tiler @@ -95,6 +95,26 @@ def make_tiled_tma_atom( :return: A TMA Copy Atom associated with the TMA tensor :rtype: Tuple[atom.CopyAtom, Tensor] """ + smem_rank = core.rank(smem_layout_) + tiler_rank = core.rank(cta_tiler) + assert smem_rank == tiler_rank or smem_rank == tiler_rank + 1, ( + f"smem_layout must be non-staged (rank(smem_layout) == rank(cta_tiler)) " + f"or staged (rank(smem_layout) == rank(cta_tiler) + 1)" + ) + + # Set the smem_layout on the operation for later retrieval + op.smem_layout = ( + smem_layout_.value + if isinstance(smem_layout_, core._ComposedLayout) + else smem_layout_ + ) + + # Slice the smem_layout if it is staged + if smem_rank == tiler_rank + 1: + smem_layout = core.select(smem_layout_, mode=list(range(tiler_rank))) + else: + smem_layout = smem_layout_ + cta_v_map = core.composition( core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip), cta_tiler, @@ -105,22 +125,21 @@ def make_tiled_tma_atom( if isinstance(smem_layout, core._ComposedLayout): smem_layout = smem_layout.value - # Set the smem_layout on the operation for later retrieval - op.smem_layout = ( - smem_layout.value - if isinstance(smem_layout, core._ComposedLayout) - else smem_layout - ) - tma_format = None if internal_type is not None: if not isinstance(internal_type, NumericMeta): raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") - use_unpack = (internal_type.width == 8 and - isinstance(gmem_tensor.element_type, NumericMeta) and - gmem_tensor.element_type.width < 8) - internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type + use_unpack = ( + internal_type.width == 8 + and isinstance(gmem_tensor.element_type, NumericMeta) + and gmem_tensor.element_type.width < 8 + ) + internal_mlir_type = ( + gmem_tensor.element_type.mlir_type + if use_unpack + else internal_type.mlir_type + ) tma_format = _cute_nvgpu_ir.TmaDataFormat( _cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack) ) @@ -380,14 +399,3 @@ def fence_tma_desc_release(*, loc=None, ip=None) -> None: loc=loc, ip=ip, ) - - -@dsl_user_op -@deprecated("`group_bulk_copy_modes` is deprecated, use `group_modes` instead") -def group_bulk_copy_modes(src: Tensor, dst: Tensor, loc=None, ip=None) -> Tuple: - """ - Copy async bulk need group mode 0, acquiring whole tensor for bulk copy - """ - mSrc = core.group_modes(src, 0, core.rank(src)) - mDst = core.group_modes(dst, 0, core.rank(dst)) - return (mSrc, mDst) diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py index 6c786be5..19249057 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py @@ -17,7 +17,6 @@ import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir from .. import core, atom from ..typing import Shape, Layout, ComposedLayout, Tensor, Numeric, NumericMeta -from ...impl_utils import check_type_in from .cpasync.copy import ( CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SNonExecTrait, @@ -96,13 +95,6 @@ def make_tiled_tma_atom_A( """ - check_type_in( - op, - [CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp], - "op", - "make_tiled_tma_atom_A", - ) - # Set the smem_layout on the operation for later retrieval op.smem_layout = ( smem_layout.value @@ -136,10 +128,16 @@ def make_tiled_tma_atom_A( if not isinstance(internal_type, NumericMeta): raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") - use_unpack = (internal_type.width == 8 and - isinstance(gmem_tensor.element_type, NumericMeta) and - gmem_tensor.element_type.width < 8) - internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type + use_unpack = ( + internal_type.width == 8 + and isinstance(gmem_tensor.element_type, NumericMeta) + and gmem_tensor.element_type.width < 8 + ) + internal_mlir_type = ( + gmem_tensor.element_type.mlir_type + if use_unpack + else internal_type.mlir_type + ) tma_format = _cute_nvgpu_ir.TmaDataFormat( _cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack) ) @@ -224,13 +222,6 @@ def make_tiled_tma_atom_B( """ - check_type_in( - op, - [CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp], - "op", - "make_tiled_tma_atom_B", - ) - # Set the smem_layout on the operation for later retrieval op.smem_layout = ( smem_layout.value @@ -264,10 +255,16 @@ def make_tiled_tma_atom_B( if not isinstance(internal_type, NumericMeta): raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") - use_unpack = (internal_type.width == 8 and - isinstance(gmem_tensor.element_type, NumericMeta) and - gmem_tensor.element_type.width < 8) - internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type + use_unpack = ( + internal_type.width == 8 + and isinstance(gmem_tensor.element_type, NumericMeta) + and gmem_tensor.element_type.width < 8 + ) + internal_mlir_type = ( + gmem_tensor.element_type.mlir_type + if use_unpack + else internal_type.mlir_type + ) tma_format = _cute_nvgpu_ir.TmaDataFormat( _cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack) ) diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py index 251d54a4..74d1ce7f 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py @@ -19,6 +19,7 @@ __all__ = [ # copy.py # "Repetition", + "TmemLoadRedOp", "Pack", "Unpack", "Ld16x64bOp", @@ -60,4 +61,5 @@ __all__ = [ "make_tmem_copy", "make_s2t_copy", "get_s2t_smem_desc_tensor", + "make_umma_smem_desc", ] diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py index 76b2198b..0f165af6 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py @@ -26,6 +26,22 @@ from ...typing import Numeric from .mma import CtaGroup +class TmemLoadRedOp(enum.Enum): + """ + An enumeration for the possible reduce operations for TMEM load operations. + """ + + MAX = _cute_nvgpu_ir.TmemLoadRedOp.max + MAXABS = _cute_nvgpu_ir.TmemLoadRedOp.maxabs + MIN = _cute_nvgpu_ir.TmemLoadRedOp.min + MINABS = _cute_nvgpu_ir.TmemLoadRedOp.minabs + + def __str__(self) -> str: + return f"{self.__class__.__name__}.{self.name}" + + def __repr__(self) -> str: + return f"<{self.__class__.__name__}.{self.name}>" + class Repetition(enum.Enum): """ An enumeration for the number of repetitions of a given TMEM copy within the instruction. @@ -390,6 +406,97 @@ class Ld32x32bTrait(Trait): pass +@dataclass(frozen=True) +class LdRed16x32bx2Op(_LdBase): + """ + 16x32bx2 TMEM load Reduce Operation. + + See the `PTX documentation `__. + This Operation corresponds to the ``.red`` and ``.16x32bx2`` qualifiers. + """ + + redOp: TmemLoadRedOp = TmemLoadRedOp.MAX + nan: bool = False + half_split_off: int = 0 + + def _make_trait( + self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs + ) -> "LdRed16x32bx2Trait": + """ + Create a trait object for the 16x32bx2 TMEM load Reduce operation. + + :param copy_internal_type: The data type for the copy operation + :type copy_internal_type: Type[Numeric] + :param loc: MLIR location information for debugging, defaults to None + :type loc: optional + :param ip: MLIR insertion point for code generation, defaults to None + :type ip: optional + :param kwargs: Additional keyword arguments + :type kwargs: dict + :return: A trait object for this load operation + :rtype: LdRed16x32bx2Trait + """ + ty = _cute_nvgpu_ir.CopyAtomSM10xTmemLoadRedType.get( + copy_internal_type.mlir_type, + 16, + 32, + self.repeat.value, + self.redOp.value, + ir.UnitAttr.get() if self.nan else None, + ir.IntegerAttr.get(ir.IntegerType.get_signless(32), self.half_split_off), + ) + return LdRed16x32bx2Trait(make_atom(ty, loc=loc, ip=ip)) + + +class LdRed16x32bx2Trait(Trait): + pass + + +@dataclass(frozen=True) +class LdRed32x32bOp(_LdBase): + """ + 32x32b TMEM load Reduce Operation. + + See the `PTX documentation `__. + This Operation corresponds to the ``red`` and ``.32x32`` qualifiers. + """ + + redOp: TmemLoadRedOp = TmemLoadRedOp.MAX + nan: bool = False + + def _make_trait( + self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs + ) -> "LdRed32x32bTrait": + """ + Create a trait object for the 32x32b TMEM load Reduce operation. + + :param copy_internal_type: The data type for the copy operation + :type copy_internal_type: Type[Numeric] + :param loc: MLIR location information for debugging, defaults to None + :type loc: optional + :param ip: MLIR insertion point for code generation, defaults to None + :type ip: optional + :param kwargs: Additional keyword arguments + :type kwargs: dict + :return: A trait object for this load operation + :rtype: LdRed32x32bTrait + """ + ty = _cute_nvgpu_ir.CopyAtomSM10xTmemLoadRedType.get( + copy_internal_type.mlir_type, + 32, + 32, + self.repeat.value, + self.redOp.value, + ir.UnitAttr.get() if self.nan else None, + None, + ) + return LdRed32x32bTrait(make_atom(ty, loc=loc, ip=ip)) + + +class LdRed32x32bTrait(Trait): + pass + + @dataclass(frozen=True) class _StBase(CopyOp): """ diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py index b70a6ea0..760f05d3 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py @@ -9,14 +9,16 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from typing import overload, Type, Tuple, Union +from typing import overload, Type, Tuple, Union, Optional from cutlass.cutlass_dsl import dsl_user_op +from cutlass._mlir import ir import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir -from cutlass._mlir.dialects import nvvm +from cutlass._mlir.dialects import nvvm, builtin from ...typing import ( + Pointer, Shape, IntTuple, Layout, @@ -27,6 +29,7 @@ from ...typing import ( NumericMeta, Int16, Int32, + Int64, ) from ... import core from ...tensor import recast_tensor @@ -102,17 +105,27 @@ def make_smem_layout_atom( SmemLayoutAtomKind.MN_SW128_32B, ): # M/N-major layout - outer = core.make_layout( - (num_contiguous_elems, 8), stride=(1, num_contiguous_elems), loc=loc, ip=ip + return core.make_composed_layout( + sw, + 0, + core.make_layout( + (num_contiguous_elems, 8), stride=(1, num_contiguous_elems) + ), + loc=loc, + ip=ip, ) else: # K-major layout - outer = core.make_layout( - (8, num_contiguous_elems), stride=(num_contiguous_elems, 1), loc=loc, ip=ip + return core.make_composed_layout( + sw, + 0, + core.make_layout( + (8, num_contiguous_elems), stride=(num_contiguous_elems, 1) + ), + loc=loc, + ip=ip, ) - return core.make_composed_layout(sw, 0, outer, loc=loc, ip=ip) - @overload def tile_to_mma_shape( @@ -190,14 +203,27 @@ def commit( mbar_ptr = mbar_ptr.llvm_ptr if mask is not None: mask = Int16(mask).ir_value(loc=loc, ip=ip) - nvvm.tcgen05_commit_arrive( - mbar_ptr, multicast_mask=mask, group=group, loc=loc, ip=ip - ) + nvvm.tcgen05_commit(mbar_ptr, multicast_mask=mask, group=group, loc=loc, ip=ip) else: - nvvm.tcgen05_commit_arrive(mbar_ptr, group=group, loc=loc, ip=ip) + nvvm.tcgen05_commit(mbar_ptr, group=group, loc=loc, ip=ip) return +@dsl_user_op +def int_to_smem_descriptor(i, *, loc=None, ip=None) -> ir.Value: + desc_type = _cute_nvgpu_ir.SmemDescType.get() + return builtin.unrealized_conversion_cast( + [desc_type], [Int64(i).ir_value(loc=loc, ip=ip)], loc=loc, ip=ip + ) + + +@dsl_user_op +def smem_descriptor_to_int(desc: ir.Value, *, loc=None, ip=None) -> Int64: + return Int64( + builtin.unrealized_conversion_cast([Int64.mlir_type], [desc], loc=loc, ip=ip) + ) + + #################################################################################################### # # Helper functions for Copies @@ -324,3 +350,55 @@ def get_s2t_smem_desc_tensor( atom._trait.value, smem_tensor.value, loc=loc, ip=ip ) return smem_desc_tensor + + +def make_umma_smem_desc( + src: Pointer, + layout: Layout, + major: str, + next_src: Optional[Pointer] = None, + *, + loc=None, + ip=None, +): + """ + Construct shared memory descriptor for UMMA. + + The `make_umma_smem_desc` operation accepts an input cute.ptr (optionally a nextSrc + pointer for the second buffer in a circular buffer scheme), alongside a cute.layout + and a major attr, then constructs the shared memory descriptor and returns it. + The layout must be describing the buffer pointed to by the input pointer and the + iterator must carry valid swizzle information. + + There are 5 supported swizzle variants: + - S<0, 4, 3> | SWIZZLE_NONE + - S<1, 4, 3> | SWIZZLE_32B + - S<2, 4, 3> | SWIZZLE_64B + - S<3, 4, 3> | SWIZZLE_128B + - S<2, 5, 2> | SWIZZLE_128B_BASE32B + + The cute.ptr must carry shared address space and must be aligned to 16B. + + :param src: The source pointer to shared memory + :type src: Pointer + :param layout: The layout describing the buffer + :type layout: Layout + :param major: The major mode attribute + :type major: str + :param next_src: Optional next source pointer for circular buffer scheme + :type next_src: Optional[Pointer] + :return: The shared memory descriptor + :rtype: SmemDescType + """ + src = src.value + if next_src is not None: + next_src = next_src.value + + return _cute_nvgpu_ir.make_umma_smem_desc( + src=src, + layout=layout.type.attribute, + major=major, + next_src=next_src, + loc=loc, + ip=ip, + ) diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py index 705f366f..f81bd21b 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py @@ -20,7 +20,7 @@ import cutlass._mlir.dialects.cute as _cute_ir import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir from cutlass._mlir import ir -from ..common import OpError +from ..common import OpError, normalize_field_to_ir_name from ... import core, atom from ...core import _pack_shape, rank, depth from ...typing import ( @@ -141,6 +141,7 @@ class Field(enum.Enum): return self.value + # Base class for all tcgen05 MMA Ops with syntax `tcgen05.mma.cta_group.kind` used to factor out some internal code @dataclass(frozen=True) class MmaOp(Tcgen05MmaOp): @@ -268,26 +269,30 @@ class MmaTraits(Trait): admissible_fields = [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B] def set(self, field, value, *, loc=None, ip=None) -> None: - if field not in self.admissible_fields: - raise ValueError( - f"expects field to be one of {self.admissible_fields}, but got {field}" + field_ir = normalize_field_to_ir_name(field, self.admissible_fields) + bool_val = Boolean(value).ir_value(loc=loc, ip=ip) + try: + self.value = _cute_nvgpu_ir.atom_set_value( + self.value, field_ir, bool_val, loc=loc, ip=ip + ) + except (TypeError, AttributeError): + # Legacy fallback + attr = ir.Attribute.parse(f"#cute_nvgpu.atom_mma_field_sm100<{field_ir}>") + self.value = _cute_nvgpu_ir.atom_set_value( + self.value, attr, bool_val, loc=loc, ip=ip ) - field_name = f"#cute_nvgpu.atom_mma_field_sm100<{field._to_ir_field_name()}>" - attr = ir.Attribute.parse(field_name) - self.value = _cute_nvgpu_ir.atom_set_value( - self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip - ) def get(self, field, *, loc=None, ip=None) -> Any: - if field not in self.admissible_fields: - raise ValueError( - f"expects field to be one of {self.admissible_fields}, but got {field}" + field_ir = normalize_field_to_ir_name(field, self.admissible_fields) + try: + return _cute_nvgpu_ir.atom_get_value( + Boolean.mlir_type, self.value, field_ir, loc=loc, ip=ip + ) + except (TypeError, AttributeError): + attr = ir.Attribute.parse(f"#cute_nvgpu.atom_mma_field_sm100<{field_ir}>") + return _cute_nvgpu_ir.atom_get_value( + Boolean.mlir_type, self.value, attr, loc=loc, ip=ip ) - field_name = f"#cute_nvgpu.atom_mma_field_sm100<{field._to_ir_field_name()}>" - attr = ir.Attribute.parse(field_name) - return _cute_nvgpu_ir.atom_get_value( - Boolean.mlir_type, self.value, attr, loc=loc, ip=ip - ) # Base class for all tcgen05 BlockScaled MMA Ops with syntax `tcgen05.mma.cta_group.kind.block_scale` used to factor out some internal code @@ -420,33 +425,58 @@ class BlockScaledMmaTraits(Trait): ] def set(self, field, value, *, loc=None, ip=None) -> None: - if field not in self.admissible_fields: - raise ValueError( - f"expects field to be one of {self.admissible_fields}, but got {field}" - ) - if field in [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B]: - value = Boolean(value).ir_value(loc=loc, ip=ip) - elif field in [Field.SFA, Field.SFB]: + field_ir = normalize_field_to_ir_name(field, self.admissible_fields) + # Derive boolean/pointer IR names from enum values, no hard-coded strings. + bool_field_ir = { + f._to_ir_field_name() + for f in self.admissible_fields + if f in (Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B) + } + ptr_field_ir = { + f._to_ir_field_name() + for f in self.admissible_fields + if f in (Field.SFA, Field.SFB) + } + # Coerce value based on field kind + if field_ir in bool_field_ir: + val = Boolean(value).ir_value(loc=loc, ip=ip) + elif field_ir in ptr_field_ir: if not isinstance(value, Pointer): raise ValueError( - f"expects value to be a pointer for {field}, but got {type(value).__name__}" + f"expects value to be a pointer for {field_ir}, but got {type(value).__name__}" ) - value = value.value - - field_name = f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field._to_ir_field_name()}>" - attr = ir.Attribute.parse(field_name) - self.value = _cute_nvgpu_ir.atom_set_value( - self.value, attr, value, loc=loc, ip=ip - ) + val = value.value + else: + raise ValueError(f"unsupported field: {field_ir}") + try: + self.value = _cute_nvgpu_ir.atom_set_value( + self.value, field_ir, val, loc=loc, ip=ip + ) + except (TypeError, AttributeError): + attr = ir.Attribute.parse( + f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field_ir}>" + ) + self.value = _cute_nvgpu_ir.atom_set_value( + self.value, attr, val, loc=loc, ip=ip + ) def get(self, field, *, loc=None, ip=None) -> Any: - if field not in [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B]: - raise ValueError(f"the get method for {field} is not supported") - field_name = f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field._to_ir_field_name()}>" - attr = ir.Attribute.parse(field_name) - return _cute_nvgpu_ir.atom_get_value( - Boolean.mlir_type, self.value, attr, loc=loc, ip=ip - ) + # Only boolean-returning fields supported for get. Derive from admissible_fields. + gettable_fields = [ + f for f in self.admissible_fields if f not in (Field.SFA, Field.SFB) + ] + field_ir = normalize_field_to_ir_name(field, gettable_fields) + try: + return _cute_nvgpu_ir.atom_get_value( + Boolean.mlir_type, self.value, field_ir, loc=loc, ip=ip + ) + except (TypeError, AttributeError): + attr = ir.Attribute.parse( + f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field_ir}>" + ) + return _cute_nvgpu_ir.atom_get_value( + Boolean.mlir_type, self.value, attr, loc=loc, ip=ip + ) # @@ -802,6 +832,7 @@ class MmaFP8Trait(MmaTraits): pass + # # MXF8F6F4 MMA # @@ -946,7 +977,7 @@ class MmaMXF4Op(BlockScaledMmaOp): f"but got {self.shape_mnk[2]}", ) - def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF8Trait": + def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4Trait": shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip) ty = _cute_nvgpu_ir.MmaAtomSM100UMMABlockScaledType.get( shape_mnk.type.attribute, @@ -1039,7 +1070,7 @@ class MmaMXF4NVF4Op(BlockScaledMmaOp): f"but got {self.shape_mnk[2]}", ) - def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF8Trait": + def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4NVF4Trait": shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip) ty = _cute_nvgpu_ir.MmaAtomSM100UMMABlockScaledType.get( shape_mnk.type.attribute, @@ -1077,6 +1108,181 @@ class MmaMXF4NVF4Trait(BlockScaledMmaTraits): pass +# +# SM103 MXF4 MMA +# + + +@dataclass(frozen=True) +class SM103MmaMXF4Op(BlockScaledMmaOp): + """ + SM103 MXF4 tcgen05 BlockScaled MMA Operation. + + See the `PTX documentation `__. + This Operation corresponds to the ``.kind::mxf4`` qualifier. + This Operation is for SM103. + """ + + descriptive_name = "tcgen05 SM103 MXF4 BlockScaled MMA Operation" + + def __init__( + self, + instruction_shape: Shape, + cta_group: CtaGroup, + a_src: OperandSource, + ) -> None: + super().__init__( + Float4E2M1FN, + Float4E2M1FN, + Float32, + Float8E8M0FNU, + 32, + instruction_shape, + cta_group, + a_src, + OperandMajorMode.K, + OperandMajorMode.K, + ) + self._verify() + + def _verify(self) -> None: + # Instruction shape verification + instruction_k = 96 + if rank(self.shape_mnk) == 2: + object.__setattr__(self, "shape_mnk", (*self.shape_mnk, instruction_k)) + if self.shape_mnk[2] != instruction_k: + raise OpError( + self, + f"expects the instruction extent in the K-mode to be {instruction_k}, " + f"but got {self.shape_mnk[2]}", + ) + + def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4Trait": + shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip) + ty = _cute_nvgpu_ir.MmaAtomSM100UMMABlockScaledType.get( + shape_mnk.type.attribute, + self.cta_group.value, + self.a_major_mode._to_ir(), + self.b_major_mode._to_ir(), + self.a_dtype.mlir_type, + self.b_dtype.mlir_type, + self.acc_dtype.mlir_type, + self.sf_dtype.mlir_type, + self.a_src._to_ir(), + self.sf_vec_size, + 1030, + ) + return MmaMXF4Trait( + make_atom( + ty, + ( + Boolean(False).ir_value(loc=loc, ip=ip), + Boolean(False).ir_value(loc=loc, ip=ip), + Boolean(False).ir_value(loc=loc, ip=ip), + core.make_ptr( + self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + ).value, + core.make_ptr( + self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + ).value, + ), + loc=loc, + ip=ip, + ) + ) + + +# +# SM103 MXF4NVF4 MMA +# + + +@dataclass(frozen=True) +class SM103MmaMXF4NVF4Op(BlockScaledMmaOp): + """ + SM103 MXF4NVF4 tcgen05 BlockScaled MMA Operation. + + See the `PTX documentation `__. + This Operation corresponds to the ``.kind::mxf4nvf4`` qualifier. + This Operation is for SM103. + """ + + descriptive_name = "tcgen05 SM103 MXF4NVF4 BlockScaled MMA Operation" + + def __init__( + self, + sf_dtype: Type[Numeric], + instruction_shape: Shape, + cta_group: CtaGroup, + a_src: OperandSource, + ) -> None: + super().__init__( + Float4E2M1FN, + Float4E2M1FN, + Float32, + sf_dtype, + 16, + instruction_shape, + cta_group, + a_src, + OperandMajorMode.K, + OperandMajorMode.K, + ) + self._verify() + + def _verify(self) -> None: + # Scale Factor data type verification + if self.sf_dtype not in [Float8E8M0FNU, Float8E4M3FN]: + raise OpError( + self, + "expects the 'sf_dtype' Op parameter to be one of Float8E8M0FNU", + ) + # Instruction shape verification + instruction_k = 96 + if rank(self.shape_mnk) == 2: + object.__setattr__(self, "shape_mnk", (*self.shape_mnk, instruction_k)) + if self.shape_mnk[2] != instruction_k: + raise OpError( + self, + f"expects the instruction extent in the K-mode to be {instruction_k}, " + f"but got {self.shape_mnk[2]}", + ) + + def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4NVF4Trait": + shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip) + ty = _cute_nvgpu_ir.MmaAtomSM100UMMABlockScaledType.get( + shape_mnk.type.attribute, + self.cta_group.value, + self.a_major_mode._to_ir(), + self.b_major_mode._to_ir(), + self.a_dtype.mlir_type, + self.b_dtype.mlir_type, + self.acc_dtype.mlir_type, + self.sf_dtype.mlir_type, + self.a_src._to_ir(), + self.sf_vec_size, + 1030, + ) + return MmaMXF4NVF4Trait( + make_atom( + ty, + ( + Boolean(False).ir_value(loc=loc, ip=ip), + Boolean(False).ir_value(loc=loc, ip=ip), + Boolean(False).ir_value(loc=loc, ip=ip), + core.make_ptr( + self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + ).value, + core.make_ptr( + self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + ).value, + ), + loc=loc, + ip=ip, + ) + ) + + #################################################################################################### # # SMEM layout atoms diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py index 09cde138..baff4839 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py @@ -82,6 +82,7 @@ class LdMatrix8x8x16bOp(BaseOp): class LdMatrix8x8x16bTrait(Trait): pass + @dataclass(frozen=True) class LdMatrix8x16x8bOp(BaseOp): """ @@ -102,15 +103,20 @@ class LdMatrix8x16x8bOp(BaseOp): self, "expects the 'num_matrices' Op parameter to be one of [1,2,4]", ) - if self.unpack_bits not in [4, 6]: - raise OpError(self, "Op unpack bits must be 4 or 6") + if self.unpack_bits not in [None, 4, 6]: + raise OpError(self, "Op unpack bits must be 4 or 6 or None") def _make_trait( self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs ) -> "LdMatrix8x16x8bTrait": - mode = _pack_shape((8, 16), loc=loc, ip=ip) - sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8 - if self.unpack_bits == 6: + # LdMatrix8x16x8b without unpacking doesn't exist + # but is equivalent to LdMatrix8x8x16b + mode_n = 8 if self.unpack_bits is None else 16 + mode = _pack_shape((8, mode_n), loc=loc, ip=ip) + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u16 + if self.unpack_bits == 4: + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8 + elif self.unpack_bits == 6: sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u6x16p32to8 ty = _cute_nvgpu_ir.CopyAtomLdsmType.get( copy_internal_type.mlir_type, @@ -125,11 +131,12 @@ class LdMatrix8x16x8bOp(BaseOp): class LdMatrix8x16x8bTrait(Trait): pass + @dataclass(frozen=True) class LdMatrix16x8x8bOp(BaseOp): """ 16x8 8b ``ldmatrix`` Operation with transpose - + There is no direct PTX correspondance to this Op. This actually lowers to ldmatrix with the ``.m16n16`` qualifier and additional address and value permutations to match stmatrix.m16n8.trans. @@ -166,6 +173,7 @@ class LdMatrix16x8x8bOp(BaseOp): ) return LdMatrix16x8x8bTrait(make_atom(ty, loc=loc, ip=ip)) + class LdMatrix16x8x8bTrait(Trait): pass @@ -176,7 +184,7 @@ class LdMatrix16x16x8bOp(BaseOp): 16x16 ``ldmatrix`` Operation with transpose and optional unpacking to 8b container. Packed source container is 16x4b elements with 64b padding or 16x6b elements with 32b padding (total 128b per 16 elements) - + See the `PTX documentation `__. This operation corresponds to the ``.m16n16`` and the ``.b4x16_p64``,``.b6x16_p32``,``.b8`` qualifiers. """ diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py index 145d9a88..781128b6 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py @@ -15,7 +15,7 @@ from typing import Type, Any import enum from cutlass import cute from cutlass.base_dsl.arch import Arch -from cutlass.cutlass_dsl import CuTeDSL +from cutlass.cutlass_dsl import BaseDSL from ..common import OpError @@ -134,7 +134,7 @@ class MmaSM120BlockScaledOp(MmaOp): def __post_init__(self) -> None: # Verify arch - arch = CuTeDSL._get_dsl().get_arch_enum() + arch = BaseDSL._get_dsl().get_arch_enum() if not arch == Arch.sm_120a: raise OpError( self, @@ -174,6 +174,7 @@ class MmaSM120BlockScaledOp(MmaOp): self, "expects the 'sf_vec_size' Op parameter to be 16 or 32", ) + def __str__(self) -> str: return ( "warp-level MXF4/MXF4NVF4 MMA Operation" @@ -190,6 +191,7 @@ class MmaSM120BlockScaledOp(MmaOp): def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None): pass + class Field(enum.Enum): """ An enumeration for the fields of the MMA Atom that can be modified at runtime. diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py index 3555790e..f3c0ecb6 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py @@ -15,12 +15,13 @@ from typing import Type, Any from cutlass.base_dsl.arch import Arch from cutlass.cutlass_dsl import BaseDSL, T +from typing_extensions import deprecated import cutlass._mlir.dialects.cute as _cute_ir import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir from cutlass._mlir import ir -from ..common import OpError +from ..common import OpError, normalize_field_to_ir_name from ...core import _pack_shape, rank, depth from ...typing import ( Shape, @@ -208,27 +209,44 @@ class MmaOp(WarpGroupMmaOp): class MmaTraits(Trait): admissible_fields = [Field.ACCUMULATE] + def _normalize_field_name(self, field: Any) -> str: + """ + Normalize a field specifier (enum or string) into the IR logical field name. + Accepted inputs: + - Field.ACCUMULATE + - "accum_c" + """ + return normalize_field_to_ir_name(field, self.admissible_fields) + def set(self, field, value, *, loc=None, ip=None) -> None: - if field not in self.admissible_fields: - raise ValueError( - f"invalid field, must be {Field.ACCUMULATE}, but got {field}" + field_ir_name = self._normalize_field_name(field) + # Prefer the newer builder that accepts a logical field name, but keep + # a fallback for legacy attribute-based construction to avoid breaking changes. + bool_val = Boolean(value).ir_value(loc=loc, ip=ip) + try: + self.value = _cute_nvgpu_ir.atom_set_value( + self.value, field_ir_name, bool_val, loc=loc, ip=ip + ) + except (TypeError, AttributeError): + # Legacy path: construct the per-arch field attribute explicitly + attr_asm = f"#cute_nvgpu.atom_mma_field_sm90<{field_ir_name}>" + attr = ir.Attribute.parse(attr_asm) + self.value = _cute_nvgpu_ir.atom_set_value( + self.value, attr, bool_val, loc=loc, ip=ip ) - field_name = f"#cute_nvgpu.atom_mma_field_sm90<{field._to_ir_field_name()}>" - attr = ir.Attribute.parse(field_name) - self.value = _cute_nvgpu_ir.atom_set_value( - self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip - ) def get(self, field, *, loc=None, ip=None) -> Any: - if field not in self.admissible_fields: - raise ValueError( - f"invalid field, must be {Field.ACCUMULATE}, but got {field}" + field_ir_name = self._normalize_field_name(field) + try: + return _cute_nvgpu_ir.atom_get_value( + Boolean.mlir_type, self.value, field_ir_name, loc=loc, ip=ip + ) + except (TypeError, AttributeError): + attr_asm = f"#cute_nvgpu.atom_mma_field_sm90<{field_ir_name}>" + attr = ir.Attribute.parse(attr_asm) + return _cute_nvgpu_ir.atom_get_value( + Boolean.mlir_type, self.value, attr, loc=loc, ip=ip ) - field_name = f"#cute_nvgpu.atom_mma_field_sm90<{field._to_ir_field_name()}>" - attr = ir.Attribute.parse(field_name) - return _cute_nvgpu_ir.atom_get_value( - Boolean.mlir_type, self.value, attr, loc=loc, ip=ip - ) @dataclass(frozen=True) diff --git a/python/CuTeDSL/cutlass/cute/runtime.py b/python/CuTeDSL/cutlass/cute/runtime.py index 4019f1b6..fbff4dff 100644 --- a/python/CuTeDSL/cutlass/cute/runtime.py +++ b/python/CuTeDSL/cutlass/cute/runtime.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # Use of this software is governed by the terms and conditions of the diff --git a/python/CuTeDSL/cutlass/cute/tensor.py b/python/CuTeDSL/cutlass/cute/tensor.py index a8f0ff7b..bf7fba17 100644 --- a/python/CuTeDSL/cutlass/cute/tensor.py +++ b/python/CuTeDSL/cutlass/cute/tensor.py @@ -20,6 +20,7 @@ from cutlass.cutlass_dsl import ( T, cutlass_arith, _binary_op_type_promote, + MLIR_DYNAMIC, BaseDSL, ) from cutlass._mlir import ir @@ -75,35 +76,8 @@ from .core import ( recast_layout, ) -from .typing import ( - IntTuple, - Coord, - Shape, - Stride, - Pointer, - Layout, - ComposedLayout, - Tensor, - AddressSpace, - is_integer, - is_int_tuple, - as_numeric, -) -from .typing import ( - Numeric, - Integer, - Boolean, - Int4, - Uint8, - Int8, - Int32, - Float4E2M1FN, - Float16, - Float32, - BFloat16, -) from .tuple import transform_leaf, product, product_like, flatten_to_tuple -from .arch import cvt_i8_bf16_intrinsic, cvt_i4_bf16_intrinsic, cvt_f4e2m1_f16_intrinsic +from .arch import cvt_i8_bf16_intrinsic, cvt_i4_bf16_intrinsic __all__ = [ @@ -439,10 +413,9 @@ class _Tensor(Tensor): return _cute_ir.get_layout(self.value, loc=loc, ip=ip) @property - @dsl_user_op @lru_cache_ir() - def shape(self, *, loc=None, ip=None) -> Shape: - return self.layout.shape_method(loc=loc, ip=ip) + def shape(self) -> Shape: + return self.layout.shape @property @lru_cache_ir() @@ -480,12 +453,23 @@ class _Tensor(Tensor): raise ValueError(f"{self} doesn't have memspace") @dsl_user_op - def load(self, *, loc=None, ip=None) -> "TensorSSA": + def load( + self, + *, + mask: Optional["TensorSSA"] = None, + pass_thru: Optional["TensorSSA"] = None, + loc=None, + ip=None, + ) -> "TensorSSA": """Load tensor elements as a vector. Loads all elements of the tensor into a vector representation, assuming the tensor has a static shape and is in a memory space that supports load operations. + :param mask: Mask vector, defaults to None + :type mask: Optional[TensorSSA] + :param pass_thru: Pass through vector, defaults to None + :type pass_thru: Optional[TensorSSA] :param loc: Source location for MLIR operation tracking, defaults to None :type loc: Optional[Location] :param ip: Insertion point for MLIR operation, defaults to None @@ -501,9 +485,15 @@ class _Tensor(Tensor): if not is_static(self.shape): raise ValueError("dynamic layout doesn't support load") - self._check_can_load_store() + self._check_can_load_store(vectorized=True) - res_vect = _cute_ir.memref_load_vec(self.value, loc=loc, ip=ip) + mask_val = None if mask is None else mask.ir_value(loc=loc, ip=ip) + pass_thru_val = ( + None if pass_thru is None else self._cvt_to_dest(pass_thru, loc=loc, ip=ip) + ) + res_vect = _cute_ir.memref_load_vec( + self.value, mask=mask_val, pass_thru=pass_thru_val, loc=loc, ip=ip + ) if self.element_type is Boolean: assert res_vect.type.element_type == T.i8(), ( f"Boolean tensor must be stored as i8 in memory, but got {res_vect.type.element_type}" @@ -515,7 +505,14 @@ class _Tensor(Tensor): return TensorSSA(res_vect, self.shape, self.element_type) @dsl_user_op - def store(self, data: "TensorSSA", *, loc=None, ip=None): + def store( + self, + data: "TensorSSA", + *, + mask: Optional["TensorSSA"] = None, + loc=None, + ip=None, + ): """Store vector data into tensor. Stores vector data into the tensor, assuming matching shapes and a memory space @@ -523,6 +520,8 @@ class _Tensor(Tensor): :param data: Vector data to store into tensor :type data: TensorSSA + :param mask: Mask vector, defaults to None + :type mask: Optional[TensorSSA] :param loc: Source location for MLIR operation tracking, defaults to None :type loc: Optional[Location] :param ip: Insertion point for MLIR operation, defaults to None @@ -538,7 +537,7 @@ class _Tensor(Tensor): if not is_static(self.shape): raise ValueError("Dynamic layout doesn't support vectorized store") - self._check_can_load_store() + self._check_can_load_store(vectorized=True) n_elems = size(self.shape, loc=loc, ip=ip) if n_elems != size(data.shape, loc=loc, ip=ip): @@ -556,7 +555,11 @@ class _Tensor(Tensor): # Implicit upcast to wider type new_data = self._cvt_to_dest(data, loc=loc, ip=ip) - return _cute_ir.memref_store_vec(new_data, self.value, loc=loc, ip=ip) + mask_val = None if mask is None else mask.ir_value(loc=loc, ip=ip) + + return _cute_ir.memref_store_vec( + new_data, self.value, mask=mask_val, loc=loc, ip=ip + ) @dsl_user_op def fill(self, value: Numeric, *, loc=None, ip=None) -> None: @@ -585,7 +588,7 @@ class _Tensor(Tensor): # Fill tensor with constant value tensor.fill(0.5) # All elements become 0.5 """ - self._check_can_load_store() + self._check_can_load_store(vectorized=True) sz = size(self, loc=loc, ip=ip) if type(sz) is not int: @@ -599,7 +602,7 @@ class _Tensor(Tensor): ) self.store(vect_val, loc=loc, ip=ip) - def _check_can_load_store(self): + def _check_can_load_store(self, vectorized: bool = False): if not isinstance(self.type, _cute_ir.MemRefType) or self.memspace not in ( AddressSpace.rmem, AddressSpace.smem, @@ -608,9 +611,9 @@ class _Tensor(Tensor): ): raise ValueError(f"{self} doesn't support load and store") - if self.type.is_swizzled: + if vectorized and isinstance(self.layout, ComposedLayout): raise NotImplementedError( - f"load & store swizzled memory is not supported yet: {self}" + "vectorized load/store on tensor with composed layout is not supported yet" ) def _check_can_dereference(self): @@ -1038,8 +1041,10 @@ def print_tensor( signed = tensor.element_type.signed else: signed = False - else: + elif isinstance(tensor.type, _cute_ir.CoordTensorType): signed = True + else: + raise ValueError(f"unsupported tensor type for print_tensor, got {tensor.type}") _cute_ir.print_view(tensor.value, verbose=verbose, is_signed=signed, loc=loc, ip=ip) @@ -1750,7 +1755,8 @@ class TensorSSA(cutlass_arith.ArithValue): idx = crd2idx(crd, self._layout, loc=loc, ip=ip) assert not isinstance(idx, tuple), "index must be scalar" idx_val = as_numeric(idx).ir_value(loc=loc, ip=ip) - res_val = vector.extractelement(self, position=idx_val, loc=loc, ip=ip) + idx_val = arith.index_cast(T.index(), idx_val, loc=loc, ip=ip) + res_val = vector.extract(self, [idx_val], [MLIR_DYNAMIC], loc=loc, ip=ip) return self.dtype(res_val) if not is_static(crd): @@ -1817,16 +1823,7 @@ class TensorSSA(cutlass_arith.ArithValue): # maybe downcast can lose signedness src = self.maybe_downcast().with_signedness(self.signed) if src_dtype.is_float and dtype.is_float: - if src_dtype == Float4E2M1FN and dtype in (Float16, Float32): - res_vect = cvt_f4e2m1_f16_intrinsic( - src, size(self.shape), loc=loc, ip=ip - ) - if dtype == Float32: - res_vect = cutlass_arith.cvtf( - res_vect, dtype.mlir_type, loc=loc, ip=ip - ) - else: - res_vect = cutlass_arith.cvtf(src, dtype.mlir_type, loc=loc, ip=ip) + res_vect = cutlass_arith.cvtf(src, dtype.mlir_type, loc=loc, ip=ip) elif src_dtype.is_float and issubclass(dtype, Integer): res_vect = cutlass_arith.fptoi( src, dtype.signed, dtype.mlir_type, loc=loc, ip=ip diff --git a/python/CuTeDSL/cutlass/cute/testing.py b/python/CuTeDSL/cutlass/cute/testing.py index 9b99209f..3cdcd84c 100644 --- a/python/CuTeDSL/cutlass/cute/testing.py +++ b/python/CuTeDSL/cutlass/cute/testing.py @@ -20,15 +20,12 @@ from typing import Type, Union, Callable, Optional, Dict, List, Any import cuda.bindings.driver as cuda_driver import cuda.bindings.runtime as cuda_runtime -import cutlass -import cutlass.base_dsl.jit_executor -import cutlass.cutlass_dsl.cuda_jit_executor from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op, const_expr from .typing import Numeric, Int8, Boolean, Tensor, Layout, Shape from . import nvgpu -from .core import recast_layout, make_layout, composition, get, rank, size, zipped_divide +from .core import recast_layout, make_layout, composition, get, rank, size from .tuple import elem_less from .tensor import ( make_rmem_tensor, @@ -39,6 +36,7 @@ from .tensor import ( ) from .atom import make_copy_atom from .algorithm import copy +from .core import zipped_divide from .runtime import from_dlpack from cutlass._mlir.dialects import builtin, cf, nvvm, vector @@ -76,7 +74,7 @@ class _CompileTimeAssertion(Assertion): def __init__( self, - tensor: _Tensor, + tensor: Tensor, num_assertions: int = 1, msgs=None, device=None, @@ -849,7 +847,9 @@ def get_workspace_count( :return: Number of workspaces needed :rtype: int """ - num_l2_cache_bytes = cutlass.utils.HardwareInfo().get_l2_cache_size_in_bytes() + from cutlass.utils import HardwareInfo + + num_l2_cache_bytes = HardwareInfo().get_l2_cache_size_in_bytes() num_workspaces = (num_l2_cache_bytes * 3) // one_workspace_bytes + 1 num_iters = warmup_iterations + iterations return num_iters if num_iters < num_workspaces else num_workspaces diff --git a/python/CuTeDSL/cutlass/cute/typing.py b/python/CuTeDSL/cutlass/cute/typing.py index 80a084da..fdb60443 100644 --- a/python/CuTeDSL/cutlass/cute/typing.py +++ b/python/CuTeDSL/cutlass/cute/typing.py @@ -12,7 +12,6 @@ from abc import ABC, abstractmethod import ctypes from typing import ForwardRef, Tuple, Union, Any, Type, List, Optional, Literal -from functools import lru_cache from cutlass.base_dsl.typing import * @@ -28,9 +27,13 @@ class SymInt: def __init__(self, width: Literal[32, 64] = 32, *, divisibility=1): if width not in [32, 64]: raise ValueError(f"Unsupported width: {width}") + self._width = width self._divisibility = divisibility + def __hash__(self): + return hash((self._width, self._divisibility)) + @property def width(self): return self._width @@ -80,6 +83,7 @@ class SymInt: else: assert False, f"Unsupported width: {self.width}" return self + def sym_int(width: Literal[32, 64] = 32, *, divisibility=1) -> SymInt: return SymInt(width, divisibility=divisibility) @@ -403,6 +407,4 @@ __all__ = [ "XTuple", "is_integer", "is_int_tuple", - "Pointer", - "Tensor", ] diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py b/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py index 690546c6..077d66b7 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py @@ -53,7 +53,6 @@ from ..base_dsl.compiler import ( KeepCUBIN, KeepPTX, GPUArch, - LinkLibraries, EnableTVMFFI, ) from ..base_dsl.runtime.jit_arg_adapters import * diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py index 2c5e8a6e..72f57b8f 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py @@ -58,8 +58,6 @@ class CudaDialectJitModule: for library in self.cuda_library: cuda_runtime.cudaLibraryUnload(library) self.cuda_library.clear() - except Exception as e: - pass finally: self._unloaded = True diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py index 83196af2..136c4fda 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py @@ -22,6 +22,7 @@ from typing import ( List, Tuple, Sequence, + Iterable, ForwardRef, Any, get_origin, @@ -34,7 +35,6 @@ from dataclasses import is_dataclass, fields from math import ceil from itertools import chain from pathlib import Path -from collections.abc import Sequence import builtins import ctypes import hashlib @@ -65,10 +65,14 @@ from cutlass._mlir.dialects import ( from cutlass._mlir.dialects._ods_common import ( get_op_result_or_op_results as _get_op_result_or_op_results, ) + +from cutlass._mlir.dialects import lir as cutlass_lir + from cutlass._mlir.extras import types as T # Helpers from ..base_dsl._mlir_helpers import arith as cutlass_arith +from ..base_dsl._mlir_helpers import lru_cache_ir from ..base_dsl._mlir_helpers.op import dsl_user_op from ..base_dsl._mlir_helpers.arith import const @@ -94,6 +98,7 @@ from .cutlass_ast_decorators import ( _loop_execute_range_dynamic, _if_execute_dynamic, _while_execute_dynamic, + _ifexp_execute_dynamic, ) from ..base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry @@ -275,6 +280,11 @@ class CutlassBaseDSL(BaseDSL): log().info(f"self: {self}") log().info(f"Entering GPU module for {self.name}") log().info(f"GPU module: {self.gpu_module}") + if not self.gpu_module: + raise DSLRuntimeError( + f"GPU module is not set, probably compilation of a kernel from different DSL decorator", + suggestion=f"Use the same DSL decorator to build the GPU module, DSL: {type(self).__name__}", + ) return ir.InsertionPoint(self.gpu_module.bodyRegion.blocks[0]) @staticmethod @@ -290,9 +300,9 @@ class CutlassBaseDSL(BaseDSL): ) def _generate_kernel_attrs(self, config: BaseDSL.LaunchConfig) -> dict: - assert isinstance(config, BaseDSL.LaunchConfig), ( - f"Expect LaunchConfig for @kernel, but got {type(config)}" - ) + assert isinstance( + config, BaseDSL.LaunchConfig + ), f"Expect LaunchConfig for @kernel, but got {type(config)}" ret = {} if config.has_max_number_threads(): @@ -381,16 +391,7 @@ class CutlassBaseDSL(BaseDSL): ) from e files.append((giant_dso_name, so_path, so_size)) - def handle_import_error(exc): - """Handle errors during package walking, ignoring ImportError and NotImplementedError.""" - if isinstance(exc, (ImportError, NotImplementedError)): - log().info(f"Skipping module due to {type(exc).__name__}: {exc}") - else: - log().warning(f"Unexpected error during package walk: {exc}") - - for lib in pkgutil.walk_packages( - [dsl_path], prefix="cutlass.", onerror=handle_import_error - ): + for lib in pkgutil.walk_packages([dsl_path], prefix="cutlass."): spec = lib.module_finder.find_spec(lib.name) if not spec or not spec.origin: continue @@ -624,12 +625,7 @@ class CutlassBaseDSL(BaseDSL): loc=None, ip=None, ): - # set to 3 for PDL, cluster size, and cooperative - max_num_attributes = 3 - - if preferred_cluster_size_x is not None: - max_num_attributes += 1 - + max_num_attributes = 17 launch_config_type = cuda_dialect.LaunchConfigType.get(max_num_attributes) if len(stream) == 0: @@ -653,8 +649,6 @@ class CutlassBaseDSL(BaseDSL): cfg = cuda_dialect.launch_cfg_create( # Launch config type launch_config_type, - # Max num of attributes the launch config can hold - # set to 3 for PDL, cluster size, and cooperative ir.IntegerAttr.get(ir.IntegerType.get_signless(32), max_num_attributes), block_size_x, block_size_y, @@ -793,15 +787,15 @@ class CutlassBaseDSL(BaseDSL): requiredArgs = kwargs.get("requiredArgs", None) loc = kwargs.get("loc", None) assert kernelSym is not None, "kernelSym being None is not expected!" - assert requiredArgs is not None, ( - "requiredArgs being None is not expected!" - ) - assert kernelOperands is not None, ( - "kernelOperands being None is not expected!" - ) - assert isinstance(requiredArgs.config, BaseDSL.LaunchConfig), ( - f"Expect LaunchConfig for @kernel, but got {type(requiredArgs.config)}" - ) + assert ( + requiredArgs is not None + ), "requiredArgs being None is not expected!" + assert ( + kernelOperands is not None + ), "kernelOperands being None is not expected!" + assert isinstance( + requiredArgs.config, BaseDSL.LaunchConfig + ), f"Expect LaunchConfig for @kernel, but got {type(requiredArgs.config)}" cfg = requiredArgs.config @@ -822,6 +816,7 @@ class CutlassBaseDSL(BaseDSL): if not isinstance(cfg.async_deps, (list, tuple)): async_deps = [cfg.async_deps] + # Prepare launch kwargs launch_kwargs = {} if cfg.has_fallback_cluster: @@ -1091,6 +1086,58 @@ class CuTeDSL(CutlassBaseDSL): return cuda_dialect.ReturnOp([], loc=loc, ip=ip) +# ============================================================================= +# CuteExperimental DSL Class +# ============================================================================= + + +class CuteExperimentalDSL(CutlassBaseDSL): + def __init__(self): + name = "CUTE_EXPERIMENTAL_DSL" + compiler_provider = compiler.Compiler(passmanager, execution_engine) + pass_sm_arch_name = "cubin-chip" + + super().__init__(name, compiler_provider, pass_sm_arch_name, preprocess=True) + + def _get_pipeline(self, pipeline): + if pipeline == None: + return "builtin.module(gpu.module(lir-to-cute{enable-cuda-dialect enable-lir-func-finalization=false}), lir-func-finalization{enable-cuda-dialect=true}, cute-to-nvvm{check-inline-asm=false cubin-format=bin enable-cuda-dialect})" + return pipeline + + @staticmethod + def generate_func_op(arg_types, arg_attrs, kernel_name, loc=None): + func_op = cutlass_lir.FuncOp( + ir.StringAttr.get(kernel_name), + ir.TypeAttr.get(ir.FunctionType.get(arg_types, [])), + loc=loc, + ) + func_op.attributes["cu_attrs"] = ir.DictAttr.get( + { + str( + cuda_dialect.CUFunctionAttribute.non_portable_cluster_size_allowed + ): ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 1), + str( + cuda_dialect.CUFunctionAttribute.max_dynamic_shared_size_bytes + ): cuda_dialect.DevMaxSharedMemoryOptinAttr.get(), + } + ) + # Monkey patch FuncOp to add an add_entry_block method, if not already defined. + if not hasattr(func_op, "add_entry_block"): + + def add_entry_block(arg_locs): + if len(func_op.body.blocks) != 0: + raise RuntimeError("The function already has an entry block.") + func_op.body.blocks.append(*arg_types) + return func_op.body.blocks[0] + + func_op.add_entry_block = add_entry_block + return func_op + + @staticmethod + def generate_func_ret_op(loc=None, ip=None): + return cutlass_lir.ReturnOp([]) + + # ============================================================================= # KernelLauncher # ============================================================================= @@ -1331,9 +1378,9 @@ def to_index(value): if is_dynamic_expression(value): if isinstance(value, Numeric): value = value.ir_value() - assert ir.IntegerType.isinstance(value.type), ( - f"expects integer type, but got {value.type}" - ) + assert ir.IntegerType.isinstance( + value.type + ), f"expects integer type, but got {value.type}" res = arith.index_cast(T.index(), value) else: res = const(int(value), ty=T.index()) @@ -1379,7 +1426,7 @@ def _validate_iter_args_structure(iter_args, ir_values): def _minmax(op, *args, loc=None, ip=None): """Computes the minimum or maximum value from the provided arguments.""" - from ..base_dsl.typing import _binary_op_type_promote + from ..base_dsl.typing import _binary_op, _binary_op_type_promote # AST Traversal doesn't support early exit in if executor x = None @@ -1830,7 +1877,7 @@ def for_generate( def _createI32Attr(value): if not isinstance(value, int): - raise DSLRuntimeError("value must be int.") + raise DSLRuntimeError(f"value must be int.") return ir.IntegerAttr.get(ir.IntegerType.get_signless(32), value) ir_iter_args = extract_mlir_values(iter_args) if iter_args is not None else None @@ -1951,7 +1998,9 @@ def if_generate( # Collect MLIR results. mlir_results = _get_op_result_or_op_results(if_op) - if not isinstance(mlir_results, list): + if not isinstance(mlir_results, list) and not isinstance( + mlir_results, ir.OpResultList + ): mlir_results = [mlir_results] # Wrap the results with their DSL types. @@ -2245,6 +2294,7 @@ executor.set_functions( any_executor=any_, all_executor=all_, builtin_redirector=_builtin_redirector, + ifexp_dynamic=_ifexp_execute_dynamic, ) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py index 5c032d1c..214eb710 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py @@ -374,7 +374,7 @@ def _loop_execute_range_dynamic( for i, d in enumerate(dyn_yield_ops) ) raise DSLRuntimeError( - f"Failed to create scf.ForOp \n\t\tstart={start_}: type : {type(start_)}" + f"Failed to create dynamic for loop \n\t\tstart={start_}: type : {type(start_)}" f"\n\t\tstop={stop_}: type : {type(stop_)}\n\t\tstep={step_}: type : {type(step_)}" f", \n\tdyn_yield_ops:\n{yield_ops}" ) from e @@ -461,7 +461,7 @@ def _if_execute_dynamic( ) except Exception as e: raise DSLRuntimeError( - f"Failed to create scf.IfOp \n\t\tpred={pred_}: type : {type(pred_)}" + f"Failed to create dynamic if \n\t\tpred={pred_}: type : {type(pred_)}" ) from e return if_op @@ -550,7 +550,7 @@ def _while_execute_dynamic( for i, d in enumerate(dyn_yield_ops) ) raise DSLRuntimeError( - f"Failed to create scf.WhileOp with yield_ops:\n{yield_ops}" + f"Failed to create dynamic while loop with yield_ops:\n{yield_ops}" ) from e def before_block_builder( @@ -643,3 +643,121 @@ def _while_execute_dynamic( before_block_builder: before_block_terminator }, # Only customize the before block ) + + +def _ifexp_execute_dynamic( + pred: "ir.Value", + block_args: tuple, + then_block: Callable, + else_block: Callable, +): + """ + Dynamically execute a Python inline if-expression (ternary) as a runtime-dispatched control flow op. + + This function builds an SCF (Structured Control Flow) `if` operation in the IR, using the given + predicate and block functions for the 'then' and 'else' branches, and infers the result types + from the return signature of those blocks. It ensures that both branches return values of the same + tree structure and types, so that the IR op can properly yield their results. + + Parameters + ---------- + pred : ir.Value + The predicate value (a boolean IR value) that determines which branch is executed. + block_args : tuple + The block arguments that are passed to the then and else blocks. + then_block : Callable + A Python function that executes the 'then' branch and returns the result(s). This will be + executed if `pred` evaluates to True. + else_block : Callable + A Python function that executes the 'else' branch and returns the result(s). This will be + executed if `pred` evaluates to False. + + Returns + ------- + list + The evaluated result(s) of the selected branch, in a standardized (possibly list-wrapped) format. + + Raises + ------ + DSLRuntimeError + If the 'then' and 'else' blocks return values of different tree structures or types, + or if IR construction fails. + + Notes + ----- + This function is a low-level implementation intended for use by the AST transformation machinery, + and not for direct user invocation. It acts as the backend for transformed Python inline if-expressions. + """ + # Infer result types by running both branches with dummy arguments in a temporary region + execution_region = scf.ExecuteRegionOp(result=[]) + execution_region.region.blocks.append() + + result_types = [] + mix_iter_args = [] + + with ir.InsertionPoint(execution_region.region.blocks[0]): + # Call the then block and unpack its results to IR values and tree structure + then_results = ScfGenerator._normalize_region_result_to_list( + then_block(*block_args) + ) + ir_values, then_tree = cutlass_dsl.unpack_to_irvalue(then_results, "ifexp", 0) + + # Call the else block and unpack its results to IR values and tree structure + else_results = ScfGenerator._normalize_region_result_to_list( + else_block(*block_args) + ) + _, else_tree = cutlass_dsl.unpack_to_irvalue(else_results, "ifexp", 0) + + # Check that both branches are structurally and type compatible + if check_tree_equal(then_tree, else_tree) != -1: + raise DSLRuntimeError( + "Then and else blocks of ifexp return different types" + ) + + # Collect result types for the SCF IfOp + result_types.extend([arg.type for arg in ir_values]) + mix_iter_args.extend(then_results) + + # Set up a generator for SCF op creation + scf_gen = ScfGenerator() + + # Function to create the IfOp with correct predicate and result types + def create_if_op(_): + pred_ = Boolean(pred) + try: + if_op = scf.IfOp( + pred_.ir_value(), + hasElse=True, + results_=result_types, + ) + except Exception as e: + raise DSLRuntimeError( + f"Failed to create dynamic if-expression \n\t\tpred={pred_}: type : {type(pred_)}" + ) from e + return if_op + + # SCF region builder for then block + def then_builder(*args): + # Just call the then_block as no arguments are passed to it + return then_block(*block_args) + + # SCF region builder for else block + def else_builder(*args): + return else_block(*block_args) + + # Prepare the list of region builders for the SCF IfOp: first for "then", then for "else" + region_builders = [then_builder, else_builder] + + ret = scf_gen.scf_execute_dynamic( + op_type_name="if", + mix_iter_args=mix_iter_args, + full_write_args_count=0, + mix_iter_arg_names=["unknown" for _ in mix_iter_args], + create_op_func=create_if_op, + region_builders=region_builders, + ) + + # Clean up: Remove the temporary execution region from the IR graph + execution_region.operation.erase() + + return ret diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py index a7c1e7e2..c8a53047 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py @@ -21,7 +21,6 @@ from cutlass._mlir.dialects import llvm from cutlass._mlir._mlir_libs._cutlass_ir import _aot_support from cutlass.cutlass_dsl.cuda_jit_executor import CudaDialectJitCompiledFunction from cutlass.base_dsl.common import DSLRuntimeError -from cutlass.base_dsl.jit_executor import ExecutionArgs from typing import Optional, Callable import tvm_ffi @@ -130,16 +129,13 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): cuda_global_state_ptr = self.address_of( self.cuda_global_state_symbol, self.ptr_type ) - - cuda_init_ptr = context.builder.get_or_load_global_func_ptr_from_text( - current_block, "cuda_init" - ) - cuda_load_to_device_ptr = context.builder.get_or_load_global_func_ptr_from_text( - current_block, "cuda_load_to_device" - ) - set_error_ptr = context.builder.get_or_load_global_func_ptr_from_text( - current_block, "TVMFFIErrorSetRaisedFromCStr" - ) + cuda_init_ptr = self.address_of("cuda_init", self.ptr_type) + cuda_load_to_device_ptr = self.address_of( + "cuda_load_to_device", self.ptr_type + ) + set_error_ptr = self.address_of( + "TVMFFIErrorSetRaisedFromCStr", self.ptr_type + ) with ir.InsertionPoint(current_block): # Call the callback function with the loaded ptr value @@ -211,6 +207,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): global_dtors = llvm.mlir_global_dtors( dtors=[], priorities=[], + data=[], ) else: # use the existing global destructors @@ -223,6 +220,9 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): global_dtors.attributes["priorities"] += [ ir.IntegerAttr.get(self.i32_type, 65535) ] # the default priority + global_dtors.attributes["data"] += [ + ir.FlatSymbolRefAttr.get(unload_func_wrapper_symbol) + ] # the data will not be used, but we need to pass something to satisfy the llvm.mlir.global_dtors op return current_block @@ -255,7 +255,8 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): current_device: Optional[ir.Value], target_device: Optional[ir.Value], ) -> ir.Block: - """Set the CUDA device index if it differs from the target device.""" + """Set the CUDA device index if it differs from the target device. + """ # If either device is None, no switching needed if current_device is None: assert target_device is None @@ -273,7 +274,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): self.cond_br( cond=devices_differ, true_block=switch_device_block, - false_block=continuation_block, + false_block=continuation_block ) # Switch device block: call cudaSetDevice @@ -287,9 +288,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): ) # Check for errors and branch to continuation - switch_device_block = self.check_cuda_error( - result, switch_device_block, context - ) + switch_device_block = self.check_cuda_error(result, switch_device_block, context) with ir.InsertionPoint(switch_device_block): self.br(continuation_block) @@ -320,9 +319,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): op_bundle_sizes=[], op_bundle_operands=[], ) - current_block = self.check_cuda_error( - get_device_result, current_block, context - ) + current_block = self.check_cuda_error(get_device_result, current_block, context) # Load the current device index from the alloca with ir.InsertionPoint(current_block): @@ -354,6 +351,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): return current_block + def find_cuda_device_index_from_params(self, context: CallContext): """Find the CUDA device index from tensor parameters.""" for param in context.params: @@ -365,9 +363,12 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): return None def create_shared_cuda_error_block( - self, current_block: ir.Block, context: CallContext + self, + current_block: ir.Block, + context: CallContext ) -> ir.Block: - """Create a shared error handling block for all CUDA errors.""" + """Create a shared error handling block for all CUDA errors. + """ # Create the shared error block after the current block (setup phase) # This block will be branched to from multiple error checking sites # It accepts the error code as a block argument @@ -397,9 +398,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): current_block = self.append_unload_to_global_dtors(current_block, context) # Create shared CUDA error handling block after the setup blocks # This reduces code duplication - all CUDA errors branch to this single block - self.cuda_error_handle_block = self.create_shared_cuda_error_block( - current_block, context - ) + self.cuda_error_handle_block = self.create_shared_cuda_error_block(current_block, context) # setup device index, will be set around the call to the target function self.cuda_device_index = self.find_cuda_device_index_from_params(context) current_block = super().__call__(current_block, context) @@ -458,6 +457,10 @@ class TVMFFIJitCompiledFunctionBase(CudaDialectJitCompiledFunction): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + # use direct call to the tvm_ffi.Function.__call__ + # to avoid most of python overhead + __call__ = tvm_ffi.Function.__call__ + def to(self, device=None): """TVM FFI function itself is already support all devices.""" return self diff --git a/python/CuTeDSL/cutlass/jax/types.py b/python/CuTeDSL/cutlass/jax/types.py index e6f84ac2..0fa414ed 100644 --- a/python/CuTeDSL/cutlass/jax/types.py +++ b/python/CuTeDSL/cutlass/jax/types.py @@ -285,6 +285,7 @@ class JaxArrayValue(JaxArray): llvm.PointerType.get(), shape_array, [], + no_wrap_flags=0, raw_constant_indices=ir.DenseI32ArrayAttr.get([i]), elem_type=i64, loc=loc, diff --git a/python/CuTeDSL/cutlass/pipeline/helpers.py b/python/CuTeDSL/cutlass/pipeline/helpers.py index 20479c15..ab4c7e57 100644 --- a/python/CuTeDSL/cutlass/pipeline/helpers.py +++ b/python/CuTeDSL/cutlass/pipeline/helpers.py @@ -10,22 +10,14 @@ # is strictly prohibited. import enum +import inspect from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Optional, Union import warnings import cutlass.cute as cute -from cutlass.cutlass_dsl import ( - Boolean, - Int32, - Int64, - if_generate, - dsl_user_op, - dsl_user_op, -) -from cutlass._mlir.dialects import llvm -import cutlass._mlir.dialects.cute as _cute_ir +from cutlass.cutlass_dsl import Boolean, Int32, if_generate, dsl_user_op ############################################################################## @@ -111,7 +103,6 @@ class PipelineOp(enum.Enum): # Async load without TMA AsyncLoad = enum.auto() - def _get_pipeline_op(type_str): return PipelineOp(type_str) @@ -336,7 +327,9 @@ class MbarrierArray(SyncObject): def arrive_and_expect_tx_with_dst( self, index: int, tx_count: int, dst: Optional[int] = None, *, loc=None, ip=None ) -> None: - cute.arch.mbarrier_arrive_and_expect_tx(self.get_barrier(index), tx_count, dst, loc=loc, ip=ip) + cute.arch.mbarrier_arrive_and_expect_tx( + self.get_barrier(index, loc=loc, ip=ip), tx_count, dst, loc=loc, ip=ip + ) @dsl_user_op def try_wait(self, index: int, phase: int, *, loc=None, ip=None) -> Boolean: @@ -386,6 +379,14 @@ class MbarrierArray(SyncObject): ) +# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator +MbarrierArray.__init__.__signature__ = inspect.Signature( + [ + inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD), + ] +) + + ############################################################################## # NamedBarrier class ############################################################################## @@ -429,14 +430,11 @@ class NamedBarrier(SyncObject): """ The unaligned flavor of arrive can be used with an arbitrary number of threads in the CTA. """ - llvm.inline_asm( - None, - [Int32(self.barrier_id).ir_value(), Int32(self.num_threads).ir_value()], - "barrier.arrive $0, $1;", - "r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, + cute.arch.barrier_arrive( + barrier_id=self.barrier_id, + number_of_threads=self.num_threads, + loc=loc, + ip=ip, ) @dsl_user_op @@ -453,15 +451,13 @@ class NamedBarrier(SyncObject): ) self.arrive_and_wait(loc=loc, ip=ip) - def wait_unaligned(self) -> None: - llvm.inline_asm( - None, - [Int32(self.barrier_id).ir_value(), Int32(self.num_threads).ir_value()], - "barrier.sync $0, $1;", - "r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, + @dsl_user_op + def wait_unaligned(self, *, loc=None, ip=None) -> None: + cute.arch.barrier( + barrier_id=self.barrier_id, + number_of_threads=self.num_threads, + loc=loc, + ip=ip, ) @dsl_user_op @@ -751,18 +747,6 @@ def agent_sync(group: Agent, is_relaxed: bool = False, *, loc=None, ip=None): ) -def _mbarrier_i64_to_ptr(val: Int64) -> cute.Pointer: - """ - Converts a smem pointer of type Int64 to cute.Pointer with 8B alignment - """ - return cute.make_ptr( - Int64, - val.ir_value(), - mem_space=_cute_ir.AddressSpace.smem, - assumed_align=8, - ) - - # NamedBarrier free functions @dsl_user_op def arrive(barrier_id: int, num_threads: int, *, loc=None, ip=None): @@ -780,19 +764,13 @@ def arrive_unaligned(barrier_id: int, num_threads: int, *, loc=None, ip=None): """ The unaligned flavor of arrive can be used with an arbitrary number of threads in the CTA. """ - llvm.inline_asm( - None, - [Int32(barrier_id).ir_value(), Int32(num_threads).ir_value()], - "barrier.arrive $0, $1;", - "r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, + cute.arch.barrier_arrive( + barrier_id=barrier_id, number_of_threads=num_threads, loc=loc, ip=ip ) @dsl_user_op -def wait(barrier_id: int, num_threads: int): +def wait(*, loc=None, ip=None): """ NamedBarriers do not have a standalone wait like mbarriers, only an arrive_and_wait. If synchronizing two warps in a producer/consumer pairing, the arrive count would be @@ -811,14 +789,8 @@ def wait_unaligned(barrier_id: int, num_threads: int, *, loc=None, ip=None): warnings.warn( "NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()." ) - llvm.inline_asm( - None, - [Int32(barrier_id).ir_value(), Int32(num_threads).ir_value()], - "barrier.sync $0, $1;", - "r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, + cute.arch.barrier( + barrier_id=barrier_id, number_of_threads=num_threads, loc=loc, ip=ip ) diff --git a/python/CuTeDSL/cutlass/pipeline/sm100.py b/python/CuTeDSL/cutlass/pipeline/sm100.py index 76a17b24..b804d2a5 100644 --- a/python/CuTeDSL/cutlass/pipeline/sm100.py +++ b/python/CuTeDSL/cutlass/pipeline/sm100.py @@ -238,7 +238,10 @@ class PipelineTmaUmma(PipelineAsync): if not defer_sync: cute.arch.mbarrier_init_fence() - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: + if ( + cta_layout_vmnk is None + or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1 + ): agent_sync(Agent.ThreadBlock) else: agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) @@ -449,7 +452,10 @@ class PipelineAsyncUmma(PipelineAsync): if not defer_sync: cute.arch.mbarrier_init_fence() - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: + if ( + cta_layout_vmnk is None + or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1 + ): agent_sync(Agent.ThreadBlock) else: agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) @@ -587,7 +593,10 @@ class PipelineUmmaAsync(PipelineAsync): if not defer_sync: cute.arch.mbarrier_init_fence() - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: + if ( + cta_layout_vmnk is None + or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1 + ): agent_sync(Agent.ThreadBlock) else: agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) diff --git a/python/CuTeDSL/cutlass/pipeline/sm90.py b/python/CuTeDSL/cutlass/pipeline/sm90.py index 385f8fa1..ccbb9ee7 100644 --- a/python/CuTeDSL/cutlass/pipeline/sm90.py +++ b/python/CuTeDSL/cutlass/pipeline/sm90.py @@ -308,7 +308,7 @@ class PipelineAsync: @dataclass(frozen=True) class PipelineCpAsync(PipelineAsync): """ - PipelineCpAsync is used for CpAsync producers and AsyncThread consumers (e.g. Hopper non-TMA mainloops). + PipelineCpAsync is used for CpAsync producers and AsyncThread consumers """ @staticmethod @@ -656,6 +656,7 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): ) if not defer_sync: + cute.arch.mbarrier_init_fence() if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: agent_sync(Agent.ThreadBlock) else: diff --git a/python/CuTeDSL/cutlass/utils/__init__.py b/python/CuTeDSL/cutlass/utils/__init__.py index adddfc98..a0909c57 100644 --- a/python/CuTeDSL/cutlass/utils/__init__.py +++ b/python/CuTeDSL/cutlass/utils/__init__.py @@ -70,7 +70,6 @@ from .tmem_allocator import TmemAllocator, get_num_tmem_alloc_cols from .layout import LayoutEnum -from . import gemm from . import distributed from .mixed_input_helpers import ( @@ -99,11 +98,17 @@ from .mixed_input_helpers import ( store_transformed_a, ) +from . import gemm + from . import hopper_helpers as sm90 from . import blackwell_helpers as sm100 - from .print_latex import print_latex, print_latex_tv +from .tensor_helpers import ( + is_fp8_dtype, + create_cute_tensor_for_fp8, +) + __all__ = [ "get_smem_capacity_in_bytes", "SmemAllocator", @@ -135,6 +140,7 @@ __all__ = [ "get_divisibility", "epilogue_tma_store", "epilogue", + "create_tensor_a", "compute_epilogue_tile_shape", "get_smem_store_op", "get_tmem_load_op", @@ -145,10 +151,12 @@ __all__ = [ "make_blockscaled_trivial_tiled_mma", "sm90", "sm100", - "print_latex", - "print_latex_tv", "gemm", - "distributed", "ClcDynamicPersistentTileSchedulerParams", "ClcDynamicPersistentTileScheduler", + "print_latex", + "print_latex_tv", + "is_fp8_dtype", + "create_cute_tensor_for_fp8", + "distributed", ] diff --git a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py index f6fa0f84..c0d146c6 100644 --- a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py +++ b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py @@ -658,7 +658,11 @@ def make_smem_layout_a( :rtype: Union[cute.Layout, cute.ComposedLayout] """ - is_k_major = (tiled_mma.op.a_major_mode == OperandMajorMode.K) if is_k_major is None else is_k_major + is_k_major = ( + (tiled_mma.op.a_major_mode == OperandMajorMode.K) + if is_k_major is None + else is_k_major + ) a_major_mode = OperandMajorMode.K if is_k_major else OperandMajorMode.MN a_smem_shape = tiled_mma.partition_shape_A( cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip), loc=loc, ip=ip @@ -712,7 +716,11 @@ def make_smem_layout_b( :rtype: Union[cute.Layout, cute.ComposedLayout] """ - is_k_major = (tiled_mma.op.b_major_mode == OperandMajorMode.K) if is_k_major is None else is_k_major + is_k_major = ( + (tiled_mma.op.b_major_mode == OperandMajorMode.K) + if is_k_major is None + else is_k_major + ) b_major_mode = OperandMajorMode.K if is_k_major else OperandMajorMode.MN b_smem_shape = tiled_mma.partition_shape_B( cute.dice(mma_tiler_mnk, (None, 1, 1), loc=loc, ip=ip), loc=loc, ip=ip diff --git a/python/CuTeDSL/cutlass/utils/blockscaled_layout.py b/python/CuTeDSL/cutlass/utils/blockscaled_layout.py index 49078db5..b3461416 100644 --- a/python/CuTeDSL/cutlass/utils/blockscaled_layout.py +++ b/python/CuTeDSL/cutlass/utils/blockscaled_layout.py @@ -84,6 +84,38 @@ def tile_atom_to_shape_SF( return sf_layout +@dsl_user_op +def make_smem_layout_sf( + tile_shape: cute.Tile, + sf_vec_size: int, + num_stages: int, + *, + loc=None, + ip=None, +) -> cute.Layout: + """ + A helper function to get dynamic SFA/SFB layout by filling dynamic A/B shape to the scale factor atom layout. + + :param Shape: The shape of the A/B tensor + :param sf_vec_size: Scale factor vector size + :param num_stages: Number of stages + + :return: The layout of the SFA/SFB tensor + :rtype: cute.Layout + """ + + smem_layout = cute.tile_to_shape( + BlockScaledBasicChunk(sf_vec_size).layout, tile_shape, (2, 1) + ) + smem_layout_staged = cute.append( + smem_layout, + cute.make_layout( + num_stages, stride=cute.cosize(cute.filter_zeros(smem_layout)) + ), + ) + return smem_layout_staged + + @dsl_user_op def make_smem_layout_sfa( tiled_mma: cute.TiledMma, @@ -214,6 +246,176 @@ def make_smem_layout_sfb( return sfb_smem_layout_staged +@dsl_user_op +def sm120_make_smem_layout_sfa( + tiled_mma: cute.TiledMma, + tile_shape_mnk: cute.Tile, + sf_vec_size: int, + num_stages: int, + *, + loc=None, + ip=None, +) -> cute.Layout: + """ + Make smem layout for SFA based on: + 1. BlockScaledBasicChunk + 2. MMA tiler shape + 3. Scale factor vector size + 4. Number of stages + + :param tiled_mma: The tiled MMA + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The mma tiler shape + :type mma_tiler_mnk: cute.Tile + :param sf_vec_size: The scale factor vector size + :type sf_vec_size: int + :param num_stages: The number of stages + :type num_stages: int + + :return: Smem layout for SFA + :rtype: cute.Layout + """ + + assert sf_vec_size == 16 or sf_vec_size == 32, "sf_vec_size must be 16 or 32" + + blk_mn = 128 + blk_sf = 4 + blk_elems = blk_mn * blk_sf + mma_nsf = tiled_mma.shape_mnk[2] // sf_vec_size + + mn_basic_block_shape = (32, 4) + mn_basic_block_stride = (16, 4) + k_basic_block_shape = (sf_vec_size, mma_nsf) + k_basic_block_stride = (0, 1) + + assert tile_shape_mnk[0] % blk_mn == 0, ( + "tile_shape_mnk[0] must be divisible by blk_mn" + ) + + sSFA_shapeM = (mn_basic_block_shape, tile_shape_mnk[0] // blk_mn) + sSF_strideM = (mn_basic_block_stride, blk_elems) + + assert tile_shape_mnk[2] % (blk_sf * mma_nsf) == 0, ( + "tile_shape_mnk[2] must be divisible by blk_sf * mma_nsf" + ) + + sSFA_shapeK = ( + k_basic_block_shape, + blk_sf // mma_nsf, + tile_shape_mnk[2] // sf_vec_size // blk_sf, + ) + sSF_strideK = ( + k_basic_block_stride, + mma_nsf, + tile_shape_mnk[0] // blk_mn * blk_elems, + ) + + sSFA_shape = (sSFA_shapeM, sSFA_shapeK) + sSFA_stride = (sSF_strideM, sSF_strideK) + + smem_layout = cute.make_layout(sSFA_shape, stride=sSFA_stride) + + # (((Atom_Inst_M, Rest_M),(Atom_Inst_K, Rest_K)), MMA_M, MMA_K, STAGE) + sfa_smem_layout_staged = cute.append( + smem_layout, + cute.make_layout( + num_stages, stride=cute.cosize(cute.filter_zeros(smem_layout)) + ), + ) + + return sfa_smem_layout_staged + + +@dsl_user_op +def sm120_make_smem_layout_sfb( + tiled_mma: cute.TiledMma, + tile_shape_mnk: cute.Tile, + sf_vec_size: int, + num_stages: int, + *, + loc=None, + ip=None, +) -> cute.Layout: + """ + Make smem layout for SFB based on: + 1. BlockScaledBasicChunk + 2. MMA tiler shape + 3. Scale factor vector size + 4. Number of stages + + :param tiled_mma: The tiled MMA + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The mma tiler shape + :type mma_tiler_mnk: cute.Tile + :param sf_vec_size: The scale factor vector size + :type sf_vec_size: int + :param num_stages: The number of stages + :type num_stages: int + + :return: Smem layout for SFA + :rtype: cute.Layout + """ + + # A single indivisible block will hold 4 scale factors of 128 rows/columns (A/B matrix). + # 4 is chosen to make consecutive 32bits of data to have scale factors for only a single row(col). + blk_mn = 128 + blk_sf = 4 + blk_elems = blk_mn * blk_sf + + assert sf_vec_size == 16 or sf_vec_size == 32, "sf_vec_size must be 16 or 32" + + assert tile_shape_mnk[1] % blk_mn == 0, ( + "tile_shape_mnk[1] must be divisible by blk_mn" + ) + + assert tile_shape_mnk[2] % sf_vec_size == 0, ( + "tile_shape_mnk[2] must be divisible by sf_vec_size" + ) + + mma_nsf = tiled_mma.shape_mnk[2] // sf_vec_size + + mn_basic_block_shape = (32, 4) + mn_basic_block_stride = (16, 4) + k_basic_block_shape = (sf_vec_size, mma_nsf) + k_basic_block_stride = (0, 1) + + assert tile_shape_mnk[1] % blk_mn == 0, ( + "tile_shape_mnk[1] must be divisible by blk_mn" + ) + + sSFA_shapeN = (mn_basic_block_shape, tile_shape_mnk[1] // blk_mn) + sSF_strideN = (mn_basic_block_stride, blk_elems) + + assert tile_shape_mnk[2] % (blk_sf * mma_nsf) == 0, ( + "tile_shape_mnk[2] must be divisible by blk_sf * mma_nsf" + ) + + sSFA_shapeK = ( + k_basic_block_shape, + blk_sf // mma_nsf, + tile_shape_mnk[2] // sf_vec_size // blk_sf, + ) + sSF_strideK = ( + k_basic_block_stride, + mma_nsf, + tile_shape_mnk[1] // blk_mn * blk_elems, + ) + + sSFA_shape = (sSFA_shapeN, sSFA_shapeK) + sSFA_stride = (sSF_strideN, sSF_strideK) + + smem_layout = cute.make_layout(sSFA_shape, stride=sSFA_stride) + + # (((Atom_Inst_M, Rest_M),(Atom_Inst_K, Rest_K)), MMA_M, MMA_K, STAGE) + sfb_smem_layout_staged = cute.append( + smem_layout, + cute.make_layout( + num_stages, stride=cute.cosize(cute.filter_zeros(smem_layout)) + ), + ) + + return sfb_smem_layout_staged + @dsl_user_op def make_tmem_layout_sfa( diff --git a/python/CuTeDSL/cutlass/utils/distributed.py b/python/CuTeDSL/cutlass/utils/distributed.py index 36c458f1..01952814 100644 --- a/python/CuTeDSL/cutlass/utils/distributed.py +++ b/python/CuTeDSL/cutlass/utils/distributed.py @@ -71,6 +71,7 @@ def ld_bypass(input_tensor: cute.Tensor): @dsl_user_op def multimem_red_release_gpu_add1( lock_ptr: Pointer, + *, loc=None, ip=None, ) -> None: @@ -89,6 +90,7 @@ def multimem_red_release_gpu_add1( @dsl_user_op def multimem_red_release_sys_add1( lock_ptr: Pointer, + *, loc=None, ip=None, ) -> None: @@ -285,7 +287,6 @@ def spin_lock_atom_cas_relaxed_wait( ip=ip, ) - ######################################################## # Multimem Load & Store ######################################################## diff --git a/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py b/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py index 9741c08a..f18a4ff3 100644 --- a/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py +++ b/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py @@ -27,6 +27,7 @@ from cutlass.utils.static_persistent_tile_scheduler import ( ) import cutlass.cute as cute + class ClcDynamicPersistentTileSchedulerParams: """A class to represent parameters for a dynamic persistent tile scheduler. @@ -98,6 +99,7 @@ class ClcDynamicPersistentTileSchedulerParams: ) return problem_ceiling_cta_mnl + class ClcDynamicPersistentTileScheduler: """A scheduler for dynamic persistent tile execution in CUTLASS/CuTe kernels. @@ -127,7 +129,7 @@ class ClcDynamicPersistentTileScheduler: :param num_tiles_executed: Counter for executed tiles. :type num_tiles_executed: Int32 :param clc_response_ptr: Pointer of the clc rsponse. - :type clc_response_ptr: Tuple[Integer, Integer, Integer, Integer] + :type clc_response_ptr: cute.Pointer :param block_idx: The block index. :type block_idx: Tuple[Integer, Integer, Integer] """ @@ -236,14 +238,17 @@ class ClcDynamicPersistentTileScheduler: @dsl_user_op def work_tile_info_from_clc_response( - self, result_addr: Int32, *, loc=None, ip=None + self, result_addr: cute.Pointer, *, loc=None, ip=None ) -> WorkTileInfo: """ Simulates parsing CLC response data in Python. result_addr: 16-byte response data (simulating shared memory access) """ m_idx, n_idx, l_idx, vld = cute.arch.clc_response(result_addr, loc=loc, ip=ip) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) cta_idx_in_cluster, cta_idy_in_cluster, _ = self.cta_id_in_cluster cur_tile_coord = (m_idx + cta_idx_in_cluster, n_idx + cta_idy_in_cluster, l_idx) return WorkTileInfo(cur_tile_coord, vld) diff --git a/python/CuTeDSL/cutlass/utils/gemm/__init__.py b/python/CuTeDSL/cutlass/utils/gemm/__init__.py index d9d89020..e6dcbbe1 100644 --- a/python/CuTeDSL/cutlass/utils/gemm/__init__.py +++ b/python/CuTeDSL/cutlass/utils/gemm/__init__.py @@ -7,7 +7,7 @@ # # Any use, reproduction, disclosure, or distribution of this software # and related documentation outside the scope permitted by the EULA -# is strictly prohibited +# is strictly prohibited. from . import sm100 diff --git a/python/CuTeDSL/cutlass/utils/gemm/sm100.py b/python/CuTeDSL/cutlass/utils/gemm/sm100.py index d23fcfba..b6f151aa 100644 --- a/python/CuTeDSL/cutlass/utils/gemm/sm100.py +++ b/python/CuTeDSL/cutlass/utils/gemm/sm100.py @@ -14,9 +14,6 @@ import cutlass.cute as cute from cutlass.cutlass_dsl import Int32, Boolean, Constexpr, const_expr import cutlass.pipeline as pipeline from cutlass.utils.static_persistent_tile_scheduler import StaticPersistentTileScheduler -from cutlass.utils.dynamic_persistent_tile_scheduler import ( - ClcDynamicPersistentTileScheduler, -) from cutlass.utils.blackwell_helpers import get_tmem_load_op, get_smem_store_op from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.nvgpu.common import CacheEvictionPriority @@ -161,8 +158,6 @@ def epilogue_tma_store( gemm_kernel, epi_tidx: Int32, warp_idx: Int32, - acc_pipeline: pipeline.PipelineAsync, - tiled_mma: cute.TiledMma, tma_atom_c: cute.CopyAtom, # Input of epilogue tCtAcc_base: cute.Tensor, @@ -171,11 +166,13 @@ def epilogue_tma_store( # Output of epilogue tCgC_base: cute.Tensor, epi_tile: cute.Tile, - tile_sched: Union[StaticPersistentTileScheduler, ClcDynamicPersistentTileScheduler], + num_tiles_executed: Int32, epilogue_op: Constexpr, - clc_pipeline: Union[pipeline.PipelineClcFetchAsync, None] = None, - clc_consumer_state: Union[pipeline.PipelineState, None] = None, -) -> None: + mma_tile_coord_mnl: Tuple[Int32, Int32, Int32], + acc_consumer_state: pipeline.PipelineState, + acc_pipeline: pipeline.PipelineAsync, + c_pipeline: pipeline.PipelineTmaStore, +) -> pipeline.PipelineState: # Layout transformation for tCgC_base # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, TILE_M, TILE_N, TILE_K) # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), TILE_M, TILE_N, TILE_K) @@ -207,142 +204,97 @@ def epilogue_tma_store( cute.group_modes(tCgC_epi, 0, 2), ) - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, gemm_kernel.num_acc_stage - ) - - # Threads/warps participating in tma store pipeline - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, - 32 * len(gemm_kernel.epilogue_warp_id), - ) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=gemm_kernel.num_c_stage, producer_group=c_producer_group - ) - epilog_sync_barrier = pipeline.NamedBarrier( barrier_id=gemm_kernel.epilog_sync_bar_id, num_threads=32 * len(gemm_kernel.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_N) + 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 = num_tiles_executed * subtile_cnt + for subtile_idx in 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) # - # Slice to per mma tile index + # Convert to C type # - # ((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_N) - tTR_tAcc = tTR_tAcc_base[ - (None, None, None, None, None, acc_consumer_state.index) - ] + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(gemm_kernel.c_dtype)) + tRS_rC.store(acc_vec) # - # Wait for accumulator buffer full + # Store C to shared memory # - 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 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(gemm_kernel.c_dtype)) - tRS_rC.store(acc_vec) - - # - # Store C to shared memory - # - c_buffer = (num_prev_subtiles + subtile_idx) % gemm_kernel.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("async.shared", space="cta") - epilog_sync_barrier.arrive_and_wait() - - # - # TMA store C to global memory - # - if warp_idx == gemm_kernel.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() - epilog_sync_barrier.arrive_and_wait() - + c_buffer = (num_prev_subtiles + subtile_idx) % gemm_kernel.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("async.shared", space="cta") epilog_sync_barrier.arrive_and_wait() # - # Async arrive accumulator buffer empty + # TMA store C to global memory # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() + if warp_idx == gemm_kernel.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() + epilog_sync_barrier.arrive_and_wait() - # - # Advance to next tile - # - # Check if tile_sched is StaticPersistentTileScheduler or any subclass inheriting from it - if const_expr(isinstance(tile_sched, StaticPersistentTileScheduler)): - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() - elif const_expr(isinstance(tile_sched, ClcDynamicPersistentTileScheduler)): - clc_pipeline.consumer_wait(clc_consumer_state) - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() - else: - # Not match - pass + epilog_sync_barrier.arrive_and_wait() - # Wait for C store complete - c_pipeline.producer_tail() + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + return acc_consumer_state @cute.jit def epilogue( gemm_kernel, epi_tidx: Int32, - acc_pipeline: pipeline.PipelineAsync, - tiled_mma: cute.TiledMma, tCtAcc_base: cute.Tensor, tCgC_base: cute.Tensor, epi_tile: cute.Tile, - tile_sched: Union[StaticPersistentTileScheduler, ClcDynamicPersistentTileScheduler], epilogue_op: Constexpr, - tmem_dealloc_barrier: pipeline.NamedBarrier, + mma_tile_coord_mnl: Tuple[Int32, Int32, Int32], + acc_consumer_state: pipeline.PipelineState, + acc_pipeline: pipeline.PipelineAsync, tCcC_base: cute.Tensor = None, mC_mnl: cute.Tensor = None, - clc_pipeline: Union[pipeline.PipelineClcFetchAsync, None] = None, - clc_consumer_state: Union[pipeline.PipelineState, None] = None, -) -> None: + overlapping_accum: Constexpr = False, +) -> pipeline.PipelineState: """ Epilogue function that stores accumulator results directly to global memory. Used when TMA store is not enabled. @@ -351,32 +303,26 @@ def epilogue( :type gemm_kernel: Any :param epi_tidx: Thread index in epilogue warp groups :type epi_tidx: Int32 - :param acc_pipeline: Accumulator pipeline for async operations - :type acc_pipeline: pipeline.PipelineAsync - :param tiled_mma: The tiled MMA configuration - :type tiled_mma: cute.TiledMma :param tCtAcc_base: Base accumulator tensor in tensor memory :type tCtAcc_base: cute.Tensor :param tCgC_base: The global memory tensor C to be copied and partitioned :type tCgC_base: cute.Tensor :param epi_tile: Epilogue tile configuration :type epi_tile: cute.Tile - :param tile_sched: Tile scheduler for persistent scheduling - :type tile_sched: StaticPersistentTileScheduler :param epilogue_op: Optional elementwise operation to apply :type epilogue_op: Constexpr - :param tmem_dealloc_barrier: Barrier for tensor memory deallocation - :type tmem_dealloc_barrier: pipeline.NamedBarrier - :param alignment_bytes: Alignment bytes for global memory store - :type alignment_bytes: int + :param mma_tile_coord_mnl: MMA tile coordinates (M, N, L) + :type mma_tile_coord_mnl: Tuple[Int32, Int32, Int32] + :param acc_consumer_state: Accumulator consumer pipeline state + :type acc_consumer_state: pipeline.PipelineState + :param acc_pipeline: Accumulator pipeline for async operations + :type acc_pipeline: pipeline.PipelineAsync :param tCcC_base: Identity/coordinate tensor C :type tCcC_base: cute.Tensor :param mC_mnl: Global memory tensor C (full tensor for predicate computation) :type mC_mnl: cute.Tensor - :param clc_pipeline: Pipeline for dynamic persistent tile scheduling - :type clc_pipeline: Union[pipeline.PipelineClcFetchAsync, None] - :param clc_consumer_state: Consumer state for dynamic persistent tile scheduling - :type clc_consumer_state: Union[pipeline.PipelineState, None] + :param overlapping_accum: Whether to use overlapping accumulator + :type overlapping_accum: Constexpr """ # Layout transformation for tCgC_base @@ -434,29 +380,21 @@ def epilogue( cC_epi = cute.flat_divide(tCcC, epi_tile) tTR_cC_partitioned = thr_copy_t2r.partition_D(cC_epi) - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, gemm_kernel.num_acc_stage - ) - - work_tile = tile_sched.initial_work_tile_info() - while work_tile.is_valid_tile: - # - # Pre-advance to next tile - # - if const_expr(isinstance(tile_sched, StaticPersistentTileScheduler)): - tile_sched.advance_to_next_work() - next_work_tile = tile_sched.get_current_work() - - # Get tile coord from current work tile - 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], + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + *mma_tile_coord_mnl, ) + ] + if const_expr(use_predication): # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_gC = tTR_gC_partitioned[ + tTR_cC = tTR_cC_partitioned[ ( None, None, @@ -466,88 +404,89 @@ def epilogue( *mma_tile_coord_mnl, ) ] - if const_expr(use_predication): - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_cC = tTR_cC_partitioned[ - ( - None, - None, - None, - None, - None, - *mma_tile_coord_mnl, - ) - ] - tTR_cC = cute.group_modes(tTR_cC, 3, cute.rank(tTR_cC)) + tTR_cC = cute.group_modes(tTR_cC, 3, cute.rank(tTR_cC)) - # 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) - ] + # Get accumulator stage index + if const_expr(overlapping_accum): + acc_stage_index = acc_consumer_state.phase + reverse_subtile = acc_stage_index == 0 + else: + acc_stage_index = acc_consumer_state.index + + # 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_stage_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)) + 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]) + for subtile_idx in range(subtile_cnt): + # Compute the actual subtile index + real_subtile_idx = subtile_idx + if const_expr(overlapping_accum): + if reverse_subtile: + real_subtile_idx = subtile_cnt - 1 - subtile_idx + # + # Get the destination and coordinate slices for this subtile + # + tTR_gC_subtile = tTR_gC[(None, None, None, real_subtile_idx)] + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) # - # Wait for accumulator buffer full + # Async arrive accumulator buffer empty # - acc_pipeline.consumer_wait(acc_consumer_state) - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - 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]) - for subtile_idx in range(subtile_cnt): - # - # Get the destination and coordinate slices for this subtile - # - tTR_gC_subtile = tTR_gC[(None, None, None, subtile_idx)] - # - # 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) - # Async arrive accumulator buffer empty - # Release early for perf + if const_expr(overlapping_accum): + # Early release when overlapping: release after processing the + # overlapping region (SF columns) so they can be reused + if subtile_idx == gemm_kernel.iter_acc_early_release_in_epilogue: + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + else: + # Release early for perf at the last subtile if subtile_idx == subtile_cnt - 1: with cute.arch.elect_one(): acc_pipeline.consumer_release(acc_consumer_state) acc_consumer_state.advance() - # - # Convert to C type - # - acc_vec = tTR_rAcc.load() - acc_vec = epilogue_op(acc_vec.to(gemm_kernel.c_dtype)) - tTR_rC.store(acc_vec) + # + # Convert to C type + # + acc_vec = tTR_rAcc.load() + acc_vec = epilogue_op(acc_vec.to(gemm_kernel.c_dtype)) + tTR_rC.store(acc_vec) - if const_expr(use_predication): - # compute predicate - tTR_cC_subtile = tTR_cC[(None, None, None, subtile_idx)] - pred_C_shape = (1, *tTR_cC_subtile.shape[1:]) - pred_C = cute.make_rmem_tensor(pred_C_shape, Boolean) - for m_idx in range(tTR_cC_subtile.shape[1]): - for n_idx in range(tTR_cC_subtile.shape[2]): - vector_first_coord = tTR_cC_subtile[(0, m_idx, n_idx)] - pred_C[(0, m_idx, n_idx)] = cute.elem_less( - vector_first_coord, mC_mnl.shape - ) - # Store C to global memory with predication - cute.copy(simt_atom, tTR_rC, tTR_gC_subtile, pred=pred_C) - else: - # Store C directly to global memory - cute.copy(simt_atom, tTR_rC, tTR_gC_subtile) + if const_expr(use_predication): + # compute predicate + tTR_cC_subtile = tTR_cC[(None, None, None, real_subtile_idx)] + pred_C_shape = (1, *tTR_cC_subtile.shape[1:]) + pred_C = cute.make_rmem_tensor(pred_C_shape, Boolean) + for m_idx in range(tTR_cC_subtile.shape[1]): + for n_idx in range(tTR_cC_subtile.shape[2]): + vector_first_coord = tTR_cC_subtile[(0, m_idx, n_idx)] + pred_C[(0, m_idx, n_idx)] = cute.elem_less( + vector_first_coord, mC_mnl.shape + ) + # Store C to global memory with predication + cute.copy(simt_atom, tTR_rC, tTR_gC_subtile, pred=pred_C) + else: + # Store C directly to global memory + cute.copy(simt_atom, tTR_rC, tTR_gC_subtile) - if const_expr(isinstance(tile_sched, StaticPersistentTileScheduler)): - work_tile = next_work_tile - elif const_expr(isinstance(tile_sched, ClcDynamicPersistentTileScheduler)): - clc_pipeline.consumer_wait(clc_consumer_state) - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() - - # Synchronize before TMEM dealloc (done by the caller) - tmem_dealloc_barrier.arrive_and_wait() + return acc_consumer_state @cute.jit @@ -908,4 +847,3 @@ def epilogue_release_flag( # Synchronize before TMEM dealloc (done by the caller) tmem_dealloc_barrier.arrive_and_wait() - diff --git a/python/CuTeDSL/cutlass/utils/hardware_info.py b/python/CuTeDSL/cutlass/utils/hardware_info.py index 1edd861c..68f05235 100644 --- a/python/CuTeDSL/cutlass/utils/hardware_info.py +++ b/python/CuTeDSL/cutlass/utils/hardware_info.py @@ -41,7 +41,23 @@ class HardwareInfo: self.driver_version = self._checkCudaErrors(driver.cuDriverGetVersion()) # Getting the max active clusters for a given cluster size - def get_max_active_clusters(self, cluster_size: int) -> int: + def get_max_active_clusters( + self, cluster_size: int, stream: driver.CUstream = None + ) -> int: + """ + Get the maximum number of active clusters for a given cluster size. + + When a stream from a green context is provided, the occupancy calculation + will reflect the reduced SM partition of the green context. + + :param cluster_size: Number of blocks per cluster (must be between 1 and 32) + :type cluster_size: int + :param stream: Optional CUDA stream handle. If provided (especially from a green context), + the occupancy calculation reflects the stream's SM partition. + :type stream: driver.CUstream, optional + :return: Maximum number of active clusters + :rtype: int + """ if self._cuda_driver_version_lt(11, 8): raise RuntimeError( "CUDA Driver version < 11.8, cannot get _max_active_clusters" @@ -94,6 +110,13 @@ class HardwareInfo: launch_config.blockDimY = 1 launch_config.blockDimZ = 1 launch_config.sharedMemBytes = max_dynamic_shared_memory + + # IMPORTANT: Set the stream for green context support + # When hStream is set, cuOccupancyMaxActiveClusters will use the context + # associated with that stream, which includes the green context's SM partition + if stream is not None: + launch_config.hStream = stream + launch_config.numAttrs = 1 # max possible cluster size is 32 cluster_dims_attr = driver.CUlaunchAttribute() @@ -178,13 +201,17 @@ class HardwareInfo: # Create a temporary directory for dumping artifacts with tempfile.TemporaryDirectory() as temp_dir: # keep-cubin will keep the cubin in the artifacts - compiled_func = cute.compile(self._host_function, options=f"--dump-dir={temp_dir} --keep-cubin") + compiled_func = cute.compile( + self._host_function, options=f"--dump-dir={temp_dir} --keep-cubin" + ) # Get the CUBIN from artifacts cubin_data = compiled_func.artifacts.CUBIN cuda_library = self._checkCudaErrors( driver.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0) ) # Enumerate kernels from the library - kernels = self._checkCudaErrors(driver.cuLibraryEnumerateKernels(1, cuda_library)) + kernels = self._checkCudaErrors( + driver.cuLibraryEnumerateKernels(1, cuda_library) + ) # Get the function from the kernel return self._checkCudaErrors(driver.cuKernelGetFunction(kernels[0])) diff --git a/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py b/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py index 41ab59bd..477fcb94 100644 --- a/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py +++ b/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py @@ -513,9 +513,9 @@ def get_smem_layout_scale( cute.size(mma_tiler[2]) % cute.size(smem_layout_scale_per_stage.outer[1]) == 0 ), "smem_layout_scale_per_stage must evenly divide tile k shape." # Shared memory buffer for scale must be at least 128B to satisfy TMA requirement - assert ( - cute.size_in_bytes(a_scale_dtype, smem_layout_scale_per_stage) >= 128 - ), "smem size for scale must be at least 128B" + assert cute.size_in_bytes(a_scale_dtype, smem_layout_scale_per_stage) >= 128, ( + "smem size for scale must be at least 128B" + ) # Scale layout in smem with multiple stages smem_layout_scale = cute.append( smem_layout_scale_per_stage, @@ -972,10 +972,9 @@ def cvt_tensor_a( for int4-to-bf16 conversion. """ from cutlass import CUDA_VERSION - # shuffle is supported since CUDA 13.1 shuffle_supported = True - if CUDA_VERSION.major < 13 or (CUDA_VERSION == 13 and CUDA_VERSION.minor < 1): + if CUDA_VERSION.major < 13 or (CUDA_VERSION.major == 13 and CUDA_VERSION.minor < 1): shuffle_supported = False shuffle = shuffle and shuffle_supported rst = src.load() diff --git a/python/CuTeDSL/cutlass/utils/smem_allocator.py b/python/CuTeDSL/cutlass/utils/smem_allocator.py index 2e1e3f89..01afb380 100644 --- a/python/CuTeDSL/cutlass/utils/smem_allocator.py +++ b/python/CuTeDSL/cutlass/utils/smem_allocator.py @@ -80,7 +80,7 @@ class SmemAllocator: GPU compute capability. :param compute_capability: The compute capability string (e.g. "70", "75", "80") - :type compute_capability: str + :type compute_capability: Optional[str] :return: The shared memory capacity in bytes :rtype: int :raises ValueError: If the compute capability is not supported diff --git a/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py b/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py index 925ff522..7b836532 100644 --- a/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py +++ b/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py @@ -9,6 +9,7 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. +import inspect from typing import Tuple from cutlass.cutlass_dsl import ( @@ -325,6 +326,14 @@ class PersistentTileSchedulerParams: return (*self.cluster_shape_mn, num_persistent_clusters) +# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator +PersistentTileSchedulerParams.__init__.__signature__ = inspect.Signature( + [ + inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD), + ] +) + + class StaticPersistentTileScheduler: """A scheduler for static persistent tile execution in CUTLASS/CuTe kernels. diff --git a/python/CuTeDSL/cutlass/utils/tensor_helpers.py b/python/CuTeDSL/cutlass/utils/tensor_helpers.py new file mode 100644 index 00000000..d7dd93fb --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/tensor_helpers.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +"""Utility functions for tensor creation and type handling.""" + +from typing import Type, Optional + +# Import only the specific types needed to avoid circular import with cutlass module +from cutlass.cute.typing import Float8E5M2, Float8E4M3FN, TFloat32, Numeric +from cutlass.cute.runtime import from_dlpack + + +def is_fp8_dtype(dtype: Type[Numeric]) -> bool: + """Check if dtype is a float8 type that doesn't support dlpack. + params dtype: The cutlass numeric type to check + type dtype: Type[cutlass.Numeric] + return: True if the dtype is Float8E5M2 or Float8E4M3FN, False otherwise + """ + return dtype in {Float8E5M2, Float8E4M3FN} + + +def create_cute_tensor_for_fp8( + storage_tensor, + dtype: Type[Numeric], + leading_dim: int, + source_f32_tensor=None, +): + """Create cute tensor, handling float8 types that don't support dlpack. + + For float8 types, the storage_tensor should be uint8 (for DLPack compatibility). + The source_f32_tensor provides the actual float32 values to convert to fp8. + + params storage_tensor: Tensor for DLPack (uint8 for fp8, otherwise the actual dtype) + params dtype: Target cutlass dtype + params leading_dim: Leading dimension for dynamic layout + paramas source_f32_tensor: Float32 source data for fp8 conversion (required for fp8) + return: A cute tensor with the appropriate dtype and layout + """ + import cutlass.torch as cutlass_torch + + cute_tensor = from_dlpack( + storage_tensor, assumed_align=16, force_tf32=dtype == TFloat32 + ) + # For float8 types, set element_type explicitly since storage is uint8 + if is_fp8_dtype(dtype): + cute_tensor.element_type = dtype + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + # For float8 types, convert data from float32 using GPU kernel + if is_fp8_dtype(dtype): + if source_f32_tensor is None: + raise ValueError("source_f32_tensor is required for fp8 types") + cute_tensor = cutlass_torch.convert_cute_tensor( + source_f32_tensor, cute_tensor, dtype, is_dynamic_layout=True + ) + return cute_tensor diff --git a/python/CuTeDSL/cutlass/utils/tmem_allocator.py b/python/CuTeDSL/cutlass/utils/tmem_allocator.py index 925f7005..37eef568 100644 --- a/python/CuTeDSL/cutlass/utils/tmem_allocator.py +++ b/python/CuTeDSL/cutlass/utils/tmem_allocator.py @@ -9,8 +9,8 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. +from math import log2, ceil from typing import Optional, Type, Union, List -from math import ceil, log2 import inspect from cutlass import const_expr @@ -29,7 +29,7 @@ from cutlass.cute.arch import get_max_tmem_alloc_cols, get_min_tmem_alloc_cols class TmemAllocator: - """A class for managing tensor memory allocation. + """A class for managing tensor memory allocation on GPUs. This class manages allocation/deallocation of tensor memory, including the mbarrier synchronization for two cta use case. @@ -81,6 +81,7 @@ class TmemAllocator: two_cta_tmem_dealloc_mbar_ptr: Optional[cute.Pointer] = None, *, arch: str = "sm_100", + dealloc_mbarrier_initialized: bool = False, loc=None, ip=None, ): @@ -126,7 +127,7 @@ class TmemAllocator: self._max_tmem_columns = get_max_tmem_alloc_cols(arch) # Init tmem dealloc mbarrier if two cta - if const_expr(self._is_two_cta): + if not dealloc_mbarrier_initialized and const_expr(self._is_two_cta): self._init_dealloc_mbarrier(loc=loc, ip=ip) def __extract_mlir_values__(self) -> list[ir.Value]: @@ -158,7 +159,8 @@ class TmemAllocator: self._is_two_cta, self._num_allocated_columns, new_two_cta_tmem_dealloc_mbar_ptr, - arch=self._arch, + arch=self._arch, # Preserve the architecture parameter + dealloc_mbarrier_initialized=True, ) @cute.jit diff --git a/python/CuTeDSL/prep_editable_install.py b/python/CuTeDSL/prep_editable_install.py index ac7d258f..b9f869c1 100644 --- a/python/CuTeDSL/prep_editable_install.py +++ b/python/CuTeDSL/prep_editable_install.py @@ -40,27 +40,6 @@ class CutlassDSLSetupError(Exception): pass -def get_package_spec(requirements_path: Optional[Path] = None) -> str: - """ - Return the pip requirement spec for nvidia-cutlass-dsl from requirements.txt. - - If anything goes wrong (file not found, parse failure, line missing), - return PACKAGE_NAME as a safe default. - """ - try: - req_path = requirements_path or Path(__file__).with_name("requirements.txt") - with open(req_path, "r", encoding="utf-8") as f: - for raw_line in f: - line = raw_line.strip() - if not line or line.startswith("#"): - continue - if line.lower().startswith(PACKAGE_NAME): - return line.split("#", 1)[0].strip() - except Exception: - pass - return PACKAGE_NAME - - def download_wheel(temp_dir: Path) -> Path: """ Download the nvidia-cutlass-dsl wheel to a temporary directory. @@ -74,10 +53,7 @@ def download_wheel(temp_dir: Path) -> Path: Raises: CutlassDSLSetupError: If download fails or wheel not found """ - # Resolve package spec from requirements, or fall back to PACKAGE_NAME - package_spec = get_package_spec() - - logger.info(f"Downloading {package_spec} wheel to {temp_dir}") + logger.info(f"Downloading {PACKAGE_NAME} wheel to {temp_dir}") try: subprocess.check_call( @@ -87,7 +63,7 @@ def download_wheel(temp_dir: Path) -> Path: "pip", "download", "--no-deps", - package_spec, + PACKAGE_NAME, "--dest", str(temp_dir), ], @@ -103,7 +79,7 @@ def download_wheel(temp_dir: Path) -> Path: raise CutlassDSLSetupError(error_msg) # Find the downloaded wheel file - wheel_pattern = f"*.whl" + wheel_pattern = f"{PACKAGE_NAME.replace('-', '_')}-*.whl" wheel_files = list(temp_dir.glob(wheel_pattern)) if not wheel_files: raise CutlassDSLSetupError( @@ -132,7 +108,7 @@ def extract_version_from_wheel(wheel_path: Path) -> str: # Construct version regex from package name # Wheel filename format: {package_name_with_underscores}-{version}-{python}-{abi}-{platform}.whl package_pattern = PACKAGE_NAME.replace("-", "_") - version_regex = rf"{re.escape(package_pattern)}-([^-]+)" + version_regex = rf"{re.escape(package_pattern)}-([^-]+)-" version_match = re.match(version_regex, wheel_filename) if version_match: @@ -156,7 +132,10 @@ def extract_version_from_wheel(wheel_path: Path) -> str: return dev_version else: - return "9.9.9.dev0" + raise CutlassDSLSetupError( + f"Could not parse version from wheel filename: {wheel_filename}" + ) + def extract_wheel_contents(wheel_path: Path, extract_dir: Path) -> None: """ diff --git a/python/CuTeDSL/pyproject.toml b/python/CuTeDSL/pyproject.toml index 66af717d..5ae73380 100644 --- a/python/CuTeDSL/pyproject.toml +++ b/python/CuTeDSL/pyproject.toml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual diff --git a/python/CuTeDSL/requirements-cu13.txt b/python/CuTeDSL/requirements-cu13.txt new file mode 100644 index 00000000..b49a23e7 --- /dev/null +++ b/python/CuTeDSL/requirements-cu13.txt @@ -0,0 +1,3 @@ +# Use `pip install -r requirements-cu13.txt` with the present file to install a +# wheel consistent with the present state of the github repository +nvidia-cutlass-dsl[cu13]==4.4.0 diff --git a/python/CuTeDSL/requirements.txt b/python/CuTeDSL/requirements.txt index 4c4fb948..bba4fe6d 100644 --- a/python/CuTeDSL/requirements.txt +++ b/python/CuTeDSL/requirements.txt @@ -1,3 +1,3 @@ # Use `pip install -r requirements.txt` with the present file to install a # wheel consistent with the present state of the github repository -nvidia-cutlass-dsl==4.4.0.dev0 +nvidia-cutlass-dsl==4.4.0 diff --git a/python/CuTeDSL/setup.sh b/python/CuTeDSL/setup.sh new file mode 100755 index 00000000..5428ece9 --- /dev/null +++ b/python/CuTeDSL/setup.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +################################################################################################# +# +# Copyright (c) 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. +# +################################################################################################# + +set -e + +# Get the directory where this script is located +SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" + +# Default to requirements.txt +REQUIREMENTS_FILE="requirements.txt" + +# Parse command line arguments +if [ $# -gt 0 ]; then + case "$1" in + --cu12) + REQUIREMENTS_FILE="requirements.txt" + echo "Installing CUDA 12 requirements..." + ;; + --cu13) + REQUIREMENTS_FILE="requirements-cu13.txt" + echo "Installing CUDA 13 requirements..." + ;; + --help|-h) + echo "Usage: $0 [--cu12|--cu13]" + echo " --cu12 Install requirements for CUDA 12 (default)" + echo " --cu13 Install requirements for CUDA 13" + exit 0 + ;; + *) + echo "Error: Unknown argument '$1'" + echo "Usage: $0 [--cu12|--cu13]" + exit 1 + ;; + esac +else + echo "Installing default requirements (CUDA 12)..." +fi + +# Check if requirements file exists +REQUIREMENTS_PATH="${SCRIPT_DIR}/${REQUIREMENTS_FILE}" +if [ ! -f "$REQUIREMENTS_PATH" ]; then + echo "Error: Requirements file not found: $REQUIREMENTS_PATH" + exit 1 +fi + +# Uninstall previous version of the CUTLASS DSL +echo "Trying to uninstall previous version of the CUTLASS DSL..." +pip uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu13 -y + +# Install requirements +echo "Installing from: $REQUIREMENTS_FILE" +pip install -r "$REQUIREMENTS_PATH" + +echo "Installation complete!" diff --git a/python/cutlass_library/heuristics_provider.py b/python/cutlass_library/heuristics_provider.py index e2d60437..aefd576e 100644 --- a/python/cutlass_library/heuristics_provider.py +++ b/python/cutlass_library/heuristics_provider.py @@ -54,6 +54,7 @@ class MatmulHeuristics: def __init__(self, gpu = None): import nvMatmulHeuristics + import inspect self.mmh_lib = nvMatmulHeuristics self.gpu = gpu @@ -62,13 +63,63 @@ class MatmulHeuristics: else: nvmmhInterfaceEx = self.mmh_lib.NvMatmulHeuristicsInterfaceEx - self.lh = nvmmhInterfaceEx( + # nvidia-matmul-heuristics 0.1.0.28 changed the API: + # - Constructor: removed 'load_discovery_implicitly' and 'gpu' params + # - GPU: now set via createHardwareDescriptor() + setHardwarePredefinedGpu() + # - setBackendValueProperty renamed to setBackendPropertyValue (simpler signature) + # - getEx: added hardware_descriptor parameter + init_params = set(inspect.signature(self.mmh_lib.NvMatmulHeuristicsInterfaceEx.__init__).parameters.keys()) + self._legacy_api = 'load_discovery_implicitly' in init_params + + init_kwargs = dict( backend=self.mmh_lib.NvMatmulHeuristicsTarget["CUTLASS3"], flags=self.mmh_lib.NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING, - load_discovery_implicitly=True, - gpu=self.mmh_lib.NvMatmulHeuristicsNvidiaGpu[self.gpu] if self.gpu else None ) + + if self._legacy_api: + # <= 0.1.0.27 + init_kwargs['gpu'] = self.mmh_lib.NvMatmulHeuristicsNvidiaGpu[self.gpu] if self.gpu else None + init_kwargs['load_discovery_implicitly'] = True + + self.lh = nvmmhInterfaceEx(**init_kwargs) + + # >= 0.1.0.28: gpu is set via hardware descriptor after construction, + # and passed to getEx() calls + self.hw_desc = None + if not self._legacy_api and self.gpu: + self.hw_desc = self.lh.createHardwareDescriptor() + if self.hw_desc is None: + raise RuntimeError("Failed to create hardware descriptor for GPU: " + self.gpu) + self.lh.setHardwarePredefinedGpu(self.hw_desc, self.mmh_lib.NvMatmulHeuristicsNvidiaGpu[self.gpu]) + self.backend = self.lh.createBackend(self.mmh_lib.NvMatmulHeuristicsTarget["CUTLASS3"]) + + if not self._legacy_api: + lh = self.lh + original_del = type(lh).__del__ + + def _safe_del(self_lh): + try: + original_del(self_lh) + except Exception: + pass + + type(lh).__del__ = _safe_del + + def __del__(self): + """Clean up resources in correct order before the library's __del__ runs.""" + try: + if hasattr(self, 'backend') and self.backend: + self.lh.destroyBackend(self.backend) + self.backend = None + if hasattr(self, 'hw_desc') and self.hw_desc: + self.lh.destroyHardwareDescriptor(self.hw_desc) + self.hw_desc = None + # Null out the handle so the library's __del__ skips nvMatmulHeuristicsDestroy + if hasattr(self, 'lh') and self.lh and hasattr(self.lh, 'handle'): + self.lh.handle = None + except Exception: + pass def _layout_from_cutlass(self, layouts): assert(len(layouts)==3) @@ -98,41 +149,45 @@ class MatmulHeuristics: else: return a_c + dtype_to_cublas[dtype_b] + dtype_to_cublas[dtype_c] + dtype_to_cublas[dtype_compute] + dtype_to_cublas[dtype_d] + def _set_backend_property(self, property, value): + """Compat wrapper: setBackendValueProperty (<=0.1.0.27) vs setBackendPropertyValue (>=0.1.0.28)""" + if self._legacy_api: + c_val = ctypes.c_int(value) + self.lh.setBackendValueProperty( + self.backend, property, + ctypes.byref(c_val), ctypes.sizeof(c_val) + ) + else: + self.lh.setBackendPropertyValue(self.backend, property, value) + def set_cta_div_n(self, div_n): - cta_n_div_requirement = ctypes.c_int(div_n) - self.lh.setBackendValueProperty( - self.backend, - self.mmh_lib.NvMatmulHeuristicsBackendProperty.CTA_TILE_N_DIV_REQUIREMENT, - ctypes.byref(cta_n_div_requirement), - ctypes.sizeof(cta_n_div_requirement) - ) + self._set_backend_property( + self.mmh_lib.NvMatmulHeuristicsBackendProperty.CTA_TILE_N_DIV_REQUIREMENT, div_n) def set_cta_div_m(self, div_m): - cta_m_div_requirement = ctypes.c_int(div_m) - self.lh.setBackendValueProperty( - self.backend, - self.mmh_lib.NvMatmulHeuristicsBackendProperty.CTA_TILE_M_DIV_REQUIREMENT, - ctypes.byref(cta_m_div_requirement), - ctypes.sizeof(cta_m_div_requirement) - ) + self._set_backend_property( + self.mmh_lib.NvMatmulHeuristicsBackendProperty.CTA_TILE_M_DIV_REQUIREMENT, div_m) def get_configs(self, m, n, k, batch_count, dtypes, layouts, align_a, align_b, voidC=False, use_fast_acc=True, count=1): - if use_fast_acc: - disable_fast_acc_for_fp8 = ctypes.c_int(0) - else: - disable_fast_acc_for_fp8 = ctypes.c_int(1) - self.lh.setBackendValueProperty( - self.backend, + self._set_backend_property( self.mmh_lib.NvMatmulHeuristicsBackendProperty.DISABLE_FAST_ACC_FOR_FP8, - ctypes.byref(disable_fast_acc_for_fp8), - ctypes.sizeof(disable_fast_acc_for_fp8) + 0 if use_fast_acc else 1 ) precision = self._precision_from_cutlass_dtypes(dtypes) layout = self._layout_from_cutlass(layouts) - matmul_problem = self.lh.makeNvMatmulHeuristicsProblem(m, n, k, layout, batch_count) - configs = self.lh.getEx(matmul_problem, count, self.backend, precision=precision) + if self._legacy_api: + matmul_problem = self.lh.makeNvMatmulHeuristicsProblem(m, n, k, layout, batch_count) + else: + # >= 0.1.0.28: takes (m,n,k) as a tuple + matmul_problem = self.lh.makeNvMatmulHeuristicsProblem((m, n, k), layout, batch_count) + + getEx_kwargs = dict(precision=precision) + if not self._legacy_api: + # >= 0.1.0.28: pass hardware descriptor to getEx + getEx_kwargs['hardware_descriptor'] = self.hw_desc + configs = self.lh.getEx(matmul_problem, count, self.backend, **getEx_kwargs) ret = [] for c in configs: diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt b/test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt index 70f6be42..4e57a807 100644 --- a/test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt +++ b/test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt @@ -70,3 +70,4 @@ cutlass_test_unit_gemm_device_add_executable( endif() add_subdirectory(narrow_precision) +add_subdirectory(extra_tests) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/CMakeLists.txt b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/CMakeLists.txt new file mode 100644 index 00000000..38b7124e --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/CMakeLists.txt @@ -0,0 +1,95 @@ +# Copyright (c) 2024 - 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. + +if (CUTLASS_NVCC_ARCHS MATCHES 100a) +add_custom_target( + cutlass_test_unit_gemm_device_sm100_tensorop_extra + DEPENDS + cutlass_test_unit_gemm_device_tensorop_runtime_datatype_sm100 + cutlass_test_unit_gemm_device_tensorop_epilogue_fusion_sm100 + cutlass_test_unit_gemm_device_fp8_tensorop_epilogue_fusion_sm100 + cutlass_test_unit_gemm_device_sm100_dense_and_bs_gemm_stage + cutlass_test_unit_gemm_device_tensorop_stride_batch_alpha_beta_sm100 +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_tensorop_runtime_datatype_sm100 + + # No batching of source to control compiler memory usage + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu + sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu + sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu + sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_tensorop_epilogue_fusion_sm100 + + # No batching of source to control compiler memory usage + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu + sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_fp8_tensorop_epilogue_fusion_sm100 + + # No batching of source to control compiler memory usage + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu + sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu + sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_sm100_dense_and_bs_gemm_stage + + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_f8_f8_f32_void_f8_stage.cu + sm100_gemm_f32_f32_f32_void_f32_stage.cu +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_tensorop_stride_batch_alpha_beta_sm100 + + # No batching of source to control compiler memory usage + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu +) +endif() diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu new file mode 100644 index 00000000..c2827b64 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu @@ -0,0 +1,310 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/thread/activation.h" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/// A Row B Col +TEST(SM100_Device_Gemm_f16t_f16n_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = void; + using ElementD = float; + using ElementCompute = float; + using ElementAccumulator = float; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::RowMajor; + using MmaTileShape_MNK = Shape<_256,_128,_128>; + using ClusterShape_MNK = Shape<_4,_4,_1>; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC, 16, + ElementD, GmemLayoutC, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, GmemLayoutA, 8, + ElementB, GmemLayoutB, 8, + ElementAccumulator, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + +/// A Col B Row +TEST(SM100_Device_Gemm_f16n_f16t_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = void; + using ElementD = float; + using ElementCompute = float; + using ElementAccumulator = float; + using GmemLayoutA = cutlass::layout::ColumnMajor; + using GmemLayoutB = cutlass::layout::RowMajor; + using GmemLayoutC = cutlass::layout::RowMajor; + using MmaTileShape_MNK = Shape<_256,_128,_128>; + using ClusterShape_MNK = Shape<_4,_4,_1>; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC, 16, + ElementD, GmemLayoutC, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, GmemLayoutA, 8, + ElementB, GmemLayoutB, 8, + ElementAccumulator, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + +/// A Row B Row +TEST(SM100_Device_Gemm_f16t_f16t_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = void; + using ElementD = float; + using ElementCompute = float; + using ElementAccumulator = float; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::RowMajor; + using GmemLayoutC = cutlass::layout::RowMajor; + using MmaTileShape_MNK = Shape<_256,_128,_128>; + using ClusterShape_MNK = Shape<_4,_4,_1>; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC, 16, + ElementD, GmemLayoutC, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, GmemLayoutA, 8, + ElementB, GmemLayoutB, 8, + ElementAccumulator, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + +/// A Col B Col +TEST(SM100_Device_Gemm_f16n_f16n_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = void; + using ElementD = float; + using ElementCompute = float; + using ElementAccumulator = float; + using GmemLayoutA = cutlass::layout::ColumnMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::RowMajor; + using MmaTileShape_MNK = Shape<_256,_128,_128>; + using ClusterShape_MNK = Shape<_4,_4,_1>; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC, 16, + ElementD, GmemLayoutC, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, GmemLayoutA, 8, + ElementB, GmemLayoutB, 8, + ElementAccumulator, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_bf16t_bf16t_bf32_void_f32n_tensor_op, 128x256x64_1x2x1) { + using ElementA = cutlass::bfloat16_t; + using LayoutA = cutlass::layout::RowMajor; + using ElementB = cutlass::bfloat16_t; + using LayoutB = cutlass::layout::RowMajor; + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape = Shape<_128,_128,_64>; + using ClusterShape = Shape<_1,_2,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, LayoutC, 8, + float, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.0); + EXPECT_TRUE(pass); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_b2b.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_b2b.cu new file mode 100644 index 00000000..1cd1ab3a --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_b2b.cu @@ -0,0 +1,257 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide SM100 back-to-back GEMM interface +*/ + + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/epilogue/thread/activation.h" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x_b2b.hpp" + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensor_op_gmma_f32_b2b, 128x64x64_1x1x1_128x256x64_1x2x1) { + + using CollectiveEpilogue0 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop0 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel0 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop0, + CollectiveEpilogue0>; + + using CollectiveEpilogue1 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop1 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel1 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop1, + CollectiveEpilogue1>; + + using Gemm0 = cutlass::gemm::device::GemmUniversalAdapter; + using Gemm1 = cutlass::gemm::device::GemmUniversalAdapter; + + EXPECT_TRUE((test::gemm::device::TestAllB2B())); +} + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensor_op_gmma_f32_b2b, 128x64x64_1x2x1_128x256x64_2x2x1) { + + using CollectiveEpilogue0 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop0 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel0 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop0, + CollectiveEpilogue0>; + + using CollectiveEpilogue1 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop1 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel1 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop1, + CollectiveEpilogue1>; + + using Gemm0 = cutlass::gemm::device::GemmUniversalAdapter; + using Gemm1 = cutlass::gemm::device::GemmUniversalAdapter; + + EXPECT_TRUE((test::gemm::device::TestAllB2B())); +} + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensor_op_gmma_f32_b2b, 128x128x64_2x2x1_256x128x64_2x1x1) { + + using CollectiveEpilogue0 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop0 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel0 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop0, + CollectiveEpilogue0>; + + using CollectiveEpilogue1 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop1 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel1 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop1, + CollectiveEpilogue1>; + + using Gemm0 = cutlass::gemm::device::GemmUniversalAdapter; + using Gemm1 = cutlass::gemm::device::GemmUniversalAdapter; + + EXPECT_TRUE((test::gemm::device::TestAllB2B())); +} + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_stream_k.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_stream_k.cu new file mode 100644 index 00000000..61ee4a83 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_stream_k.cu @@ -0,0 +1,235 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 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. + * + **************************************************************************************************/ +/*! \file + \brief Tests for device-wide GEMM interface with stream-K scheduling +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/kernel/tile_scheduler.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +using namespace cute; + +TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 128x256x64_1x2x1) { + using ElementA = cutlass::half_t; + using LayoutA = cutlass::layout::RowMajor; + using ElementB = cutlass::half_t; + using LayoutB = cutlass::layout::RowMajor; + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_128,_64>; + using ClusterShape_MNK = Shape<_1,_2,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Testbed = Testbed3x; + bool result = TestSmall(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048}); + EXPECT_TRUE(result); +} + +TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 256x128x64_2x1x1) { + using ElementA = cutlass::half_t; + using LayoutA = cutlass::layout::RowMajor; + using ElementB = cutlass::half_t; + using LayoutB = cutlass::layout::RowMajor; + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_128,_64>; + using ClusterShape_MNK = Shape<_2,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Testbed = Testbed3x; + bool result = TestSmall(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048}); + EXPECT_TRUE(result); +} + +TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 256x256x64_2x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::RowMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_128,_64>; + using ClusterShape_MNK = Shape<_2,_2,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Testbed = Testbed3x; + bool result = TestSmall(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048}); + EXPECT_TRUE(result); +} + +/////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_stream_k, 256x128x64_2x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_64,_64>; + using ClusterShape_MNK = Shape<_2,_4,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Testbed = Testbed3x; + bool result = TestSmall(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048}); + EXPECT_TRUE(result); +} + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_swap_ab_bias_relu.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_swap_ab_bias_relu.cu new file mode 100644 index 00000000..872aafea --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_swap_ab_bias_relu.cu @@ -0,0 +1,168 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x64 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensorop_1cta_f32_bias_relu_aux, 128x256x64_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using ElementA = cutlass::half_t; + using ElementB = cutlass::half_t; + using ElementC = cutlass::half_t; + using ElementD = cutlass::half_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAux = cutlass::uint1b_t; + using MmaTileShape = cute::Shape<_128,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBiasEltActAux< + LayoutC, cutlass::epilogue::thread::Clamp, ElementD, ElementCompute, ElementAux, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensorop_2cta_f32_bias_gelu_aux, 128x1024x64_2x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::half_t; + using ElementB = cutlass::half_t; + using ElementC = cutlass::half_t; + using ElementD = cutlass::half_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementAux = cutlass::half_t; + using ElementBias = cutlass::half_t; + using MmaTileShape = Shape<_128,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBiasEltActAux< + LayoutC, cutlass::epilogue::thread::GELU, ElementD, ElementCompute, ElementAux, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestSmallFusion(2, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_tensor_op_f32.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_tensor_op_f32.cu new file mode 100644 index 00000000..d514ba32 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_tensor_op_f32.cu @@ -0,0 +1,96 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/thread/activation.h" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_f16t_f16t_f32_void_f16n_tensor_op, 128x256x64_1x2x1) { + using ElementA = cutlass::half_t; + using LayoutA = cutlass::layout::RowMajor; + using ElementB = cutlass::half_t; + using LayoutB = cutlass::layout::RowMajor; + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_128,_64>; + using ClusterShape_MNK = Shape<_1,_2,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.0); + EXPECT_TRUE(pass); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_void_f16_stage.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_void_f16_stage.cu new file mode 100644 index 00000000..651a25c9 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_void_f16_stage.cu @@ -0,0 +1,233 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Test for sm100 dense gemm stage +*/ +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/numeric_types.h" +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +namespace cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x128x64_1x1x1_0_tnt_align8_1sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::half_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x256x64_1x1x1_0_tnt_align8_1sm_stage4 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::half_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x128x64_2x1x1_0_tnt_align8_2sm_stage8 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::half_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x256x64_2x1x1_0_tnt_align8_2sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::half_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +TEST(cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x128x64_1x1x1_0_tnt_align8_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x128x64_1x1x1_0_tnt_align8_1sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x256x64_1x1x1_0_tnt_align8_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x256x64_1x1x1_0_tnt_align8_1sm_stage4::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 4); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x128x64_2x1x1_0_tnt_align8_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x128x64_2x1x1_0_tnt_align8_2sm_stage8::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 8); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x256x64_2x1x1_0_tnt_align8_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x256x64_2x1x1_0_tnt_align8_2sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f32_f32_f32_void_f32_stage.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f32_f32_f32_void_f32_stage.cu new file mode 100644 index 00000000..ae7681f8 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f32_f32_f32_void_f32_stage.cu @@ -0,0 +1,234 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Test for sm100 dense gemm stage +*/ + +#include "../../../../common/cutlass_unit_test.h" +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/numeric_types.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +namespace cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x128x32_1x1x1_0_tnt_align4_1sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + float, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::ColumnMajor, 4, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x256x32_1x1x1_0_tnt_align4_1sm_stage4 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + float, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::ColumnMajor, 4, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x128x32_2x1x1_0_tnt_align4_2sm_stage8 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + float, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::ColumnMajor, 4, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x256x32_2x1x1_0_tnt_align4_2sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + float, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::ColumnMajor, 4, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +TEST(cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x128x32_1x1x1_0_tnt_align4_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x128x32_1x1x1_0_tnt_align4_1sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x256x32_1x1x1_0_tnt_align4_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x256x32_1x1x1_0_tnt_align4_1sm_stage4::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 4); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x128x32_2x1x1_0_tnt_align4_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x128x32_2x1x1_0_tnt_align4_2sm_stage8::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 8); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x256x32_2x1x1_0_tnt_align4_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x256x32_2x1x1_0_tnt_align4_2sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu new file mode 100644 index 00000000..8397d294 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu @@ -0,0 +1,154 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_e2m1t_e2m1n_f32t_tensorop_2sm_f32_runtime_datatype, 512x512x128_4x4x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::RowMajor, 128, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E2M1, cute::UMMA::MXF8F6F4Format::E2M1); + EXPECT_TRUE(pass); + +} + + +TEST(SM100_Device_Gemm_e2m1t_e2m1n_f32t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::RowMajor, 128, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E2M1, cute::UMMA::MXF8F6F4Format::E2M1); + EXPECT_TRUE(pass); +} + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu new file mode 100644 index 00000000..af8316fe --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu @@ -0,0 +1,154 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_e3m2t_e2m3n_f32t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float6_t, cutlass::layout::RowMajor, 128, + cutlass::type_erased_dynamic_float6_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E3M2, cute::UMMA::MXF8F6F4Format::E2M3); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e3m2t_e2m3n_f32t_tensorop_1sm_f32_runtime_datatype, 512x512x128_4x4x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float6_t, cutlass::layout::RowMajor, 128, + cutlass::type_erased_dynamic_float6_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E3M2, cute::UMMA::MXF8F6F4Format::E2M3); + EXPECT_TRUE(pass); + +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu new file mode 100644 index 00000000..97e3442c --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu @@ -0,0 +1,107 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_e4m3t_e2m1n_f32t_tensorop_2sm_f32_runtime_datatype, 256x128x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized2Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E2M1); + EXPECT_TRUE(pass); + +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f32_void_f8_stage.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f32_void_f8_stage.cu new file mode 100644 index 00000000..7f091314 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f32_void_f8_stage.cu @@ -0,0 +1,234 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Test for sm100 gemm stage +*/ + +#include "../../../../common/cutlass_unit_test.h" +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/numeric_types.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +namespace cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x128x128_1x1x1_0_tnt_align16_1sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x256x128_1x1x1_0_tnt_align16_1sm_stage4 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveoutEpi, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x128x128_2x1x1_0_tnt_align16_2sm_stage9 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x256x128_2x1x1_0_tnt_align16_2sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +TEST(cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x128x128_1x1x1_0_tnt_align16_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x128x128_1x1x1_0_tnt_align16_1sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} + +TEST(cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x256x128_1x1x1_0_tnt_align16_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x256x128_1x1x1_0_tnt_align16_1sm_stage4::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 4); +} + +TEST(cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x128x128_2x1x1_0_tnt_align16_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x128x128_2x1x1_0_tnt_align16_2sm_stage9::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 9); +} + +TEST(cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x256x128_2x1x1_0_tnt_align16_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x256x128_2x1x1_0_tnt_align16_2sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu new file mode 100644 index 00000000..17428cb2 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu @@ -0,0 +1,319 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 128x128x128 ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_gelu, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ScaledGELU_taylor, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_gelu, 64x128x128_1x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_64,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ScaledGELU_taylor, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_gelu, 256x128x128_2x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ScaledGELU_taylor, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_gelu, 512x512x128_4x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ScaledGELU_taylor, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_dgelu, 512x512x128_4x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombDeEltAct< + LayoutC, cutlass::epilogue::thread::dGELU, ElementD, ElementCompute, ElementD>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu new file mode 100644 index 00000000..c801f6d9 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu @@ -0,0 +1,328 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 128x128x128 ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu_amax_aux, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::uint1b_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3t_tensorop_1sm_f32_colbias_relu_amax_aux, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::uint1b_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerColBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_gelu_amax_aux, 64x128x128_1x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::float_e4m3_t; + using MmaTileShape = cute::Shape<_64,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::GELU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_gelu_amax_aux, 256x128x128_2x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::float_e4m3_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::GELU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_gelu_amax_aux, 512x512x128_4x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::float_e4m3_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::GELU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu new file mode 100644 index 00000000..3bdb269d --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu @@ -0,0 +1,365 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 128x128x128 ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu_beta0, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3t_tensorop_1sm_f32_colbias_relu, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerColBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu, 64x128x128_1x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_64,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_relu, 256x128x128_2x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_relu, 512x512x128_4x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu new file mode 100644 index 00000000..b2146f53 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu @@ -0,0 +1,295 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_e5m2t_e4m3n_e4m3t_tensorop_2sm_f32_runtime_datatype, 256x128x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + cutlass::float_e4m3_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E4M3); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e5m2t_e4m3n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + cutlass::float_e4m3_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E4M3); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e4m3t_e5m2n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + cutlass::float_e4m3_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E5M2); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + cutlass::float_e4m3_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E4M3); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e5m2t_e5m2n_e5m2t_tensorop_2sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e5m2_t, cutlass::layout::RowMajor, 16, + cutlass::float_e5m2_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e5m2_t, + float, + cutlass::float_e5m2_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E5M2); + EXPECT_TRUE(pass); + +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu new file mode 100644 index 00000000..bb90bf87 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu @@ -0,0 +1,219 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////// Test Batch alpha and beta ////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1cta_s32_batch_alpha_beta, 128x64x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + + using FusionOperation = cutlass::epilogue::fusion::LinearCombination< + ElementD, + ElementCompute, + ElementC, + ElementBias + >; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); // beta is [1.0, 2.0] + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu_batch_alpha_beta, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); // beta is [0.5, 1.5] + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu__batch_alpha_beta0, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, -1.0); // beta is [-1.0, 0.0] + EXPECT_TRUE(pass); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu new file mode 100644 index 00000000..c5a1c46d --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu @@ -0,0 +1,278 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(CUTLASS_SM100_FAMILY_ARCHS_ENABLED)) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x128 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_bias_relu, 128x64x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLu, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x128 4x2x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_bias_relu, 512x128x128_4x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLu, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 64x256x128 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s32n_tensorop_1cta_s32_bias_relu, 64x256x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int32_t; + using ElementD = int32_t; + using ElementAccumulator = int32_t; + using ElementCompute = int32_t; + using ElementBias = int32_t; + using MmaTileShape = cute::Shape<_64,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLu, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 64x256x128 2x4x1 TMEM 2x2 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_2cta_s32_bias_relu, 128x1024x128_2x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = Shape<_128,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLu, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} + +#endif // #if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(CUTLASS_SM100_FAMILY_ARCHS_ENABLED)) + diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu new file mode 100644 index 00000000..7ac363d0 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu @@ -0,0 +1,343 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 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. + * + **************************************************************************************************/ + +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(CUTLASS_SM100_FAMILY_ARCHS_ENABLED)) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////// Test Vector alpha and vector beta ////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x128 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_vector_alpha_beta, 128x64x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x128 4x2x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_vector_alpha_beta, 512x128x128_4x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 64x256x128 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_vector_alpha_beta, 64x256x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_64,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 64x256x128 2x4x1 TMEM 2x2 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_vector_alpha_beta, 128x1024x128_2x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_64,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +//////////////////////////////// Dynamic vector/scalar broadcast /////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_dynamic_vector_alpha_beta, 128x64x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + constexpr bool force_legacy_epilogue = false; + constexpr bool apply_alignment_offset = false; + // non-batched host scalar, beta 0 + EXPECT_TRUE((TestSmallFusion(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_HOST, VectorScale::DISABLED))); + // non-batched host scalar, beta 1 + EXPECT_TRUE((TestSmallFusion(1.0, 1.0, CheckEquality::EXACT, ScalarLoc::ON_HOST, VectorScale::DISABLED))); + // batched device scalar, beta 0 + EXPECT_TRUE((TestSmallFusion(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::DISABLED))); + // batched device scalar, beta 1 + EXPECT_TRUE((TestSmallFusion(1.0, 1.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::DISABLED))); +} + +#endif // #if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(CUTLASS_SM100_FAMILY_ARCHS_ENABLED))